black&white_to_color/black&white_to_color.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using TensorFlow backend.\n"
     ]
    }
   ],
   "source": [
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from keras.layers import Input, Dense, Convolution2D, MaxPooling2D, UpSampling2D, Reshape, Flatten\n",
    "from keras.models import Model\n",
    "from IPython.display import SVG\n",
    "from keras.utils.visualize_util import model_to_dot\n",
    "from keras.callbacks import ModelCheckpoint\n",
    "\n",
    "import os\n",
    "import random\n",
    "import keras\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ROW = 80\n",
    "COL = 80\n",
    "CHANNELS = 3\n",
    "\n",
    "TRAIN_DIR = 'cat_dog_train/'\n",
    "TEST_DIR = 'cat_dog_test/'\n",
    "\n",
    "train_images = [TRAIN_DIR+i for i in os.listdir(TRAIN_DIR)] # use this for full dataset\n",
    "\n",
    "def readImg(imgFile):\n",
    "    colorImg = cv2.imread(imgFile, cv2.IMREAD_COLOR)\n",
    "    colorImg = cv2.cvtColor(colorImg, cv2.COLOR_BGR2RGB)\n",
    "    colorImg = cv2.resize(colorImg, (ROW, COL))/255.0\n",
    "    \n",
    "    greyImg = cv2.imread(imgFile, cv2.IMREAD_GRAYSCALE)\n",
    "    greyImg = cv2.cvtColor(greyImg,cv2.COLOR_GRAY2RGB)\n",
    "    greyImg = cv2.resize(greyImg, (ROW, COL))/255.0\n",
    "    \n",
    "    return greyImg,colorImg\n",
    "    \n",
    "    \n",
    "\n",
    "def generateDate(imgFiles):\n",
    "    count = len(imgFiles)\n",
    "    dataX = np.ndarray((count, ROW, COL,CHANNELS), dtype=float)\n",
    "    dataY = np.ndarray((count, ROW, COL,CHANNELS), dtype=float)\n",
    "    for i, image_file in enumerate(imgFiles):\n",
    "        gImg,cImg = readImg(image_file)\n",
    "        dataX[i] = gImg\n",
    "        dataY[i] = cImg  \n",
    "    return dataX,dataY\n",
    "\n",
    "random.shuffle(train_images)\n",
    "\n",
    "\n",
    "import math\n",
    "def chunked(iterable, n):\n",
    "    chunksize = int(math.ceil(len(iterable) / n))\n",
    "    return (iterable[i * chunksize:i * chunksize + chunksize]\n",
    "            for i in range(n))\n",
    "\n",
    "BATCHNUMBER= 20\n",
    "\n",
    "chuckedTrainList = list(chunked(train_images,BATCHNUMBER))\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "ename": "SyntaxError",
     "evalue": "invalid syntax (<ipython-input-22-691912db50a0>, line 1)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;36m  File \u001b[0;32m\"<ipython-input-22-691912db50a0>\"\u001b[0;36m, line \u001b[0;32m1\u001b[0m\n\u001b[0;31m    BaseLevel = ROW////2\u001b[0m\n\u001b[0m                      ^\u001b[0m\n\u001b[0;31mSyntaxError\u001b[0m\u001b[0;31m:\u001b[0m invalid syntax\n"
     ]
    }
   ],
   "source": [
    "BaseLevel = ROW////2\n",
    "input_img = Input(shape=(ROW,COL,CHANNELS))\n",
    "x = Convolution2D(256, 5, 5, activation='relu', border_mode='same')(input_img)\n",
    "x = MaxPooling2D((2, 2), border_mode='same')(x)\n",
    "x = Convolution2D(128, 5, 5, activation='relu', border_mode='same')(x)\n",
    "x = MaxPooling2D((2, 2), border_mode='same')(x)\n",
    "x = Flatten()(x)\n",
    "\n",
    "encoded = Dense(2000)(x)\n",
    "oneD = Dense(BaseLevel*BaseLevel*128)(encoded)\n",
    "fold = Reshape((BaseLevel,BaseLevel,128))(oneD)\n",
    "\n",
    "x = UpSampling2D((2, 2))(fold)\n",
    "x = Convolution2D(128, 5, 5, activation='relu', border_mode='same')(x)\n",
    "x = UpSampling2D((2, 2))(x)\n",
    "decoded = Convolution2D(3, 5, 5, activation='sigmoid', border_mode='same')(x)\n",
    "\n",
    "autoencoder = Model(input_img, decoded)\n",
    "autoencoder.compile(optimizer='adam', loss='binary_crossentropy')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "collapsed": false,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "checkpoint = ModelCheckpoint(filepath=\"/home/yhfy2006/machineLearningInPython/AutoEncoder/logs/weights.hdf5\", verbose=1, save_best_only=True)\n",
    "earlyStopping=keras.callbacks.EarlyStopping(monitor='val_loss', patience=0, verbose=0, mode='auto')\n",
    "\n",
    "cbks = [checkpoint,earlyStopping]\n",
    "\n",
    "\n",
    "digitalImg = dict()\n",
    "epoch = 500\n",
    "\n",
    "for ii in range(epoch):\n",
    "    i = ii%len(chuckedTrainList)\n",
    "    print(\"epoch =====>\"+str(ii))\n",
    "    imgfiles = chuckedTrainList[i]\n",
    "    if str(i) in digitalImg:\n",
    "        dataX = digitalImg[str(i)][0]\n",
    "        dataY = digitalImg[str(i)][1]\n",
    "    else:\n",
    "        dataX,dataY = generateDate(imgfiles)\n",
    "        digitalImg[str(i)] = (dataX,dataY)\n",
    "        \n",
    "    autoencoder.fit(dataX, dataY,\n",
    "                nb_epoch=1,\n",
    "                batch_size=50,\n",
    "                shuffle=True,\n",
    "                verbose=2,\n",
    "                validation_split=0.3,\n",
    "                callbacks=cbks)\n",
    "        \n",
    "\n",
    "    \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#autoencoder.load_weights(\"logs/weights_best_20170301.hdf5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "collapsed": false,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "test_images =  [TEST_DIR+i for i in os.listdir(TEST_DIR)]\n",
    "test_images =  test_images[80:92]\n",
    "\n",
    "testX,testY = generateDate(test_images)\n",
    "\n",
    "predict = autoencoder.predict(testX)\n",
    "\n",
    "for i,imageFile in enumerate(test_images):\n",
    "    plt.imshow(testX[i])\n",
    "    plt.show()\n",
    "    \n",
    "    plt.imshow(predict[i].astype(float))\n",
    "    plt.show()\n",
    "    \n",
    "    plt.imshow(testY[i])\n",
    "    plt.show()\n",
    "\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 144,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "____________________________________________________________________________________________________\n",
      "Layer (type)                     Output Shape          Param #     Connected to                     \n",
      "====================================================================================================\n",
      "input_21 (InputLayer)            (None, 500, 500, 3)   0                                            \n",
      "____________________________________________________________________________________________________\n",
      "convolution2d_133 (Convolution2D (None, 500, 500, 16)  448         input_21[0][0]                   \n",
      "____________________________________________________________________________________________________\n",
      "maxpooling2d_58 (MaxPooling2D)   (None, 250, 250, 16)  0           convolution2d_133[0][0]          \n",
      "____________________________________________________________________________________________________\n",
      "convolution2d_134 (Convolution2D (None, 250, 250, 8)   1160        maxpooling2d_58[0][0]            \n",
      "____________________________________________________________________________________________________\n",
      "maxpooling2d_59 (MaxPooling2D)   (None, 125, 125, 8)   0           convolution2d_134[0][0]          \n",
      "____________________________________________________________________________________________________\n",
      "convolution2d_135 (Convolution2D (None, 125, 125, 8)   584         maxpooling2d_59[0][0]            \n",
      "____________________________________________________________________________________________________\n",
      "maxpooling2d_60 (MaxPooling2D)   (None, 63, 63, 8)     0           convolution2d_135[0][0]          \n",
      "____________________________________________________________________________________________________\n",
      "convolution2d_136 (Convolution2D (None, 63, 63, 8)     584         maxpooling2d_60[0][0]            \n",
      "____________________________________________________________________________________________________\n",
      "upsampling2d_58 (UpSampling2D)   (None, 126, 126, 8)   0           convolution2d_136[0][0]          \n",
      "____________________________________________________________________________________________________\n",
      "convolution2d_137 (Convolution2D (None, 126, 126, 8)   584         upsampling2d_58[0][0]            \n",
      "____________________________________________________________________________________________________\n",
      "upsampling2d_59 (UpSampling2D)   (None, 252, 252, 8)   0           convolution2d_137[0][0]          \n",
      "____________________________________________________________________________________________________\n",
      "convolution2d_138 (Convolution2D (None, 250, 250, 16)  1168        upsampling2d_59[0][0]            \n",
      "____________________________________________________________________________________________________\n",
      "upsampling2d_60 (UpSampling2D)   (None, 500, 500, 16)  0           convolution2d_138[0][0]          \n",
      "____________________________________________________________________________________________________\n",
      "convolution2d_139 (Convolution2D (None, 500, 500, 3)   435         upsampling2d_60[0][0]            \n",
      "====================================================================================================\n",
      "Total params: 4,963\n",
      "Trainable params: 4,963\n",
      "Non-trainable params: 0\n",
      "____________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "autoencoder.summary()"
   ]
  }
 ],
 "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
}