ni-mvi_2_train.ipynb 13.8 KB
Newer Older
Matej Choma's avatar
Matej Choma committed
1 2 3 4
{
 "cells": [
  {
   "cell_type": "code",
Matej Choma's avatar
Matej Choma committed
5
   "execution_count": 23,
Matej Choma's avatar
Matej Choma committed
6 7 8 9 10 11 12 13 14 15 16 17 18 19 20
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.utils.data import random_split\n",
    "\n",
    "import pytorch_lightning as pl\n",
    "from pytorch_lightning.callbacks import ModelCheckpoint\n",
    "from pytorch_lightning.callbacks.early_stopping import EarlyStopping\n",
    "from pytorch_lightning.loggers import TensorBoardLogger\n",
    "\n",
    "from core import datasets\n",
    "from core.utils import visualization\n",
Matej Choma's avatar
Matej Choma committed
21
    "from core.utils.helpers import torch_scale, dump_json\n",
Matej Choma's avatar
Matej Choma committed
22
    "from core.metrics import ssim, psnr, mse\n",
Matej Choma's avatar
Matej Choma committed
23 24 25
    "\n",
    "from core.base import BaseModule\n",
    "from core.modules import UnetBT, BaseUpsampler\n",
Matej Choma's avatar
Matej Choma committed
26
    "from core.module_unet import UnetCascade\n",
Matej Choma's avatar
Matej Choma committed
27
    "from core.module_gan import SRResNet, SRGAN\n",
Matej Choma's avatar
Matej Choma committed
28
    "\n",
Matej Choma's avatar
Matej Choma committed
29 30 31 32 33 34 35 36 37 38 39
    "import numpy as np\n",
    "\n",
    "generalconfig = {\n",
    "    'val_loss': {\n",
    "        'mse': nn.MSELoss(),\n",
    "        'ssim': ssim,\n",
    "        'psnr': psnr\n",
    "    },\n",
    "    \n",
    "    'range_base': (0, 255),\n",
    "}"
Matej Choma's avatar
Matej Choma committed
40 41 42
   ]
  },
  {
Matej Choma's avatar
Matej Choma committed
43
   "cell_type": "markdown",
Matej Choma's avatar
Matej Choma committed
44
   "metadata": {},
Matej Choma's avatar
Matej Choma committed
45 46 47
   "source": [
    "### Dataset Init"
   ]
Matej Choma's avatar
Matej Choma committed
48 49 50
  },
  {
   "cell_type": "code",
Matej Choma's avatar
Matej Choma committed
51
   "execution_count": 24,
Matej Choma's avatar
Matej Choma committed
52
   "metadata": {},
Matej Choma's avatar
Matej Choma committed
53 54
   "outputs": [],
   "source": [
Matej Choma's avatar
Matej Choma committed
55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77
    "#######################################\n",
    "#   DIV2K MILD UNKNOWN\n",
    "#######################################\n",
    "\n",
    "# dataconfig = {\n",
    "#     'batch_size': 16,\n",
    "#     'train_kwargs': {\n",
    "#         'dir_hr': 'data/datasets/DIV2K/train_HR_480/',\n",
    "#         'dir_lr': 'data/datasets/DIV2K/train_LR_120_unknown/',\n",
    "#         'crop': (128, 128)\n",
    "#     },\n",
    "#     'valid_kwargs': {\n",
    "#         'dir_hr': 'data/datasets/DIV2K/valid_HR_1000/',\n",
    "#         'dir_lr': 'data/datasets/DIV2K/valid_LR_250_unknown/'\n",
    "#     }\n",
    "# }\n",
    "# LOGDIR = 'lightning_logs/div2k_mild/'\n",
    "\n",
    "#######################################\n",
    "#   DIV2K BICUBIC\n",
    "#######################################\n",
    "\n",
    "# dataconfig = {\n",
Matej Choma's avatar
Matej Choma committed
78
    "#     'batch_size': 16,\n",
Matej Choma's avatar
Matej Choma committed
79
    "#     'train_kwargs': {\n",
Matej Choma's avatar
Matej Choma committed
80 81
    "#         'dir_hr': 'data/datasets/DIV2K/train_HR_480/',\n",
    "#         'dir_lr': 'data/datasets/DIV2K/train_LR_120_bicubic/',\n",
Matej Choma's avatar
Matej Choma committed
82 83 84
    "#         'crop': (128, 128)\n",
    "#     },\n",
    "#     'valid_kwargs': {\n",
Matej Choma's avatar
Matej Choma committed
85 86
    "#         'dir_hr': 'data/datasets/DIV2K/valid_HR_1000/',\n",
    "#         'dir_lr': 'data/datasets/DIV2K/valid_LR_250_bicubic/'\n",
Matej Choma's avatar
Matej Choma committed
87 88
    "#     }\n",
    "# }\n",
Matej Choma's avatar
Matej Choma committed
89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112
    "# LOGDIR = 'lightning_logs/div2k_bicubic/'\n",
    "\n",
    "#######################################\n",
    "#   RAIN\n",
    "#######################################\n",
    "\n",
    "dataconfig = {\n",
    "    'batch_size': 16,\n",
    "    'train_kwargs': {\n",
    "        'dir_hr': 'data/datasets/rain/train_128_HR/',\n",
    "        'dir_lr': 'data/datasets/rain/train_32_LR/',\n",
    "#         'crop': (512, 512)\n",
    "#         'crop': (128, 128)\n",
    "    },\n",
    "    'valid_kwargs': {\n",
    "        'dir_hr': 'data/datasets/rain/valid_HR/',\n",
    "        'dir_lr': 'data/datasets/rain/valid_LR/'\n",
    "    },\n",
    "    'test_kwargs': {\n",
    "        'dir_hr': 'data/datasets/rain/test_HR/',\n",
    "        'dir_lr': 'data/datasets/rain/test_LR/'\n",
    "    }\n",
    "}\n",
    "LOGDIR = 'lightning_logs/rain_128/'\n",
Matej Choma's avatar
Matej Choma committed
113
    "\n",
Matej Choma's avatar
Matej Choma committed
114
    "#######################################\n",
Matej Choma's avatar
Matej Choma committed
115
    "\n",
Matej Choma's avatar
Matej Choma committed
116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132
    "# init datasets\n",
    "d_train = datasets.SRDataset(**dataconfig['train_kwargs'])\n",
    "d_val = datasets.SRDataset(**dataconfig['valid_kwargs'])\n",
    "\n",
    "train_loader = DataLoader(d_train, batch_size=dataconfig['batch_size'], num_workers=4)\n",
    "val_loader = DataLoader(d_val, batch_size=4, num_workers=4)\n",
    "\n",
    "if 'test_kwargs' in dataconfig:\n",
    "    d_test = datasets.SRDataset(**dataconfig['test_kwargs'])\n",
    "    test_loader = DataLoader(d_test, batch_size=4, num_workers=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model Init"
Matej Choma's avatar
Matej Choma committed
133 134 135 136
   ]
  },
  {
   "cell_type": "code",
Matej Choma's avatar
Matej Choma committed
137
   "execution_count": 3,
Matej Choma's avatar
Matej Choma committed
138
   "metadata": {},
Matej Choma's avatar
Matej Choma committed
139 140 141
   "outputs": [],
   "source": [
    "#######################################\n",
Matej Choma's avatar
Matej Choma committed
142
    "#   BASE UPSAMPLER\n",
Matej Choma's avatar
Matej Choma committed
143
    "#######################################\n",
Matej Choma's avatar
Matej Choma committed
144
    "NAME = 'base_upsampler'\n",
Matej Choma's avatar
Matej Choma committed
145 146
    "HPARAMS = {\n",
    "    'lr_scale': 4,\n",
Matej Choma's avatar
Matej Choma committed
147
    "    'upsample_mode': 'bicubic',\n",
Matej Choma's avatar
Matej Choma committed
148 149 150 151 152
    "    \n",
    "    'loss': nn.L1Loss(),\n",
    "    \n",
    "    'range_in': (0, 1),\n",
    "    'range_out': (0, 1),\n",
Matej Choma's avatar
Matej Choma committed
153
    "    **generalconfig\n",
Matej Choma's avatar
Matej Choma committed
154 155 156
    "}\n",
    "\n",
    "# logger\n",
Matej Choma's avatar
Matej Choma committed
157
    "logger = TensorBoardLogger(LOGDIR, name=NAME, default_hp_metric=False, version=0)\n",
Matej Choma's avatar
Matej Choma committed
158 159
    "\n",
    "# model\n",
Matej Choma's avatar
Matej Choma committed
160
    "model = BaseUpsampler(**HPARAMS)\n",
Matej Choma's avatar
Matej Choma committed
161 162 163 164 165
    "model.float() ;"
   ]
  },
  {
   "cell_type": "code",
Matej Choma's avatar
Matej Choma committed
166
   "execution_count": 25,
Matej Choma's avatar
Matej Choma committed
167
   "metadata": {},
Matej Choma's avatar
Matej Choma committed
168 169 170
   "outputs": [],
   "source": [
    "#######################################\n",
Matej Choma's avatar
Matej Choma committed
171
    "#   U-Net Cascade\n",
Matej Choma's avatar
Matej Choma committed
172
    "#######################################\n",
Matej Choma's avatar
Matej Choma committed
173
    "NAME = 'unetcascade'\n",
Matej Choma's avatar
Matej Choma committed
174
    "HPARAMS = {\n",
Matej Choma's avatar
Matej Choma committed
175 176 177
    "    'n_casc': 3,\n",
    "    'rcab_feat': 128,\n",
    "    'bn': False,\n",
Matej Choma's avatar
Matej Choma committed
178 179
    "    'lr_scale': 4,\n",
    "    \n",
Matej Choma's avatar
Matej Choma committed
180 181
    "    'loss': nn.MSELoss(),\n",
    "\n",
Matej Choma's avatar
Matej Choma committed
182
    "    'range_in': (0, 1),\n",
Matej Choma's avatar
Matej Choma committed
183 184
    "    'range_out': (-1, 1),\n",
    "    **generalconfig\n",
Matej Choma's avatar
Matej Choma committed
185 186 187
    "}\n",
    "\n",
    "# logger\n",
Matej Choma's avatar
Matej Choma committed
188
    "logger = TensorBoardLogger(LOGDIR, name=NAME, default_hp_metric=False)\n",
Matej Choma's avatar
Matej Choma committed
189 190
    "\n",
    "# model\n",
Matej Choma's avatar
Matej Choma committed
191
    "model = UnetCascade(**HPARAMS)\n",
Matej Choma's avatar
Matej Choma committed
192 193 194 195 196
    "model.float() ;"
   ]
  },
  {
   "cell_type": "code",
Matej Choma's avatar
Matej Choma committed
197
   "execution_count": 28,
Matej Choma's avatar
Matej Choma committed
198
   "metadata": {},
Matej Choma's avatar
Matej Choma committed
199 200 201
   "outputs": [],
   "source": [
    "#######################################\n",
Matej Choma's avatar
Matej Choma committed
202
    "#   SRResNet\n",
Matej Choma's avatar
Matej Choma committed
203
    "#######################################\n",
Matej Choma's avatar
Matej Choma committed
204
    "NAME = 'srresnet'\n",
Matej Choma's avatar
Matej Choma committed
205 206 207
    "HPARAMS = {\n",
    "    'lr_scale': 4,\n",
    "    \n",
Matej Choma's avatar
Matej Choma committed
208 209
    "    'loss': nn.MSELoss(),\n",
    "\n",
Matej Choma's avatar
Matej Choma committed
210
    "    'range_in': (0, 1),\n",
Matej Choma's avatar
Matej Choma committed
211 212
    "    'range_out': (-1, 1),\n",
    "    **generalconfig\n",
Matej Choma's avatar
Matej Choma committed
213 214 215
    "}\n",
    "\n",
    "# logger\n",
Matej Choma's avatar
Matej Choma committed
216
    "logger = TensorBoardLogger(LOGDIR, name=NAME, default_hp_metric=False)\n",
Matej Choma's avatar
Matej Choma committed
217 218
    "\n",
    "# model\n",
Matej Choma's avatar
Matej Choma committed
219
    "model = SRResNet(**HPARAMS)\n",
Matej Choma's avatar
Matej Choma committed
220 221 222 223 224
    "model.float() ;"
   ]
  },
  {
   "cell_type": "code",
Matej Choma's avatar
Matej Choma committed
225
   "execution_count": 30,
Matej Choma's avatar
Matej Choma committed
226 227 228 229
   "metadata": {},
   "outputs": [],
   "source": [
    "#######################################\n",
Matej Choma's avatar
Matej Choma committed
230
    "#   SRGAN\n",
Matej Choma's avatar
Matej Choma committed
231
    "#######################################\n",
Matej Choma's avatar
Matej Choma committed
232
    "NAME = 'srgan'\n",
Matej Choma's avatar
Matej Choma committed
233
    "HPARAMS = {\n",
Matej Choma's avatar
Matej Choma committed
234
    "    'W': 128,\n",
Matej Choma's avatar
Matej Choma committed
235
    "    'gen_init': 'lightning_logs/div2k_bicubic/srresnet/version_0/checkpoints/epoch=9.ckpt',\n",
Matej Choma's avatar
Matej Choma committed
236 237
    "    \n",
    "    'range_in': (0, 1),\n",
Matej Choma's avatar
Matej Choma committed
238 239
    "    'range_out': (-1, 1),\n",
    "    **generalconfig\n",
Matej Choma's avatar
Matej Choma committed
240 241 242
    "}\n",
    "\n",
    "# logger\n",
Matej Choma's avatar
Matej Choma committed
243
    "logger = TensorBoardLogger(LOGDIR, name=NAME, default_hp_metric=False, version=1)\n",
Matej Choma's avatar
Matej Choma committed
244 245
    "\n",
    "# model\n",
Matej Choma's avatar
Matej Choma committed
246
    "model = SRGAN(**HPARAMS)\n",
Matej Choma's avatar
Matej Choma committed
247 248 249
    "model.float() ;"
   ]
  },
Matej Choma's avatar
Matej Choma committed
250 251 252 253 254 255 256
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
Matej Choma's avatar
Matej Choma committed
257 258
  {
   "cell_type": "code",
Matej Choma's avatar
Matej Choma committed
259 260
   "execution_count": null,
   "metadata": {},
Matej Choma's avatar
Matej Choma committed
261
   "outputs": [],
Matej Choma's avatar
Matej Choma committed
262
   "source": [
Matej Choma's avatar
Matej Choma committed
263 264 265 266 267 268 269 270 271
    "# checkpoints\n",
    "checkpoint_callback = ModelCheckpoint(\n",
    "    monitor='val_loss',\n",
    "    filename='{epoch:02d}-{val_loss:.4f}',\n",
    "    save_top_k=3,\n",
    "    mode='min',\n",
    ")\n",
    "\n",
    "# training\n",
Matej Choma's avatar
Matej Choma committed
272 273 274
    "trainer = pl.Trainer(\n",
    "    gpus=1,\n",
    "    auto_select_gpus=True,\n",
Matej Choma's avatar
Matej Choma committed
275
    "    callbacks=[checkpoint_callback],\n",
Matej Choma's avatar
Matej Choma committed
276
    "    logger=logger,\n",
Matej Choma's avatar
Matej Choma committed
277
    "    progress_bar_refresh_rate=10,\n",
Matej Choma's avatar
Matej Choma committed
278 279 280 281 282
    "    max_epochs=50,\n",
    "#     resume_from_checkpoint='lightning_logs/rain/unetcascade/version_2/checkpoints/epoch=29.ckpt'\n",
    "#     resume_from_checkpoint='lightning_logs/div2k_bicubic/srresnet/version_0/checkpoints/epoch=9.ckpt'\n",
    "#     resume_from_checkpoint='lightning_logs/div2k_bicubic/unetcascade/version_0/checkpoints/epoch=9.ckpt'\n",
    "#     resume_from_checkpoint='lightning_logs/rain/unetcascade/version_1/checkpoints/epoch=19.ckpt'\n",
Matej Choma's avatar
Matej Choma committed
283
    ")\n",
Matej Choma's avatar
Matej Choma committed
284 285 286 287 288 289 290
    "trainer.fit(model, train_loader, val_loader)\n",
    "\n",
    "# log after\n",
    "dump_json(dataconfig, trainer.logger.log_dir + '/dataconfig.json')\n",
    "trainer.save_checkpoint(trainer.logger.log_dir + f'/checkpoints/epoch={trainer.current_epoch}.ckpt')"
   ]
  },
Matej Choma's avatar
Matej Choma committed
291 292 293 294 295 296 297 298 299 300
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "dump_json(dataconfig, trainer.logger.log_dir + '/dataconfig.json')\n",
    "trainer.save_checkpoint(trainer.logger.log_dir + f'/checkpoints/epoch={trainer.current_epoch}.ckpt')"
   ]
  },
