LOL_prepare_data.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "from sklearn.model_selection import train_test_split\n",
    "import random\n",
    "\n",
    "random.seed(98052)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## split for train and test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# input directories\n",
    "data_dir = '<raw data directory>'\n",
    "sub_folders = ['contourbufferstrips', 'fieldborders', 'filterstrips', 'riparian', 'terraces', 'waterways', 'wsb']\n",
    "output_dir = '<augmented data directory>'\n",
    "os.makedirs(output_dir, exist_ok =True)\n",
    "\n",
    "RATIO = 0.3 # 70% as training, 30% as testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# split data into train/jpg, train/polygon, test/jpg, test/polygon folders\n",
    "for sub_folder in sub_folders:\n",
    "    \n",
    "    image_dir = os.path.join(data_dir, sub_folder)\n",
    "\n",
    "    raw_jpg = os.path.join(image_dir, 'jpg')\n",
    "    raw_poly = os.path.join(image_dir, 'polygon')\n",
    "\n",
    "    file_names = next(os.walk(raw_jpg))[2]\n",
    "    train, test = train_test_split(file_names, test_size=RATIO, random_state = 98052)\n",
    "\n",
    "    test_dir = os.path.join(data_dir, sub_folder+'_test')\n",
    "    os.makedirs(test_dir, exist_ok = True)\n",
    "\n",
    "    test_jpg = os.path.join(test_dir, 'jpg')\n",
    "    os.makedirs(test_jpg, exist_ok = True)\n",
    "    test_poly = os.path.join(test_dir, 'polygon')\n",
    "    os.makedirs(test_poly, exist_ok = True)\n",
    "\n",
    "    # move 30% dataset to the test folers\n",
    "    for f_test in test:\n",
    "        path1_jpg = os.path.join(raw_jpg, f_test)\n",
    "        path2_jpg = os.path.join(test_jpg, f_test)\n",
    "        shutil.move(path1_jpg, path2_jpg)   \n",
    "\n",
    "        fn_img, ext = os.path.splitext(os.path.basename(f_test))\n",
    "        for name in next(os.walk(raw_poly))[2]:\n",
    "            if name.startswith(fn_img):        \n",
    "                path1_poly = os.path.join(raw_poly, name)\n",
    "                path2_poly = os.path.join(test_poly, name)\n",
    "                shutil.move(path1_poly, path2_poly)\n",
    "                \n",
    "    if bool(set(next(os.walk(raw_jpg))[2]) & set(next(os.walk(test_jpg))[2])):\n",
    "        print (\"there is overlap between training jpg and test jpg\")\n",
    "    else:\n",
    "        pass\n",
    "    \n",
    "    if bool(set(next(os.walk(raw_poly))[2]) & set(next(os.walk(test_poly))[2])):\n",
    "        print (\"there is overlap between training polygon and test polygon\")\n",
    "    else:\n",
    "        pass\n",
    "    \n",
    "    # rename raw image folder to training image folder\n",
    "    os.rename(image_dir, image_dir+'_train')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## merge for train and test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "subcategory_folders = ['waterways', 'fieldborders', 'terraces', 'wsb']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# create corresponding folders\n",
    "input_train = os.path.join(output_dir, 'train')\n",
    "input_test = os.path.join(output_dir, 'test')\n",
    "\n",
    "os.makedirs(input_train, exist_ok = True)\n",
    "os.makedirs(input_test, exist_ok = True)\n",
    "\n",
    "input_train_jpg = os.path.join(input_train, 'jpg')\n",
    "os.makedirs(input_train_jpg, exist_ok = True)\n",
    "input_train_poly = os.path.join(input_train, 'polygon')\n",
    "os.makedirs(input_train_poly, exist_ok = True)\n",
    "\n",
    "input_test_jpg = os.path.join(input_test, 'jpg')\n",
    "os.makedirs(input_test_jpg, exist_ok = True)\n",
    "input_test_poly = os.path.join(input_test, 'polygon')\n",
    "os.makedirs(input_test_poly, exist_ok = True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# merge training dataset\n",
    "\n",
    "for folder in subcategory_folders:\n",
    "    folder_dir = folder+'_train'\n",
    "    \n",
    "    image_dir = os.path.join(data_dir, folder_dir)\n",
    "    raw_jpg = os.path.join(image_dir, 'jpg')\n",
    "    raw_poly = os.path.join(image_dir, 'polygon')\n",
    "    \n",
    "    for jpg in next(os.walk(raw_jpg))[2]:\n",
    "        path1_jpg = os.path.join(raw_jpg, jpg)\n",
    "        path2_jpg = os.path.join(input_train_jpg, jpg)\n",
    "        shutil.copy(path1_jpg, path2_jpg)   \n",
    "    \n",
    "    for poly in next(os.walk(raw_poly))[2]:\n",
    "        path1_poly = os.path.join(raw_poly, poly)\n",
    "        path2_poly = os.path.join(input_train_poly, poly)\n",
    "        shutil.copy(path1_poly, path2_poly)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# merge testing dataset\n",
    "\n",
    "for folder in subcategory_folders:\n",
    "    folder_dir = folder+'_test'\n",
    "    \n",
    "    image_dir = os.path.join(data_dir, folder_dir)\n",
    "    raw_jpg = os.path.join(image_dir, 'jpg')\n",
    "    raw_poly = os.path.join(image_dir, 'polygon')\n",
    "    \n",
    "    for jpg in next(os.walk(raw_jpg))[2]:\n",
    "        path1_jpg = os.path.join(raw_jpg, jpg)\n",
    "        path2_jpg = os.path.join(input_test_jpg, jpg)\n",
    "        shutil.copy(path1_jpg, path2_jpg)   \n",
    "    \n",
    "    for poly in next(os.walk(raw_poly))[2]:\n",
    "        path1_poly = os.path.join(raw_poly, poly)\n",
    "        path2_poly = os.path.join(input_test_poly, poly)\n",
    "        shutil.copy(path1_poly, path2_poly)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "if bool(set(next(os.walk(input_train_poly))[2]) & set(next(os.walk(input_test_poly))[2])):\n",
    "    print (\"there is overlap between training jpg and test jpg\")\n",
    "else:\n",
    "    pass\n",
    "\n",
    "if bool(set(next(os.walk(input_train_jpg))[2]) & set(next(os.walk(input_test_jpg))[2])):\n",
    "    print (\"there is overlap between training polygon and test polygon\")\n",
    "else:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## data aug"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import cv2\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "import imgaug as ia\n",
    "import matplotlib.pyplot as plt\n",
    "from imgaug import augmenters as iaa\n",
    "from matplotlib.pyplot import imshow\n",
    "import matplotlib.gridspec as gridspec\n",
    "%matplotlib inline "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### define augmentation directories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "image_location = os.path.join(output_dir, 'train/jpg')\n",
    "poly_location = os.path.join(output_dir, 'train/polygon')\n",
    "\n",
    "category_names = ['contourbufferstrips', 'fieldborders', 'filterstrips', 'riparian', 'terraces', 'waterways', 'wsb']\n",
    "\n",
    "aug_dir = os.path.join(output_dir, 'train_aug')\n",
    "os.makedirs(aug_dir, exist_ok =True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### helper functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def load_image(image):\n",
    "    img = Image.open(image)\n",
    "    data = np.asarray(img, dtype = 'uint8')\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def load_mask(poly):\n",
    "    img = Image.open(poly).convert('L')\n",
    "    data = np.asarray(img, dtype = 'uint8')\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def visualize_imgs(dim_x, dim_y, images):\n",
    "    gs = gridspec.GridSpec(dim_x, dim_y, top=1., bottom=0., right=1., left=0., hspace=0.,\n",
    "        wspace=0.) # the size of grid will be adjusted to the number of augmented images\n",
    "\n",
    "    for i in range(len(images)):\n",
    "        ax = plt.subplot(gs[i])\n",
    "        ax.imshow(images[i])\n",
    "        ax.set_xticks([])\n",
    "        ax.set_yticks([])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def seq_functions(aug_method):\n",
    "     \n",
    "    seq_fliplr = iaa.Sequential([\n",
    "        iaa.Fliplr(1.0), # horizontal flips\n",
    "    ], random_order=False) \n",
    "    \n",
    "    seq_flipud = iaa.Sequential([\n",
    "        iaa.Flipud(1.0), # vertical flips\n",
    "    ], random_order=False) \n",
    "    \n",
    "    seq_flip = iaa.Sequential([\n",
    "        iaa.Fliplr(1.0), # horizontal flips\n",
    "        iaa.Flipud(1.0), # vertical flips\n",
    "    ], random_order=False) \n",
    "        \n",
    "    seq_rotate45 = iaa.Sequential([\n",
    "            iaa.Affine(\n",
    "            rotate=45,\n",
    "        )\n",
    "    ], random_order=False) \n",
    "    \n",
    "    seq_rotate90 = iaa.Sequential([\n",
    "            iaa.Affine(\n",
    "            rotate=90,\n",
    "        )\n",
    "    ], random_order=False) \n",
    "    \n",
    "    seq_rotate135 = iaa.Sequential([\n",
    "            iaa.Affine(\n",
    "            rotate=135,\n",
    "        )\n",
    "    ], random_order=False) \n",
    "    \n",
    "    seq_dict = {'fliplr': seq_fliplr, 'flipud': seq_flipud, 'flip': seq_flip, \\\n",
    "                'rotate45': seq_rotate45, 'rotate90': seq_rotate90, 'rotate135': seq_rotate135,}\n",
    "    \n",
    "    seq = seq_dict[aug_method]\n",
    "    \n",
    "    return seq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def aug_mask_image(_mask, image, num_aug, operation_name):\n",
    "        \n",
    "    masks1 = np.zeros((num_aug, _mask.shape[0], _mask.shape[1]))\n",
    "    masks2 = np.zeros((num_aug, _mask.shape[0], _mask.shape[1]))\n",
    "    masks3 = np.zeros((num_aug, _mask.shape[0], _mask.shape[1]))\n",
    "    masks4 = np.zeros((num_aug, _mask.shape[0], _mask.shape[1]))\n",
    "    masks5 = np.zeros((num_aug, _mask.shape[0], _mask.shape[1]))\n",
    "    masks6 = np.zeros((num_aug, _mask.shape[0], _mask.shape[1]))\n",
    "    masks7 = np.zeros((num_aug, _mask.shape[0], _mask.shape[1]))\n",
    "\n",
    "    _img = image \n",
    "    images = np.zeros((num_aug, _img.shape[0], _img.shape[1], 3), 'uint8') # dtype uint8.\n",
    "\n",
    "    for i in range(num_aug):\n",
    "        masks1[i] = _mask[:,:,0]\n",
    "        masks2[i] = _mask[:,:,1]\n",
    "        masks3[i] = _mask[:,:,2]\n",
    "        masks4[i] = _mask[:,:,3]\n",
    "        masks5[i] = _mask[:,:,4]\n",
    "        masks6[i] = _mask[:,:,5]\n",
    "        masks7[i] = _mask[:,:,6]\n",
    "        \n",
    "        images[i] = _img\n",
    "    \n",
    "    seq = seq_functions(operation_name)\n",
    "\n",
    "    seq_det = seq.to_deterministic() # set the random parameters to derterministic\n",
    "    \n",
    "    images_aug = seq_det.augment_images(images)\n",
    "    \n",
    "    masks1_aug = seq_det.augment_images(masks1)\n",
    "    masks2_aug = seq_det.augment_images(masks2)\n",
    "    masks3_aug = seq_det.augment_images(masks3)\n",
    "    masks4_aug = seq_det.augment_images(masks4)\n",
    "    masks5_aug = seq_det.augment_images(masks5)\n",
    "    masks6_aug = seq_det.augment_images(masks6)\n",
    "    masks7_aug = seq_det.augment_images(masks7)\n",
    "    \n",
    "    return images_aug, masks1_aug, masks2_aug, masks3_aug, masks4_aug, masks5_aug, masks6_aug, masks7_aug,"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### data aug operations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "operations = ['fliplr', 'flipud', 'flip', 'rotate90']\n",
    "\n",
    "for operation_name in operations:\n",
    "\n",
    "    sub_operation = os.path.join(aug_dir, operation_name)\n",
    "    os.makedirs(sub_operation, exist_ok = True)\n",
    "    \n",
    "    for image_name in next(os.walk(image_location))[2]:\n",
    "        image = os.path.join(image_location, image_name)\n",
    "\n",
    "        num_aug = 1\n",
    "        _img = load_image(image)\n",
    "\n",
    "        mask = np.zeros([256, 256, len(category_names)])\n",
    "        fn_img, ext = os.path.splitext(os.path.basename(image_name))\n",
    "\n",
    "        mask_names = []\n",
    "        for i in range(len(category_names)):\n",
    "            if fn_img.split('_')[0] == category_names[i]:\n",
    "                poly = os.path.join(poly_location, image_name)\n",
    "            else:\n",
    "                poly = os.path.join(poly_location, fn_img + '_' + category_names[i] + '.jpg')\n",
    "            mask[:,:,i:i+1] = np.expand_dims(load_mask(poly), axis = 2)\n",
    "            mask_names.append(poly)\n",
    "\n",
    "\n",
    "        images_aug, masks1_aug, masks2_aug, masks3_aug, masks4_aug, masks5_aug, masks6_aug, masks7_aug = \\\n",
    "        aug_mask_image(mask, _img, num_aug, operation_name)\n",
    "\n",
    "        masks = [masks1_aug, masks2_aug, masks3_aug, masks4_aug, masks5_aug, masks6_aug, masks7_aug]\n",
    "        \n",
    "        output_jpg_path = os.path.join(sub_operation, 'jpg')\n",
    "        output_poly_path = os.path.join(sub_operation, 'polygon')\n",
    "        os.makedirs(output_jpg_path, exist_ok = True)\n",
    "        os.makedirs(output_poly_path, exist_ok = True)\n",
    "        \n",
    "        for i in range(num_aug):\n",
    "            fn1,fn2, fn3, ex = image_name.split('_')\n",
    "            file_name = fn1 + '_' + fn2 + '_' + fn3 + '_' + operation_name + str(i) + '_' + ex\n",
    "            output_name = os.path.join(output_jpg_path, file_name)\n",
    "            cv2.imwrite(output_name, images_aug[i])\n",
    "\n",
    "        for i in range(len(mask_names)):\n",
    "            fn_img, ext = os.path.splitext(os.path.basename(mask_names[i]))\n",
    "            mask_array = masks[i]\n",
    "\n",
    "            for j in range(num_aug):\n",
    "                if fn_img.split('_')[-1] == \"merged\":\n",
    "                    fn1,fn2, fn3, ex = mask_names[i].split('/')[-1].split('_')\n",
    "                    file_name = fn1 + '_' + fn2 + '_' + fn3 + '_' + operation_name + str(j) + '_' + ex\n",
    "                else:\n",
    "                    fn1,fn2, fn3, fn4, ex = mask_names[i].split('/')[-1].split('_')\n",
    "                    file_name = fn1 + '_' + fn2 + '_' + fn3 + '_' +  operation_name + str(j) + '_' + fn4 + '_'+ ex\n",
    "                output_name = os.path.join(output_poly_path, file_name)\n",
    "    \n",
    "                cv2.imwrite(output_name, mask_array[j])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### merge augmentaion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# merge training dataset\n",
    "\n",
    "for folder_dir in operations:\n",
    "    \n",
    "    image_dir = os.path.join(aug_dir, folder_dir)\n",
    "    raw_jpg = os.path.join(image_dir, 'jpg')\n",
    "    raw_poly = os.path.join(image_dir, 'polygon')\n",
    "    \n",
    "    output_jpg = os.path.join(aug_dir, 'jpg')\n",
    "    output_poly = os.path.join(aug_dir, 'polygon')\n",
    "    os.makedirs(output_jpg, exist_ok =True)\n",
    "    os.makedirs(output_poly, exist_ok =True)\n",
    "    \n",
    "    for jpg in next(os.walk(raw_jpg))[2]:\n",
    "        path1_jpg = os.path.join(raw_jpg, jpg)\n",
    "        shutil.copy(path1_jpg, output_jpg)   \n",
    "    \n",
    "    for poly in next(os.walk(raw_poly))[2]:\n",
    "        path1_poly = os.path.join(raw_poly, poly)\n",
    "        shutil.copy(path1_poly, output_poly)"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python [default]",
   "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}