train_video_super_resolution (REDS4 sample dataset).ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using TensorFlow backend.\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "from keras.optimizers import Adam\n",
    "from EDVR_arch import EDVR, charbonnier_penalty\n",
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "path = './dataset/'\n",
    "images = os.listdir(path+'GT')\n",
    "nframes = 5\n",
    "center = nframes//2\n",
    "batch_size = 1 "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create train data \n",
    "\n",
    "N=2\n",
    "#### X is frames[t-N : t+N]\n",
    "\n",
    "#### y is frames[t]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((11, 5, 180, 320, 3), (11, 720, 1280, 3))"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X = []\n",
    "y = []\n",
    "for k, img in enumerate(images):\n",
    "    y.append(cv2.imread(os.path.join(path, 'GT', img)))\n",
    "    X.append(cv2.imread(os.path.join(path, 'blur_bicubic', img)))\n",
    "X = np.stack(X)\n",
    "y = np.stack(y)\n",
    "\n",
    "X_trn = []\n",
    "for i in range(len(images)):\n",
    "    next_frames = X[i:i+center+1]\n",
    "    if i<center:\n",
    "        prev_frames = X[:i]\n",
    "    else:\n",
    "        prev_frames = X[i-center:i]\n",
    "    \n",
    "    to_fill = nframes - next_frames.shape[0] - prev_frames.shape[0]\n",
    "    if to_fill:\n",
    "        if len(prev_frames) and i<nframes:\n",
    "            pad_x = np.repeat(prev_frames[0][None], to_fill, axis=0)\n",
    "            xx = np.concatenate((pad_x, prev_frames, next_frames))\n",
    "        else:\n",
    "            if i>nframes:\n",
    "                pad_x = np.repeat(next_frames[-1][None], to_fill, axis=0)\n",
    "                xx = np.concatenate((prev_frames, next_frames, pad_x))\n",
    "            else:\n",
    "                pad_x = np.repeat(next_frames[0][None], to_fill, axis=0)\n",
    "                xx = np.concatenate((pad_x, prev_frames, next_frames))\n",
    "    else:\n",
    "        xx = np.concatenate((prev_frames, next_frames))\n",
    "    X_trn.append(xx)\n",
    "X_trn = np.stack(X_trn)\n",
    "X_trn.shape, y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "False\n"
     ]
    }
   ],
   "source": [
    "# Check if ground truth has same dimensions as input images\n",
    "HR_in = False if np.prod(X_trn.shape[2:]) < np.prod(y.shape[1:]) else True\n",
    "print(HR_in)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /usr/local/lib/python3.5/dist-packages/tensorflow/python/framework/op_def_library.py:263: colocate_with (from tensorflow.python.framework.ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Colocations handled automatically by placer.\n"
     ]
    }
   ],
   "source": [
    "VideoSuperResolution = EDVR(inp_shape=X_trn.shape[2:],\n",
    "                            nf=64, nframes=nframes,\n",
    "                            groups=8, front_RBs=5,\n",
    "                            back_RBs=10, center=None,\n",
    "                            predeblur=True, HR_in=HR_in)\n",
    "model = VideoSuperResolution.get_EDVR_model()\n",
    "optimizer = Adam(lr=4e-4, beta_1=.9, beta_2=0.999)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /usr/local/lib/python3.5/dist-packages/tensorflow/python/ops/math_ops.py:3066: to_int32 (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use tf.cast instead.\n",
      "Train on 9 samples, validate on 2 samples\n",
      "Epoch 1/5\n"
     ]
    }
   ],
   "source": [
    "model.compile(optimizer=optimizer, loss=charbonnier_penalty)\n",
    "model.fit(x=X_trn/255, y=y/255, batch_size=batch_size,\n",
    "          epochs=5, validation_split=0.1, shuffle=True)"
   ]
  }
 ],
 "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}