Deconvnet_CNN.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import random\n",
    "import tensorflow as tf\n",
    "import time\n",
    "import tarfile\n",
    "import numpy as np\n",
    "import cv2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class DeconvNet:\n",
    "    \n",
    "    def __init__(self, checkpoint_dir='./checkpoints/'):\n",
    "        self.saver = tf.train.Saver(max_to_keep = 5, keep_checkpoint_every_n_hours =1)\n",
    "        config = tf.ConfigProto(log_device_placement=True, allow_soft_placement=True)\n",
    "        self.session = tf.Session(config = config)\n",
    "        self.session.run(tf.global_variables_initializer())\n",
    "        self.checkpoint_dir = checkpoint_dir\n",
    "        \n",
    "    def weight_variable(self, shape):\n",
    "        \"\"\"\n",
    "        Create a Weight tensor\n",
    "        argument:\n",
    "            shape : The shape of Weight\n",
    "        \"\"\"\n",
    "        initial = tf.truncated_normal(shape, stddev=0.1)\n",
    "        return tf.Variable(initial)\n",
    "    def bias_variable(self, shape):\n",
    "        \"\"\"\n",
    "        Create a Bias tensor\n",
    "        argument:\n",
    "            shape : The shape of Bias\n",
    "        \"\"\"\n",
    "        initial = tf.truncated_normal(shape, stddev=0.1)\n",
    "        return tf.Variable(initial)\n",
    "    \n",
    "    def conv_layer(self, x, W_shape, b_shape, name, padding='SAME'):\n",
    "        \"\"\"\n",
    "        Create a convolutional layer\n",
    "        argument:\n",
    "            x: input tensor, shape = [m,h,w,c]\n",
    "            W_shape : The shape of filter\n",
    "            b_shape : The shape of bias\n",
    "            name : layer's name\n",
    "            padding : padding , default=SAME\n",
    "            \n",
    "        detail:\n",
    "            stride : all using stride = [1,1,1,1]\n",
    "            activation funtion : all using relu\n",
    "        \"\"\"\n",
    "        W = self.weight_variable(W_shape)\n",
    "        b = self.bias_variable(b_shape)\n",
    "        return tf.nn.relu(tf.nn.conv2d(x, W, stride=[1,1,1,1], padding=padding) + b)\n",
    "        \n",
    "    def pooling_layer(self, x):\n",
    "        \"\"\"\n",
    "        Pooling layer\n",
    "        argument:\n",
    "            x: input tensor\n",
    "        return:\n",
    "            pooling_output, pooling_argmax\n",
    "        detail:\n",
    "            pooling_argmax: Store the origin position of the max value\n",
    "        \"\"\"\n",
    "        return tf.nn.max_pool_with_argmax(x, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')\n",
    "    \n",
    "    def deconv_layer(self, x, W_shape, b_shape, name, padding='SAME'):\n",
    "        \"\"\"\n",
    "        conv2d_transpose(value, filter, output_shape, strides, padding=\"SAME\", data_format=\"NHWC\", name=None)\n",
    "        Deconvolutional layer\n",
    "        argument:\n",
    "            x : input tensor\n",
    "            W_shape: the filter shape that is same as the filter coming\n",
    "            b_shape: the bias shape that is same as the bias coming\n",
    "            name: layer's name\n",
    "        \"\"\"\n",
    "        W = self.weight_variable(W_shape)\n",
    "        b = self.bias_variable([b_shape])\n",
    "\n",
    "        x_shape = tf.shape(x)\n",
    "        out_shape = tf.stack([x_shape[0], x_shape[1], x_shape[2], W_shape[2]])\n",
    "\n",
    "        return tf.nn.conv2d_transpose(x, W, out_shape, [1, 1, 1, 1], padding=padding) + b\n",
    "        \n",
    "    def unravel_argmax(self, argmax, shape):\n",
    "        \"\"\"\n",
    "        I don't get it\n",
    "        I get it a little bit LOL 3/5\n",
    "        \"\"\"\n",
    "        output_list = []\n",
    "        output_list.append(argmax // (shape[2] * shape[3]))\n",
    "        output_list.append(argmax % (shape[2] * shape[3]) // shape[3])\n",
    "        return tf.stack(output_list)\n",
    "\n",
    "    def unpool_layer2x2(self, x, raveled_argmax, out_shape):\n",
    "        \"\"\"\n",
    "        I don't get it\n",
    "        \"\"\"\n",
    "        argmax = self.unravel_argmax(raveled_argmax, tf.to_int64(out_shape))\n",
    "        output = tf.zeros([out_shape[1], out_shape[2], out_shape[3]])\n",
    "\n",
    "        height = tf.shape(output)[0]\n",
    "        width = tf.shape(output)[1]\n",
    "        channels = tf.shape(output)[2]\n",
    "\n",
    "        t1 = tf.to_int64(tf.range(channels))\n",
    "        t1 = tf.tile(t1, [((width + 1) // 2) * ((height + 1) // 2)])\n",
    "        t1 = tf.reshape(t1, [-1, channels])\n",
    "        t1 = tf.transpose(t1, perm=[1, 0])\n",
    "        t1 = tf.reshape(t1, [channels, (height + 1) // 2, (width + 1) // 2, 1])\n",
    "\n",
    "        t2 = tf.squeeze(argmax)\n",
    "        t2 = tf.stack((t2[0], t2[1]), axis=0)\n",
    "        t2 = tf.transpose(t2, perm=[3, 1, 2, 0])\n",
    "\n",
    "        t = tf.concat([t2, t1], 3)\n",
    "        indices = tf.reshape(t, [((height + 1) // 2) * ((width + 1) // 2) * channels, 3])\n",
    "\n",
    "        x1 = tf.squeeze(x)\n",
    "        x1 = tf.reshape(x1, [-1, channels])\n",
    "        x1 = tf.transpose(x1, perm=[1, 0])\n",
    "        values = tf.reshape(x1, [-1])\n",
    "\n",
    "        delta = tf.SparseTensor(indices, values, tf.to_int64(tf.shape(output)))\n",
    "        return tf.expand_dims(tf.sparse_tensor_to_dense(tf.sparse_reorder(delta)), 0)\n",
    "    \n",
    "    def build(self):\n",
    "        \"\"\"\n",
    "        Build up the model : DevconvNet \n",
    "        placeholder:\n",
    "            self.x : input image\n",
    "            self.y : labeled image\n",
    "            self.rate : learning rate \n",
    "        \"\"\"\n",
    "        self.x = tf.placeholder(tf.float32, shape=(1, None, None, 3))\n",
    "        self.y = tf.placeholder(tf.int64, shape=(1, None, None))\n",
    "        expected = tf.expand_dims(self.y, -1)\n",
    "        self.rate = tf.placeholder(tf.float32, shape=[])\n",
    "\n",
    "        conv_1_1 = self.conv_layer(self.x, [3, 3, 3, 64], 64, 'conv_1_1')\n",
    "        conv_1_2 = self.conv_layer(conv_1_1, [3, 3, 64, 64], 64, 'conv_1_2')\n",
    "\n",
    "        pool_1, pool_1_argmax = self.pool_layer(conv_1_2)\n",
    "\n",
    "        conv_2_1 = self.conv_layer(pool_1, [3, 3, 64, 128], 128, 'conv_2_1')\n",
    "        conv_2_2 = self.conv_layer(conv_2_1, [3, 3, 128, 128], 128, 'conv_2_2')\n",
    "\n",
    "        pool_2, pool_2_argmax = self.pool_layer(conv_2_2)\n",
    "\n",
    "        conv_3_1 = self.conv_layer(pool_2, [3, 3, 128, 256], 256, 'conv_3_1')\n",
    "        conv_3_2 = self.conv_layer(conv_3_1, [3, 3, 256, 256], 256, 'conv_3_2')\n",
    "        conv_3_3 = self.conv_layer(conv_3_2, [3, 3, 256, 256], 256, 'conv_3_3')\n",
    "\n",
    "        pool_3, pool_3_argmax = self.pool_layer(conv_3_3)\n",
    "\n",
    "        conv_4_1 = self.conv_layer(pool_3, [3, 3, 256, 512], 512, 'conv_4_1')\n",
    "        conv_4_2 = self.conv_layer(conv_4_1, [3, 3, 512, 512], 512, 'conv_4_2')\n",
    "        conv_4_3 = self.conv_layer(conv_4_2, [3, 3, 512, 512], 512, 'conv_4_3')\n",
    "\n",
    "        pool_4, pool_4_argmax = self.pool_layer(conv_4_3)\n",
    "\n",
    "        conv_5_1 = self.conv_layer(pool_4, [3, 3, 512, 512], 512, 'conv_5_1')\n",
    "        conv_5_2 = self.conv_layer(conv_5_1, [3, 3, 512, 512], 512, 'conv_5_2')\n",
    "        conv_5_3 = self.conv_layer(conv_5_2, [3, 3, 512, 512], 512, 'conv_5_3')\n",
    "\n",
    "        pool_5, pool_5_argmax = self.pool_layer(conv_5_3)\n",
    "\n",
    "        fc_6 = self.conv_layer(pool_5, [7, 7, 512, 4096], 4096, 'fc_6')\n",
    "        fc_7 = self.conv_layer(fc_6, [1, 1, 4096, 4096], 4096, 'fc_7')\n",
    "\n",
    "        deconv_fc_6 = self.deconv_layer(fc_7, [7, 7, 512, 4096], 512, 'fc6_deconv')\n",
    "\n",
    "        unpool_5 = self.unpool_layer2x2(deconv_fc_6, pool_5_argmax, tf.shape(conv_5_3))\n",
    "\n",
    "        deconv_5_3 = self.deconv_layer(unpool_5, [3, 3, 512, 512], 512, 'deconv_5_3')\n",
    "        deconv_5_2 = self.deconv_layer(deconv_5_3, [3, 3, 512, 512], 512, 'deconv_5_2')\n",
    "        deconv_5_1 = self.deconv_layer(deconv_5_2, [3, 3, 512, 512], 512, 'deconv_5_1')\n",
    "\n",
    "        unpool_4 = self.unpool_layer2x2(deconv_5_1, pool_4_argmax, tf.shape(conv_4_3))\n",
    "\n",
    "        deconv_4_3 = self.deconv_layer(unpool_4, [3, 3, 512, 512], 512, 'deconv_4_3')\n",
    "        deconv_4_2 = self.deconv_layer(deconv_4_3, [3, 3, 512, 512], 512, 'deconv_4_2')\n",
    "        deconv_4_1 = self.deconv_layer(deconv_4_2, [3, 3, 256, 512], 256, 'deconv_4_1')\n",
    "\n",
    "        unpool_3 = self.unpool_layer2x2(deconv_4_1, pool_3_argmax, tf.shape(conv_3_3))\n",
    "\n",
    "        deconv_3_3 = self.deconv_layer(unpool_3, [3, 3, 256, 256], 256, 'deconv_3_3')\n",
    "        deconv_3_2 = self.deconv_layer(deconv_3_3, [3, 3, 256, 256], 256, 'deconv_3_2')\n",
    "        deconv_3_1 = self.deconv_layer(deconv_3_2, [3, 3, 128, 256], 128, 'deconv_3_1')\n",
    "\n",
    "        unpool_2 = self.unpool_layer2x2(deconv_3_1, pool_2_argmax, tf.shape(conv_2_2))\n",
    "\n",
    "        deconv_2_2 = self.deconv_layer(unpool_2, [3, 3, 128, 128], 128, 'deconv_2_2')\n",
    "        deconv_2_1 = self.deconv_layer(deconv_2_2, [3, 3, 64, 128], 64, 'deconv_2_1')\n",
    "\n",
    "        unpool_1 = self.unpool_layer2x2(deconv_2_1, pool_1_argmax, tf.shape(conv_1_2))\n",
    "\n",
    "        deconv_1_2 = self.deconv_layer(unpool_1, [3, 3, 64, 64], 64, 'deconv_1_2')\n",
    "        deconv_1_1 = self.deconv_layer(deconv_1_2, [3, 3, 32, 64], 32, 'deconv_1_1')\n",
    "\n",
    "        score_1 = self.deconv_layer(deconv_1_1, [1, 1, 21, 32], 21, 'score_1')\n",
    "\n",
    "        logits = tf.reshape(score_1, (-1, 21))\n",
    "        cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=tf.reshape(expected, [-1]), logits=logits, name='x_entropy')\n",
    "        self.loss = tf.reduce_mean(cross_entropy, name='x_entropy_mean')\n",
    "\n",
    "        self.train_step = tf.train.AdamOptimizer(self.rate).minimize(self.loss)\n",
    "\n",
    "        self.prediction = tf.argmax(tf.reshape(tf.nn.softmax(logits), tf.shape(score_1)), dimension=3)\n",
    "        self.accuracy = tf.reduce_sum(tf.pow(self.prediction - expected, 2))\n",
    "    \n",
    "    def train(self, train_stage=1, training_steps=5, restore_session=False, learning_rate=1e-6):\n",
    "        if restore_session:\n",
    "            step_start = restore_session()\n",
    "        else:\n",
    "            step_start = 0\n",
    "\n",
    "        if train_stage == 1:\n",
    "            \"\"\"\n",
    "            feed in train data\n",
    "            \"\"\"\n",
    "            trainset = None\n",
    "        else:\n",
    "            \"\"\"\n",
    "            feed in dev data\n",
    "            \"\"\"\n",
    "            trainset = None\n",
    "\n",
    "        for i in range(step_start, step_start+training_steps):\n",
    "            \n",
    "            # pick random line from file\n",
    "            random_line = random.choice(trainset)\n",
    "            image_file = random_line.split(' ')[0]\n",
    "            ground_truth_file = random_line.split(' ')[1]\n",
    "            image = np.float32(cv2.imread('data' + image_file))\n",
    "            ground_truth = cv2.imread('data' + ground_truth_file[:-1], cv2.IMREAD_GRAYSCALE)\n",
    "            #\n",
    "            # norm to 21 classes [0-20] (see paper)\n",
    "            ground_truth = (ground_truth / 255) * 20\n",
    "            print('run train step: '+str(i))\n",
    "            start = time.time()\n",
    "            self.train_step.run(session=self.session, feed_dict={self.x: [image], self.y: [ground_truth], self.rate: learning_rate})\n",
    "\n",
    "            if i % 10000 == 0:\n",
    "                print('step {} finished in {:.2f} s with loss of {:.6f}'.format(i, time.time() - start, self.loss.eval(session=self.session, feed_dict={self.x: [image], self.y: [ground_truth]})))\n",
    "                self.saver.save(self.session, self.checkpoint_dir+'model', global_step=i)\n",
    "                print('Model {} saved'.format(i))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "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.5.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}