ResNet_DeepUD.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow.python.training.moving_averages import assign_moving_average\n",
    "import numpy as np\n",
    "import os, pdb\n",
    "import cv2\n",
    "import numpy as np\n",
    "import random as rn\n",
    "import tensorflow as tf\n",
    "import threading\n",
    "import time\n",
    "from sklearn import metrics\n",
    "import utils\n",
    "global n_classes\n",
    "\n",
    "\n",
    "n_classes = 2\n",
    "\n",
    "\n",
    "def activation(x,name=\"activation\"):\n",
    "    return tf.nn.swish(x, name=name)\n",
    "    \n",
    "def conv2d(name, l_input, w, b, s, p):\n",
    "    l_input = tf.nn.conv2d(l_input, w, strides=[1,s,s,1], padding=p, name=name)\n",
    "    l_input = l_input+b\n",
    "\n",
    "    return l_input\n",
    "\n",
    "def batchnorm(conv, isTraining, name='bn'):\n",
    "    return tf.layers.batch_normalization(conv, training=isTraining, name=\"bn\"+name)\n",
    "\n",
    "def initializer(in_filters, out_filters, name):\n",
    "    w1 = tf.get_variable(name+\"W\", [3, 3, in_filters, out_filters], initializer=tf.truncated_normal_initializer())\n",
    "    b1 = tf.get_variable(name+\"B\", [out_filters], initializer=tf.truncated_normal_initializer())\n",
    "    return w1, b1\n",
    "  \n",
    "def residual_block(in_x, in_filters, out_filters, stride, isDownSampled, name, isTraining):\n",
    "    global ema_gp\n",
    "    # first convolution layer\n",
    "    if isDownSampled:\n",
    "      in_x = tf.nn.avg_pool(in_x, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='VALID')\n",
    "      \n",
    "    x = batchnorm(in_x, isTraining, name=name+'FirstBn')\n",
    "    x = activation(x)\n",
    "    w1, b1 = initializer(in_filters, in_filters, name+\"first_res\")\n",
    "    x = conv2d(name+'r1', x, w1, b1, 1, \"SAME\")\n",
    "\n",
    "    # second convolution layer\n",
    "    x = batchnorm(x, isTraining, name=name+'SecondBn')\n",
    "    x = activation(x)\n",
    "    w2, b2 = initializer(in_filters, out_filters, name+\"Second_res\")\n",
    "    x = conv2d(name+'r2', x, w2, b2, 1, \"SAME\")\n",
    "    \n",
    "    if in_filters != out_filters:\n",
    "        difference = out_filters - in_filters\n",
    "        left_pad = difference // 2\n",
    "        right_pad = difference - left_pad\n",
    "        identity = tf.pad(in_x, [[0, 0], [0, 0], [0, 0], [left_pad, right_pad]])\n",
    "        return x + identity\n",
    "    else:\n",
    "        return in_x + x\n",
    "\n",
    "      \n",
    "def ResNet(_X, isTraining):\n",
    "    global n_classes\n",
    "    w1 = tf.get_variable(\"initW\", [7, 7, 3, 96], initializer=tf.truncated_normal_initializer())\n",
    "    b1 = tf.get_variable(\"initB\", [96], initializer=tf.truncated_normal_initializer())\n",
    "    x = conv2d('conv1', _X, w1, b1, 4, \"VALID\")\n",
    "    \n",
    "    filters_num = [96,128,256]\n",
    "    block_num = [2,2,2]\n",
    "    l_cnt = 1\n",
    "    for i in range(len(filters_num)):\n",
    "      for j in range(block_num[i]):\n",
    "          \n",
    "          if ((j==block_num[i]-1) & (i<len(filters_num)-1)):\n",
    "            x = residual_block(x, filters_num[i], filters_num[i+1], 2, True, 'ResidualBlock%d'%(l_cnt), isTraining)\n",
    "            print('[L-%d] Build %dth connection layer %d from %d to %d channels' % (l_cnt, i, j, filters_num[i], filters_num[i+1]))\n",
    "          else:\n",
    "            x = residual_block(x, filters_num[i], filters_num[i], 1, False, 'ResidualBlock%d'%(l_cnt), isTraining)\n",
    "            print('[L-%d] Build %dth residual block %d with %d channels' % (l_cnt,i, j, filters_num[i]))\n",
    "          l_cnt +=1\n",
    "\n",
    "    saliency = x\n",
    "    x_shape = x.get_shape().as_list()\n",
    "    dense1 = x_shape[1]*x_shape[2]*x_shape[3]\n",
    "    W = tf.get_variable(\"featW\", [dense1, 128], initializer=tf.truncated_normal_initializer())\n",
    "    b = tf.get_variable(\"featB\", [128], initializer=tf.truncated_normal_initializer())\n",
    "    dense1 = tf.reshape(x, [-1, dense1])\n",
    "    feat = tf.nn.softmax(tf.matmul(dense1, W) + b)\n",
    "    \n",
    "    with tf.variable_scope('Final'):\n",
    "        x = batchnorm(x, isTraining, name='FinalBn')\n",
    "        x = activation(x)\n",
    "        wo, bo=initializer(filters_num[-1], n_classes, \"FinalOutput\")\n",
    "        x = conv2d('final', x, wo, bo, 1, \"SAME\")\n",
    "\n",
    "\n",
    "        x=tf.reduce_mean(x, [1, 2])\n",
    "\n",
    "        W = tf.get_variable(\"FinalW\", [n_classes, n_classes], initializer=tf.truncated_normal_initializer())\n",
    "        b = tf.get_variable(\"FinalB\", [n_classes], initializer=tf.truncated_normal_initializer())\n",
    "\n",
    "        out = tf.matmul(x, W) + b\n",
    "                            \n",
    "\n",
    "    return out, feat, saliency\n",
    "\n",
    "\n",
    "#==========================================================================\n",
    "#=============Reading data in multithreading manner========================\n",
    "#==========================================================================\n",
    "def read_labeled_image_list(image_list_file, training_img_dir):\n",
    "    f = open(image_list_file, 'r')\n",
    "    filenames = []\n",
    "    labels = []\n",
    "\n",
    "    for line in f:\n",
    "        filename, label = line[:-1].split(' ')\n",
    "        filename = training_img_dir+filename\n",
    "        filenames.append(filename)\n",
    "        labels.append(int(label))\n",
    "        \n",
    "    return filenames, labels\n",
    "    \n",
    "    \n",
    "def read_images_from_disk(input_queue, size1=64):\n",
    "    label = input_queue[1]\n",
    "    fn=input_queue[0]\n",
    "    file_contents = tf.read_file(input_queue[0])\n",
    "    example = tf.image.decode_jpeg(file_contents, channels=3)\n",
    "    \n",
    "    #example = tf.image.decode_png(file_contents, channels=3, name=\"dataset_image\") # png fo rlfw\n",
    "    example=tf.image.resize_images(example, [size1,size1])\n",
    "    return example, label, fn\n",
    "    \n",
    "def setup_inputs(sess, filenames, training_img_dir, image_size=64, crop_size=64, isTest=False, batch_size=128):\n",
    "    \n",
    "    # Read each image file\n",
    "    image_list, label_list = read_labeled_image_list(filenames, training_img_dir)\n",
    "\n",
    "    images = tf.cast(image_list, tf.string)\n",
    "    labels = tf.cast(label_list, tf.int64)\n",
    "     # Makes an input queue\n",
    "    if isTest is False:\n",
    "        isShuffle = True\n",
    "        numThr = 4\n",
    "    else:\n",
    "        isShuffle = False\n",
    "        numThr = 1\n",
    "        \n",
    "    input_queue = tf.train.slice_input_producer([images, labels], shuffle=isShuffle)\n",
    "    image, y,fn = read_images_from_disk(input_queue)\n",
    "\n",
    "    channels = 3\n",
    "    image.set_shape([None, None, channels])\n",
    "        \n",
    "    # Crop and other random augmentations\n",
    "    if isTest is False:\n",
    "        image = tf.image.random_flip_left_right(image)\n",
    "        image = tf.image.random_saturation(image, .95, 1.05)\n",
    "        image = tf.image.random_brightness(image, .05)\n",
    "        image = tf.image.random_contrast(image, .95, 1.05)\n",
    "        \n",
    "    image = tf.cast(image, tf.float32)/255.0\n",
    "    \n",
    "    image, y,fn = tf.train.batch([image, y, fn], batch_size=batch_size, capacity=batch_size*3, num_threads=numThr, name='labels_and_images')\n",
    "\n",
    "    tf.train.start_queue_runners(sess=sess)\n",
    "\n",
    "    return image, y, fn, len(label_list)\n",
    "\n",
    "#==========================================================================\n",
    "#=============Reading data in multithreading manner========================\n",
    "#==========================================================================\n",
    "def read_labeled_image_list2(image_list_file, training_img_dir):\n",
    "    f = open(image_list_file, 'r')\n",
    "    filenames = []\n",
    "    filenames2 = []\n",
    "    labels = []\n",
    "\n",
    "    for line in f:\n",
    "        filename, fn2, label = line[:-1].split(' ')\n",
    "        filename = training_img_dir+filename\n",
    "        filenames.append(filename)\n",
    "        filename = training_img_dir+fn2\n",
    "        filenames2.append(filename)\n",
    "        labels.append(int(label))\n",
    "        \n",
    "    return filenames, filenames2, labels\n",
    "    \n",
    "    \n",
    "def read_images_from_disk2(input_queue, size1=64):\n",
    "    label = input_queue[2]\n",
    "    fn=input_queue[0]\n",
    "    file_contents = tf.read_file(input_queue[0])\n",
    "    file_contents2 = tf.read_file(input_queue[1])\n",
    "    example = tf.image.decode_jpeg(file_contents, channels=3)\n",
    "    example2 = tf.image.decode_jpeg(file_contents2, channels=3)\n",
    "    \n",
    "    #example = tf.image.decode_png(file_contents, channels=3, name=\"dataset_image\") # png fo rlfw\n",
    "    example=tf.image.resize_images(example, [size1,size1])\n",
    "    example2=tf.image.resize_images(example2, [size1,size1])\n",
    "    return example, example2, label, fn\n",
    "    \n",
    "def setup_inputs2(sess, filenames, training_img_dir, image_size=64, crop_size=64, isTest=False, batch_size=128):\n",
    "    \n",
    "    # Read each image file\n",
    "    image_list, image2_list, label_list = read_labeled_image_list2(filenames, training_img_dir)\n",
    "\n",
    "    images = tf.cast(image_list, tf.string)\n",
    "    images2 = tf.cast(image2_list, tf.string)\n",
    "    labels = tf.cast(label_list, tf.int64)\n",
    "     # Makes an input queue\n",
    "    if isTest is False:\n",
    "        isShuffle = True\n",
    "        numThr = 4\n",
    "    else:\n",
    "        isShuffle = False\n",
    "        numThr = 1\n",
    "        \n",
    "    input_queue = tf.train.slice_input_producer([images, images2, labels], shuffle=isShuffle)\n",
    "    image, image2, y,fn = read_images_from_disk2(input_queue)\n",
    "\n",
    "    channels = 3\n",
    "    image.set_shape([None, None, channels])\n",
    "    image2.set_shape([None, None, channels])\n",
    "        \n",
    "    # Crop and other random augmentations\n",
    "    if isTest is False:\n",
    "        image = tf.image.random_flip_left_right(image)\n",
    "        image = tf.image.random_saturation(image, .95, 1.05)\n",
    "        image = tf.image.random_brightness(image, .05)\n",
    "        image = tf.image.random_contrast(image, .95, 1.05)\n",
    "        image2 = tf.image.random_flip_left_right(image2)\n",
    "        image2 = tf.image.random_saturation(image2, .95, 1.05)\n",
    "        image2 = tf.image.random_brightness(image2, .05)\n",
    "        image2 = tf.image.random_contrast(image2, .95, 1.05)\n",
    "        \n",
    "\n",
    "    image = tf.cast(image, tf.float32)/255.0\n",
    "    image2 = tf.cast(image2, tf.float32)/255.0\n",
    "    \n",
    "    image, image2, y,fn = tf.train.batch([image, image2, y, fn], batch_size=batch_size, capacity=batch_size*3, num_threads=numThr, name='labels_and_images')\n",
    "\n",
    "    tf.train.start_queue_runners(sess=sess)\n",
    "\n",
    "    return image, image2, y, fn, len(label_list)\n",
    "'''\n",
    "Main Program:\n",
    "Fake image detection based on siamese network (with constrastive loss)\n",
    "'''\n",
    "batch_size = 32\n",
    "display_step = 80\n",
    "learning_rate = tf.placeholder(tf.float32)      # Learning rate to be fed\n",
    "lr = 1e-3                         # Learning rate start\n",
    "tst = tf.placeholder(tf.bool)\n",
    "iter = tf.placeholder(tf.int32)\n",
    "print('GO!!')\n",
    "\n",
    "\n",
    "# In[ ]:\n",
    "\n",
    "\n",
    "# Setup the tensorflow...\n",
    "config = tf.ConfigProto()\n",
    "config.gpu_options.allow_growth = True\n",
    "sess = tf.Session(config=config)\n",
    "\n",
    "print(\"Preparing the training & validation data...\")\n",
    "# pairwise_progressGAN.txt ==> PGGAN is excluded in the training set\n",
    "# pairwise_wgan.txt        ==> WGAN is excluded in the training set\n",
    "sia_data, sia_data2, sia_labels, sialist1, slen1 = setup_inputs2(sess, \"data/pairwise_progressGAN.txt\", \"./\", batch_size=batch_size)\n",
    "# For Level-2 learning : Classifier learning\n",
    "train_data, train_labels, filelist1, glen1 = setup_inputs(sess, \"data/trainprogressGAN.txt\", \"ㄡ/\", batch_size=batch_size)\n",
    "# For Level-2 test: valprogressGAN.txt means that the PGGAN is used to verify the system performance.\n",
    "val_data, val_labels, filelist2, tlen1 = setup_inputs(sess, \"data/valprogressGAN.txt\", \"./\", batch_size=1000,isTest=True)\n",
    "print(\"Found %d pairs, %d training images, and %d validation images...\" % (slen1, glen1, tlen1))\n",
    "\n",
    "max_iter = glen1*15  # How many epochs we want to used to train...\n",
    "print(\"Preparing the training model with learning rate = %.5f...\" % (lr))\n",
    "\n",
    "# Make a model\n",
    "with tf.variable_scope(\"ResNet\") as scope:\n",
    "    #=================Level-1 training based on constrative learning===================\n",
    "    _, feat1,_ = ResNet(sia_data, True)\n",
    "    scope.reuse_variables()\n",
    "    _, feat2,_ = ResNet(sia_data2, True)\n",
    "    scope.reuse_variables()\n",
    "    #==================================================================================\n",
    "    #=================Level-2 training based on corss-entropy==========================\n",
    "    pred, _,_ = ResNet(train_data, True)\n",
    "    scope.reuse_variables()\n",
    "    valpred, _, saliency = ResNet(val_data, False)\n",
    "    #==================================================================================\n",
    "\n",
    "\n",
    "#==================Set up the constrative loss===============================\n",
    "with tf.name_scope('ContrastiveLoss'):\n",
    "  margin = 0.5\n",
    "  labels_t = tf.cast(train_labels, tf.float32)\n",
    "  labels_f = tf.cast(1-train_labels, tf.float32)         # labels_ = !labels;\n",
    "  eucd2 = tf.pow(feat1- feat2, 2.0)\n",
    "  eucd2 = tf.reduce_sum(eucd2, [1])\n",
    "  eucd = tf.sqrt(eucd2+1e-10, name=\"eucd\")\n",
    "  C = tf.constant(margin, name=\"C\")\n",
    "  pos = labels_t * eucd2\n",
    "  neg = labels_f *tf.pow(tf.maximum(C- eucd, 0), 2)\n",
    "  losses = pos + neg\n",
    "  sialoss = tf.reduce_mean(losses, name=\"Contrastive_loss\")\n",
    "#=============================================================================\n",
    "#==================Set up the optimizer=======================================\n",
    "with tf.name_scope('Loss_and_Accuracy'):\n",
    "  update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)\n",
    "  with tf.control_dependencies(update_ops):\n",
    "    t_vars=tf.trainable_variables() \n",
    "    #t_vars=[var for var in t_vars if 'Final']\n",
    "    cost = tf.losses.sparse_softmax_cross_entropy(labels=train_labels, logits=pred)\n",
    "    optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(cost, var_list=t_vars)\n",
    "    sia_optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(sialoss)\n",
    "  \n",
    "  #================Set up the accuracy measurement============================\n",
    "  correct_prediction = tf.equal(tf.argmax(pred, 1), train_labels)   # Training acc.\n",
    "  accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))\n",
    "  correct_prediction2 = tf.equal(tf.argmax(valpred, 1), val_labels) # Validation acc\n",
    "  accuracy2 = tf.reduce_mean(tf.cast(correct_prediction2, tf.float32))\n",
    "\n",
    "\n",
    "#==========Summary============================\n",
    "tf.summary.scalar(\"Contrastive_loss\", sialoss)\n",
    "tf.summary.scalar('Loss', cost)\n",
    "tf.summary.scalar('Training_Accuracy', accuracy)\n",
    "tf.summary.scalar('Validation_Accuracy', accuracy2)\n",
    "\n",
    "saver = tf.train.Saver()\n",
    "init = tf.global_variables_initializer()\n",
    "sess.run(init)\n",
    "step = 0\n",
    "writer = tf.summary.FileWriter(\"logs/sia/\", sess.graph)\n",
    "summaries = tf.summary.merge_all()\n",
    "\n",
    "print(\"We are going to pretrain model susing ResNet based on contrastive loss!!!\")\n",
    "start_lr = lr\n",
    "while (step * batch_size) < max_iter:\n",
    "    epoch1=np.floor((step*batch_size)/glen1)\n",
    "    # Learning rate decay\n",
    "    if (((step*batch_size)%glen1 < batch_size) & (lr==1e-3) & (epoch1 >2)):\n",
    "        lr /= 10\n",
    "    \n",
    "    if epoch1 <=1:\n",
    "        sess.run([sia_optimizer],  feed_dict={learning_rate: lr})\n",
    "    else:\n",
    "        # Learning rate decay at level-2 training\n",
    "#         if start_lr == lr:\n",
    "#             lr = lr /10\n",
    "        sess.run([optimizer],  feed_dict={learning_rate: lr})\n",
    "        \n",
    "    if (step % 15000==1) & (step>15000):\n",
    "        save_path = saver.save(sess, \"checkpoints/tf_deepUD_model_iter\" + str(step) + \".ckpt\")\n",
    "        print(\"Model saved in file at iteration %d: %s\" % (step*batch_size,save_path))\n",
    "\n",
    "    if step>0 and step % display_step == 0:\n",
    "        # calculate the loss\n",
    "        loss, acc, summaries_string, sia_val = sess.run([cost, accuracy, summaries, sialoss])\n",
    "        print(\"Iter=%d/epoch=%d, Loss=%.6f, Contrastive loss=%.6f, Training Accuracy=%.6f, lr=%f\" % (step*batch_size, epoch1 ,loss, sia_val, acc, lr))\n",
    "        writer.add_summary(summaries_string, step)\n",
    "    \n",
    "    if step>0 and (step % (display_step*20) == 0):\n",
    "        rounds = tlen1 // 1000\n",
    "        valacc=[]\n",
    "        vis=[]\n",
    "        tis=[]\n",
    "        for k in range(rounds):\n",
    "          a2, vi, ti = sess.run([accuracy2, tf.argmax(valpred, 1), val_labels])\n",
    "          valacc.append(a2)\n",
    "          vis.append(vi)\n",
    "          tis.append(ti)\n",
    "        tis = np.reshape(np.asarray(tis), [-1])\n",
    "        vis = np.reshape(np.asarray(vis), [-1])\n",
    "        precision=metrics.precision_score(tis, vis) \n",
    "        recall=metrics.recall_score(tis, vis)\n",
    "        \n",
    "        sal, valimg = sess.run([saliency, val_data])\n",
    "        utils.batchimwrite2(sal, 'saliency_img/sal')\n",
    "        utils.batchimwrite2(valimg, 'saliency_img/img')\n",
    "\n",
    "        print(\"Iter=%d/epoch=%d, Validation Accuracy=%.6f, Precision=%.6f, Recall=%.6f\" % (step*batch_size, epoch1 , np.mean(valacc), precision, recall))\n",
    "\n",
    "  \n",
    "    step += 1\n",
    "print(\"Optimization Finished!\")\n",
    "save_path = saver.save(sess, \"checkpoints/tf_deepUD_model.ckpt\")\n",
    "print(\"Model saved in file: %s\" % save_path)\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}