{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "NOTE\n",
    "--\n",
    "\n",
    "This is a CNN version instead of PFF; for individual study; will be merged later.\n",
    "\n",
    "Shu Kong \n",
    "\n",
    "20190130"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Predictive Filter Flow for Non-Uniform Motion Blur Removal\n",
    "================\n",
    "**Author**: `Shu Kong <https://www.ics.uci.edu/~skong2/pff.html>`\n",
    "\n",
    "**Date**: Nov. 2018\n",
    "\n",
    "This is a jupyter demo for showing how to train, evaluate, and visualize the model of non-uniform motion deblur. It squeezes the following parts. Please read through, uncomment lines and run accordingly for your interest.\n",
    "\n",
    "-  defining model architecture\n",
    "-  defining the loss function\n",
    "-  training protocal\n",
    "-  evaluation protocal\n",
    "-  analysis using PCA, k-means and t-SNE.\n",
    "\n",
    "\n",
    "Others to notice: \n",
    "- Pytorch is used for this work.\n",
    "- Please manually install packages as shown below if necessary\n",
    "- For the training set, the DIV2K dataset is not uploaded due to its large size. If you want to train the model using the whole training set, please download manually, chop each image into overlapping 512x512 sub-images, and put them in the corresponding folder (read through this jupyter script to see where to put them:-) For demonstration and turtorial purpose, other small-scale training images are uploaded here."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "import packages\n",
    "------------------\n",
    "\n",
    "Some packages are installed automatically if you use Anaconda. As pytorch is used here, you are expected to install that in your machine. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function, division\n",
    "import os, random, time, copy\n",
    "from skimage import io, transform\n",
    "import numpy as np\n",
    "import os.path as path\n",
    "import scipy.io as sio\n",
    "from scipy import misc\n",
    "from scipy import ndimage, signal\n",
    "import scipy\n",
    "import pickle\n",
    "\n",
    "import math\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image\n",
    "from io import BytesIO\n",
    "from skimage import data, img_as_float\n",
    "from skimage.measure import compare_ssim as ssim\n",
    "from skimage.measure import compare_psnr as psnr\n",
    "\n",
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.optim import lr_scheduler \n",
    "import torch.nn.functional as F\n",
    "from torch.autograd import Variable\n",
    "\n",
    "import torchvision\n",
    "from torchvision import datasets, models, transforms\n",
    "\n",
    "from libs_deblur.utils.metrics import *\n",
    "from libs_deblur.utils.flow_functions import *\n",
    "from libs_deblur.models.pixel_embedding_model import *\n",
    "from libs_deblur.datasetMotionBlur import *\n",
    "from libs_deblur.trainvalGaussBlur_cnn import *\n",
    "import libs_deblur.pyblur\n",
    "import warnings # ignore warnings\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " Setup config parameters\n",
    " -----------------\n",
    " \n",
    " There are several things to setup, like which GPU to use, where to read images and save files, etc. Please read and understand this. By default, you should be able to run this script smoothly by changing nothing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# set device, which gpu to use.\n",
    "device ='cpu'\n",
    "if torch.cuda.is_available(): device='cuda:0'\n",
    "\n",
    "################## set attributes for this project/experiment ##################\n",
    "# config result folder\n",
    "exp_dir = './libs_deblur' # experiment directory, used for reading the init model\n",
    "project_name = 'demo' # name this project as \"demo\", used for saving files\n",
    "path_to_root = './libs_deblur/dataset' # where to fetch data\n",
    "\n",
    "\n",
    "batch_size = 16 # small batch size for demonstration; using larger batch size (like 56 and 64) for training\n",
    "\n",
    "embedding_dim = 16 # dimension of the learned embedding space\n",
    "kernel_size = 17 # the kernel size in the filter flow\n",
    "cropSize = [64, 64] # patch size for training the model\n",
    "sigmaMin=0.5\n",
    "sigmaMax=2\n",
    "\n",
    "lambda_norm = 0.1\n",
    "total_epoch_num = 500 # total number of epoch in training\n",
    "base_lr = 0.0005 # base learning rate\n",
    "\n",
    "torch.cuda.device_count()\n",
    "torch.cuda.empty_cache()\n",
    "\n",
    "save_dir = os.path.join(exp_dir, project_name) # where to save the log file and trained models.\n",
    "print(save_dir)    \n",
    "if not os.path.exists(save_dir): os.makedirs(save_dir)\n",
    "log_filename = os.path.join(save_dir, 'train.log')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define model architecture\n",
    "---------\n",
    "\n",
    "Here is the definition of the model architecture. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SiamesePixelEmbed(nn.Module):\n",
    "    def __init__(self, emb_dimension=64, filterSize=11, device='cpu', pretrained=False):\n",
    "        super(SiamesePixelEmbed, self).__init__()\n",
    "        self.device = device\n",
    "        self.emb_dimension = emb_dimension  \n",
    "        self.PEMbase = PixelEmbedModelResNet18(emb_dimension=self.emb_dimension, pretrained=pretrained)  \n",
    "        self.rawEmbFeature1 = 0\n",
    "        self.rawEmbFeature2 = 0        \n",
    "        self.embFeature1_to_2 = 0\n",
    "        self.embFeature1_to_2 = 0\n",
    "        self.filterSize = filterSize\n",
    "        self.filterSize2Channel = self.filterSize**2\n",
    "                \n",
    "        self.ordered_embedding = nn.Sequential(            \n",
    "            nn.Conv2d(self.emb_dimension, self.filterSize2Channel, kernel_size=3, padding=1, bias=False),\n",
    "            nn.ReLU(True),\n",
    "            nn.BatchNorm2d(self.filterSize2Channel),     \n",
    "            nn.Conv2d(self.filterSize2Channel, self.filterSize2Channel, kernel_size=3, padding=1, bias=False),\n",
    "            nn.ReLU(True),\n",
    "            nn.BatchNorm2d(self.filterSize2Channel),            \n",
    "            nn.Conv2d(self.filterSize2Channel, 1, kernel_size=3, padding=1, bias=True)\n",
    "        )\n",
    "        \n",
    "        \n",
    "    def forward(self, inputs1):        \n",
    "        self.rawEmbFeature1 = self.PEMbase.forward(inputs1)        \n",
    "        self.embFeature1_to_2 = self.ordered_embedding(self.rawEmbFeature1)        \n",
    "        #self.embFeature1_to_2 = F.softmax(self.embFeature1_to_2, 1)        \n",
    "        return self.embFeature1_to_2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define loss function\n",
    "---------\n",
    "\n",
    "Here is the L1 loss at pixel level, including how to warp the image with the predictive filter flow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LossOrderedPairReconstruction(nn.Module):\n",
    "    def __init__(self, device='cpu'):\n",
    "        super(LossOrderedPairReconstruction, self).__init__()\n",
    "        self.device = device\n",
    "        self.reconstructImage = 0\n",
    "        \n",
    "    def forward(self, image1, image2):\n",
    "        N,C,H,W = image1.size()\n",
    "        self.reconstructImage = image1\n",
    "        diff = self.reconstructImage - image2               \n",
    "        diff = torch.abs(diff)       \n",
    "        totloss = torch.sum(torch.sum(torch.sum(torch.sum(diff))))        \n",
    "        return totloss/(N*C*H*W)    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Setup dataset\n",
    "-----------\n",
    "\n",
    "Here is where to fetch training and testing dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "################## dataset ###################\n",
    "transform4Image = transforms.Compose(\n",
    "    [transforms.ToTensor(),\n",
    "     transforms.Normalize((127.,127.,127.),(127.,127.,127.)) # (mean, std)\n",
    "    ]) # (mean, std)\n",
    "\n",
    "whole_datasets = {set_name: \n",
    "                  Dataset4MotionBlur(root_dir=path_to_root,\n",
    "                                     size=cropSize, set_name=set_name, \n",
    "                                     transform=transform4Image, \n",
    "                                     sigmaMin=sigmaMin, sigmaMax=sigmaMax)\n",
    "                  for set_name in ['train', 'val']}\n",
    "\n",
    "\n",
    "dataloaders = {set_name: DataLoader(whole_datasets[set_name], \n",
    "                                    batch_size=batch_size,\n",
    "                                    shuffle=set_name=='train', \n",
    "                                    num_workers=5) # num_work can be set to batch_size\n",
    "               for set_name in ['train', 'val']}\n",
    "\n",
    "dataset_sizes = {set_name: len(whole_datasets[set_name]) for set_name in ['train', 'val']}\n",
    "print(dataset_sizes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Visualizing some training images\n",
    "-----------\n",
    "\n",
    "Visualizing some training image pairs, as well as the random blurring kernel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample = iter(dataloaders['train']).next()\n",
    "kernelList, imgList1, imgList2 = sample\n",
    "\n",
    "imgList1 = imgList1.to(device)\n",
    "imgList2 = imgList2.to(device)\n",
    "\n",
    "\n",
    "figWinNumHeight, figWinNumWidth, subwinCount = batch_size/4, 4*3, 1\n",
    "plt.figure(figsize=(15,8), dpi=88, facecolor='w', edgecolor='k') # figsize -- inch-by-inch\n",
    "plt.clf()\n",
    "\n",
    "for sampleIndex in range(batch_size):\n",
    "    # visualize blurry image\n",
    "    plt.subplot(figWinNumHeight,figWinNumWidth,subwinCount)\n",
    "    subwinCount += 1\n",
    "    noisy_img = imgList1[sampleIndex].cpu().numpy().squeeze().transpose((1,2,0))\n",
    "    noisy_img = (noisy_img+1)/2    \n",
    "    noisy_img = noisy_img.clip(0,1)        \n",
    "    plt.imshow(noisy_img, cmap='gray')    \n",
    "    plt.axis('off')\n",
    "    plt.title('blurry img')\n",
    "    \n",
    "    \n",
    "    plt.subplot(figWinNumHeight,figWinNumWidth,subwinCount)\n",
    "    subwinCount += 1\n",
    "    img = imgList2[sampleIndex].cpu().numpy().squeeze().transpose((1,2,0))\n",
    "    img = (img+1)/2\n",
    "    img = img.clip(0,1)        \n",
    "    plt.imshow(img, cmap='gray')\n",
    "    plt.axis('off')\n",
    "    plt.title('original img')\n",
    "    \n",
    "    plt.subplot(figWinNumHeight,figWinNumWidth,subwinCount)\n",
    "    subwinCount += 1\n",
    "    kernel = kernelList[sampleIndex].cpu().numpy().squeeze()\n",
    "\n",
    "    plt.imshow(kernel, cmap='gray')\n",
    "    plt.axis('off')\n",
    "    plt.title('kernel')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Training demo\n",
    "-------------\n",
    "\n",
    "Run this cell to train the model (essentially resuming the training over the model provided here). Stop it if you do not want to keep training. This is just for demonstrating how to train and how to tweak the architecture.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "################## init model ###################\n",
    "initModel = SiamesePixelEmbed(emb_dimension=embedding_dim, \n",
    "                              filterSize=kernel_size,\n",
    "                              device=device, pretrained=False)\n",
    "\n",
    "initModel.load_state_dict(torch.load(os.path.join(exp_dir,'epoch-445.paramOnly')))\n",
    "initModel.to(device);\n",
    "\n",
    "\n",
    "# decreasing the momentum in batch normalization, as the default is too large for this work.\n",
    "allLayer_dict = initModel.state_dict()\n",
    "child_counter = 0\n",
    "for child in initModel.PEMbase.children():\n",
    "    for i in range(len(child)):        \n",
    "        if 'BatchNorm2d' in str(type(child[i])): \n",
    "            child[i].momentum=0.001\n",
    "            #print(child[i])            \n",
    "\n",
    "for i in range(len(initModel.ordered_embedding)):    \n",
    "    if 'BatchNorm2d' in str(type(initModel.ordered_embedding[i])): \n",
    "            initModel.ordered_embedding[i].momentum=0.001\n",
    "            #print(initModel.ordered_embedding[i])\n",
    "            \n",
    "child_counter = 0\n",
    "for child in initModel.children():\n",
    "    #print(\" child\", child_counter, \"is:\")\n",
    "    #print(child)\n",
    "    child_counter += 1            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "################## loss function ###################\n",
    "loss_1_to_2 = LossOrderedPairReconstruction(device=device)\n",
    "\n",
    "loss_l1norm = nn.L1Loss(size_average=True)\n",
    "\n",
    "optimizer_ft = optim.Adam([{'params': initModel.PEMbase.parameters()},\n",
    "                           {'params': initModel.ordered_embedding.parameters(), 'lr': base_lr},                           \n",
    "                         ], lr=base_lr)\n",
    "\n",
    "\n",
    "# Decay LR by a factor of 0.5 every int(total_epoch_num/5) epochs\n",
    "exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=int(total_epoch_num/5), gamma=0.5)\n",
    "\n",
    "################## start training ###################\n",
    "fn = open(log_filename,'w')\n",
    "fn.write(log_filename+'\\t'+device+'\\n\\n')\n",
    "#fn.write(path.basename(__file__)+'\\n\\n')\n",
    "fn.close()\n",
    "file_to_note_bestModel = os.path.join(save_dir,'note_bestModel.log')\n",
    "fn = open(file_to_note_bestModel, 'w')\n",
    "fn.write('Record of best models on the way.\\n')\n",
    "fn.close()\n",
    "\n",
    "model_ft = train_model(initModel, dataloaders, dataset_sizes, \n",
    "                       loss_1_to_2, \n",
    "                       optimizer_ft, exp_lr_scheduler,\n",
    "                       num_epochs=total_epoch_num, \n",
    "                       work_dir=save_dir, device=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Leaving blank\n",
    "=====\n",
    "\n",
    "This is the end of the demo. Should have questions, please contact through the following\n",
    "\n",
    "### Shu Kong\n",
    "\n",
    "\n",
    "### aimerykong - at - gmail dot com\n",
    "\n",
    "\n",
    "\n",
    "If you find anything inspires you, please cite\n",
    "\n",
    "\n",
    "    @inproceedings{kong2018PPF,\n",
    "      title={Image Reconstruction with Predictive Filter Flow},\n",
    "      author={Kong, Shu and Fowlkes, Charless},\n",
    "      booktitle={arxiv},\n",
    "      year={2018}\n",
    "    }\n"
   ]
  }
 ],
 "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
