Style Data Augmentation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Style Image Augmentation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The original dataset of Chinese painting only has 100 Chinese painting from 8 categories. This notebook will augment the total images to 944, with each classes to have about 100 images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "#import libraries\n",
    "import glob\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from tensorflow.keras.preprocessing.image import ImageDataGenerator\n",
    "from matplotlib.pyplot import imread, imshow, subplots, show\n",
    "import imageio\n",
    "import tensorflow as tf\n",
    "import math\n",
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Read image file path\n",
    "filenames_n0 = glob.glob('./style_images_original/claborate/*')\n",
    "filenames_n1 = glob.glob('./style_images_original/Fuchun/*')\n",
    "filenames_n2 = glob.glob('./style_images_original/Others/*')\n",
    "filenames_n3 = glob.glob('./style_images_original/Shanshui/*')\n",
    "filenames_n4 = glob.glob('./style_images_original/Shuimo/*')\n",
    "filenames_n5 = glob.glob('./style_images_original/Thinline/*')\n",
    "filenames_n6 = glob.glob('./style_images_original/Wuguanzhou/*')\n",
    "filenames_n7 = glob.glob('./style_images_original/Xishan/*')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>filename</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>./style_images/claborate\\claborate1.jpg</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>./style_images/claborate\\claborate3.jpg</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>./style_images/claborate\\claborate4.jpg</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>./style_images/claborate\\claborate5.jpg</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>./style_images/claborate\\claborate6.jpg</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                  filename\n",
       "0  ./style_images/claborate\\claborate1.jpg\n",
       "1  ./style_images/claborate\\claborate3.jpg\n",
       "2  ./style_images/claborate\\claborate4.jpg\n",
       "3  ./style_images/claborate\\claborate5.jpg\n",
       "4  ./style_images/claborate\\claborate6.jpg"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Make a dataframe based on the filenames\n",
    "df = pd.DataFrame(filenames_n0, columns = [\"filename\"])\n",
    "df2 = pd.DataFrame(filenames_n1, columns = [\"filename\"])\n",
    "df3 = pd.DataFrame(filenames_n2, columns = [\"filename\"])\n",
    "df4 = pd.DataFrame(filenames_n3, columns = [\"filename\"])\n",
    "df5 = pd.DataFrame(filenames_n4, columns = [\"filename\"])\n",
    "df6 = pd.DataFrame(filenames_n5, columns = [\"filename\"])\n",
    "df7 = pd.DataFrame(filenames_n6, columns = [\"filename\"])\n",
    "df8 = pd.DataFrame(filenames_n7, columns = [\"filename\"])\n",
    "\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "#convert images into RGB format \n",
    "def to_rgb(df):\n",
    "    for i in df['filename']:\n",
    "        im = Image.open(i)\n",
    "        rgb_im = im.convert('RGB')\n",
    "        rgb_im.save(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Define function to randomly augment the datafram of images with rotation, width shift, height shit or zoom shift\n",
    "def random_augment(df, ite_num,path):\n",
    "    filename=''\n",
    "    dftemp = df\n",
    "    \n",
    "    dftemp['image'] = dftemp['filename'].apply(lambda x:imread(x))\n",
    "    dftemp['image'] = dftemp['image'].apply(lambda x:x.reshape((1, x.shape[0], x.shape[1], x.shape[2])))\n",
    "    \n",
    "    for i,k in enumerate(dftemp['image']):\n",
    "        for j in range(ite_num):\n",
    "            ran = np.random.randint(1,5,dtype='int') #random generate a number to decide which augmentation to run for each iteration\n",
    "            if ran == 1:\n",
    "                data_generator = ImageDataGenerator(rotation_range=90,fill_mode='reflect')\n",
    "            elif ran == 2:\n",
    "                data_generator = ImageDataGenerator(width_shift_range=0.5,fill_mode='reflect')\n",
    "            elif ran == 3:\n",
    "                data_generator = ImageDataGenerator(height_shift_range=0.5,fill_mode='reflect')\n",
    "            elif ran == 4:\n",
    "                data_generator = ImageDataGenerator(zoom_range=[0.5, 2.0],fill_mode='reflect')\n",
    "            \n",
    "            data_generator.fit(k)\n",
    "            image_iterator = data_generator.flow(k)\n",
    "            #set a unique filename for each image generated\n",
    "            filename = path + str(i)+ '_'+ str(j) + '.jpg'\n",
    "            #save image\n",
    "            imageio.imwrite(filename,image_iterator.next()[0].astype('uint8'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Convert to make sure all images have only RGB channels\n",
    "to_rgb(df)\n",
    "to_rgb(df2)\n",
    "to_rgb(df3)\n",
    "to_rgb(df4)\n",
    "to_rgb(df5)\n",
    "to_rgb(df6)\n",
    "to_rgb(df7)\n",
    "to_rgb(df8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Make each category to have at least 100 dataset after augmentation\n",
    "min_num = 100"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run Augmentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "#df1\n",
    "ite_num = math.ceil(100/len(df))\n",
    "random_augment(df,ite_num,'./style_images/claborate/') "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "#df2\n",
    "ite_num = math.ceil(100/len(df2))\n",
    "random_augment(df2,ite_num,'./style_images/Fuchun/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "#df3\n",
    "ite_num = math.ceil(100/len(df3))\n",
    "random_augment(df3,ite_num,'./style_images/Others/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "#df4\n",
    "ite_num = math.ceil(100/len(df4))\n",
    "random_augment(df4,ite_num,'./style_images/Shanshui/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "#df5\n",
    "ite_num = math.ceil(100/len(df5))\n",
    "random_augment(df5,ite_num,'./style_images/Shuimo/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "#df6\n",
    "ite_num = math.ceil(100/len(df6))\n",
    "random_augment(df6,ite_num,'./style_images/Thinline/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "#df7\n",
    "ite_num = math.ceil(100/len(df7))\n",
    "random_augment(df7,ite_num,'./style_images/Wuguanzhou/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "#df8\n",
    "ite_num = math.ceil(100/len(df8))\n",
    "random_augment(df8,ite_num,'./style_images/Xishan/')"
   ]
  }
 ],
 "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}