{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.gridspec as gridspec\n",
    "import cv2\n",
    "import os\n",
    "import numpy as np\n",
    "import glob\n",
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def random_shearing(img, num, border):\n",
    "    rows = img.shape[0]\n",
    "    cols = img.shape[1]\n",
    "    if num == 0:\n",
    "        pts1 = np.float32([[5,5],[20,5],[2,20]])\n",
    "        pts2 = np.float32([[10,10],[20,5],[5,25]])\n",
    "    elif num == 1:\n",
    "        pts1 = np.float32([[5,5],[15,5],[2,20]])\n",
    "        pts2 = np.float32([[5,10],[10,10],[5,25]])\n",
    "    elif num == 2:\n",
    "        pts1 = np.float32([[5,5],[15,5],[5,20]])\n",
    "        pts2 = np.float32([[5,10],[10,10],[5,25]])\n",
    "    elif num == 3:\n",
    "        pts1 = np.float32([[5,5],[10,5],[2,20]])\n",
    "        pts2 = np.float32([[5,10],[10,10],[5,25]])\n",
    "    elif num == 4:\n",
    "        pts1 = np.float32([[5,5],[10,5],[2,20]])\n",
    "        pts2 = np.float32([[5,10],[10,10],[5,30]])\n",
    "    else:\n",
    "        pts1 = np.float32([[5,5],[10,5],[10,20]])\n",
    "        pts2 = np.float32([[5,10],[10,10],[5,30]])\n",
    "    M = cv2.getAffineTransform(pts1,pts2)\n",
    "    return cv2.warpAffine(img, M, (cols,rows), borderValue=border)\n",
    "\n",
    "def random_rotation(img, degree, border):\n",
    "    rows = img.shape[0]\n",
    "    cols = img.shape[1]\n",
    "    M = cv2.getRotationMatrix2D((cols/2,rows/2),degree,1)\n",
    "    return cv2.warpAffine(img,M,(cols,rows), borderValue=border)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def transform_image(img, skt, ang_range, shear_range, trans_range):\n",
    "    '''\n",
    "    This function transforms images to generate new images.\n",
    "    The function takes in following arguments,\n",
    "    1- Image\n",
    "    2- ang_range: Range of angles for rotation\n",
    "    3- shear_range: Range of values to apply affine transform to\n",
    "    4- trans_range: Range of values to apply translations over.\n",
    "\n",
    "    A Random uniform distribution is used to generate different parameters for transformation\n",
    "\n",
    "    '''\n",
    "    # Rotation\n",
    "\n",
    "    ang_rot = np.random.uniform(ang_range)-ang_range/2\n",
    "    rows,cols,ch = img.shape    \n",
    "    Rot_M = cv2.getRotationMatrix2D((cols/2,rows/2),ang_rot,1)\n",
    "\n",
    "    # Translation\n",
    "    tr_x = trans_range*np.random.uniform()-trans_range/2\n",
    "    tr_y = trans_range*np.random.uniform()-trans_range/2\n",
    "    Trans_M = np.float32([[1,0,tr_x],[0,1,tr_y]])\n",
    "\n",
    "    # Shear\n",
    "    pts1 = np.float32([[5,5],[20,5],[5,20]])\n",
    "\n",
    "    pt1 = 5+shear_range*np.random.uniform()-shear_range/2\n",
    "    pt2 = 20+shear_range*np.random.uniform()-shear_range/2\n",
    "\n",
    "    pts2 = np.float32([[pt1,5],[pt2,pt1],[5,pt2]])\n",
    "\n",
    "    shear_M = cv2.getAffineTransform(pts1,pts2)\n",
    "\n",
    "    # Border\n",
    "    idx = 0\n",
    "    border_img = tuple([int(img[idx][0][0]), int(img[idx][0][1]), int(img[idx][0][2])])\n",
    "    border_skt = tuple([int(skt[0][0][0]), int(skt[0][0][1]), int(skt[0][0][2])])\n",
    "    \n",
    "    img = cv2.warpAffine(img,Rot_M,(cols,rows), borderValue=border_img)\n",
    "    img = cv2.warpAffine(img,Trans_M,(cols,rows), borderValue=border_img)\n",
    "    img = cv2.warpAffine(img,shear_M,(cols,rows), borderValue=border_img)\n",
    "    \n",
    "    skt = cv2.warpAffine(skt,Rot_M,(cols,rows), borderValue=border_skt)\n",
    "    skt = cv2.warpAffine(skt,Trans_M,(cols,rows), borderValue=border_skt)\n",
    "    skt = cv2.warpAffine(skt,shear_M,(cols,rows), borderValue=border_skt)\n",
    "\n",
    "    return img, skt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "sketch_dir = 'Dataset/Augmented sketch/'\n",
    "photo_dir = 'Dataset/Augmented photo/'\n",
    "\n",
    "if not os.path.exists(sketch_dir):\n",
    "    os.mkdir(sketch_dir)\n",
    "\n",
    "if not os.path.exists(photo_dir):\n",
    "    os.mkdir(photo_dir)\n",
    "\n",
    "p_filenames = glob.glob('Dataset/CUHK/Training photo/*')\n",
    "s_filenames = glob.glob('Dataset/CUHK/Training sketch/*')\n",
    "\n",
    "counter = 0\n",
    "for i in range(len(p_filenames)):\n",
    "    im = cv2.imread(p_filenames[i])\n",
    "    sk = cv2.imread(s_filenames[i])\n",
    "\n",
    "    for j in range(200):\n",
    "        img, skt = transform_image(im, sk, 40, 10, 10)\n",
    "\n",
    "        cv2.imwrite(photo_dir + str(counter) + '.jpg', img)\n",
    "        cv2.imwrite(sketch_dir + str(counter) + '.jpg', skt)\n",
    "\n",
    "        counter += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}