Inference_unet_generated_imgs.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES']='0'\n",
    "\n",
    "from fastai import *\n",
    "from fastai.vision import *\n",
    "from fastai.vision.gan import *\n",
    "from ArNet.generators import *\n",
    "from ArNet.critics import *\n",
    "from ArNet.dataset import *\n",
    "from ArNet.loss import *\n",
    "from ArNet.save import *\n",
    "from ArNet.fid_loss import *\n",
    "from ArNet.ssim import *\n",
    "from ArNet.metrics import *\n",
    "\n",
    "import torchvision\n",
    "import geffnet # efficient/ mobile net\n",
    "from PIL import Image as PilImage\n",
    "import cv2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def toEven(sz):\n",
    "    tempSz = [sz[0], sz[1]]\n",
    "    if sz[0]%2 != 0:\n",
    "        tempSz[0] += 1\n",
    "    if sz[1]%2 != 0:\n",
    "        tempSz[1] += 1\n",
    "    return tempSz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = Path('./dataset/')\n",
    "\n",
    "path_fullRes = path/'DIV2K_train_HR'\n",
    "\n",
    "path_lowRes_512 = path/'DIV2K_train_LR_512_QF20'\n",
    "path_lowRes_Full = path/'DIV2K_train_LR_Full_QF20'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = geffnet.mobilenetv3_rw\n",
    "loss_func = lpips_loss()\n",
    "\n",
    "bs=1\n",
    "\n",
    "data_gen = get_DIV2k_data_QF(path_lowRes_Full, path_fullRes, bs=bs, sz=1024)\n",
    "\n",
    "learn_gen = gen_learner_wide(data=data_gen,\n",
    "                             gen_loss=loss_func,\n",
    "                             arch = model,\n",
    "                             nf_factor=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "weights = \"/data/students_home/fmameli/repos/Artifact_Removal_GAN/models/unet_wideNf2_superRes_mobilenetV3_GAN_best\"\n",
    "learn_gen.load(weights, with_opt=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn_gen.export(\"/data/students_home/fmameli/repos/Artifact_Removal_GAN/models/standard.pkl\" )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = geffnet.mobilenetv3_small_minimal_100\n",
    "loss_func = lpips_loss()\n",
    "\n",
    "bs=1\n",
    "\n",
    "data_gen = get_DIV2k_data_QF(path_lowRes_Full, path_fullRes, bs=bs, sz=512)\n",
    "\n",
    "learn_gen = gen_learner_wide(data=data_gen,\n",
    "                             gen_loss=loss_func,\n",
    "                             arch = model,\n",
    "                             nf_factor=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "weights = \"/data/students_home/fmameli/repos/Artifact_Removal_GAN/wandb/run-20200515_110129-unet_wideNf2_superRes_mobilenetMinimal_1k_P64px_VGG_SuperRes_05-15_13_01/bestmodel\"\n",
    "learn_gen.load(weights, with_opt=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn_gen.export(\"/data/students_home/fmameli/repos/Artifact_Removal_GAN/models/minimal.pkl\" )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Minimal model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "root_model_path = Path(\"/data/students_home/fmameli/repos/Artifact_Removal_GAN/models/\")\n",
    "exported_model_minimal  =Path(\"/data/students_home/fmameli/repos/Artifact_Removal_GAN/models/minimal.pkl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "testFrameList = ImageImageList.from_folder(path/'DIV2K_train_LR_512_QF20', presort=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn_min = load_learner(path=root_model_path, file=exported_model_minimal, test=testFrameList[:24], tfm_y=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn_min.data.batch_size = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "i = 803\n",
    "id_img = str(i).zfill(4)\n",
    "\n",
    "img_low = open_image(\"dataset/DIV2K_train_LR_512_QF20/\" + id_img + \".jpg\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n 1 -r 100\n",
    "p,img_hr,b = learn_min.predict(img_low)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n 1 -r 4\n",
    "p = learn_min.get_preds(ds_type=DatasetType.Test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Standard GAN model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "root_model_path = Path(\"/data/students_home/fmameli/repos/Artifact_Removal_GAN/models/\")\n",
    "exported_model_standard =Path(\"/data/students_home/fmameli/repos/Artifact_Removal_GAN/models/standard.pkl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "testFrameList = ImageImageList.from_folder(path/'DIV2K_train_LR_512_QF20', presort=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn_std = load_learner(path=root_model_path, file=exported_model_standard, test=testFrameList[:24], tfm_y=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn_std.data.batch_size = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n 1 -r 100\n",
    "p,img_hr,b = learn_std.predict(img_low)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n 1 -r 4\n",
    "p = learn_std.get_preds(ds_type=DatasetType.Test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Image processing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def toEven(sz):\n",
    "    tempSz = [sz[0], sz[1]]\n",
    "    if sz[0]%2 != 0:\n",
    "        tempSz[0] += 1\n",
    "    if sz[1]%2 != 0:\n",
    "        tempSz[1] += 1\n",
    "    return tempSz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "weights = \"/data/students_home/fmameli/repos/SuperRes/models/unet_wideNf2_superRes_mobilenetV3_GAN_best\"\n",
    "learn_gen.load(weights, with_opt=False)\n",
    "\n",
    "for i in range(1, 901):\n",
    "    id_img = str(i).zfill(4)\n",
    "\n",
    "    img_low = open_image(\"dataset/DIV2K_train_LR_1024_QF20/\" + id_img + \".jpg\")\n",
    "    size=toEven(img_low.size)\n",
    "    data_gen = get_DIV2k_data_QF(path_lowRes_Full, path_fullRes, bs=bs, sz=size)\n",
    "\n",
    "    learn_gen.data = data_gen\n",
    "    \n",
    "    p,img_hr,b = learn_gen.predict(img_low)\n",
    "    p.save(\"output_imgs/\" + id_img + \"_GAN.png\")\n",
    "    print(\"dataset/DIV2K_train_LR_1024_QF20/\" + id_img + \".jpg\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "weights = \"/data/students_home/fmameli/repos/SuperRes/models/unet_superRes_mobilenetV3_MSE_gen_512px_2\"\n",
    "learn_gen.load(weights, with_opt=False)\n",
    "\n",
    "for i in range(1, 901):\n",
    "    id_img = str(i).zfill(4)\n",
    "\n",
    "    img_low = open_image(\"dataset/DIV2K_train_LR_1024_QF20/\" + id_img + \".jpg\")\n",
    "    size=toEven(img_low.size)\n",
    "    data_gen = get_DIV2k_data_QF(path_lowRes_Full, path_fullRes, bs=bs, sz=size)\n",
    "\n",
    "    learn_gen.data = data_gen\n",
    "    \n",
    "    p,img_hr,b = learn_gen.predict(img_low)\n",
    "    p.save(\"output_imgs/\" + id_img + \"_MSE.png\")\n",
    "    print(\"dataset/DIV2K_train_LR_1024_QF20/\" + id_img + \".jpg\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "weights = \"/data/students_home/fmameli/repos/SuperRes/models/unet_superRes_mobilenetV3_SSIM_gen_512px_0\"\n",
    "learn_gen.load(weights, with_opt=False)\n",
    "\n",
    "for i in range(1, 901):\n",
    "    id_img = str(i).zfill(4)\n",
    "\n",
    "    img_low = open_image(\"dataset/DIV2K_train_LR_1024_QF20/\" + id_img + \".jpg\")\n",
    "    size=toEven(img_low.size)\n",
    "    data_gen = get_DIV2k_data_QF(path_lowRes_Full, path_fullRes, bs=bs, sz=size)\n",
    "\n",
    "    learn_gen.data = data_gen\n",
    "    \n",
    "    p,img_hr,b = learn_gen.predict(img_low)\n",
    "    p.save(\"output_imgs/\" + id_img + \"_SSIM.png\")\n",
    "    print(\"dataset/DIV2K_train_LR_1024_QF20/\" + id_img + \".jpg\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "weights = \"/data/students_home/fmameli/repos/SuperRes/models/unet_wideNf2_superRes_mobilenetV3_GAN_best\"\n",
    "learn_gen.load(weights, with_opt=False)\n",
    "\n",
    "p_gan,img_hr,b = learn_gen.predict(img_low)\n",
    "p_gan"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "weights = \"/data/students_home/fmameli/repos/SuperRes/models/unet_superRes_mobilenetV3_MSE_gen_512px_2\"\n",
    "learn_gen.load(weights, with_opt=False)\n",
    "\n",
    "p_mse,img_hr,b = learn_gen.predict(img_low)\n",
    "p_mse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "weights = \"/data/students_home/fmameli/repos/SuperRes/models/unet_superRes_mobilenetV3_SSIM_gen_512px_0\"\n",
    "learn_gen.load(weights, with_opt=False)\n",
    "\n",
    "p_ssim,img_hr,b = learn_gen.predict(img_low)\n",
    "p_ssim"
   ]
  }
 ],
 "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.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}