See-In-The-Dark.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 0. Set up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.optim as optim\n",
    "from torch.optim import lr_scheduler\n",
    "import torch.nn as nn\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import rawpy\n",
    "from tqdm import tqdm as pbar\n",
    "import copy\n",
    "from livelossplot import PlotLosses\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn\n",
    "seaborn.set()\n",
    "import scipy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path = 'dataset'\n",
    "# np.random.seed(0)\n",
    "# torch.manual_seed(0)\n",
    "# torch.backends.cudnn.deterministic = True\n",
    "# torch.backends.cudnn.benchmark = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Preprocess raw data from camera sensor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](figures/3a.png)\n",
    "\n",
    "Pack raw Bayer sensor data into 4 channels (R-G-B-G). By doing this also reduces resolution by factor of 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.1 Pack raw is used for input "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pack_raw(raw):\n",
    "    \"\"\"\n",
    "    Input: object returned from rawpy.imread()\n",
    "    Output: numpy array in shape (1424, 2128, 4)\n",
    "    \"\"\"\n",
    "    \n",
    "    im = raw.raw_image_visible.astype(np.float32) # shape of (2848, 4256)\n",
    "    im = np.maximum(im - 512, 0) / (16383 - 512) #subtract the black level\n",
    "    im = np.expand_dims(im, axis=2) # shape of (2848, 4256, 1)\n",
    "\n",
    "    img_shape = im.shape # (H, W, 1)\n",
    "    H = img_shape[0]\n",
    "    W = img_shape[1]\n",
    "    \n",
    "    # Pack into 4 channels\n",
    "    red = im[0:H:2,0:W:2,:]\n",
    "    green_1 = im[0:H:2,1:W:2,:]\n",
    "    blue = im[1:H:2,1:W:2,:]\n",
    "    green_2 = im[1:H:2,0:W:2,:]\n",
    "    \n",
    "    # Final shape: (1424, 2128, 4)\n",
    "    out = np.concatenate((red, green_1, blue, green_2), axis=2)\n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# x_img = rawpy.imread(data_path + '/Sony/short/00001_00_0.04s.ARW')\n",
    "# x_img = pack_raw(x_img)\n",
    "# x_img.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.2 Post process is used for ground true"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def post_process(raw):\n",
    "    \"\"\"\n",
    "    Input: object returned from rawpy.imgread()\n",
    "    Output: numpy array in shape (2848, 4256, 3)\n",
    "    \"\"\"\n",
    "    max_output = 65535.0\n",
    "    im = raw.postprocess(use_camera_wb=True, no_auto_bright=True, output_bps=16)\n",
    "    im = np.float32(im / max_output)\n",
    "    return im"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# y_img = rawpy.imread(data_path + '/Sony/long/00001_00_10s.ARW')\n",
    "# y_img = post_process(y_img)\n",
    "# y_img.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.3 Batch process all data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Files' name explanation**\n",
    "\n",
    "The file lists are provided. In each row, there are a short-exposed image path, the corresponding long-exposed image path, camera ISO and F number. \n",
    "Note that multiple short-exposed images may correspond to the same long-exposed image.\n",
    "\n",
    "The file name contains the image information. For example, in \"10019_00_0.033s.RAF\":\n",
    "- the first digit \"1\" means it is from the test set (\"0\" for training set and \"2\" for validation set)\n",
    "- 0019\" is the image ID\n",
    "- the following \"00\" is the number in the sequence/burst\n",
    "- \"0.033s\" is the exposure time 1/30 seconds.\n",
    "\n",
    "There are some misalignment with the ground-truth for image 10034, 10045, 10172. I've removed those images for quantitative results, but they still can be used for qualitative evaluations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_file_list(file_list):\n",
    "    data = pd.read_csv(data_path + file_list, sep=\" \", header = None, names = ['X', 'Y', 'ISO', 'F-stop'])\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train_list = read_file_list('/Sony_train_list.txt')\n",
    "# train_list.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def batch_process_raw(data, hide_progree=False):\n",
    "    \"\"\"\n",
    "    Input: Pandas dataframe returned from read_file_list\n",
    "    Output: a dictionary of \n",
    "            X : amplified numpy array\n",
    "            Y : numpy array\n",
    "            X_Y_map: numpy array of indexes of corresponding pair of X and Y\n",
    "    \"\"\"\n",
    "    \n",
    "    # Multiple Xs can have the same Y    \n",
    "    m_x = len(data)\n",
    "    m_y = data['Y'].nunique()\n",
    "    \n",
    "    X = np.zeros((m_x, 1424, 2128, 4), dtype=np.float32)\n",
    "    Y = np.zeros((m_y, 2848, 4256, 3), dtype=np.float32)\n",
    "   \n",
    "    # Mapping of of X to Y\n",
    "    X_map = []\n",
    "    Y_map = []\n",
    "    \n",
    "    for i in pbar(range(m_x), disable=hide_progree):\n",
    "        x_path = data.iloc[i][0][1:] # remove the \".\" in the name\n",
    "        y_path = data.iloc[i][1][1:] # remove the \".\" in the name\n",
    "        \n",
    "        # Shutter speed is in the file name\n",
    "        x_shutter_speed = x_path.split('_')[-1].split('s.')[0]\n",
    "        y_shutter_speed = y_path.split('_')[-1].split('s.')[0]\n",
    "        amp_ratio = float(y_shutter_speed)/float(x_shutter_speed)\n",
    "        \n",
    "        X[i] = pack_raw(rawpy.imread(data_path + x_path)) * amp_ratio\n",
    "    \n",
    "    for i in pbar(range(m_y), disable=hide_progree):\n",
    "        current_y = data['Y'].unique()[i]\n",
    "        \n",
    "        y_path = current_y[1:]\n",
    "        Y[i] = post_process(rawpy.imread(data_path + y_path))\n",
    "        \n",
    "        # Maping of X to Y\n",
    "        X_map_temp = data['Y'][data['Y']==current_y].index.tolist()\n",
    "        Y_map_temp = [i]*len(X_map_temp)\n",
    "        X_map += X_map_temp\n",
    "        Y_map += Y_map_temp\n",
    "    \n",
    "    X_Y_map = np.array((X_map, Y_map), dtype=np.int32).T\n",
    "    dataset = {'X':X, 'Y':Y, 'X_Y_map':X_Y_map}\n",
    "    \n",
    "    return dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train_dataset = batch_process_raw(train_list.head(10), True)\n",
    "# print(\"Shape of X_train:\", train_dataset['X'].shape)\n",
    "# print(\"Shape of Y_train:\", train_dataset['Y'].shape)\n",
    "# print(\"Shape of X_Y_map_train:\", train_dataset['X_Y_map'].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Data augmentation\n",
    "Random crop, flip, and tranpose data, then amplify the result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def numpy_to_torch(image):\n",
    "    \"\"\"\n",
    "    Input: numpy array (H x W x C)\n",
    "    Output: torch tensory (C x H x W)\n",
    "    \"\"\"\n",
    "    image = image.transpose((2, 0, 1))\n",
    "    torch_tensor = torch.from_numpy(image)\n",
    "    return torch_tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def augment_data(x_input, y_output, ps):\n",
    "    \"\"\"\n",
    "    Input: numpy arrays with shape (H x W x C), patch_size = integer\n",
    "    Output: X: augmented torch tensor with shape (C x atch_size x patch_size)\n",
    "            Y: augmented numpy arrays with shape (Cx 2*patch_size x 2*patch_size)\n",
    "    \"\"\"\n",
    "    \n",
    "    # Random crop\n",
    "    H = x_input.shape[0]\n",
    "    W = x_input.shape[1]\n",
    "    xx = np.random.randint(0, W-ps)\n",
    "    yy = np.random.randint(0, H-ps)\n",
    "    x_patch = x_input[yy:yy+ps, xx:xx+ps,:]\n",
    "    y_patch = y_output[yy*2:yy*2+ps*2, xx*2:xx*2+ps*2,:]\n",
    "\n",
    "    # Random flip first axis\n",
    "    if np.random.randint(2, size=1)[0] == 1:\n",
    "        x_patch = np.flip(x_patch, axis=0)\n",
    "        y_patch = np.flip(y_patch, axis=0)\n",
    "    \n",
    "    # Random flip second axis\n",
    "    if np.random.randint(2, size=1)[0] == 1:\n",
    "        x_patch = np.flip(x_patch, axis=1)\n",
    "        y_patch = np.flip(y_patch, axis=1)\n",
    "    \n",
    "    # Random transpose\n",
    "    if np.random.randint(2, size=1)[0] == 1:\n",
    "        x_patch = np.transpose(x_patch, (1, 0, 2))\n",
    "        y_patch = np.transpose(y_patch, (1, 0, 2))\n",
    "    \n",
    "    # Clip saturated value\n",
    "    x_patch = np.clip(x_patch, a_min=0.0, a_max=1.0)\n",
    "    y_patch = np.clip(y_patch, a_min=0.0, a_max=1.0)\n",
    "        \n",
    "    return numpy_to_torch(x_patch), numpy_to_torch(y_patch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# x_aug, y_aug = augment_data(X_train[0], Y_train[0], 512)\n",
    "# print(\"Shape of X_aug:\", x_aug.shape)\n",
    "# print(\"Shape of Y_aug:\", y_aug.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. Make batches of image patches for training, validation and testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_batch(dataset, image_indexes, patch_size):\n",
    "    \"\"\"\n",
    "    Prepare a batch for training\n",
    "    Input:  a dictionary of X, Y, X_Y_map (returned from batch_process_raw())\n",
    "            image_indexces: a subset size m of random permuatation of X_Y_map \n",
    "    Output: X torch tensor of shape (m, 4, patch_size, patch_size)\n",
    "            Y torch tensor of shape (m, 3, 2*patch_size, 2*patch_size)\n",
    "    \"\"\"\n",
    "    \n",
    "    X = dataset['X']\n",
    "    Y = dataset['Y']\n",
    "    \n",
    "    m = len(image_indexes)\n",
    "    \n",
    "    X_patches = torch.zeros(m, 4, patch_size, patch_size, dtype=torch.float32, device=device)\n",
    "    Y_patches = torch.zeros(m, 3, 2*patch_size, 2*patch_size, dtype=torch.float32, device=device)\n",
    "    \n",
    "    for i in range(m):\n",
    "        x_index, y_index = image_indexes[i]\n",
    "        X_patches[i], Y_patches[i] = augment_data(X[x_index], Y[y_index], patch_size)\n",
    "    \n",
    "    return X_patches, Y_patches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_batch_test(dataset, image_indexes):\n",
    "    \"\"\"\n",
    "    Prepare a batch (full res) for testing\n",
    "    \"\"\"\n",
    "    X = dataset['X']\n",
    "    Y = dataset['Y']\n",
    "    \n",
    "    m = len(image_indexes)\n",
    "    \n",
    "    X_images = torch.zeros(m, 4, 1424, 2128, dtype=torch.float32, device=device)\n",
    "    Y_images = torch.zeros(m, 3, 2848, 4256, dtype=torch.float32, device=device)\n",
    "    \n",
    "    for i in range(m):\n",
    "        x_index, y_index = image_indexes[i]\n",
    "        X_images[i] = numpy_to_torch(np.clip(X[x_index], a_min=0.0, a_max=1.0))\n",
    "        Y_images[i] = numpy_to_torch(np.clip(Y[y_index], a_min=0.0, a_max=1.0))\n",
    "    \n",
    "    return X_images, Y_images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# batch_size = 2\n",
    "# random_orders = np.random.permutation(train_dataset['X_Y_map'])\n",
    "# splitted_random_orders = np.array_split(random_orders, range(batch_size, len(random_orders), batch_size))\n",
    "\n",
    "# first_batch_indexes = splitted_random_orders\n",
    "# last_batch_indexes = splitted_random_orders\n",
    "\n",
    "# x_batch, y_batch = make_batch(train_dataset, first_batch_indexes, 512)\n",
    "# print('Shape of first batch X and Y:', x_batch.shape, y_batch.shape)\n",
    "\n",
    "# x_batch, y_batch = make_batch(train_dataset, last_batch_indexes, 512)\n",
    "# print('Shape of last batch X and Y:', x_batch.shape, y_batch.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4. Model architecture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DoubleConv(nn.Module):\n",
    "    #  Conv -> BN -> LReLU -> Conv -> BN -> LReLU\n",
    "    def __init__(self, in_ch, out_ch):\n",
    "        super().__init__()\n",
    "        self.f = nn.Sequential(\n",
    "            nn.Conv2d(in_ch, out_ch, 3, padding=1),\n",
    "            nn.BatchNorm2d(out_ch),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            nn.Conv2d(out_ch, out_ch, 3, padding=1),\n",
    "            nn.BatchNorm2d(out_ch),\n",
    "            nn.LeakyReLU(0.2, inplace=True),)\n",
    "    def forward(self, x):\n",
    "        x = self.f(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "class Down(nn.Module):\n",
    "    def __init__(self, in_ch, out_ch):\n",
    "        super().__init__()\n",
    "        self.f = nn.Sequential(\n",
    "            nn.MaxPool2d(2),\n",
    "            DoubleConv(in_ch, out_ch),)\n",
    "    def forward(self, x):\n",
    "        x = self.f(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "class Up(nn.Module):\n",
    "    # upsample and concat\n",
    "    def __init__(self, in_ch, out_ch):\n",
    "        super().__init__()\n",
    "        self.upsample = nn.ConvTranspose2d(in_ch, in_ch//2, 2, stride=2)\n",
    "        self.conv = DoubleConv(in_ch, out_ch)\n",
    "    def forward(self, x1, x2):\n",
    "        x1 = self.upsample(x1)\n",
    "        x = torch.cat([x2, x1], dim=1)\n",
    "        x = self.conv(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "class OutConv(nn.Module):\n",
    "    def __init__(self, in_ch, out_ch):\n",
    "        super(OutConv, self).__init__()\n",
    "        self.f = nn.Conv2d(in_ch, out_ch, 1)\n",
    "    def forward(self, x):\n",
    "        x = self.f(x)\n",
    "        return x\n",
    "\n",
    "class Unet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.inc = DoubleConv(4, 32)\n",
    "        self.d1 = Down(32, 64)\n",
    "        self.d2 = Down(64, 128)\n",
    "        self.d3 = Down(128, 256)\n",
    "        self.d4 = Down(256, 512)\n",
    "\n",
    "        self.u1 = Up(512, 256)\n",
    "        self.u2 = Up(256, 128)\n",
    "        self.u3 = Up(128, 64)\n",
    "        self.u4 = Up(64, 32)\n",
    "        self.outc = OutConv(32, 12)\n",
    "        self.pixel_shuffle = nn.PixelShuffle(2)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x1 = self.inc(x)\n",
    "        x2 = self.d1(x1)\n",
    "        x3 = self.d2(x2)\n",
    "        x4 = self.d3(x3)\n",
    "        x5 = self.d4(x4)\n",
    "        x = self.u1(x5, x4)\n",
    "        x = self.u2(x, x3)\n",
    "        x = self.u3(x, x2)\n",
    "        x = self.u4(x, x1)\n",
    "        x = self.outc(x)\n",
    "        x = self.pixel_shuffle(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 5. Traing and testing code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_psnr(target, output):\n",
    "    \"\"\"\n",
    "    Calculate Peak Signal To Noise Ratio\n",
    "    Input: torch tensor of shape (m, C, H, W)\n",
    "    Output: average of PSTR for that batch\n",
    "    \"\"\"\n",
    "    m, C, H, W = target.shape\n",
    "    sum_psnr = 0 \n",
    "    \n",
    "    for i in range(m):\n",
    "        output[i] = torch.clamp(output[i], min=0.0, max=1.0)\n",
    "        mse = torch.sum((target[i] - output[i])**2)/(C*H*W)\n",
    "        psnr =  -10*torch.log10(mse)\n",
    "        sum_psnr += psnr\n",
    "        \n",
    "    return sum_psnr/m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(model, train_dataset, val_dataset, optimizer, scheduler, check_point, batch_size, num_epochs):\n",
    "    liveloss = PlotLosses()\n",
    "    criterion = nn.L1Loss()\n",
    "    m_train = len(train_dataset['X_Y_map'])\n",
    "    best_psnr = 0.0\n",
    "    best_model_weights = copy.deepcopy(model.state_dict())\n",
    "    \n",
    "    for epoch in pbar(range(num_epochs)):\n",
    "        plot_logs = {}\n",
    "        logs = []\n",
    "        scheduler.step()\n",
    "        \n",
    "        # Each epoch has a training and validation phase\n",
    "        for phase in ['train', 'validation']:\n",
    "            psnr_epoch = 0\n",
    "            \n",
    "            if phase == 'train':\n",
    "                model.train()\n",
    "\n",
    "                # Shuffle training set\n",
    "                random_orders = np.random.permutation(train_dataset['X_Y_map'])\n",
    "                # Split the training set in into batches\n",
    "                splitted_random_orders = np.array_split(random_orders, range(batch_size, len(random_orders), batch_size))\n",
    "                \n",
    "                # Iterate over data\n",
    "                for a_batch_index in splitted_random_orders:\n",
    "                    image, target = make_batch(train_dataset, a_batch_index, patch_size=512)\n",
    "                    \n",
    "                    # Zero gradient\n",
    "                    optimizer.zero_grad()\n",
    "                    \n",
    "                    # Forward pass\n",
    "                    y_hat = model(image)\n",
    "                    \n",
    "                    # Calculate loss\n",
    "                    psnr_batch = calculate_psnr(target.detach(), y_hat.detach()).item()\n",
    "                    loss = criterion(target, y_hat)\n",
    "                    psnr_epoch += psnr_batch * image.size(0)\n",
    "                    \n",
    "                    # Backward pass\n",
    "                    loss.backward()\n",
    "                    optimizer.step()\n",
    "                    \n",
    "                # Update logs\n",
    "                psnr_epoch = psnr_epoch / m_train\n",
    "                plot_logs['PSNR'] = psnr_epoch\n",
    "                logs.append(psnr_epoch)\n",
    "                                    \n",
    "            else:    \n",
    "                val_psnr_epoch = test_model(model, val_dataset)\n",
    "                \n",
    "                # Update logs\n",
    "                plot_logs['val_PSNR'] = val_psnr_epoch\n",
    "                logs.append(val_psnr_epoch)\n",
    "                \n",
    "                # Save best model\n",
    "                if val_psnr_epoch > best_psnr:\n",
    "                    best_psnr = val_psnr_epoch\n",
    "                    best_model_weights = copy.deepcopy(model.state_dict())\n",
    "                    \n",
    "                # Check point\n",
    "                if epoch%check_point==0:\n",
    "                    torch.save(best_model_weights, 'trained_model/best_model.pt')\n",
    "\n",
    "        # Update live plot every epoch\n",
    "        liveloss.update(plot_logs)\n",
    "        liveloss.draw()\n",
    "        \n",
    "        # Write to log file every epoch\n",
    "        # Epoch - Best Val PSNR - Train  PSNR - Val PSNR\n",
    "        f = open(\"trained_model/training_log.txt\", \"a\")\n",
    "        f.write(\"\\n{:4d} \\t{:.5f} \\t{:.5f} \\t{:.5f}\".format(epoch, best_psnr, logs[0], logs[1]))\n",
    "        f.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_model(model, dataset, hide_progress=True):\n",
    "    model.eval()\n",
    "    m_test = len(dataset['X_Y_map'])\n",
    "    test_psnr = 0\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        # Iterate over data\n",
    "        for i in pbar(dataset['X_Y_map'], disable=hide_progress):\n",
    "            image, target = make_batch_test(dataset, np.expand_dims(i, 0))\n",
    "\n",
    "            # Forward pass\n",
    "            y_hat = model(image)\n",
    "\n",
    "            # Calculate loss\n",
    "            test_psnr_batch = calculate_psnr(target, y_hat).item()\n",
    "            test_psnr += test_psnr_batch\n",
    "            \n",
    "    return test_psnr / m_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def display_an_example(model, image_list, dataset, index):\n",
    "    \"\"\"\n",
    "    Display a single example\n",
    "    \"\"\"\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        image, ground_truth = make_batch_test(dataset, np.expand_dims(dataset['X_Y_map'][index], 0))\n",
    "        y_hat = model(image)\n",
    "        y_hat = torch.clamp(y_hat, min=0.0, max=1.0)\n",
    "    \n",
    "        # Convert from torch tensor to numpy\n",
    "        y_hat = y_hat.squeeze().cpu().numpy().transpose((1, 2, 0))\n",
    "        ground_truth = ground_truth.squeeze().cpu().numpy().transpose((1, 2, 0))\n",
    "    \n",
    "    x_path = image_list.iloc[index][0][1:] # remove the \".\" in the name   \n",
    "    image_to_display = post_process(rawpy.imread(data_path + x_path))\n",
    "    fig=plt.figure(figsize=(30, 10))\n",
    "    \n",
    "    fig.add_subplot(1, 3, 1)\n",
    "    plt.imshow(image_to_display, vmin=0, vmax=1)\n",
    "    plt.title('Original image')\n",
    "    plt.axis('off')\n",
    "    plt.grid(b=None)\n",
    "    \n",
    "    fig.add_subplot(1, 3, 2)\n",
    "    plt.imshow(y_hat, vmin=0, vmax=1)\n",
    "    plt.title('Denoised by model')\n",
    "    plt.axis('off')\n",
    "    plt.grid(b=None)\n",
    "    \n",
    "    fig.add_subplot(1, 3, 3)\n",
    "    plt.imshow(ground_truth, vmin=0, vmax=1)\n",
    "    plt.title('Ground Truth')\n",
    "    plt.axis('off')\n",
    "    plt.grid(b=None)\n",
    "\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "def display_custom_image(model, image_path, amp_ratio, render=False):\n",
    "    model.eval()\n",
    "        \n",
    "    orig_image = post_process(rawpy.imread(image_path))\n",
    "    \n",
    "    fig=plt.figure(figsize=(20, 10))\n",
    "    fig.add_subplot(1, 2, 1)\n",
    "    plt.imshow(orig_image, vmin=0, vmax=1)\n",
    "    plt.title('Original image')\n",
    "    plt.axis('off')\n",
    "    plt.grid(b=None)\n",
    "    \n",
    "    image = pack_raw(rawpy.imread(image_path)) * amp_ratio\n",
    "    image = numpy_to_torch(np.clip(image, a_min=0.0, a_max=1.0)).unsqueeze(0)\n",
    "    image = image.to(device)\n",
    "    with torch.no_grad():\n",
    "        y_hat = model(image)\n",
    "        y_hat = torch.clamp(y_hat, min=0.0, max=1.0)\n",
    "    image = y_hat.squeeze().cpu().numpy().transpose((1, 2, 0))\n",
    "        \n",
    "    fig.add_subplot(1, 2, 2)\n",
    "    plt.imshow(image, vmin=0, vmax=1)\n",
    "    plt.title('Denoised by model')\n",
    "    plt.axis('off')\n",
    "    plt.grid(b=None)\n",
    "    \n",
    "    if render:\n",
    "        scipy.misc.toimage(image * 255, high=255, low=0, cmin=0, cmax=255).save('custom_images/processed.png')\n",
    "\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 6. Put everything together"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using cuda:3 to train\n"
     ]
    }
   ],
   "source": [
    "# Train on cuda if available\n",
    "device = torch.device('cuda:3' if torch.cuda.is_available() else 'cpu')\n",
    "print('Using', device, 'to train')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Train dataset\n",
    "# train_list = read_file_list('/Sony_train_list.txt')\n",
    "# train_dataset = batch_process_raw(train_list)\n",
    "\n",
    "# # Validation dataset\n",
    "# val_list = read_file_list('/Sony_val_list.txt')\n",
    "# val_dataset = batch_process_raw(val_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Inialize and load model\n",
    "# my_model = Unet()\n",
    "# my_model.load_state_dict(torch.load('trained_model/best_model.pt'))\n",
    "# my_model = my_model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Initialize optimizer\n",
    "# optimizer = optim.Adam(my_model.parameters(), lr=1e-5, amsgrad=True)\n",
    "# scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=[1000], gamma=0.1)\n",
    "\n",
    "# # Train model\n",
    "# train_model(my_model, train_dataset, val_dataset, optimizer, scheduler, check_point=10, batch_size=32, num_epochs=1500)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 562/562 [01:39<00:00,  5.02it/s]\n",
      "100%|██████████| 47/47 [00:43<00:00,  1.08it/s]\n"
     ]
    }
   ],
   "source": [
    "# Test dataset\n",
    "test_list = read_file_list('/Sony_test_list.txt')\n",
    "test_dataset = batch_process_raw(test_list)\n",
    "\n",
    "# Inialize and load model\n",
    "my_model = Unet()\n",
    "my_model.load_state_dict(torch.load('trained_model/best_model.pt'))\n",
    "my_model = my_model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 562/562 [02:28<00:00,  3.79it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Peak Signal Noise Ratio on test dataset 28.83\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "score = test_model(my_model, test_dataset, hide_progress=False)\n",
    "print('Peak Signal Noise Ratio on test dataset {:.2f}'.format(score))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test custom image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "# display_an_example(my_model, test_list, test_dataset, 65)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "# display_custom_image(my_model, 'custom_images/image_1.arw', 8)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:see-in-dark]",
   "language": "python",
   "name": "conda-env-see-in-dark-py"
  },
  "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}