model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import os\n",
    "import random\n",
    "import tensorflow.contrib.slim as slim\n",
    "import time\n",
    "import logging\n",
    "import numpy as np\n",
    "import pickle\n",
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "logger = logging.getLogger('Training begin')\n",
    "logger.setLevel(logging.INFO)\n",
    "# formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')\n",
    "ch = logging.StreamHandler()\n",
    "ch.setLevel(logging.INFO)\n",
    "logger.addHandler(ch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "tf.app.flags.DEFINE_boolean('random_flip_up_down', False, \"Whether to random flip up down\")\n",
    "tf.app.flags.DEFINE_boolean('random_brightness', True, \"whether to adjust brightness\")\n",
    "tf.app.flags.DEFINE_boolean('random_contrast', True, \"whether to random constrast\")\n",
    "\n",
    "tf.app.flags.DEFINE_integer('charset_size', 3755, \"Choose the first `charset_size` character to conduct our experiment.\")\n",
    "tf.app.flags.DEFINE_integer('image_size', 64, \"Needs to provide same value as in training.\")\n",
    "tf.app.flags.DEFINE_boolean('gray', True, \"whether to change the rbg to gray\")\n",
    "tf.app.flags.DEFINE_integer('max_steps', 12002, 'the max training steps ')\n",
    "tf.app.flags.DEFINE_integer('eval_steps', 50, \"the step num to eval\")\n",
    "tf.app.flags.DEFINE_integer('save_steps', 2000, \"the steps to save\")\n",
    "\n",
    "tf.app.flags.DEFINE_string('checkpoint_dir', './output/checkpoint/', 'the checkpoint dir')\n",
    "tf.app.flags.DEFINE_string('train_data_dir', './output/train/', 'the train dataset dir')\n",
    "tf.app.flags.DEFINE_string('test_data_dir', './output/test/', 'the test dataset dir')\n",
    "tf.app.flags.DEFINE_string('log_dir', './log', 'the logging dir')\n",
    "\n",
    "tf.app.flags.DEFINE_boolean('restore', False, 'whether to restore from checkpoint')\n",
    "tf.app.flags.DEFINE_boolean('epoch', 1, 'Number of epoches')\n",
    "tf.app.flags.DEFINE_boolean('batch_size', 128, 'Validation batch size')\n",
    "tf.app.flags.DEFINE_string('mode', 'train', 'Running mode. One of {\"train\", \"valid\", \"test\"}')\n",
    "FLAGS = tf.app.flags.FLAGS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "check_d = \"./output/checkpoint\" \n",
    "if not os.path.exists(check_d):\n",
    "    os.mkdir(check_d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train\n",
      "Begin training\n",
      "./output/train/03755\n"
     ]
    }
   ],
   "source": [
    "class DataIterator:\n",
    "    def __init__(self, data_dir):\n",
    "        # Set FLAGS.charset_size to a small value if available computation power is limited.\n",
    "        truncate_path = data_dir + ('%05d' % FLAGS.charset_size)\n",
    "        print(truncate_path)\n",
    "        self.image_names = []\n",
    "        for root, sub_folder, file_list in os.walk(data_dir):\n",
    "            if root < truncate_path:\n",
    "                self.image_names += [os.path.join(root, file_path) for file_path in file_list]\n",
    "        random.shuffle(self.image_names)\n",
    "        self.labels = [int(file_name[len(data_dir):].split(os.sep)[0]) for file_name in self.image_names]\n",
    "\n",
    "    @property\n",
    "    def size(self):\n",
    "        return len(self.labels)\n",
    "\n",
    "    @staticmethod\n",
    "    def data_augmentation(images):\n",
    "        if FLAGS.random_flip_up_down:\n",
    "            images = tf.image.random_flip_up_down(images)\n",
    "        if FLAGS.random_brightness:\n",
    "            images = tf.image.random_brightness(images, max_delta=0.3)\n",
    "        if FLAGS.random_contrast:\n",
    "            images = tf.image.random_contrast(images, 0.8, 1.2)\n",
    "        return images\n",
    "\n",
    "    def input_pipeline(self, batch_size, num_epochs=None, aug=False):\n",
    "        images_tensor = tf.convert_to_tensor(self.image_names, dtype=tf.string)\n",
    "        labels_tensor = tf.convert_to_tensor(self.labels, dtype=tf.int64)\n",
    "        input_queue = tf.train.slice_input_producer([images_tensor, labels_tensor], num_epochs=num_epochs)\n",
    "\n",
    "        labels = input_queue[1]\n",
    "        images_content = tf.read_file(input_queue[0])\n",
    "        images = tf.image.convert_image_dtype(tf.image.decode_png(images_content, channels=1), tf.float32)\n",
    "        if aug:\n",
    "            images = self.data_augmentation(images)\n",
    "        new_size = tf.constant([FLAGS.image_size, FLAGS.image_size], dtype=tf.int32)\n",
    "        images = tf.image.resize_images(images, new_size)\n",
    "        image_batch, label_batch = tf.train.shuffle_batch([images, labels], batch_size=batch_size, capacity=50000,\n",
    "                                                          min_after_dequeue=10000)\n",
    "        return image_batch, label_batch\n",
    "\n",
    "\n",
    "def build_graph(top_k):\n",
    "    # with tf.device('/cpu:0'):\n",
    "    keep_prob = tf.placeholder(dtype=tf.float32, shape=[], name='keep_prob')\n",
    "    images = tf.placeholder(dtype=tf.float32, shape=[None, 64, 64, 1], name='image_batch')\n",
    "    labels = tf.placeholder(dtype=tf.int64, shape=[None], name='label_batch')\n",
    "\n",
    "    conv_1 = slim.conv2d(images, 64, [3, 3], 1, padding='SAME', scope='conv1')\n",
    "    max_pool_1 = slim.max_pool2d(conv_1, [2, 2], [2, 2], padding='SAME')\n",
    "    conv_2 = slim.conv2d(max_pool_1, 128, [3, 3], padding='SAME', scope='conv2')\n",
    "    max_pool_2 = slim.max_pool2d(conv_2, [2, 2], [2, 2], padding='SAME')\n",
    "    conv_3 = slim.conv2d(max_pool_2, 256, [3, 3], padding='SAME', scope='conv3')\n",
    "    max_pool_3 = slim.max_pool2d(conv_3, [2, 2], [2, 2], padding='SAME')\n",
    "\n",
    "    flatten = slim.flatten(max_pool_3)\n",
    "    fc1 = slim.fully_connected(slim.dropout(flatten, keep_prob), 1024, activation_fn=tf.nn.tanh, scope='fc1')\n",
    "    logits = slim.fully_connected(slim.dropout(fc1, keep_prob), FLAGS.charset_size, activation_fn=None, scope='fc2')\n",
    "        # logits = slim.fully_connected(flatten, FLAGS.charset_size, activation_fn=None, reuse=reuse, scope='fc')\n",
    "    loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits, labels=labels))\n",
    "    accuracy = tf.reduce_mean(tf.cast(tf.equal(tf.argmax(logits, 1), labels), tf.float32))\n",
    "\n",
    "    global_step = tf.get_variable(\"step\", [], initializer=tf.constant_initializer(0.0), trainable=False)\n",
    "    rate = tf.train.exponential_decay(2e-4, global_step, decay_steps=2000, decay_rate=0.97, staircase=True)\n",
    "    train_op = tf.train.AdamOptimizer(learning_rate=rate).minimize(loss, global_step=global_step)\n",
    "    probabilities = tf.nn.softmax(logits)\n",
    "\n",
    "    tf.summary.scalar('loss', loss)\n",
    "    tf.summary.scalar('accuracy', accuracy)\n",
    "    merged_summary_op = tf.summary.merge_all()\n",
    "    predicted_val_top_k, predicted_index_top_k = tf.nn.top_k(probabilities, k=top_k)\n",
    "    accuracy_in_top_k = tf.reduce_mean(tf.cast(tf.nn.in_top_k(probabilities, labels, top_k), tf.float32))\n",
    "\n",
    "    return {'images': images,\n",
    "            'labels': labels,\n",
    "            'keep_prob': keep_prob,\n",
    "            'top_k': top_k,\n",
    "            'global_step': global_step,\n",
    "            'train_op': train_op,\n",
    "            'loss': loss,\n",
    "            'accuracy': accuracy,\n",
    "            'accuracy_top_k': accuracy_in_top_k,\n",
    "            'merged_summary_op': merged_summary_op,\n",
    "            'predicted_distribution': probabilities,\n",
    "            'predicted_index_top_k': predicted_index_top_k,\n",
    "            'predicted_val_top_k': predicted_val_top_k}\n",
    "\n",
    "\n",
    "def train():\n",
    "    print('Begin training')\n",
    "    train_feeder = DataIterator(data_dir='./output/train/')\n",
    "    test_feeder = DataIterator(data_dir='./output/test/')\n",
    "    with tf.Session() as sess:\n",
    "        train_images, train_labels = train_feeder.input_pipeline(batch_size=FLAGS.batch_size, aug=True)\n",
    "        test_images, test_labels = test_feeder.input_pipeline(batch_size=FLAGS.batch_size)\n",
    "        graph = build_graph(top_k=1)\n",
    "        sess.run(tf.global_variables_initializer())\n",
    "        coord = tf.train.Coordinator()\n",
    "        threads = tf.train.start_queue_runners(sess=sess, coord=coord)\n",
    "        saver = tf.train.Saver()\n",
    "\n",
    "        train_writer = tf.summary.FileWriter(FLAGS.log_dir + '/train', sess.graph)\n",
    "        test_writer = tf.summary.FileWriter(FLAGS.log_dir + '/val')\n",
    "        start_step = 0\n",
    "        if FLAGS.restore:\n",
    "            ckpt = tf.train.latest_checkpoint(FLAGS.checkpoint_dir)\n",
    "            if ckpt:\n",
    "                saver.restore(sess, ckpt)\n",
    "                print(\"restore from the checkpoint {0}\".format(ckpt))\n",
    "                start_step += int(ckpt.split('-')[-1])\n",
    "\n",
    "        logger.info(':::Training Start:::')\n",
    "        try:\n",
    "            while not coord.should_stop():\n",
    "                start_time = time.time()\n",
    "                train_images_batch, train_labels_batch = sess.run([train_images, train_labels])\n",
    "                feed_dict = {graph['images']: train_images_batch,\n",
    "                             graph['labels']: train_labels_batch,\n",
    "                             graph['keep_prob']: 0.8}\n",
    "                _, loss_val, train_summary, step = sess.run(\n",
    "                    [graph['train_op'], graph['loss'], graph['merged_summary_op'], graph['global_step']],\n",
    "                    feed_dict=feed_dict)\n",
    "                train_writer.add_summary(train_summary, step)\n",
    "                end_time = time.time()\n",
    "                logger.info(\"the step {0} takes {1} loss {2}\".format(step, end_time - start_time, loss_val))\n",
    "                if step > FLAGS.max_steps:\n",
    "                    break\n",
    "                if step % FLAGS.eval_steps == 1:\n",
    "                    test_images_batch, test_labels_batch = sess.run([test_images, test_labels])\n",
    "                    feed_dict = {graph['images']: test_images_batch,\n",
    "                                 graph['labels']: test_labels_batch,\n",
    "                                 graph['keep_prob']: 1.0}\n",
    "                    accuracy_test, test_summary = sess.run(\n",
    "                        [graph['accuracy'], graph['merged_summary_op']],\n",
    "                        feed_dict=feed_dict)\n",
    "                    test_writer.add_summary(test_summary, step)\n",
    "                    logger.info('===============Eval a batch=======================')\n",
    "                    logger.info('the step {0} test accuracy: {1}'\n",
    "                                .format(step, accuracy_test))\n",
    "                    logger.info('===============Eval a batch=======================')\n",
    "                if step % FLAGS.save_steps == 1:\n",
    "                    logger.info('Save the ckpt of {0}'.format(step))\n",
    "                    saver.save(sess, os.path.join(FLAGS.checkpoint_dir, 'my-model'),\n",
    "                               global_step=graph['global_step'])\n",
    "        except tf.errors.OutOfRangeError:\n",
    "            logger.info('==================Train Finished================')\n",
    "            saver.save(sess, os.path.join(FLAGS.checkpoint_dir, 'my-model'), global_step=graph['global_step'])\n",
    "        finally:\n",
    "            coord.request_stop()\n",
    "        coord.join(threads)\n",
    "\n",
    "\n",
    "def validation():\n",
    "    print('validation')\n",
    "    test_feeder = DataIterator(data_dir='./output/test/')\n",
    "\n",
    "    final_predict_val = []\n",
    "    final_predict_index = []\n",
    "    groundtruth = []\n",
    "\n",
    "    with tf.Session() as sess:\n",
    "        test_images, test_labels = test_feeder.input_pipeline(batch_size=FLAGS.batch_size, num_epochs=1)\n",
    "        graph = build_graph(3)\n",
    "\n",
    "        sess.run(tf.global_variables_initializer())\n",
    "        sess.run(tf.local_variables_initializer())  # initialize test_feeder's inside state\n",
    "\n",
    "        coord = tf.train.Coordinator()\n",
    "        threads = tf.train.start_queue_runners(sess=sess, coord=coord)\n",
    "\n",
    "        saver = tf.train.Saver()\n",
    "        ckpt = tf.train.latest_checkpoint(FLAGS.checkpoint_dir)\n",
    "        if ckpt:\n",
    "            saver.restore(sess, ckpt)\n",
    "            print(\"restore from the checkpoint {0}\".format(ckpt))\n",
    "\n",
    "        logger.info(':::Start validation:::')\n",
    "        try:\n",
    "            i = 0\n",
    "            acc_top_1, acc_top_k = 0.0, 0.0\n",
    "            while not coord.should_stop():\n",
    "                i += 1\n",
    "                start_time = time.time()\n",
    "                test_images_batch, test_labels_batch = sess.run([test_images, test_labels])\n",
    "                feed_dict = {graph['images']: test_images_batch,\n",
    "                             graph['labels']: test_labels_batch,\n",
    "                             graph['keep_prob']: 1.0}\n",
    "                batch_labels, probs, indices, acc_1, acc_k = sess.run([graph['labels'],\n",
    "                                                                       graph['predicted_val_top_k'],\n",
    "                                                                       graph['predicted_index_top_k'],\n",
    "                                                                       graph['accuracy'],\n",
    "                                                                       graph['accuracy_top_k']], feed_dict=feed_dict)\n",
    "                final_predict_val += probs.tolist()\n",
    "                final_predict_index += indices.tolist()\n",
    "                groundtruth += batch_labels.tolist()\n",
    "                acc_top_1 += acc_1\n",
    "                acc_top_k += acc_k\n",
    "                end_time = time.time()\n",
    "                logger.info(\"the batch {0} takes {1} seconds, accuracy = {2}(top_1) {3}(top_k)\"\n",
    "                            .format(i, end_time - start_time, acc_1, acc_k))\n",
    "\n",
    "        except tf.errors.OutOfRangeError:\n",
    "            logger.info('==================Validation Finished================')\n",
    "            acc_top_1 = acc_top_1 * FLAGS.batch_size / test_feeder.size\n",
    "            acc_top_k = acc_top_k * FLAGS.batch_size / test_feeder.size\n",
    "            logger.info('top 1 accuracy {0} top k accuracy {1}'.format(acc_top_1, acc_top_k))\n",
    "        finally:\n",
    "            coord.request_stop()\n",
    "        coord.join(threads)\n",
    "    return {'prob': final_predict_val, 'indices': final_predict_index, 'groundtruth': groundtruth}\n",
    "\n",
    "\n",
    "def inference(image):\n",
    "    print('inference')\n",
    "    temp_image = Image.open(image).convert('L')\n",
    "    temp_image = temp_image.resize((FLAGS.image_size, FLAGS.image_size), Image.ANTIALIAS)\n",
    "    temp_image = np.asarray(temp_image) / 255.0\n",
    "    temp_image = temp_image.reshape([-1, 64, 64, 1])\n",
    "    with tf.Session() as sess:\n",
    "        logger.info('========start inference============')\n",
    "        # images = tf.placeholder(dtype=tf.float32, shape=[None, 64, 64, 1])\n",
    "        # Pass a shadow label 0. This label will not affect the computation graph.\n",
    "        graph = build_graph(top_k=3)\n",
    "        saver = tf.train.Saver()\n",
    "        ckpt = tf.train.latest_checkpoint(FLAGS.checkpoint_dir)\n",
    "        if ckpt:\n",
    "            saver.restore(sess, ckpt)\n",
    "        predict_val, predict_index = sess.run([graph['predicted_val_top_k'], graph['predicted_index_top_k']],\n",
    "                                              feed_dict={graph['images']: temp_image, graph['keep_prob']: 1.0})\n",
    "    return predict_val, predict_index\n",
    "\n",
    "\n",
    "def main(_):\n",
    "    print(FLAGS.mode)\n",
    "    if FLAGS.mode == \"train\":\n",
    "        train()\n",
    "    elif FLAGS.mode == 'validation':\n",
    "        dct = validation()\n",
    "        result_file = 'result.dict'\n",
    "        logger.info('Write result into {0}'.format(result_file))\n",
    "        with open(result_file, 'wb') as f:\n",
    "            pickle.dump(dct, f)\n",
    "        logger.info('Write file ends')\n",
    "    elif FLAGS.mode == 'inference':\n",
    "        image_path = './output/test/00190/13320.png'\n",
    "        final_predict_val, final_predict_index = inference(image_path)\n",
    "        logger.info('the result info label {0} predict index {1} predict_val {2}'.format(190, final_predict_index,\n",
    "                                                                                         final_predict_val))\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    tf.app.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}