Dataset_Generate_Critic_Imgs.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES']='1'\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_preds(dl, path_gen, learn):\n",
    "    i=0\n",
    "    names = dl.dataset.items\n",
    "    \n",
    "    for b in dl:\n",
    "        preds = learn.pred_batch(batch=b, reconstruct=True, ds_type=DatasetType.Valid)\n",
    "        for o in preds:\n",
    "            o.save(path_gen/names[i].name)\n",
    "            print(path_gen/names[i].name)\n",
    "            i += 1\n",
    "            \n",
    "\n",
    "def save_gen_images(data_gen, path_gen, learn):\n",
    "    if path_gen.exists(): shutil.rmtree(path_gen)\n",
    "    path_gen.mkdir(exist_ok=True)\n",
    "    save_preds(data_gen.fix_dl, path_gen, learn)\n",
    "    save_preds(data_gen.valid_dl, path_gen, learn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_DIV2k_data_Input(pLow, pFull, bs: int, sz: int):\n",
    "    \"\"\"Given the path of low resolution images\n",
    "       returns a databunch\n",
    "    \"\"\"\n",
    "    src = ImageImageList.from_folder(pLow, presort=True).split_by_idxs(\n",
    "        train_idx=list(range(0, 800)), valid_idx=list(range(800, 900)))\n",
    "\n",
    "    data = (src.label_from_func(\n",
    "        lambda x: pFull/(x.name.replace(\".jpg\", \".png\"))\n",
    "    ).transform(\n",
    "        size=sz,\n",
    "        tfm_y=True,\n",
    "    ).databunch(bs=bs, num_workers=8, no_check=True)\n",
    "        .normalize(imagenet_stats, do_y=True))\n",
    "    data.c = 3\n",
    "    return data"
   ]
  },
  {
   "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'"
   ]
  },
  {
   "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_dummy_databunch(1, 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": "markdown",
   "metadata": {},
   "source": [
    "# Generator dataset with exported model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "weights = \"/data/students_home/fmameli/repos/Artifact_Removal_GAN/models/unet_wideNf2_mobileV3_DivFlickr1k_P64px_SuperRes_gen_3\"\n",
    "learn_gen.load(weights, with_opt=False)\n",
    "\n",
    "learn_gen.export(\"/data/students_home/fmameli/repos/Artifact_Removal_GAN/models/std_path.pkl\" )"
   ]
  },
  {
   "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/std_patch.pkl\")\n",
    "\n",
    "learn_std = load_learner(path=root_model_path, file=exported_model_standard)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(1, 900):\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_dummy_databunch(1, size)\n",
    "\n",
    "    learn_std.data = data_gen\n",
    "    \n",
    "    p,img_hr,b = learn_std.predict(img_low)\n",
    "    p.save(\"dataset/MobilenetV3_Patch_GEN/\" + id_img + \"_LPIPS.png\")\n",
    "    print(\"dataset/MobilenetV3_Patch_GEN/\" + id_img + \".png\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Input imgs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "proj_id = 'unet_superRes_mobilenetV3_Input'\n",
    "\n",
    "path_input = path/(proj_id + '_image_gen')\n",
    "path_input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "bs=1\n",
    "sz=512\n",
    "nf_factor = 2\n",
    "\n",
    "data_gen = get_DIV2k_data_Input(path_lowRes_512, path_fullRes, bs=bs, sz=sz)\n",
    "\n",
    "if path_input.exists(): shutil.rmtree(path_input)\n",
    "path_input.mkdir(exist_ok=True)\n",
    "i=0\n",
    "names = data_gen.fix_dl.dataset.items\n",
    "for img in data_gen.fix_dl.dataset:\n",
    "    img[0].save(path_input/names[i].name)\n",
    "    print(path_input/names[i].name)\n",
    "    i += 1\n",
    "\n",
    "i=0\n",
    "names = data_gen.valid_dl.dataset.items\n",
    "for img in data_gen.valid_dl.dataset:\n",
    "    img[0].save(path_input/names[i].name)\n",
    "    print(path_input/names[i].name)\n",
    "    i += 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SSIM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "proj_id = 'unet_superRes_mobilenetV3_SSIM'\n",
    "\n",
    "path_ssim = path/(proj_id + '_image_gen')\n",
    "path_ssim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = geffnet.mobilenetv3_rw\n",
    "loss_func = SSIM()\n",
    "\n",
    "bs=1\n",
    "sz=512\n",
    "nf_factor = 2\n",
    "\n",
    "data_gen = get_DIV2k_data_QF(path_lowRes_512, path_fullRes, bs=bs, sz=sz)\n",
    "\n",
    "learn_gen = gen_learner_wide(data=data_gen,\n",
    "                                 gen_loss=loss_func,\n",
    "                                 arch = model,\n",
    "                                 nf_factor=nf_factor)\n",
    "\n",
    "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",
    "save_gen_images(data_gen, path_ssim, learn_gen)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MSE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "proj_id = 'unet_superRes_mobilenetV3_MSE'\n",
    "\n",
    "path_mse = path/(proj_id + '_image_gen')\n",
    "path_mse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = geffnet.mobilenetv3_rw\n",
    "loss_func = nn.MSELoss()\n",
    "\n",
    "bs=1\n",
    "sz=512\n",
    "nf_factor = 2\n",
    "\n",
    "data_gen = get_DIV2k_data_QF(path_lowRes_512, path_fullRes, bs=bs, sz=sz)\n",
    "\n",
    "learn_gen = gen_learner_wide(data=data_gen,\n",
    "                                 gen_loss=loss_func,\n",
    "                                 arch = model,\n",
    "                                 nf_factor=nf_factor)\n",
    "\n",
    "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",
    "save_gen_images(data_gen, path_mse, learn_gen)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LPIPS Patch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "proj_id = 'unet_superRes_mobilenetV3_LPIPS_Patch'\n",
    "\n",
    "path_lpips_patch = path/(proj_id + '_image_gen')\n",
    "path_lpips_patch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = geffnet.mobilenetv3_rw\n",
    "loss_func = lpips_loss()\n",
    "\n",
    "bs=2\n",
    "sz=512\n",
    "nf_factor = 2\n",
    "\n",
    "data_gen = get_DIV2k_data_QF(path_lowRes_512, path_fullRes, bs=bs, sz=sz)\n",
    "\n",
    "learn_gen = gen_learner_wide(data=data_gen,\n",
    "                                 gen_loss=loss_func,\n",
    "                                 arch = model,\n",
    "                                 nf_factor=nf_factor)\n",
    "\n",
    "learn_gen.load(\"/data/students_home/fmameli/repos/SuperRes/models/unet_wideNf2_superRes_mobilenetV3_Patches64px_gen_64px_0\")\n",
    "\n",
    "save_gen_images(data_gen, path_lpips_patch, learn_gen)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LPIPS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "proj_id = 'unet_superRes_mobilenetV3_LPIPS'\n",
    "\n",
    "path_lpips = path/(proj_id + '_image_gen')\n",
    "path_lpips"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = geffnet.mobilenetv3_rw\n",
    "loss_func = lpips_loss()\n",
    "\n",
    "bs=2\n",
    "sz=512\n",
    "nf_factor = 2\n",
    "\n",
    "data_gen = get_DIV2k_data_QF(path_lowRes_512, path_fullRes, bs=bs, sz=sz)\n",
    "\n",
    "learn_gen = gen_learner_wide(data=data_gen,\n",
    "                                 gen_loss=loss_func,\n",
    "                                 arch = model,\n",
    "                                 nf_factor=nf_factor)\n",
    "\n",
    "weights = \"/data/students_home/fmameli/repos/SuperRes/models/unet_superRes_mobilenetV3_LPIPS_Tuned_gen_512px_0\"\n",
    "learn_gen.load(weights, with_opt=False)\n",
    "\n",
    "save_gen_images(data_gen, path_lpips, learn_gen)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# GAN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "proj_id = 'unet_superRes_mobilenetV3_GAN'\n",
    "\n",
    "path_gan = path/(proj_id + '_image_gen')\n",
    "path_gan"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = geffnet.mobilenetv3_rw\n",
    "loss_func = lpips_loss()\n",
    "\n",
    "bs=2\n",
    "sz=512\n",
    "nf_factor = 2\n",
    "\n",
    "data_gen = get_DIV2k_data_QF(path_lowRes_512, path_fullRes, bs=bs, sz=sz)\n",
    "\n",
    "learn_gen = gen_learner_wide(data=data_gen,\n",
    "                                 gen_loss=loss_func,\n",
    "                                 arch = model,\n",
    "                                 nf_factor=nf_factor)\n",
    "\n",
    "weights = \"/data/students_home/fmameli/repos/SuperRes/models/unet_superRes_mobilenetV3_LPIPS_Tuned_gen_512px_0\"\n",
    "learn_gen.load(weights, with_opt=False)\n",
    "\n",
    "save_gen_images(data_gen, path_gan, learn_gen)"
   ]
  }
 ],
 "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}