Application.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import skimage as sk\n",
    "from skimage import io, util\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import math, sys, os\n",
    "from residual_model_resdnet import *\n",
    "from MMNet_TBPTT import *\n",
    "import glob\n",
    "from subprocess import call, check_output\n",
    "from tqdm import *\n",
    "from subprocess import Popen, PIPE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rescale_to_255f(img):\n",
    "    dtype = img.dtype\n",
    "    if dtype == np.uint16:\n",
    "        img = img / 2**16 * 2**8\n",
    "    elif dtype == np.uint8:\n",
    "        img = img.astype(np.float32)\n",
    "    return img\n",
    "\n",
    "def check_pattern(img_path):\n",
    "    p = Popen(['dcraw','-i','-v',img_path], stdin=PIPE, stdout=PIPE, stderr=PIPE)\n",
    "    output, err = p.communicate()\n",
    "    output = str(output)\n",
    "    cfa = output.split('Filter pattern:')[1][1:5]\n",
    "    if 'RGGB' in cfa or 'Fujifilm' in output:\n",
    "        rollx , rolly = 0, 0\n",
    "    elif 'GBRG' in cfa:\n",
    "        rollx , rolly = -1, 0\n",
    "    elif 'GRBG' in cfa:\n",
    "        rollx , rolly = 0, -1\n",
    "    else:\n",
    "        raise NotImplementedError\n",
    "    return rollx , rolly\n",
    "\n",
    "\n",
    "def linrgb_to_srgb(img):\n",
    "    \"\"\" Convert linRGB color space to sRGB \n",
    "        https://en.wikipedia.org/wiki/SRGB\n",
    "    \"\"\"\n",
    "    assert img.dtype in [np.float32, np.float64] \n",
    "    img = img.copy()\n",
    "    mask = img < 0.0031308\n",
    "    img[~mask] = (img[~mask]**(1/2.4))*(1.055) - 0.055\n",
    "    img[mask] = img[mask] * 12.92\n",
    "    return img\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define HyperParameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "args_noise_estimation = True # wheter to estimate noise or not\n",
    "args_init = True # wheter to initialize the input with bilinear\n",
    "args_use_gpu = True \n",
    "args_block_size = (512, 512)\n",
    "args_model = 'pretrained_models/bayer_noisy/' # model path\n",
    "# Define folder with RAW images\n",
    "args_img_folder = '/home/datasets/raise/' # folder of RAW images\n",
    "args_output_folder = 'output/' # save results to folder\n",
    "args_type = '.png' # image type to save as\n",
    "if 'xtrans' in args_model:\n",
    "    args_pattern = 'xtrans'\n",
    "else:\n",
    "    args_pattern = 'RGGB'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_params = torch.load(args_model+'model_best.pth')\n",
    "\n",
    "model = ResNet_Den(BasicBlock, model_params[2], weightnorm=True)\n",
    "mmnet = MMNet(model, max_iter=model_params[1])\n",
    "for param in mmnet.parameters():\n",
    "    param.requires_grad = False\n",
    "\n",
    "mmnet.load_state_dict(model_params[0])\n",
    "if args_use_gpu:\n",
    "    mmnet = mmnet.cuda()\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Process Images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(args_output_folder):\n",
    "    os.makedirs(args_output_folder)\n",
    "    \n",
    "filepaths_img = glob.glob(args_img_folder+'*')\n",
    "filepaths_img.sort()\n",
    "filepaths_img = np.random.choice(filepaths_img, 50, replace=False)\n",
    "cnt = 0\n",
    "for img_path in tqdm(filepaths_img):\n",
    "    try:\n",
    "        if cnt > 50: \n",
    "            break\n",
    "        else:\n",
    "            cnt += 1\n",
    "        print('Processing ', img_path)\n",
    "        call([\"dcraw\",\"-j\",\"-d\",\"-T\",\"-4\",\"-w\", \"+M\", img_path])\n",
    "        # convert to RGGB CFA \n",
    "        rollx, rolly = check_pattern(img_path)\n",
    "        img_path = img_path.split(\".\")\n",
    "        img_path[-1] = '.tiff'\n",
    "        img_path = \"\".join(img_path)\n",
    "        img = io.imread(img_path)\n",
    "\n",
    "        img = np.roll(img, rollx,rolly)\n",
    "        res = rescale_to_255f(img)\n",
    "        \n",
    "        # pad according to block size\n",
    "        if res.shape[0] % args_block_size[0] != 0:\n",
    "            mod = args_block_size[0]- res.shape[0] % args_block_size[0]\n",
    "            res = np.pad(res, ((0,mod),(0,0)), 'constant')\n",
    "\n",
    "        if res.shape[1] % args_block_size[1] != 0:\n",
    "            mod = args_block_size[1]- res.shape[1] % args_block_size[1]\n",
    "            res = np.pad(res, ((0,0), (0,mod)), 'constant')\n",
    "        blocks = util.view_as_blocks(res, block_shape=args_block_size)\n",
    "\n",
    "        def process_patch(patch):\n",
    "            with torch.no_grad():\n",
    "                mmnet.eval()\n",
    "                mosaic = torch.FloatTensor(patch).float()[None]\n",
    "                # padding in order to ensure no boundary artifacts\n",
    "                mosaic = F.pad(mosaic[:,None],(8,8,8,8),'reflect')[:,0]\n",
    "                shape = mosaic[0].shape\n",
    "                mask = utils.generate_mask(shape, pattern=args_pattern)\n",
    "                M = torch.FloatTensor(mask)[None]\n",
    "                mosaic = mosaic[...,None]*M\n",
    "\n",
    "                mosaic = mosaic.permute(0,3,1,2)\n",
    "                M = M.permute(0,3,1,2)\n",
    "\n",
    "                p = Demosaic(mosaic.float(), M.float())\n",
    "                if args_use_gpu:\n",
    "                    p.cuda_()\n",
    "                xcur = mmnet.forward_all_iter(p, max_iter=mmnet.max_iter, init=args_init, noise_estimation=args_noise_estimation)\n",
    "\n",
    "                return xcur[0].cpu().data.permute(1,2,0).numpy()[8:-8,8:-8]\n",
    "\n",
    "        # demosaick image\n",
    "        block_size = blocks.shape[-2:]\n",
    "        num_blocks = blocks.shape[:2]\n",
    "        original_size = (blocks.shape[0] * blocks.shape[2], blocks.shape[1] * blocks.shape[3])\n",
    "        final_img = np.zeros((original_size[0], original_size[1],3), dtype=np.float32)\n",
    "        for i in range(num_blocks[0]):\n",
    "            for j in range(num_blocks[1]):\n",
    "                patch_result = process_patch(blocks[i,j])\n",
    "                final_img[i*block_size[0]:(i+1)*block_size[0], j*block_size[1]:(j+1)*block_size[1]] = patch_result\n",
    "        \n",
    "        final_img = final_img[:img.shape[0],:img.shape[1]]\n",
    "        final_img = np.roll(final_img, -rollx, -rolly)\n",
    "        # remove intermmidiate .tiff image\n",
    "        call([\"rm\", img_path])\n",
    "        img_path = img_path.replace(args_img_folder, args_output_folder)\n",
    "        # save the linRGB image\n",
    "        io.imsave(img_path.replace('.tiff', args_type),final_img.astype(np.uint8))\n",
    "        # save the sRGB image\n",
    "        srgb = linrgb_to_srgb(final_img/255)\n",
    "        io.imsave(img_path.replace('.tiff', '_srgb'+args_type),srgb.clip(0,1))\n",
    "    except Exception as e:\n",
    "        print(e)"
   ]
  }
 ],
 "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}