Video_Magnification.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "lQW-nsqB2dEe"
   },
   "outputs": [],
   "source": [
    "# Clone the official github repo here\n",
    "!git clone https://github.com/12dmodel/deep_motion_mag.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "AzFD2Nvx5Icd"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Video_Magnification_backup.ipynb  Video_Magnification.ipynb\n"
     ]
    }
   ],
   "source": [
    "# Go into the  corresponding folder\n",
    "cd deep_motion_mag"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "7zxYWImU403o"
   },
   "outputs": [],
   "source": [
    "# Install all the requirements\n",
    "!pip install -r requirements.txt\n",
    "!ls\n",
    "# Warning:: \n",
    "# Since this repository works on all the older versions of library after this step you might need to consider restarting the runtime for proper installation.\n",
    "# if you restart runtime your all the above updates would be lost so please run once all cells from cell number 3 otherwise you'll stuck in error."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "_ba92YWV5KgH"
   },
   "outputs": [],
   "source": [
    "# Download the pretrained model into the folder after confirming that you are in the repo's folder.\n",
    "!wget https://people.csail.mit.edu/tiam/deepmag/data.zip\n",
    "\n",
    "# Unzip the Folder\n",
    "!unzip data.zip\n",
    "\n",
    "# Delete the zip.\n",
    "!rm -r data.zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "mdSCKXzH5dCR"
   },
   "outputs": [],
   "source": [
    "# For testing purpose you have to follow link on the github repo, we are using the same baby video to show the working.\n",
    "!wget https://people.csail.mit.edu/mrub/evm/video/baby.mp4\n",
    "\n",
    "# We have to extract the frames to the corresponding folder.\n",
    "# 1. Create the corresponding folder (Last 'baby' should be replaced by your video name)\n",
    "!mkdir /content/deep_motion_mag/data/vids/baby \n",
    "\n",
    "# 2. Extract the frames to that corresponding folder. ('ffmpeg -i source -f image2 destination')\n",
    "!ffmpeg -i /content/deep_motion_mag/baby.mp4 -f image2 /content/deep_motion_mag/data/vids/baby/%06d.png\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "-thz2VeK9d_R"
   },
   "outputs": [],
   "source": [
    "''' Since avconv has been removed and we have only ffmpeg, we have to change the magnet.py file to default to ffmpeg to remove errors.'''\n",
    "# Remove the magnet.py file\n",
    "!rm magnet.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "9xEVvGgb87uP"
   },
   "outputs": [],
   "source": [
    "# Save the new magnet .py file from the repo just with a change that default library is ffmpeg.\n",
    "%%writefile magnet.py\n",
    "from __future__ import print_function\n",
    "\n",
    "import os\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import cv2\n",
    "import time\n",
    "\n",
    "from glob import glob\n",
    "from scipy.signal import firwin, butter\n",
    "from functools import partial\n",
    "from tqdm import tqdm, trange\n",
    "from subprocess import call\n",
    "\n",
    "from modules import L1_loss\n",
    "from modules import res_encoder, res_decoder, res_manipulator\n",
    "from modules import residual_block, conv2d\n",
    "from utils import load_train_data, mkdir, imread, save_images\n",
    "from preprocessor import preprocess_image, preproc_color\n",
    "from data_loader import read_and_decode_3frames\n",
    "\n",
    "# Change here if you use ffmpeg.\n",
    "DEFAULT_VIDEO_CONVERTER = 'ffmpeg'\n",
    "\n",
    "\n",
    "class MagNet3Frames(object):\n",
    "\n",
    "    def __init__(self, sess, name, arch_config):\n",
    "        self.sess = sess\n",
    "        self.exp_name = name\n",
    "        self.is_graph_built = False\n",
    "        self.n_channels = arch_config[\"n_channels\"]\n",
    "        self.arch_config = arch_config\n",
    "        self.encoder_dims = arch_config[\"ynet_3frames\"][\"enc_dims\"]\n",
    "        self.num_enc_resblk = arch_config[\"ynet_3frames\"][\"num_enc_resblk\"]\n",
    "        self.num_man_resblk = arch_config[\"ynet_3frames\"][\"num_man_resblk\"]\n",
    "        self.num_man_conv = arch_config[\"ynet_3frames\"][\"num_man_conv\"]\n",
    "        self.num_man_aft_conv = arch_config[\"ynet_3frames\"][\"num_man_aft_conv\"]\n",
    "        self.num_dec_resblk = arch_config[\"ynet_3frames\"][\"num_dec_resblk\"]\n",
    "        self.num_texture_resblk = \\\n",
    "            arch_config[\"ynet_3frames\"][\"num_texture_resblk\"]\n",
    "        self.texture_dims = arch_config[\"ynet_3frames\"][\"texture_dims\"]\n",
    "        self.texture_downsample = \\\n",
    "            arch_config[\"ynet_3frames\"][\"texture_downsample\"]\n",
    "        self.use_texture_conv = arch_config[\"ynet_3frames\"][\"use_texture_conv\"]\n",
    "        self.shape_dims = arch_config[\"ynet_3frames\"][\"shape_dims\"]\n",
    "        self.num_shape_resblk = \\\n",
    "            arch_config[\"ynet_3frames\"][\"num_shape_resblk\"]\n",
    "        self.use_shape_conv = arch_config[\"ynet_3frames\"][\"use_shape_conv\"]\n",
    "        self.decoder_dims = self.texture_dims + self.shape_dims\n",
    "        self.probe_pt = {}\n",
    "        self.manipulator = partial(res_manipulator,\n",
    "                                   layer_dims=self.encoder_dims,\n",
    "                                   num_resblk=self.num_man_resblk,\n",
    "                                   num_conv=self.num_man_conv,\n",
    "                                   num_aft_conv=self.num_man_aft_conv,\n",
    "                                   probe_pt=self.probe_pt)\n",
    "\n",
    "    def _encoder(self, image):\n",
    "        enc = res_encoder(image,\n",
    "                          layer_dims=self.encoder_dims,\n",
    "                          num_resblk=self.num_enc_resblk)\n",
    "\n",
    "        texture_enc = enc\n",
    "        shape_enc = enc\n",
    "        # first convolution on common encoding\n",
    "        if self.use_texture_conv:\n",
    "            stride = 2 if self.texture_downsample else 1\n",
    "            texture_enc = tf.nn.relu(conv2d(texture_enc, self.texture_dims,\n",
    "                                            3, stride,\n",
    "                                            name='enc_texture_conv'))\n",
    "        else:\n",
    "            assert self.texture_dims == self.encoder_dims, \\\n",
    "                \"Texture dim ({}) must match encoder dim ({}) \" \\\n",
    "                \"if texture_conv is not used.\".format(self.texture_dims,\n",
    "                                                      self.encoder_dims)\n",
    "            assert not self.texture_downsample, \\\n",
    "                \"Must use texture_conv if texture_downsample.\"\n",
    "        if self.use_shape_conv:\n",
    "            shape_enc = tf.nn.relu(conv2d(shape_enc, self.shape_dims,\n",
    "                                          3, 1, name='enc_shape_conv'))\n",
    "        else:\n",
    "            assert self.shape_dims == self.encoder_dims, \\\n",
    "                \"Shape dim ({}) must match encoder dim ({}) \" \\\n",
    "                \"if shape_conv is not used.\".format(self.shape_dims,\n",
    "                                                    self.encoder_dims)\n",
    "\n",
    "        for i in range(self.num_texture_resblk):\n",
    "            name = 'texture_enc_{}'.format(i)\n",
    "            if i == 0:\n",
    "                # for backward compatibility\n",
    "                name = 'texture_enc'\n",
    "            texture_enc = residual_block(texture_enc, self.texture_dims, 3, 1,\n",
    "                                         name)\n",
    "        for i in range(self.num_shape_resblk):\n",
    "            name = 'shape_enc_{}'.format(i)\n",
    "            if i == 0:\n",
    "                # for backward compatibility\n",
    "                name = 'shape_enc'\n",
    "            shape_enc = residual_block(shape_enc, self.shape_dims,\n",
    "                                       3, 1, name)\n",
    "        return texture_enc, shape_enc\n",
    "\n",
    "    def _decoder(self, texture_enc, shape_enc):\n",
    "        if self.texture_downsample:\n",
    "            texture_enc = tf.image.resize_nearest_neighbor(\n",
    "                            texture_enc,\n",
    "                            tf.shape(texture_enc)[1:3] \\\n",
    "                            * 2)\n",
    "            texture_enc = tf.pad(texture_enc, [[0, 0], [1, 1], [1, 1], [0, 0]],\n",
    "                                 \"REFLECT\")\n",
    "            texture_enc = tf.nn.relu(conv2d(texture_enc, self.texture_dims,\n",
    "                                            3, 1, padding='VALID',\n",
    "                                            name='texture_upsample'))\n",
    "\n",
    "        enc = tf.concat([texture_enc, shape_enc], axis=3)\n",
    "        # Needs double the channel because we concat the two encodings.\n",
    "        return res_decoder(enc,\n",
    "                           layer_dims=self.decoder_dims,\n",
    "                           out_channels=self.n_channels,\n",
    "                           num_resblk=self.num_dec_resblk)\n",
    "\n",
    "    def image_transformer(self,\n",
    "                          image_a,\n",
    "                          image_b,\n",
    "                          amplification_factor,\n",
    "                          im_size,\n",
    "                          options,\n",
    "                          is_training,\n",
    "                          reuse=False,\n",
    "                          name='ynet_3frames'):\n",
    "        with tf.variable_scope(name, reuse=reuse):\n",
    "            with tf.variable_scope('encoder'):\n",
    "                self.texture_a, self.shape_a = self._encoder(image_a)\n",
    "            with tf.variable_scope('encoder', reuse=True):\n",
    "                self.texture_b, self.shape_b = self._encoder(image_b)\n",
    "            with tf.variable_scope('manipulator'):\n",
    "                self.out_shape_enc = self.manipulator(self.shape_a,\n",
    "                                                      self.shape_b,\n",
    "                                                      amplification_factor)\n",
    "            with tf.variable_scope('decoder'):\n",
    "                return self._decoder(self.texture_b, self.out_shape_enc)\n",
    "\n",
    "    def save(self, checkpoint_dir, step):\n",
    "        model_name = self.exp_name\n",
    "\n",
    "        if not os.path.exists(checkpoint_dir):\n",
    "            os.makedirs(checkpoint_dir)\n",
    "\n",
    "        self.saver.save(self.sess,\n",
    "                        os.path.join(checkpoint_dir, model_name),\n",
    "                        global_step=step)\n",
    "\n",
    "    def load(self, checkpoint_dir, loader=None):\n",
    "        if not loader:\n",
    "            loader = self.saver\n",
    "        print(\" [*] Reading checkpoint...\")\n",
    "        if os.path.isdir(checkpoint_dir):\n",
    "            ckpt = tf.train.get_checkpoint_state(checkpoint_dir)\n",
    "            if ckpt and ckpt.model_checkpoint_path:\n",
    "                ckpt_name = ckpt.model_checkpoint_path\n",
    "            else:\n",
    "                ckpt_name = None\n",
    "        else:\n",
    "            # load from file\n",
    "            ckpt_name = checkpoint_dir\n",
    "        if ckpt_name:\n",
    "            loader.restore(self.sess, ckpt_name)\n",
    "            print('Loaded from ckpt: ' + ckpt_name)\n",
    "            self.ckpt_name = ckpt_name\n",
    "            return True\n",
    "        else:\n",
    "            return False\n",
    "\n",
    "    def _build_feed_model(self):\n",
    "        self.test_input = tf.placeholder(tf.float32,\n",
    "                                         [None, None, None,\n",
    "                                             self.n_channels * 3],\n",
    "                                         name='test_AB_and_output')\n",
    "        self.test_amplification_factor = tf.placeholder(tf.float32,\n",
    "                                                        [None],\n",
    "                                                        name='amplification_factor')\n",
    "        self.test_image_a = self.test_input[:, :, :, :self.n_channels]\n",
    "        self.test_image_b = self.test_input[:, :, :, self.n_channels:(2 * self.n_channels)]\n",
    "        self.test_amplified_frame = self.test_input[:, :, :, (2*self.n_channels):(3 * self.n_channels)]\n",
    "        self.test_output = self.image_transformer(\n",
    "                               self.test_image_a,\n",
    "                               self.test_image_b,\n",
    "                               self.test_amplification_factor,\n",
    "                               [self.image_height, self.image_width],\n",
    "                               self.arch_config,\n",
    "                               False,\n",
    "                               False)\n",
    "        self.test_output = tf.clip_by_value(self.test_output, -1.0, 1.0)\n",
    "        self.saver = tf.train.Saver()\n",
    "        self.is_graph_built = True\n",
    "\n",
    "    def setup_for_inference(self, checkpoint_dir, image_width, image_height):\n",
    "        \"\"\"Setup model for inference.\n",
    "        Build computation graph, initialize variables, and load checkpoint.\n",
    "        \"\"\"\n",
    "        self.image_width = image_width\n",
    "        self.image_height = image_height\n",
    "        # Figure out image dimension\n",
    "        self._build_feed_model()\n",
    "        ginit_op = tf.global_variables_initializer()\n",
    "        linit_op = tf.local_variables_initializer()\n",
    "        self.sess.run([ginit_op, linit_op])\n",
    "\n",
    "        if self.load(checkpoint_dir):\n",
    "            print(\"[*] Load Success\")\n",
    "        else:\n",
    "            raise RuntimeError('MagNet: Failed to load checkpoint file.')\n",
    "        self.is_graph_built = True\n",
    "\n",
    "    def inference(self, frameA, frameB, amplification_factor):\n",
    "        \"\"\"Run Magnification on two frames.\n",
    "        Args:\n",
    "            frameA: path to first frame\n",
    "            frameB: path to second frame\n",
    "            amplification_factor: float for amplification factor\n",
    "        \"\"\"\n",
    "        in_frames = [load_train_data([frameA, frameB, frameB],\n",
    "                     gray_scale=self.n_channels==1, is_testing=True)]\n",
    "        in_frames = np.array(in_frames).astype(np.float32)\n",
    "\n",
    "        out_amp = self.sess.run(self.test_output,\n",
    "                                feed_dict={self.test_input: in_frames,\n",
    "                                           self.test_amplification_factor:\n",
    "                                           [amplification_factor]})\n",
    "        return out_amp\n",
    "\n",
    "    def run(self,\n",
    "            checkpoint_dir,\n",
    "            vid_dir,\n",
    "            frame_ext,\n",
    "            out_dir,\n",
    "            amplification_factor,\n",
    "            velocity_mag=False):\n",
    "        \"\"\"Magnify a video in the two-frames mode.\n",
    "        Args:\n",
    "            checkpoint_dir: checkpoint directory.\n",
    "            vid_dir: directory containing video frames videos are processed\n",
    "                in sorted order.\n",
    "            out_dir: directory to place output frames and resulting video.\n",
    "            amplification_factor: the amplification factor,\n",
    "                with 0 being no change.\n",
    "            velocity_mag: if True, process video in Dynamic mode.\n",
    "        \"\"\"\n",
    "        vid_name = os.path.basename(out_dir)\n",
    "        # make folder\n",
    "        mkdir(out_dir)\n",
    "        vid_frames = sorted(glob(os.path.join(vid_dir, '*.' + frame_ext)))\n",
    "        first_frame = vid_frames[0]\n",
    "        im = imread(first_frame)\n",
    "        image_height, image_width = im.shape\n",
    "        if not self.is_graph_built:\n",
    "            self.setup_for_inference(checkpoint_dir, image_width, image_height)\n",
    "        try:\n",
    "            i = int(self.ckpt_name.split('-')[-1])\n",
    "            print(\"Iteration number is {:d}\".format(i))\n",
    "            vid_name = vid_name + '_' + str(i)\n",
    "        except:\n",
    "            print(\"Cannot get iteration number\")\n",
    "        if velocity_mag:\n",
    "            print(\"Running in Dynamic mode\")\n",
    "\n",
    "        prev_frame = first_frame\n",
    "        desc = vid_name if len(vid_name) < 10 else vid_name[:10]\n",
    "        for frame in tqdm(vid_frames, desc=desc):\n",
    "            file_name = os.path.basename(frame)\n",
    "            out_amp = self.inference(prev_frame, frame, amplification_factor)\n",
    "\n",
    "            im_path = os.path.join(out_dir, file_name)\n",
    "            save_images(out_amp, [1, 1], im_path)\n",
    "            if velocity_mag:\n",
    "                prev_frame = frame\n",
    "\n",
    "        # Try to combine it into a video\n",
    "        call([DEFAULT_VIDEO_CONVERTER, '-y', '-f', 'image2', '-r', '30', '-i',\n",
    "              os.path.join(out_dir, '%06d.png'), '-c:v', 'libx264',\n",
    "              os.path.join(out_dir, vid_name + '.mp4')]\n",
    "            )\n",
    "\n",
    "    # Temporal Operations\n",
    "    def _build_IIR_filtering_graphs(self):\n",
    "        \"\"\"\n",
    "        Assume a_0 = 1\n",
    "        \"\"\"\n",
    "        self.input_image = tf.placeholder(tf.float32,\n",
    "                                          [1, self.image_height,\n",
    "                                              self.image_width,\n",
    "                                           self.n_channels],\n",
    "                                          name='input_image')\n",
    "        self.filtered_enc = tf.placeholder(tf.float32,\n",
    "                                           [1, None, None,\n",
    "                                            self.shape_dims],\n",
    "                                           name='filtered_enc')\n",
    "        self.out_texture_enc = tf.placeholder(tf.float32,\n",
    "                                              [1, None, None,\n",
    "                                               self.texture_dims],\n",
    "                                              name='out_texture_enc')\n",
    "        self.ref_shape_enc = tf.placeholder(tf.float32,\n",
    "                                            [1, None, None,\n",
    "                                             self.shape_dims],\n",
    "                                            name='ref_shape_enc')\n",
    "        self.amplification_factor = tf.placeholder(tf.float32, [None],\n",
    "                                                   name='amplification_factor')\n",
    "        with tf.variable_scope('ynet_3frames'):\n",
    "            with tf.variable_scope('encoder'):\n",
    "                self.texture_enc, self.shape_rep = \\\n",
    "                    self._encoder(self.input_image)\n",
    "            with tf.variable_scope('manipulator'):\n",
    "                # set encoder a to zero because we do temporal filtering\n",
    "                # instead of taking the difference.\n",
    "                self.out_shape_enc = self.manipulator(0.0,\n",
    "                                                      self.filtered_enc,\n",
    "                                                      self.amplification_factor)\n",
    "                self.out_shape_enc += self.ref_shape_enc - self.filtered_enc\n",
    "            with tf.variable_scope('decoder'):\n",
    "                self.output_image = tf.clip_by_value(\n",
    "                                        self._decoder(self.out_texture_enc,\n",
    "                                                      self.out_shape_enc),\n",
    "                                        -1.0, 1.0)\n",
    "\n",
    "        self.saver = tf.train.Saver()\n",
    "\n",
    "    def run_temporal(self,\n",
    "                     checkpoint_dir,\n",
    "                     vid_dir,\n",
    "                     frame_ext,\n",
    "                     out_dir,\n",
    "                     amplification_factor,\n",
    "                     fl, fh, fs,\n",
    "                     n_filter_tap,\n",
    "                     filter_type):\n",
    "        \"\"\"Magnify video with a temporal filter.\n",
    "        Args:\n",
    "            checkpoint_dir: checkpoint directory.\n",
    "            vid_dir: directory containing video frames videos are processed\n",
    "                in sorted order.\n",
    "            out_dir: directory to place output frames and resulting video.\n",
    "            amplification_factor: the amplification factor,\n",
    "                with 0 being no change.\n",
    "            fl: low cutoff frequency.\n",
    "            fh: high cutoff frequency.\n",
    "            fs: sampling rate of the video.\n",
    "            n_filter_tap: number of filter tap to use.\n",
    "            filter_type: Type of filter to use. Can be one of \"fir\",\n",
    "                \"butter\", or \"differenceOfIIR\". For \"differenceOfIIR\",\n",
    "                fl and fh specifies rl and rh coefficients as in Wadhwa et al.\n",
    "        \"\"\"\n",
    "\n",
    "        nyq = fs / 2.0\n",
    "        if filter_type == 'fir':\n",
    "            filter_b = firwin(n_filter_tap, [fl, fh], nyq=nyq, pass_zero=False)\n",
    "            filter_a = []\n",
    "        elif filter_type == 'butter':\n",
    "            filter_b, filter_a = butter(n_filter_tap, [fl/nyq, fh/nyq],\n",
    "                                        btype='bandpass')\n",
    "            filter_a = filter_a[1:]\n",
    "        elif filter_type == 'differenceOfIIR':\n",
    "            # This is a copy of what Neal did. Number of taps are ignored.\n",
    "            # Treat fl and fh as rl and rh as in Wadhwa's code.\n",
    "            # Write down the difference of difference equation in Fourier\n",
    "            # domain to proof this:\n",
    "            filter_b = [fh - fl, fl - fh]\n",
    "            filter_a = [-1.0*(2.0 - fh - fl), (1.0 - fl) * (1.0 - fh)]\n",
    "        else:\n",
    "            raise ValueError('Filter type must be either '\n",
    "                             '[\"fir\", \"butter\", \"differenceOfIIR\"] got ' + \\\n",
    "                             filter_type)\n",
    "        head, tail = os.path.split(out_dir)\n",
    "        tail = tail + '_fl{}_fh{}_fs{}_n{}_{}'.format(fl, fh, fs,\n",
    "                                                      n_filter_tap,\n",
    "                                                      filter_type)\n",
    "        out_dir = os.path.join(head, tail)\n",
    "        vid_name = os.path.basename(out_dir)\n",
    "        # make folder\n",
    "        mkdir(out_dir)\n",
    "        vid_frames = sorted(glob(os.path.join(vid_dir, '*.' + frame_ext)))\n",
    "        first_frame = vid_frames[0]\n",
    "        im = imread(first_frame)\n",
    "        image_height, image_width = im.shape\n",
    "        if not self.is_graph_built:\n",
    "            self.image_width = image_width\n",
    "            self.image_height = image_height\n",
    "            # Figure out image dimension\n",
    "            self._build_IIR_filtering_graphs()\n",
    "            ginit_op = tf.global_variables_initializer()\n",
    "            linit_op = tf.local_variables_initializer()\n",
    "            self.sess.run([ginit_op, linit_op])\n",
    "\n",
    "            if self.load(checkpoint_dir):\n",
    "                print(\"[*] Load Success\")\n",
    "            else:\n",
    "                raise RuntimeError('MagNet: Failed to load checkpoint file.')\n",
    "            self.is_graph_built = True\n",
    "        try:\n",
    "            i = int(self.ckpt_name.split('-')[-1])\n",
    "            print(\"Iteration number is {:d}\".format(i))\n",
    "            vid_name = vid_name + '_' + str(i)\n",
    "        except:\n",
    "            print(\"Cannot get iteration number\")\n",
    "\n",
    "        if len(filter_a) is not 0:\n",
    "            x_state = []\n",
    "            y_state = []\n",
    "\n",
    "            for frame in tqdm(vid_frames, desc='Applying IIR'):\n",
    "                file_name = os.path.basename(frame)\n",
    "                frame_no, _ = os.path.splitext(file_name)\n",
    "                frame_no = int(frame_no)\n",
    "                in_frames = [load_train_data([frame, frame, frame],\n",
    "                             gray_scale=self.n_channels==1, is_testing=True)]\n",
    "                in_frames = np.array(in_frames).astype(np.float32)\n",
    "\n",
    "                texture_enc, x = self.sess.run([self.texture_enc, self.shape_rep],\n",
    "                                               feed_dict={\n",
    "                                                   self.input_image:\n",
    "                                                   in_frames[:, :, :, :3],})\n",
    "                x_state.insert(0, x)\n",
    "                # set up initial condition.\n",
    "                while len(x_state) < len(filter_b):\n",
    "                    x_state.insert(0, x)\n",
    "                if len(x_state) > len(filter_b):\n",
    "                    x_state = x_state[:len(filter_b)]\n",
    "                y = np.zeros_like(x)\n",
    "                for i in range(len(x_state)):\n",
    "                    y += x_state[i] * filter_b[i]\n",
    "                for i in range(len(y_state)):\n",
    "                    y -= y_state[i] * filter_a[i]\n",
    "                # update y state\n",
    "                y_state.insert(0, y)\n",
    "                if len(y_state) > len(filter_a):\n",
    "                    y_state = y_state[:len(filter_a)]\n",
    "\n",
    "                out_amp = self.sess.run(self.output_image,\n",
    "                                        feed_dict={self.out_texture_enc:\n",
    "                                                     texture_enc,\n",
    "                                                   self.filtered_enc: y,\n",
    "                                                   self.ref_shape_enc: x,\n",
    "                                                   self.amplification_factor:\n",
    "                                                     [amplification_factor]})\n",
    "\n",
    "                im_path = os.path.join(out_dir, file_name)\n",
    "                out_amp = np.squeeze(out_amp)\n",
    "                out_amp = (127.5*(out_amp+1)).astype('uint8')\n",
    "                cv2.imwrite(im_path, cv2.cvtColor(out_amp,\n",
    "                                                  code=cv2.COLOR_RGB2BGR))\n",
    "        else:\n",
    "            # This does FIR in fourier domain. Equivalent to cyclic\n",
    "            # convolution.\n",
    "            x_state = None\n",
    "            for i, frame in tqdm(enumerate(vid_frames),\n",
    "                                 desc='Getting encoding'):\n",
    "                file_name = os.path.basename(frame)\n",
    "                in_frames = [load_train_data([frame, frame, frame],\n",
    "                                             gray_scale=self.n_channels==1, is_testing=True)]\n",
    "                in_frames = np.array(in_frames).astype(np.float32)\n",
    "\n",
    "                texture_enc, x = self.sess.run([self.texture_enc, self.shape_rep],\n",
    "                                               feed_dict={\n",
    "                                                   self.input_image:\n",
    "                                                      in_frames[:, :, :, :3],})\n",
    "                if x_state is None:\n",
    "                    x_state = np.zeros(x.shape + (len(vid_frames),),\n",
    "                                       dtype='float32')\n",
    "                x_state[:, :, :, :, i] = x\n",
    "\n",
    "            filter_fft = np.fft.fft(np.fft.ifftshift(filter_b),\n",
    "                                    n=x_state.shape[-1])\n",
    "            # Filtering\n",
    "            for i in trange(x_state.shape[1], desc=\"Applying FIR filter\"):\n",
    "                x_fft = np.fft.fft(x_state[:, i, :, :], axis=-1)\n",
    "                x_fft *= filter_fft[np.newaxis, np.newaxis, np.newaxis, :]\n",
    "                x_state[:, i, :, :] = np.fft.ifft(x_fft)\n",
    "\n",
    "            for i, frame in tqdm(enumerate(vid_frames), desc='Decoding'):\n",
    "                file_name = os.path.basename(frame)\n",
    "                frame_no, _ = os.path.splitext(file_name)\n",
    "                frame_no = int(frame_no)\n",
    "                in_frames = [load_train_data([frame, frame, frame],\n",
    "                                             gray_scale=self.n_channels==1, is_testing=True)]\n",
    "                in_frames = np.array(in_frames).astype(np.float32)\n",
    "                texture_enc, _ = self.sess.run([self.texture_enc, self.shape_rep],\n",
    "                                               feed_dict={\n",
    "                                                   self.input_image:\n",
    "                                                      in_frames[:, :, :, :3],\n",
    "                                                         })\n",
    "                out_amp = self.sess.run(self.output_image,\n",
    "                                        feed_dict={self.out_texture_enc: texture_enc,\n",
    "                                                   self.filtered_enc: x_state[:, :, :, :, i],\n",
    "                                                   self.ref_shape_enc: x,\n",
    "                                                   self.amplification_factor: [amplification_factor]})\n",
    "\n",
    "                im_path = os.path.join(out_dir, file_name)\n",
    "                out_amp = np.squeeze(out_amp)\n",
    "                out_amp = (127.5*(out_amp+1)).astype('uint8')\n",
    "                cv2.imwrite(im_path, cv2.cvtColor(out_amp,\n",
    "                                                  code=cv2.COLOR_RGB2BGR))\n",
    "            del x_state\n",
    "\n",
    "        # Try to combine it into a video\n",
    "        call([DEFAULT_VIDEO_CONVERTER, '-y', '-f', 'image2', '-r', '30', '-i',\n",
    "              os.path.join(out_dir, '%06d.png'), '-c:v', 'libx264',\n",
    "              os.path.join(out_dir, vid_name + '.mp4')]\n",
    "            )\n",
    "\n",
    "    # Training code.\n",
    "    def _build_training_graph(self, train_config):\n",
    "        self.global_step = tf.Variable(0, trainable=False)\n",
    "        filename_queue = tf.train.string_input_producer(\n",
    "                            [os.path.join(train_config[\"dataset_dir\"],\n",
    "                                          'train.tfrecords')],\n",
    "                            num_epochs=train_config[\"num_epochs\"])\n",
    "        frameA, frameB, frameC, frameAmp, amplification_factor = \\\n",
    "            read_and_decode_3frames(filename_queue,\n",
    "                                    (train_config[\"image_height\"],\n",
    "                                     train_config[\"image_width\"],\n",
    "                                     self.n_channels))\n",
    "        min_after_dequeue = 1000\n",
    "        num_threads = 16\n",
    "        capacity = min_after_dequeue + \\\n",
    "            (num_threads + 2) * train_config[\"batch_size\"]\n",
    "\n",
    "        frameA, frameB, frameC, frameAmp, amplification_factor = \\\n",
    "            tf.train.shuffle_batch([frameA,\n",
    "                                    frameB,\n",
    "                                    frameC,\n",
    "                                    frameAmp,\n",
    "                                    amplification_factor],\n",
    "                                   batch_size=train_config[\"batch_size\"],\n",
    "                                   capacity=capacity,\n",
    "                                   num_threads=num_threads,\n",
    "                                   min_after_dequeue=min_after_dequeue)\n",
    "\n",
    "        frameA = preprocess_image(frameA, train_config)\n",
    "        frameB = preprocess_image(frameB, train_config)\n",
    "        frameC = preprocess_image(frameC, train_config)\n",
    "        self.loss_function = partial(self._loss_function,\n",
    "                                     train_config=train_config)\n",
    "        self.output = self.image_transformer(frameA,\n",
    "                                             frameB,\n",
    "                                             amplification_factor,\n",
    "                                             [train_config[\"image_height\"],\n",
    "                                              train_config[\"image_width\"]],\n",
    "                                             self.arch_config, True, False)\n",
    "        self.reg_loss = tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)\n",
    "        if self.reg_loss and train_config[\"weight_decay\"] > 0.0:\n",
    "            print(\"Adding Regularization Weights.\")\n",
    "            self.loss = self.loss_function(self.output, frameAmp) + \\\n",
    "                train_config[\"weight_decay\"] * tf.add_n(self.reg_loss)\n",
    "        else:\n",
    "            print(\"No Regularization Weights.\")\n",
    "            self.loss = self.loss_function(self.output, frameAmp)\n",
    "        # Add regularization more\n",
    "        # TODO: Hardcoding the network name scope here.\n",
    "        with tf.variable_scope('ynet_3frames/encoder', reuse=True):\n",
    "            texture_c, shape_c = self._encoder(frameC)\n",
    "            self.loss = self.loss + \\\n",
    "                train_config[\"texture_loss_weight\"] * L1_loss(texture_c, self.texture_a) + \\\n",
    "                train_config[\"shape_loss_weight\"] * L1_loss(shape_c, self.shape_b)\n",
    "\n",
    "        self.loss_sum = tf.summary.scalar('train_loss', self.loss)\n",
    "        self.image_sum = tf.summary.image('train_B_OUT',\n",
    "                                          tf.concat([frameB, self.output],\n",
    "                                                    axis=2),\n",
    "                                          max_outputs=2)\n",
    "        if self.n_channels == 3:\n",
    "            self.image_comp_sum = tf.summary.image('train_GT_OUT',\n",
    "                                                   frameAmp - self.output,\n",
    "                                                   max_outputs=2)\n",
    "            self.image_orig_comp_sum = tf.summary.image('train_ORIG_OUT',\n",
    "                                                        frameA - self.output,\n",
    "                                                        max_outputs=2)\n",
    "        else:\n",
    "            self.image_comp_sum = tf.summary.image('train_GT_OUT',\n",
    "                                                   tf.concat([frameAmp,\n",
    "                                                              self.output,\n",
    "                                                              frameAmp],\n",
    "                                                             axis=3),\n",
    "                                                   max_outputs=2)\n",
    "            self.image_orig_comp_sum = tf.summary.image('train_ORIG_OUT',\n",
    "                                                        tf.concat([frameA,\n",
    "                                                                   self.output,\n",
    "                                                                   frameA],\n",
    "                                                                  axis=3),\n",
    "                                                        max_outputs=2)\n",
    "        self.saver = tf.train.Saver(max_to_keep=train_config[\"ckpt_to_keep\"])\n",
    "\n",
    "    # Loss function\n",
    "    def _loss_function(self, a, b, train_config):\n",
    "        # Use train_config to implement more advance losses.\n",
    "        with tf.variable_scope(\"loss_function\"):\n",
    "            return L1_loss(a, b) * train_config[\"l1_loss_weight\"]\n",
    "\n",
    "    def train(self, train_config):\n",
    "        # Define training graphs\n",
    "        self._build_training_graph(train_config)\n",
    "\n",
    "        self.lr = tf.train.exponential_decay(train_config[\"learning_rate\"],\n",
    "                                             self.global_step,\n",
    "                                             train_config[\"decay_steps\"],\n",
    "                                             train_config[\"lr_decay\"],\n",
    "                                             staircase=True)\n",
    "        self.optim_op = tf.train.AdamOptimizer(self.lr,\n",
    "                                               beta1=train_config[\"beta1\"]) \\\n",
    "            .minimize(self.loss,\n",
    "                      var_list=tf.trainable_variables(),\n",
    "                      global_step=self.global_step)\n",
    "\n",
    "        ginit_op = tf.global_variables_initializer()\n",
    "        linit_op = tf.local_variables_initializer()\n",
    "        self.sess.run([ginit_op, linit_op])\n",
    "\n",
    "        self.writer = tf.summary.FileWriter(train_config[\"logs_dir\"],\n",
    "                                            self.sess.graph)\n",
    "        coord = tf.train.Coordinator()\n",
    "        threads = tf.train.start_queue_runners(sess=self.sess, coord=coord)\n",
    "\n",
    "        start_time = time.time()\n",
    "        for v in tf.trainable_variables():\n",
    "            print(v)\n",
    "        if train_config[\"continue_train\"] and \\\n",
    "                self.load(train_config[\"checkpoint_dir\"]):\n",
    "            print('[*] Load Success')\n",
    "        elif train_config[\"restore_dir\"] and \\\n",
    "                self.load(train_config[\"restore_dir\"],\n",
    "                          tf.train.Saver(var_list=tf.trainable_variables())):\n",
    "            self.sess.run(self.global_step.assign(0))\n",
    "            print('[*] Restore success')\n",
    "        else:\n",
    "            print('Training from scratch.')\n",
    "        try:\n",
    "            while not coord.should_stop():\n",
    "                _, loss_sum_str = self.sess.run([self.optim_op, self.loss_sum])\n",
    "                global_step = self.sess.run(self.global_step)\n",
    "                self.writer.add_summary(loss_sum_str, global_step)\n",
    "\n",
    "                if global_step % 100 == 0:\n",
    "                    # Write image summary.\n",
    "                    img_sum_str, img_comp_str, img_orig_str = \\\n",
    "                            self.sess.run([self.image_sum,\n",
    "                                           self.image_comp_sum,\n",
    "                                           self.image_orig_comp_sum])\n",
    "                    self.writer.add_summary(img_sum_str, global_step)\n",
    "                    self.writer.add_summary(img_comp_str, global_step)\n",
    "                    self.writer.add_summary(img_orig_str, global_step)\n",
    "\n",
    "                elapsed_time = time.time() - start_time\n",
    "                print (\"Steps: %2d time: %4.4f (%4.4f steps/sec)\" % (\n",
    "                    global_step, elapsed_time,\n",
    "                    float(global_step) / elapsed_time))\n",
    "\n",
    "                if np.mod(global_step, train_config[\"save_freq\"]) == 2:\n",
    "                    self.save(train_config[\"checkpoint_dir\"], global_step)\n",
    "\n",
    "        except tf.errors.OutOfRangeError:\n",
    "            print('Done Training.')\n",
    "        finally:\n",
    "            coord.request_stop()\n",
    "            coord.join(threads)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "LTnNj_rb8XaO"
   },
   "outputs": [],
   "source": [
    "# To run the test script we are using the dynamic mode you can use any mode that is specified by the author.\n",
    "!sh run_on_test_videos.sh o3f_hmhm2_bg_qnoise_mix4_nl_n_t_ds3 baby 10 yes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "JIMJ7gqE8m23"
   },
   "outputs": [],
   "source": [
    "# The output video will be under the path /data/output/(folder name)/file name (folder name is baby in our case and file name is baby_259002.mp4 )\n",
    "# Navigating to the output Video to see the output\n",
    "cd /content/deep_motion_mag/data/output/baby"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "jBhpy_0H-Niu"
   },
   "outputs": [],
   "source": [
    "# Confirming that we are in the right folder and seeing the file name\n",
    "!ls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "cegzgiAs-u3a"
   },
   "outputs": [],
   "source": [
    "# Copy the video file to the drive so that you can manipulate them (View, Download, Share)!\n",
    "# Below command assumes that you have a folder named magnified_videos in drive, if you want to save to other location then change the destination folder.\n",
    "cp /content/deep_motion_mag/data/output/baby/baby_259002.mp4 /content/drive/My\\ Drive/magnified_videos/\n",
    "\n",
    "# Now you are done you can see your videos, into the destined folder."
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "name": "Video Magnification.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}