Matej Choma's avatar
Matej Choma committed
301 302 303 304 305
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Validate/Test"
Matej Choma's avatar
Matej Choma committed
306 307 308 309
   ]
  },
  {
   "cell_type": "code",
Matej Choma's avatar
Matej Choma committed
310
   "execution_count": 4,
Matej Choma's avatar
Matej Choma committed
311 312 313 314 315 316 317 318
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: True\n",
      "TPU available: None, using: 0 TPU cores\n",
Matej Choma's avatar
Matej Choma committed
319
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]\n"
Matej Choma's avatar
Matej Choma committed
320 321 322 323 324 325
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
Matej Choma's avatar
Matej Choma committed
326
      "Testing: 100%|██████████| 53/53 [00:36<00:00,  1.47it/s]\n",
Matej Choma's avatar
Matej Choma committed
327 328 329 330
      "--------------------------------------------------------------------------------\n",
      "DATALOADER:0 TEST RESULTS\n",
      "{}\n",
      "--------------------------------------------------------------------------------\n"
Matej Choma's avatar
Matej Choma committed
331 332 333 334 335
     ]
    },
    {
     "data": {
      "text/plain": [
Matej Choma's avatar
Matej Choma committed
336
       "[{}]"
Matej Choma's avatar
Matej Choma committed
337 338
      ]
     },
Matej Choma's avatar
Matej Choma committed
339
     "execution_count": 4,
Matej Choma's avatar
Matej Choma committed
340 341 342 343 344
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
Matej Choma's avatar
Matej Choma committed
345
    "model.test_prefix = 'test_epoch'\n",
Matej Choma's avatar
Matej Choma committed
346 347 348 349 350
    "\n",
    "trainer = pl.Trainer(\n",
    "    gpus=1,\n",
    "    auto_select_gpus=True,\n",
    "    logger=logger,\n",
Matej Choma's avatar
Matej Choma committed
351 352
    "    progress_bar_refresh_rate=10,\n",
    "#     resume_from_checkpoint='lightning_logs/rain_128/srgan/version_1/checkpoints/epoch=00-val_loss=0.0185.ckpt'\n",
Matej Choma's avatar
Matej Choma committed
353
    ")\n",
Matej Choma's avatar
Matej Choma committed
354
    "trainer.test(model, test_loader)"
Matej Choma's avatar
Matej Choma committed
355 356 357 358 359 360 361
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Explore"
Matej Choma's avatar
Matej Choma committed
362 363 364 365
   ]
  },
  {
   "cell_type": "code",
Matej Choma's avatar
Matej Choma committed
366
   "execution_count": 3,
Matej Choma's avatar
Matej Choma committed
367 368 369
   "metadata": {},
   "outputs": [],
   "source": [
Matej Choma's avatar
Matej Choma committed
370 371 372 373 374 375 376 377 378
    "import matplotlib.pyplot as plt\n",
    "from skimage.transform import resize\n",
    "\n",
    "srresnet_div2k = SRResNet.load_from_checkpoint('lightning_logs/rain_128/srresnet/version_0/checkpoints/epoch=49-val_loss=0.0096.ckpt')\n",
    "srresnet = SRResNet.load_from_checkpoint('lightning_logs/rain_128/srresnet/version_1/checkpoints/epoch=48-val_loss=0.0096.ckpt')\n",
    "unet_div2k = UnetCascade.load_from_checkpoint('lightning_logs/rain_128/unetcascade/version_0/checkpoints/epoch=29-val_loss=0.0101.ckpt')\n",
    "unet = UnetCascade.load_from_checkpoint('lightning_logs/rain_128/unetcascade/version_1/checkpoints/epoch=06-val_loss=0.0155.ckpt')\n",
    "srgan_div2k = SRGAN.load_from_checkpoint('lightning_logs/rain_128/srgan/version_0/checkpoints/epoch=03-val_loss=0.0358.ckpt')\n",
    "srgan = SRGAN.load_from_checkpoint('lightning_logs/rain_128/srgan/version_1/checkpoints/epoch=00-val_loss=0.0185.ckpt')"
Matej Choma's avatar
Matej Choma committed
379 380 381 382
   ]
  },
  {
   "cell_type": "code",
Matej Choma's avatar
Matej Choma committed
383
   "execution_count": 18,
Matej Choma's avatar
Matej Choma committed
384 385 386
   "metadata": {},
   "outputs": [],
   "source": [
Matej Choma's avatar
Matej Choma committed
387 388
    "import matplotlib as mpl\n",
    "mpl.rcParams['axes.titlesize'] = 24"
Matej Choma's avatar
Matej Choma committed
389 390 391 392
   ]
  },
  {
   "cell_type": "code",
Matej Choma's avatar
Matej Choma committed
393
   "execution_count": 4,
Matej Choma's avatar
Matej Choma committed
394
   "metadata": {},
Matej Choma's avatar
Matej Choma committed
395
   "outputs": [],
Matej Choma's avatar
Matej Choma committed
396
   "source": [
Matej Choma's avatar
Matej Choma committed
397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420
    "i, o = d_val[47]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "sl = (slice(None), slice(460, 560), slice(360, 460))\n",
    "# sl = (slice(None), slice(None), slice(None))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(30, 30)) # figsize is in inches\n",
    "\n",
    "plt.subplot(331)\n",
    "visualization.show(i[:, 115:140, 90:115], fig=False, title='Input')\n",
    "# visualization.show(i, fig=False, title='Input')\n",
Matej Choma's avatar
Matej Choma committed
421
    "\n",
Matej Choma's avatar
Matej Choma committed
422 423 424
    "plt.subplot(332)\n",
    "b = BaseUpsampler(upsample_mode='bicubic').predict(i)\n",
    "visualization.show(b[sl], fig=False, title='Bicubic')\n",
Matej Choma's avatar
Matej Choma committed
425
    "\n",
Matej Choma's avatar
Matej Choma committed
426 427
    "plt.subplot(333)\n",
    "visualization.show(o[sl], fig=False, title='Output')\n",
Matej Choma's avatar
Matej Choma committed
428
    "\n",
Matej Choma's avatar
Matej Choma committed
429 430 431
    "plt.subplot(335)\n",
    "p = srresnet_div2k.predict(i)\n",
    "visualization.show(p[sl], fig=False, title='SRResNet DIV2K')\n",
Matej Choma's avatar
Matej Choma committed
432
    "\n",
Matej Choma's avatar
Matej Choma committed
433 434 435
    "plt.subplot(338)\n",
    "p = srresnet.predict(i)\n",
    "visualization.show(p[sl], fig=False, title='SRResNet')\n",
Matej Choma's avatar
Matej Choma committed
436
    "\n",
Matej Choma's avatar
Matej Choma committed
437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455
    "plt.subplot(334)\n",
    "p = unet_div2k.predict(i)\n",
    "visualization.show(p[sl], fig=False, title='U-net DIV2K')\n",
    "\n",
    "plt.subplot(337)\n",
    "p = unet.predict(i)\n",
    "visualization.show(p[sl], fig=False, title='U-net')\n",
    "\n",
    "\n",
    "plt.subplot(336)\n",
    "p = srgan_div2k.predict(i)\n",
    "visualization.show(p[sl], fig=False, title='SRGAN DIV2K')\n",
    "\n",
    "plt.subplot(339)\n",
    "p = srgan.predict(i)\n",
    "visualization.show(p[sl], fig=False, title='SRGAN')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.savefig('imgs/pred_small.png')"
Matej Choma's avatar
Matej Choma committed
456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python fitmvi",
   "language": "python",
   "name": "fitmvi"
  },
  "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}