Ptcho_CelebA_clean.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using TensorFlow backend.\n"
     ]
    }
   ],
   "source": [
    "# IMPORTING MODULES\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import scipy.io as sci\n",
    "import matplotlib.pyplot as plt\n",
    "import tensorflow as tf\n",
    "import keras.backend as K\n",
    "import numpy as np\n",
    "from Datasets import *\n",
    "from CelebAGenerator import *\n",
    "from SVHNGenerator import *\n",
    "K.set_learning_phase(0)\n",
    "from glob import glob\n",
    "import os\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "from skimage.measure import compare_psnr, compare_ssim\n",
    "import scipy.misc as scipy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# LOADING VARIABLES\n",
    "var1 = sci.loadmat('Parameters for 56 image,15 Dia,0.65 overlap,9 grid,0.01 subsamp')\n",
    "spacing= var1['spacing'].astype('float32')\n",
    "P_op = var1['P_op'].astype('float32')\n",
    "samplingIndices = var1['samplingIndices'].astype('int32')\n",
    "pupil = var1['pupil'].astype('float32')[:,:,1]\n",
    "h = 56#var1['h'].astype('float32')\n",
    "w = 56#var1['w'].astype('float32')\n",
    "\n",
    "###############################################################################\n",
    "###################### this portion is for subsampling experiments only #######\n",
    "###############################################################################\n",
    "#var1 = sci.loadmat('Parameters for 56 image,15 Dia,0.65 overlap,9 grid,0.01 subsamp')\n",
    "#var2 = sci.loadmat('Parameters for 56 image,15 Dia,0.65 overlap,9 grid,0.02 subsamp')\n",
    "#var3 = sci.loadmat('Parameters for 56 image,15 Dia,0.65 overlap,9 grid,0.03 subsamp')\n",
    "#var4 = sci.loadmat('Parameters for 56 image,15 Dia,0.65 overlap,9 grid,0.05 subsamp')\n",
    "#var5 = sci.loadmat('Parameters for 56 image,15 Dia,0.65 overlap,9 grid,0.10 subsamp')\n",
    "#spacing= var1['spacing'].astype('float32')\n",
    "\n",
    "#P_op1 = var1['P_op'].astype('float32')\n",
    "#P_op2 = var2['P_op'].astype('float32')\n",
    "#P_op3 = var3['P_op'].astype('float32')\n",
    "#P_op4 = var4['P_op'].astype('float32')\n",
    "#P_op5 = var5['P_op'].astype('float32')\n",
    "\n",
    "\n",
    "#samplingIndices = var1['samplingIndices'].astype('int32')\n",
    "#pupil = var1['pupil'].astype('float32')[:,:,1]\n",
    "#h = 56#var1['h'].astype('float32')\n",
    "#w = 56#var1['w'].astype('float32')\n",
    "\n",
    "\n",
    "###############################################################################\n",
    "###############################################################################\n",
    "###############################################################################\n",
    "\n",
    "# HYPERPARAMETERS\n",
    "LEARNING_RATE =  0.05   #0.001\n",
    "RANDOM_RESTARTS = 1\n",
    "#NOISE_STD       = 0.005\n",
    "STEPS           = 1000\n",
    "optimizer       = tf.train.AdamOptimizer(learning_rate=LEARNING_RATE)\n",
    "\n",
    "PLOT_LOSS       = True\n",
    "SAVE_RESULTS    = True\n",
    "rand_rest       = 1 #random restarts\n",
    "latent_dim      = 50\n",
    "#L               = 64\n",
    "cam_arr         = P_op.shape[2]\n",
    "\n",
    "# FUNCTION DEFINITIONS\n",
    "def fftshift(x):\n",
    "    xx = tf.reshape(x,[L,L,-1])\n",
    "    a =xx[0:int(L/2),0:int(L/2),:]\n",
    "    b =xx[0:int(L/2),int(L/2):L,:]\n",
    "    c =xx[int(L/2):L,0:int(L/2),:]\n",
    "    d =xx[int(L/2):L,int(L/2):L,:]\n",
    "    B1 =tf.concat([d,c],axis=1)\n",
    "    B2 =tf.concat([b,a],axis=1)\n",
    "    B3 =tf.concat([B1,B2],axis=0)\n",
    "    B3 =tf.reshape(B3,[L,L,-1])\n",
    "    return B3\n",
    "\n",
    "\n",
    "def forward_map(x): # change this for subsampling remove P_op\n",
    "    h = tf.constant(L);xx=tf.to_complex64(tf.reshape(x,[L,L]));\n",
    "    X   = tf.reshape(fftshift(tf.fft2d(xx)),[h,h])\n",
    "    X1   = tf.reshape(tf.pad(X, tf.constant([[int(np.floor(spacing*(9-1)/2)),int(np.floor(spacing*(9-1)/2))],[int(np.floor(spacing*(9-1)/2)),int(np.floor(spacing*(9-1)/2))]]), \"CONSTANT\"),[-1,1])\n",
    "    ind1 = tf.reshape(tf.gather(X1,tf.constant(samplingIndices-1)),[h,h,-1])\n",
    "    ifft_mul = tf.ifft2d(tf.to_complex64(tf.multiply(tf.transpose(ind1),tf.constant(pupil,dtype=tf.complex64))))\n",
    "    abs_ifft_mul = tf.transpose((tf.abs(ifft_mul)),[1,2,0])\n",
    "    sub_samp = tf.multiply(tf.constant(P_op),abs_ifft_mul)\n",
    "    return [sub_samp]\n",
    "\n",
    "# LOADING CelebA GENERATOR\n",
    "#gen = CelebAGenerator()\n",
    "#gen.GenerateModel()\n",
    "#gen.LoadWeights()\n",
    "#G = gen.GetModels()\n",
    "#G.trainable = False\n",
    "\n",
    "DATASET         = \"mnist\"                  # \"celeba\" or \"mnist\"\n",
    "# LOADING GENERATOR\n",
    "if DATASET == \"celeba\":\n",
    "    gen = CelebAGenerator()\n",
    "    gen.GenerateModel()\n",
    "    gen.LoadWeights()\n",
    "    G = gen.GetModels()\n",
    "    L  = 64\n",
    "    \n",
    "if DATASET == \"mnist\":\n",
    "    gen = SVHNGenerator()\n",
    "    gen.GenerateModel()\n",
    "    gen.LoadWeights()\n",
    "    _, _, G = gen.GetModels()\n",
    "    L   = 56\n",
    "#G.trainable = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(56, 56)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pupil.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$\\Huge{\\textbf{ Range subsample Script (in loop)} (MNIST)}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "Orig_Path   = './original_images/mnist/*.png'\n",
    "Orig_Path   = './original_images/mnist/*.png'\n",
    "P_op_arr = [P_op1,P_op2,P_op3,P_op4,P_op5]\n",
    "\n",
    "\n",
    "for kk in range(len(X_Orig)):\n",
    "    images = []\n",
    "    im1 = X_Orig[kk]\n",
    "    \n",
    "    for j in range(len(P_op_arr)):\n",
    "        \n",
    "\n",
    "        P_op = P_op_arr[j]\n",
    "        \n",
    "        \n",
    "        if DATASET == \"celeba\":\n",
    "            gen = CelebAGenerator()\n",
    "            gen.GenerateModel()\n",
    "            gen.LoadWeights()\n",
    "            G = gen.GetModels()\n",
    "            L  = 64\n",
    "    \n",
    "        if DATASET == \"mnist\":\n",
    "            gen = SVHNGenerator()\n",
    "            gen.GenerateModel()\n",
    "            gen.LoadWeights()\n",
    "            _, _, G = gen.GetModels()\n",
    "            L   = 56\n",
    "        G.trainable = False\n",
    "\n",
    "        im = im1 \n",
    "        abs_ifft_mul = []\n",
    "\n",
    "\n",
    "        v1 = im\n",
    "        #v = (v1/255).astype('float32')\n",
    "        k = (samplingIndices-1);\n",
    "        X   = np.fft.fftshift(np.fft.fft2(v1))#/(h*w)**0.5 \n",
    "        Xx   = np.reshape(np.pad(X,int(np.floor(spacing*(9-1))/2),'constant'),[-1,1])\n",
    "        ind = np.reshape(Xx[k],[L,L,-1])\n",
    "        ifft_mul = np.fft.ifft2(ind.transpose()*pupil)\n",
    "        ifft_mul = np.transpose(ifft_mul,[1,2,0])\n",
    "        abs_ifft_mul = (P_op*np.abs(ifft_mul))\n",
    "\n",
    "        #abs_ifft_mul = np.clip((abs_ifft_mul + np.random.normal(0,0,size=abs_ifft_mul.shape)),0,1)\n",
    "\n",
    "        #plt.subplot(1,2,1)\n",
    "        #plt.imshow(im)\n",
    "        #plt.subplot(1,2,2)\n",
    "        #plt.imshow(abs_ifft_mul_display[np.int(cam_arr/2),:,:,:])\n",
    "        # FORWARD MAP\n",
    "        z_tf                = tf.Variable(tf.random_normal(shape=(rand_rest, latent_dim)))\n",
    "        Y_tf                = tf.placeholder(dtype=\"float32\", shape=(L,L,cam_arr)) #PUT Y HERE AND CHANGE SHAPE\n",
    "\n",
    "        xG_tf               = G(z_tf)[:,:,:,0]\n",
    "        #xG_tf               = (xG_tf+1)/2\n",
    "        yGpred_tf           = forward_map(xG_tf,P_op)\n",
    "        Loss_tf             = tf.reduce_mean( tf.square( Y_tf - yGpred_tf ))\n",
    "        opt                 = optimizer.minimize(Loss_tf, var_list=[z_tf])\n",
    "\n",
    "        # SOLVING LEAST SQUARES\n",
    "        # SOLVING LEAST SQUARES\n",
    "        sess = K.get_session()\n",
    "        sess.run(tf.variables_initializer([z_tf]))\n",
    "        Loss = []\n",
    "        for i in tqdm(range(STEPS)):\n",
    "            _ , loss = sess.run([opt,Loss_tf], feed_dict={Y_tf:abs_ifft_mul}) #put Y and A here too\n",
    "            if (i %50) == 0:\n",
    "                x_hat= sess.run([xG_tf])\n",
    "                x_hat = np.array(x_hat)\n",
    "                print(compare_ssim(im1,x_hat[0,0,:,:].astype('float64')))\n",
    "                plt.imshow(x_hat[0,0,:,:],cmap='gray')\n",
    "                plt.show() \n",
    "            Loss.append(np.mean(loss))\n",
    "        \n",
    "        scipy.imsave('test'+str(kk+1)+'_subsampled_'+str(j+1)+'.png',x_hat[0,0,:,:])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$\\Huge{\\textbf{ DeepPtych+ code for MNIST -- Subsampling}}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "Orig_Path   = './original_images/mnist/*.png'\n",
    "X_Orig = np.array([ imread(path) for path in glob(Orig_Path)])/255\n",
    "P_op_arr = [P_op1,P_op2,P_op3,P_op4,P_op5]\n",
    "\n",
    "\n",
    "for kk in range(len(X_Orig)):\n",
    "    images = []\n",
    "    im1 = X_Orig[kk]\n",
    "    \n",
    "    for j in range(len(P_op_arr)):\n",
    "        \n",
    "\n",
    "        P_op = P_op_arr[j]\n",
    "        \n",
    "        \n",
    "        if DATASET == \"celeba\":\n",
    "            gen = CelebAGenerator()\n",
    "            gen.GenerateModel()\n",
    "            gen.LoadWeights()\n",
    "            G = gen.GetModels()\n",
    "            L  = 64\n",
    "    \n",
    "        if DATASET == \"mnist\":\n",
    "            gen = SVHNGenerator()\n",
    "            gen.GenerateModel()\n",
    "            gen.LoadWeights()\n",
    "            _, _, G = gen.GetModels()\n",
    "            L   = 56\n",
    "        G.trainable = False\n",
    "\n",
    "        im = im1 \n",
    "        abs_ifft_mul = []\n",
    "\n",
    "\n",
    "        v1 = im\n",
    "        #v = (v1/255).astype('float32')\n",
    "        k = (samplingIndices-1);\n",
    "        X   = np.fft.fftshift(np.fft.fft2(v1))#/(h*w)**0.5 \n",
    "        Xx   = np.reshape(np.pad(X,int(np.floor(spacing*(9-1))/2),'constant'),[-1,1])\n",
    "        ind = np.reshape(Xx[k],[L,L,-1])\n",
    "        ifft_mul = np.fft.ifft2(ind.transpose()*pupil)\n",
    "        ifft_mul = np.transpose(ifft_mul,[1,2,0])\n",
    "        abs_ifft_mul = (P_op*np.abs(ifft_mul))\n",
    "\n",
    "        #abs_ifft_mul = np.clip((abs_ifft_mul + np.random.normal(0,0,size=abs_ifft_mul.shape)),0,1)\n",
    "\n",
    "        #plt.subplot(1,2,1)\n",
    "        #plt.imshow(im)\n",
    "        #plt.subplot(1,2,2)\n",
    "        #plt.imshow(abs_ifft_mul_display[np.int(cam_arr/2),:,:,:])\n",
    "        # FORWARD MAP\n",
    "        # FORWARD MAP\n",
    "        z_tf                = tf.Variable(tf.random_normal(shape=(rand_rest, latent_dim)))\n",
    "        Y_tf                = tf.placeholder(dtype=\"float32\", shape=(L,L,cam_arr)) #PUT Y HERE AND CHANGE SHAPE\n",
    "\n",
    "        xG_tf               = G(z_tf)[:,:,:,0]\n",
    "        #xG_tf               = (xG_tf+1)/2\n",
    "        yGpred_tf           = forward_map(xG_tf,P_op)\n",
    "\n",
    "\n",
    "\n",
    "        x_tf  = tf.Variable(tf.random_normal(mean = 0.5, stddev  = 0.01,shape=([1,L,L]),dtype = 'float32'),constraint=lambda t: tf.clip_by_value(t, 0, 1))\n",
    "        #x_tf = tf.clip_by_value(x_tf,0,1)\n",
    "\n",
    "        loss1 = tf.reduce_mean( tf.square( Y_tf - yGpred_tf ))\n",
    "        loss2 = tf.reduce_mean( tf.square( Y_tf - forward_map(x_tf,P_op)))\n",
    "        loss3 =  tf.reduce_mean(tf.square(x_tf - xG_tf))\n",
    "        loss4 = tf.image.total_variation(x_tf)\n",
    "\n",
    "        Loss_tf = 0.0*loss1 + 50.0*loss2 + 0.01*loss3 #+ 0.0000005*loss4\n",
    "        opt     = optimizer.minimize(Loss_tf, var_list=[z_tf,x_tf])\n",
    "\n",
    "        # SOLVING LEAST SQUARES\n",
    "        # SOLVING LEAST SQUARES\n",
    "        # SOLVING LEAST SQUARES\n",
    "        sess = K.get_session()\n",
    "        sess.run(tf.variables_initializer([z_tf,x_tf]))\n",
    "        Loss = []\n",
    "        for i in tqdm(range(STEPS)):\n",
    "            _ , loss = sess.run([opt,Loss_tf], feed_dict={Y_tf:abs_ifft_mul}) #put Y and A here too\n",
    "            if (i %50) == 0:\n",
    "                x_hat= sess.run([x_tf])\n",
    "                x_hat = np.array(x_hat)\n",
    "                print(compare_ssim(im1,x_hat[0,0,:,:].astype('float64')))\n",
    "                plt.imshow(x_hat[0,0,:,:],cmap='gray')\n",
    "                plt.show() \n",
    "            Loss.append(np.mean(loss))\n",
    "        \n",
    "        scipy.imsave('test'+str(kk+1)+'_subsampled_'+str(j+1)+'.png',x_hat[0,0,:,:])"
   ]
  }
 ],
 "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}