MCNet_UCF101.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Implementation of MCNet in Pytorch\n",
        "### Uses UCF-101 for trainig and prediction"
      ],
      "metadata": {
        "id": "vzQkYH9arWCk"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0jnFRJYm7BYO"
      },
      "outputs": [],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AVUWQtea7C8o"
      },
      "outputs": [],
      "source": [
        "import tensorflow.compat.v1 as tf\n",
        "tf.disable_v2_behavior()\n",
        "tf.compat.v1.disable_eager_execution()\n",
        "\n",
        "import cv2\n",
        "import numpy as np\n",
        "import os\n",
        "from os import listdir, makedirs, system\n",
        "from os.path import exists\n",
        "from argparse import ArgumentParser\n",
        "from joblib import Parallel, delayed\n",
        "import sys\n",
        "import time\n",
        "import imageio\n",
        "\n",
        "import scipy.io as sio\n",
        "import scipy.misc as sm\n",
        "\n",
        "import random\n",
        "\n",
        "! pip3 install imageio\n",
        "import imageio"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oFLMR0lCA6nt"
      },
      "source": [
        "# BasicConvLSTMCell\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pgTm_FOyAiLb"
      },
      "outputs": [],
      "source": [
        "class BasicConvLSTMCell(object):\n",
        "    \"\"\"Basic Conv LSTM recurrent network cell.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, shape, filter_size, num_features, forget_bias=1.0,\n",
        "                 input_size=None, state_is_tuple=False, activation=tf.nn.tanh):\n",
        "        \"\"\"Initialize the basic Conv LSTM cell.\n",
        "        Args:\n",
        "          shape: int tuple thats the height and width of the cell\n",
        "          filter_size: int tuple thats the height and width of the filter\n",
        "          num_features: int thats the depth of the cell \n",
        "          forget_bias: float, The bias added to forget gates (see above).\n",
        "          input_size: Deprecated and unused.\n",
        "          state_is_tuple: If True, accepted and returned states are 2-tuples of\n",
        "            the `c_state` and `m_state`.  If False, they are concatenated\n",
        "            along the column axis.  The latter behavior will soon be deprecated.\n",
        "          activation: Activation function of the inner states.\n",
        "        \"\"\"\n",
        "        if input_size is not None:\n",
        "            logging.warn(\"%s: The input_size parameter is deprecated.\", self)\n",
        "        self.shape = shape\n",
        "        self.filter_size = filter_size\n",
        "        self.num_features = num_features\n",
        "        self._forget_bias = forget_bias\n",
        "        self._state_is_tuple = state_is_tuple\n",
        "        self._activation = activation\n",
        "\n",
        "    @property\n",
        "    def state_size(self):\n",
        "        return (LSTMStateTuple(self._num_units, self._num_units)\n",
        "                if self._state_is_tuple else 2 * self._num_units)\n",
        "\n",
        "    @property\n",
        "    def output_size(self):\n",
        "        return self._num_units\n",
        "\n",
        "    def __call__(self, inputs, state, scope=None, reuse=False):\n",
        "        \"\"\"Long short-term memory cell (LSTM).\"\"\"\n",
        "        # \"BasicLSTMCell\"\n",
        "        with tf.variable_scope(scope or type(self).__name__, reuse=reuse):\n",
        "            # Parameters of gates are concatenated into one multiply for efficiency.\n",
        "            if self._state_is_tuple:\n",
        "                c, h = state\n",
        "            else:\n",
        "                c, h = tf.split(axis=3, num_or_size_splits=2, value=state)\n",
        "            concat = _conv_linear([inputs, h], self.filter_size,\n",
        "                                  self.num_features * 4, True)\n",
        "            # i = input_gate, j = new_input, f = forget_gate, o = output_gate\n",
        "            i, j, f, o = tf.split(axis=3, num_or_size_splits=4, value=concat)\n",
        "\n",
        "            new_c = (c * tf.nn.sigmoid(f + self._forget_bias) + tf.nn.sigmoid(i) *\n",
        "                     self._activation(j))\n",
        "            new_h = self._activation(new_c) * tf.nn.sigmoid(o)\n",
        "\n",
        "            if self._state_is_tuple:\n",
        "                new_state = LSTMStateTuple(new_c, new_h)\n",
        "            else:\n",
        "                new_state = tf.concat(axis=3, values=[new_c, new_h])\n",
        "            return new_h, new_state\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rDhOgoYyAwLh"
      },
      "outputs": [],
      "source": [
        "def _conv_linear(args, filter_size, num_features, bias,\n",
        "                 bias_start=0.0, scope=None, reuse=False):\n",
        "    \"\"\"convolution:\n",
        "    Args:\n",
        "      args: a 4D Tensor or a list of 4D, batch x n, Tensors.\n",
        "      filter_size: int tuple of filter height and width.\n",
        "      num_features: int, number of features.\n",
        "      bias_start: starting value to initialize the bias; 0 by default.\n",
        "      scope: VariableScope for the created subgraph; defaults to \"Linear\".\n",
        "      reuse: For reusing already existing weights\n",
        "    Returns:\n",
        "      A 4D Tensor with shape [batch h w num_features]\n",
        "    Raises:\n",
        "      ValueError: if some of the arguments has unspecified or wrong shape.\n",
        "    \"\"\"\n",
        "\n",
        "    # Calculate the total size of arguments on dimension 1.\n",
        "    total_arg_size_depth = 0\n",
        "    shapes = [a.get_shape().as_list() for a in args]\n",
        "    for shape in shapes:\n",
        "        if len(shape) != 4:\n",
        "            raise ValueError(\n",
        "                \"Linear is expecting 4D arguments: %s\" % str(shapes))\n",
        "        if not shape[3]:\n",
        "            raise ValueError(\n",
        "                \"Linear expects shape[4] of arguments: %s\" % str(shapes))\n",
        "        else:\n",
        "            total_arg_size_depth += shape[3]\n",
        "\n",
        "    dtype = [a.dtype for a in args][0]\n",
        "\n",
        "    # Now the computation.\n",
        "    with tf.variable_scope(scope or \"Conv\", reuse=reuse):\n",
        "        matrix = tf.get_variable(\n",
        "            \"Matrix\", [filter_size[0], filter_size[1],\n",
        "                       total_arg_size_depth, num_features], dtype=dtype)\n",
        "        if len(args) == 1:\n",
        "            res = tf.nn.conv2d(args[0], matrix, strides=[\n",
        "                               1, 1, 1, 1], padding='SAME')\n",
        "        else:\n",
        "            res = tf.nn.conv2d(tf.concat(axis=3, values=args), matrix,\n",
        "                               strides=[1, 1, 1, 1], padding='SAME')\n",
        "        if not bias:\n",
        "            return res\n",
        "        bias_term = tf.get_variable(\n",
        "            \"Bias\", [num_features],\n",
        "            dtype=dtype, initializer=tf.constant_initializer(bias_start,\n",
        "                                                             dtype=dtype)\n",
        "        )\n",
        "    return res + bias_term\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h2jWQuhrA5On"
      },
      "source": [
        "# Utils\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vDSupFgiBbA4"
      },
      "outputs": [],
      "source": [
        "def transform(image):\n",
        "    return image/127.5 - 1.\n",
        "\n",
        "\n",
        "def inverse_transform(images):\n",
        "    return (images+1.)/2.\n",
        "\n",
        "\n",
        "def save_images(images, size, image_path):\n",
        "    return imsave(inverse_transform(images)*255., size, image_path)\n",
        "\n",
        "\n",
        "def merge(images, size):\n",
        "    h, w = images.shape[1], images.shape[2]\n",
        "    img = np.zeros((h * size[0], w * size[1], 3))\n",
        "\n",
        "    for idx, image in enumerate(images):\n",
        "        i = idx % size[1]\n",
        "        j = idx // size[1]\n",
        "        img[j*h:j*h+h, i*w:i*w+w, :] = image\n",
        "\n",
        "    return img\n",
        "\n",
        "\n",
        "def imsave(images, size, path):\n",
        "    return imageio.imwrite(path, merge(images, size))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ooLd7c7IBgHp"
      },
      "outputs": [],
      "source": [
        "def get_minibatches_idx(n, minibatch_size, shuffle=False):\n",
        "    \"\"\" \n",
        "    Used to shuffle the dataset at each iteration.\n",
        "    \"\"\"\n",
        "\n",
        "    idx_list = np.arange(n, dtype=\"int32\")\n",
        "\n",
        "    if shuffle:\n",
        "        random.shuffle(idx_list)\n",
        "\n",
        "    minibatches = []\n",
        "    minibatch_start = 0\n",
        "    for i in range(n // minibatch_size):\n",
        "        minibatches.append(idx_list[minibatch_start:\n",
        "                                    minibatch_start + minibatch_size])\n",
        "        minibatch_start += minibatch_size\n",
        "\n",
        "    if (minibatch_start != n):\n",
        "        # Make a minibatch out of what is left\n",
        "        minibatches.append(idx_list[minibatch_start:])\n",
        "\n",
        "    return zip(range(len(minibatches)), minibatches)\n",
        "\n",
        "\n",
        "def draw_frame(img, is_input):\n",
        "    if img.shape[2] == 1:\n",
        "        img = np.repeat(img, [3], axis=2)\n",
        "\n",
        "    if is_input:\n",
        "        img[:2, :, 0] = img[:2, :, 2] = 0\n",
        "        img[:, :2, 0] = img[:, :2, 2] = 0\n",
        "        img[-2:, :, 0] = img[-2:, :, 2] = 0\n",
        "        img[:, -2:, 0] = img[:, -2:, 2] = 0\n",
        "        img[:2, :, 1] = 255\n",
        "        img[:, :2, 1] = 255\n",
        "        img[-2:, :, 1] = 255\n",
        "        img[:, -2:, 1] = 255\n",
        "    else:\n",
        "        img[:2, :, 0] = img[:2, :, 1] = 0\n",
        "        img[:, :2, 0] = img[:, :2, 2] = 0\n",
        "        img[-2:, :, 0] = img[-2:, :, 1] = 0\n",
        "        img[:, -2:, 0] = img[:, -2:, 1] = 0\n",
        "        img[:2, :, 2] = 255\n",
        "        img[:, :2, 2] = 255\n",
        "        img[-2:, :, 2] = 255\n",
        "        img[:, -2:, 2] = 255\n",
        "\n",
        "    return img\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oRwXvavFBk9x"
      },
      "outputs": [],
      "source": [
        "def load_kth_data(f_name, data_path, image_size, K, T):\n",
        "    flip = np.random.binomial(1, .5, 1)[0]\n",
        "    tokens = f_name.split()\n",
        "    vid_path = data_path + tokens[0] + \"_uncomp.avi\"\n",
        "    vid = imageio.get_reader(vid_path, \"ffmpeg\")\n",
        "    low = int(tokens[1])\n",
        "    high = np.min([int(tokens[2]), vid.get_length()])-K-T+1\n",
        "    if low == high:\n",
        "        stidx = 0\n",
        "    else:\n",
        "        if low >= high:\n",
        "            print(vid_path)\n",
        "        stidx = np.random.randint(low=low, high=high)\n",
        "    seq = np.zeros((image_size, image_size, K+T, 1), dtype=\"float32\")\n",
        "    for t in range(K+T):\n",
        "        img = cv2.cvtColor(cv2.resize(vid.get_data(stidx+t),\n",
        "                           (image_size, image_size)),\n",
        "                           cv2.COLOR_RGB2GRAY)\n",
        "        seq[:, :, t] = transform(img[:, :, None])\n",
        "\n",
        "    if flip == 1:\n",
        "        seq = seq[:, ::-1]\n",
        "\n",
        "    diff = np.zeros((image_size, image_size, K-1, 1), dtype=\"float32\")\n",
        "    for t in range(1, K):\n",
        "        prev = inverse_transform(seq[:, :, t-1])\n",
        "        next = inverse_transform(seq[:, :, t])\n",
        "        diff[:, :, t-1] = next.astype(\"float32\")-prev.astype(\"float32\")\n",
        "\n",
        "    return seq, diff\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SAZ7BbDAgKD4"
      },
      "outputs": [],
      "source": [
        "def load_ucf_data(f_name, data_path, trainlist, K, T):\n",
        "    flip = np.random.binomial(1, .5, 1)[0]\n",
        "    vid_path = data_path + f_name\n",
        "    img_size = [240, 320]\n",
        "\n",
        "    while True:\n",
        "        try:\n",
        "            vid = imageio.get_reader(vid_path, \"ffmpeg\")\n",
        "            low = 1\n",
        "            high = vid.get_length()-K-T+1\n",
        "            if low == high:\n",
        "                stidx = 0\n",
        "            else:\n",
        "                stidx = np.random.randint(low=low, high=high)\n",
        "            seq = np.zeros((img_size[0], img_size[1], K+T, 3),\n",
        "                           dtype=\"float32\")\n",
        "        #   print('#')\n",
        "            for t in range(K+T):\n",
        "                img = cv2.resize(vid.get_data(stidx+t),\n",
        "                                 (img_size[1], img_size[0]))[:, :, ::-1]\n",
        "                seq[:, :, t] = transform(img)\n",
        "\n",
        "            if flip == 1:\n",
        "                seq = seq[:, ::-1]\n",
        "\n",
        "            diff = np.zeros((img_size[0], img_size[1], K-1, 1),\n",
        "                            dtype=\"float32\")\n",
        "        #   print('+')\n",
        "            for t in range(1, K):\n",
        "                prev = inverse_transform(seq[:, :, t-1])*255\n",
        "                prev = cv2.cvtColor(prev.astype(\"uint8\"), cv2.COLOR_BGR2GRAY)\n",
        "                next = inverse_transform(seq[:, :, t])*255\n",
        "                next = cv2.cvtColor(next.astype(\"uint8\"), cv2.COLOR_BGR2GRAY)\n",
        "                diff[:, :, t-1, 0] = (next.astype(\"float32\") -\n",
        "                                      prev.astype(\"float32\"))/255.\n",
        "        #   print('*')\n",
        "            break\n",
        "        except Exception:\n",
        "            # In case the current video is bad load a random one\n",
        "            rep_idx = np.random.randint(low=0, high=len(trainlist))\n",
        "            f_name = trainlist[rep_idx]\n",
        "            vid_path = data_path + f_name\n",
        "            print('.', end='')\n",
        "    return seq, diff\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OEawOF5rBmtf"
      },
      "source": [
        "# Ops"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MaTX8sqwBynX"
      },
      "outputs": [],
      "source": [
        "def batch_norm(inputs, name, train=True, reuse=False):\n",
        "    return tf.layers.batch_normalization(inputs=inputs, training=train,\n",
        "                                         reuse=reuse, name=name, scale=True)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lYIEx9qmBrPr"
      },
      "outputs": [],
      "source": [
        "def conv2d(input_, output_dim,\n",
        "           k_h=5, k_w=5, d_h=2, d_w=2, stddev=0.02,\n",
        "           name=\"conv2d\", reuse=False, padding='SAME'):\n",
        "    with tf.variable_scope(name, reuse=reuse):\n",
        "        w = tf.get_variable('w', [k_h, k_w, input_.get_shape()[-1], output_dim],\n",
        "                            initializer=tf.initializers.glorot_uniform())  # tf.contrib.layers.xavier_initializer())\n",
        "        conv = tf.nn.conv2d(input_, w, strides=[\n",
        "                            1, d_h, d_w, 1], padding=padding)\n",
        "\n",
        "        biases = tf.get_variable('biases', [output_dim],\n",
        "                                 initializer=tf.constant_initializer(0.0))\n",
        "        conv = tf.reshape(tf.nn.bias_add(conv, biases), conv.get_shape())\n",
        "\n",
        "        return conv\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "z5Jz7ZMfBxdC"
      },
      "outputs": [],
      "source": [
        "def deconv2d(input_, output_shape,\n",
        "             k_h=5, k_w=5, d_h=2, d_w=2, stddev=0.02,\n",
        "             name=\"deconv2d\", reuse=False, with_w=False, padding='SAME'):\n",
        "    with tf.variable_scope(name, reuse=reuse):\n",
        "        # filter : [height, width, output_channels, in_channels]\n",
        "        w = tf.get_variable('w', [k_h, k_h, output_shape[-1],\n",
        "                                  input_.get_shape()[-1]],\n",
        "                            initializer=tf.initializers.glorot_uniform())\n",
        "        # print(d_h,d_w)\n",
        "        try:\n",
        "            deconv = tf.nn.conv2d_transpose(input_, w,\n",
        "                                            output_shape=output_shape,\n",
        "                                            strides=[1, d_h, d_w, 1],\n",
        "                                            padding=padding)\n",
        "\n",
        "        # Support for verisons of TensorFlow before 0.7.0\n",
        "        except AttributeError:\n",
        "            deconv = tf.nn.deconv2d(input_, w, output_shape=output_shape,\n",
        "                                    strides=[1, d_h, d_w, 1])\n",
        "\n",
        "        biases = tf.get_variable('biases', [output_shape[-1]],\n",
        "                                 initializer=tf.constant_initializer(0.0))\n",
        "        deconv = tf.reshape(tf.nn.bias_add(deconv, biases), deconv.get_shape())\n",
        "\n",
        "        if with_w:\n",
        "            return deconv, w, biases\n",
        "        else:\n",
        "            return deconv\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zYv9x7-lB2An"
      },
      "outputs": [],
      "source": [
        "def lrelu(x, leak=0.2, name=\"lrelu\"):\n",
        "    with tf.variable_scope(name):\n",
        "        f1 = 0.5 * (1 + leak)\n",
        "        f2 = 0.5 * (1 - leak)\n",
        "        return f1 * x + f2 * abs(x)\n",
        "\n",
        "\n",
        "def relu(x):\n",
        "    return tf.nn.relu(x)\n",
        "\n",
        "\n",
        "def tanh(x):\n",
        "    return tf.nn.tanh(x)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "V_Ni-ecmB26F"
      },
      "outputs": [],
      "source": [
        "def shape2d(a):\n",
        "    \"\"\"\n",
        "    a: a int or tuple/list of length 2\n",
        "    \"\"\"\n",
        "    if type(a) == int:\n",
        "        return [a, a]\n",
        "    if isinstance(a, (list, tuple)):\n",
        "        assert len(a) == 2\n",
        "        return list(a)\n",
        "    raise RuntimeError(\"Illegal shape: {}\".format(a))\n",
        "\n",
        "\n",
        "def shape4d(a):\n",
        "    # for use with tensorflow\n",
        "    return [1] + shape2d(a) + [1]\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PM2ZV5vRB8PD"
      },
      "outputs": [],
      "source": [
        "def UnPooling2x2ZeroFilled(x):\n",
        "    out = tf.concat(axis=3, values=[x, tf.zeros_like(x)])\n",
        "    out = tf.concat(axis=2, values=[out, tf.zeros_like(out)])\n",
        "\n",
        "    sh = x.get_shape().as_list()\n",
        "    if None not in sh[1:]:\n",
        "        out_size = [-1, sh[1] * 2, sh[2] * 2, sh[3]]\n",
        "        return tf.reshape(out, out_size)\n",
        "    else:\n",
        "        sh = tf.shape(x)\n",
        "        return tf.reshape(out, [-1, sh[1] * 2, sh[2] * 2, sh[3]])\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "X8f5MXV5B8wr"
      },
      "outputs": [],
      "source": [
        "def MaxPooling(x, shape, stride=None, padding='VALID'):\n",
        "    \"\"\"\n",
        "    MaxPooling on images.\n",
        "    :param input: NHWC tensor.\n",
        "    :param shape: int or [h, w]\n",
        "    :param stride: int or [h, w]. default to be shape.\n",
        "    :param padding: 'valid' or 'same'. default to 'valid'\n",
        "    :returns: NHWC tensor.\n",
        "    \"\"\"\n",
        "    padding = padding.upper()\n",
        "    shape = shape4d(shape)\n",
        "    if stride is None:\n",
        "        stride = shape\n",
        "    else:\n",
        "        stride = shape4d(stride)\n",
        "\n",
        "    return tf.nn.max_pool(x, ksize=shape, strides=stride, padding=padding)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fLniEKCQB-z1"
      },
      "outputs": [],
      "source": [
        "def FixedUnPooling(x, shape):\n",
        "    \"\"\"\n",
        "    Unpool the input with a fixed mat to perform kronecker product with.\n",
        "    :param input: NHWC tensor\n",
        "    :param shape: int or [h, w]\n",
        "    :returns: NHWC tensor\n",
        "    \"\"\"\n",
        "    shape = shape2d(shape)\n",
        "\n",
        "    # a faster implementation for this special case\n",
        "    return UnPooling2x2ZeroFilled(x)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "45wuS3dzCAVA"
      },
      "outputs": [],
      "source": [
        "def gdl(gen_frames, gt_frames, alpha):\n",
        "    \"\"\"\n",
        "    Calculates the sum of GDL losses between the predicted and gt frames.\n",
        "    @param gen_frames: The predicted frames at each scale.\n",
        "    @param gt_frames: The ground truth frames at each scale\n",
        "    @param alpha: The power to which each gradient term is raised.\n",
        "    @return: The GDL loss.\n",
        "    \"\"\"\n",
        "    # create filters [-1, 1] and [[1],[-1]]\n",
        "    # for diffing to the left and down respectively.\n",
        "    pos = tf.constant(np.identity(3), dtype=tf.float32)\n",
        "    neg = -1 * pos\n",
        "    # [-1, 1]\n",
        "    filter_x = tf.expand_dims(tf.stack([neg, pos]), 0)\n",
        "    # [[1],[-1]]\n",
        "    filter_y = tf.stack([tf.expand_dims(pos, 0), tf.expand_dims(neg, 0)])\n",
        "    strides = [1, 1, 1, 1]  # stride of (1, 1)\n",
        "    padding = 'SAME'\n",
        "\n",
        "    gen_dx = tf.abs(tf.nn.conv2d(\n",
        "        gen_frames, filter_x, strides, padding=padding))\n",
        "    gen_dy = tf.abs(tf.nn.conv2d(\n",
        "        gen_frames, filter_y, strides, padding=padding))\n",
        "    gt_dx = tf.abs(tf.nn.conv2d(gt_frames, filter_x, strides, padding=padding))\n",
        "    gt_dy = tf.abs(tf.nn.conv2d(gt_frames, filter_y, strides, padding=padding))\n",
        "\n",
        "    grad_diff_x = tf.abs(gt_dx - gen_dx)\n",
        "    grad_diff_y = tf.abs(gt_dy - gen_dy)\n",
        "\n",
        "    gdl_loss = tf.reduce_mean((grad_diff_x ** alpha + grad_diff_y ** alpha))\n",
        "\n",
        "    return gdl_loss"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "V4CuKIYTCCva"
      },
      "outputs": [],
      "source": [
        "def linear(input_, output_size, name, stddev=0.02, bias_start=0.0,\n",
        "           reuse=False, with_w=False):\n",
        "    shape = input_.get_shape().as_list()\n",
        "\n",
        "    with tf.variable_scope(name, reuse=reuse):\n",
        "        matrix = tf.get_variable(\"Matrix\", [shape[1], output_size], tf.float32,\n",
        "                                 tf.random_normal_initializer(stddev=stddev))\n",
        "        bias = tf.get_variable(\"bias\", [output_size],\n",
        "                               initializer=tf.constant_initializer(bias_start))\n",
        "        if with_w:\n",
        "            return tf.matmul(input_, matrix) + bias, matrix, bias\n",
        "        else:\n",
        "            return tf.matmul(input_, matrix) + bias"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c32GsEX0CHpn"
      },
      "source": [
        "# MCNet\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZxAiOkZSCEYt"
      },
      "outputs": [],
      "source": [
        "class MCNET(object):\n",
        "    def __init__(self, image_size, batch_size=32, c_dim=3,\n",
        "                 K=10, T=10, checkpoint_dir=None, is_train=True):\n",
        "\n",
        "        self.batch_size = batch_size\n",
        "        self.image_size = image_size\n",
        "        self.is_train = is_train\n",
        "\n",
        "        self.gf_dim = 64\n",
        "        self.df_dim = 64\n",
        "\n",
        "        self.c_dim = c_dim\n",
        "        self.K = K\n",
        "        self.T = T\n",
        "        self.diff_shape = [batch_size, self.image_size[0],\n",
        "                           self.image_size[1], K-1, 1]\n",
        "        self.xt_shape = [batch_size, self.image_size[0],\n",
        "                         self.image_size[1], c_dim]\n",
        "        self.target_shape = [batch_size, self.image_size[0], self.image_size[1],\n",
        "                             K+T, c_dim]\n",
        "\n",
        "        self.build_model()\n",
        "\n",
        "    def build_model(self):\n",
        "        self.diff_in = tf.placeholder(\n",
        "            tf.float32, self.diff_shape, name='diff_in')\n",
        "        self.xt = tf.placeholder(tf.float32, self.xt_shape, name='xt')\n",
        "        self.target = tf.placeholder(\n",
        "            tf.float32, self.target_shape, name='target')\n",
        "\n",
        "        cell = BasicConvLSTMCell([self.image_size[0]//8, self.image_size[1]//8],\n",
        "                                 [3, 3], 256)\n",
        "        pred = self.forward(self.diff_in, self.xt, cell)\n",
        "\n",
        "        self.G = tf.concat(axis=3, values=pred)\n",
        "        if self.is_train:\n",
        "            true_sim = inverse_transform(self.target[:, :, :, self.K:, :])\n",
        "            if self.c_dim == 1:\n",
        "                true_sim = tf.tile(true_sim, [1, 1, 1, 1, 3])\n",
        "            true_sim = tf.reshape(tf.transpose(true_sim, [0, 3, 1, 2, 4]),\n",
        "                                  [-1, self.image_size[0],\n",
        "                                   self.image_size[1], 3])\n",
        "            gen_sim = inverse_transform(self.G)\n",
        "            if self.c_dim == 1:\n",
        "                gen_sim = tf.tile(gen_sim, [1, 1, 1, 1, 3])\n",
        "            gen_sim = tf.reshape(tf.transpose(gen_sim, [0, 3, 1, 2, 4]),\n",
        "                                 [-1, self.image_size[0],\n",
        "                                  self.image_size[1], 3])\n",
        "            binput = tf.reshape(self.target[:, :, :, :self.K, :],\n",
        "                                [self.batch_size, self.image_size[0],\n",
        "                                 self.image_size[1], -1])\n",
        "            btarget = tf.reshape(self.target[:, :, :, self.K:, :],\n",
        "                                 [self.batch_size, self.image_size[0],\n",
        "                                  self.image_size[1], -1])\n",
        "            bgen = tf.reshape(self.G, [self.batch_size,\n",
        "                                       self.image_size[0],\n",
        "                                       self.image_size[1], -1])\n",
        "\n",
        "            good_data = tf.concat(axis=3, values=[binput, btarget])\n",
        "            gen_data = tf.concat(axis=3, values=[binput, bgen])\n",
        "\n",
        "            with tf.variable_scope(\"DIS\", reuse=False):\n",
        "                self.D, self.D_logits = self.discriminator(good_data)\n",
        "\n",
        "            with tf.variable_scope(\"DIS\", reuse=True):\n",
        "                self.D_, self.D_logits_ = self.discriminator(gen_data)\n",
        "\n",
        "            self.L_p = tf.reduce_mean(\n",
        "                tf.square(self.G-self.target[:, :, :, self.K:, :])\n",
        "            )\n",
        "            self.L_gdl = gdl(gen_sim, true_sim, 1.)\n",
        "            self.L_img = self.L_p + self.L_gdl\n",
        "\n",
        "            self.d_loss_real = tf.reduce_mean(\n",
        "                tf.nn.sigmoid_cross_entropy_with_logits(\n",
        "                    logits=self.D_logits, labels=tf.ones_like(self.D)\n",
        "                )\n",
        "            )\n",
        "            self.d_loss_fake = tf.reduce_mean(\n",
        "                tf.nn.sigmoid_cross_entropy_with_logits(\n",
        "                    logits=self.D_logits_, labels=tf.zeros_like(self.D_)\n",
        "                )\n",
        "            )\n",
        "            self.d_loss = self.d_loss_real + self.d_loss_fake\n",
        "            self.L_GAN = tf.reduce_mean(\n",
        "                tf.nn.sigmoid_cross_entropy_with_logits(\n",
        "                    logits=self.D_logits_, labels=tf.ones_like(self.D_)\n",
        "                )\n",
        "            )\n",
        "\n",
        "            self.loss_sum = tf.summary.scalar(\"L_img\", self.L_img)\n",
        "            self.L_p_sum = tf.summary.scalar(\"L_p\", self.L_p)\n",
        "            self.L_gdl_sum = tf.summary.scalar(\"L_gdl\", self.L_gdl)\n",
        "            self.L_GAN_sum = tf.summary.scalar(\"L_GAN\", self.L_GAN)\n",
        "            self.d_loss_sum = tf.summary.scalar(\"d_loss\", self.d_loss)\n",
        "            self.d_loss_real_sum = tf.summary.scalar(\n",
        "                \"d_loss_real\", self.d_loss_real)\n",
        "            self.d_loss_fake_sum = tf.summary.scalar(\n",
        "                \"d_loss_fake\", self.d_loss_fake)\n",
        "\n",
        "            self.t_vars = tf.trainable_variables()\n",
        "            self.g_vars = [var for var in self.t_vars if 'DIS' not in var.name]\n",
        "            self.d_vars = [var for var in self.t_vars if 'DIS' in var.name]\n",
        "            num_param = 0.0\n",
        "            for var in self.g_vars:\n",
        "                num_param += int(np.prod(var.get_shape()))\n",
        "            print(\"Number of parameters: %d\" % num_param)\n",
        "        self.saver = tf.train.Saver(max_to_keep=10)\n",
        "\n",
        "    def forward(self, diff_in, xt, cell):\n",
        "        # Initial state\n",
        "        state = tf.zeros([self.batch_size, self.image_size[0]//8,\n",
        "                          self.image_size[1]//8, 512])\n",
        "        reuse = False\n",
        "        # Encoder\n",
        "        for t in range(self.K-1):\n",
        "            enc_h, res_m = self.motion_enc(diff_in[:, :, :, t, :], reuse=reuse)\n",
        "            h_dyn, state = cell(enc_h, state, scope='lstm', reuse=reuse)\n",
        "            reuse = True\n",
        "\n",
        "        pred = []\n",
        "        # Decoder\n",
        "        for t in range(self.T):\n",
        "            if t == 0:\n",
        "                h_cont, res_c = self.content_enc(xt, reuse=False)\n",
        "                h_tp1 = self.comb_layers(h_dyn, h_cont, reuse=False)\n",
        "                res_connect = self.residual(res_m, res_c, reuse=False)\n",
        "                x_hat = self.dec_cnn(h_tp1, res_connect, reuse=False)\n",
        "            else:\n",
        "                enc_h, res_m = self.motion_enc(diff_in, reuse=True)\n",
        "                h_dyn, state = cell(enc_h, state, scope='lstm', reuse=True)\n",
        "                h_cont, res_c = self.content_enc(xt, reuse=reuse)\n",
        "                h_tp1 = self.comb_layers(h_dyn, h_cont, reuse=True)\n",
        "                res_connect = self.residual(res_m, res_c, reuse=True)\n",
        "                x_hat = self.dec_cnn(h_tp1, res_connect, reuse=True)\n",
        "\n",
        "            if self.c_dim == 3:\n",
        "                x_hat_rgb = tf.concat(axis=3,\n",
        "                                      values=[x_hat[:, :, :, 2:3], x_hat[:, :, :, 1:2],\n",
        "                                              x_hat[:, :, :, 0:1]])\n",
        "                xt_rgb = tf.concat(axis=3,\n",
        "                                   values=[xt[:, :, :, 2:3], xt[:, :, :, 1:2],\n",
        "                                           xt[:, :, :, 0:1]])\n",
        "\n",
        "                x_hat_gray = 1./255.*tf.image.rgb_to_grayscale(\n",
        "                    inverse_transform(x_hat_rgb)*255.\n",
        "                )\n",
        "                xt_gray = 1./255.*tf.image.rgb_to_grayscale(\n",
        "                    inverse_transform(xt_rgb)*255.\n",
        "                )\n",
        "            else:\n",
        "                x_hat_gray = inverse_transform(x_hat)\n",
        "                xt_gray = inverse_transform(xt)\n",
        "\n",
        "            diff_in = x_hat_gray - xt_gray\n",
        "            xt = x_hat\n",
        "            pred.append(tf.reshape(x_hat, [self.batch_size, self.image_size[0],\n",
        "                                           self.image_size[1], 1, self.c_dim]))\n",
        "\n",
        "        return pred\n",
        "\n",
        "    def motion_enc(self, diff_in, reuse):\n",
        "        res_in = []\n",
        "        conv1 = relu(conv2d(diff_in, output_dim=self.gf_dim, k_h=5, k_w=5,\n",
        "                            d_h=1, d_w=1, name='dyn_conv1',  reuse=reuse))\n",
        "        res_in.append(conv1)\n",
        "        pool1 = MaxPooling(conv1, [2, 2])\n",
        "\n",
        "        conv2 = relu(conv2d(pool1, output_dim=self.gf_dim*2, k_h=5, k_w=5,\n",
        "                            d_h=1, d_w=1, name='dyn_conv2', reuse=reuse))\n",
        "        res_in.append(conv2)\n",
        "        pool2 = MaxPooling(conv2, [2, 2])\n",
        "\n",
        "        conv3 = relu(conv2d(pool2, output_dim=self.gf_dim*4, k_h=7, k_w=7,\n",
        "                            d_h=1, d_w=1, name='dyn_conv3', reuse=reuse))\n",
        "        res_in.append(conv3)\n",
        "        pool3 = MaxPooling(conv3, [2, 2])\n",
        "        return pool3, res_in\n",
        "\n",
        "    def content_enc(self, xt, reuse):\n",
        "        res_in = []\n",
        "        conv1_1 = relu(conv2d(xt, output_dim=self.gf_dim, k_h=3, k_w=3,\n",
        "                              d_h=1, d_w=1, name='cont_conv1_1', reuse=reuse))\n",
        "        conv1_2 = relu(conv2d(conv1_1, output_dim=self.gf_dim, k_h=3, k_w=3,\n",
        "                              d_h=1, d_w=1, name='cont_conv1_2', reuse=reuse))\n",
        "        res_in.append(conv1_2)\n",
        "        pool1 = MaxPooling(conv1_2, [2, 2])\n",
        "\n",
        "        conv2_1 = relu(conv2d(pool1, output_dim=self.gf_dim*2, k_h=3, k_w=3,\n",
        "                              d_h=1, d_w=1, name='cont_conv2_1', reuse=reuse))\n",
        "        conv2_2 = relu(conv2d(conv2_1, output_dim=self.gf_dim*2, k_h=3, k_w=3,\n",
        "                              d_h=1, d_w=1, name='cont_conv2_2', reuse=reuse))\n",
        "        res_in.append(conv2_2)\n",
        "        pool2 = MaxPooling(conv2_2, [2, 2])\n",
        "\n",
        "        conv3_1 = relu(conv2d(pool2, output_dim=self.gf_dim*4, k_h=3, k_w=3,\n",
        "                              d_h=1, d_w=1, name='cont_conv3_1', reuse=reuse))\n",
        "        conv3_2 = relu(conv2d(conv3_1, output_dim=self.gf_dim*4, k_h=3, k_w=3,\n",
        "                              d_h=1, d_w=1, name='cont_conv3_2', reuse=reuse))\n",
        "        conv3_3 = relu(conv2d(conv3_2, output_dim=self.gf_dim*4, k_h=3, k_w=3,\n",
        "                              d_h=1, d_w=1, name='cont_conv3_3', reuse=reuse))\n",
        "        res_in.append(conv3_3)\n",
        "        pool3 = MaxPooling(conv3_3, [2, 2])\n",
        "        return pool3, res_in\n",
        "\n",
        "    def comb_layers(self, h_dyn, h_cont, reuse=False):\n",
        "        comb1 = relu(conv2d(tf.concat(axis=3, values=[h_dyn, h_cont]),\n",
        "                            output_dim=self.gf_dim*4, k_h=3, k_w=3,\n",
        "                            d_h=1, d_w=1, name='comb1', reuse=reuse))\n",
        "        comb2 = relu(conv2d(comb1, output_dim=self.gf_dim*2, k_h=3, k_w=3,\n",
        "                            d_h=1, d_w=1, name='comb2', reuse=reuse))\n",
        "        h_comb = relu(conv2d(comb2, output_dim=self.gf_dim*4, k_h=3, k_w=3,\n",
        "                             d_h=1, d_w=1, name='h_comb', reuse=reuse))\n",
        "        return h_comb\n",
        "\n",
        "    def residual(self, input_dyn, input_cont, reuse=False):\n",
        "        n_layers = len(input_dyn)\n",
        "        res_out = []\n",
        "        for l in range(n_layers):\n",
        "            input_ = tf.concat(axis=3, values=[input_dyn[l], input_cont[l]])\n",
        "            out_dim = input_cont[l].get_shape()[3]\n",
        "            res1 = relu(conv2d(input_, output_dim=out_dim,\n",
        "                               k_h=3, k_w=3, d_h=1, d_w=1,\n",
        "                               name='res'+str(l)+'_1', reuse=reuse))\n",
        "            res2 = conv2d(res1, output_dim=out_dim, k_h=3, k_w=3,\n",
        "                          d_h=1, d_w=1, name='res'+str(l)+'_2', reuse=reuse)\n",
        "            res_out.append(res2)\n",
        "        return res_out\n",
        "\n",
        "    def dec_cnn(self, h_comb, res_connect, reuse=False):\n",
        "        shapel3 = [self.batch_size, self.image_size[0]//4,\n",
        "                   self.image_size[1]//4, self.gf_dim*4]\n",
        "        shapeout3 = [self.batch_size, self.image_size[0]//4,\n",
        "                     self.image_size[1]//4, self.gf_dim*2]\n",
        "        depool3 = FixedUnPooling(h_comb, [2, 2])\n",
        "        deconv3_3 = relu(deconv2d(relu(tf.add(depool3, res_connect[2])),\n",
        "                                  output_shape=shapel3, k_h=3, k_w=3,\n",
        "                                  d_h=1, d_w=1, name='dec_deconv3_3', reuse=reuse))\n",
        "        deconv3_2 = relu(deconv2d(deconv3_3, output_shape=shapel3, k_h=3, k_w=3,\n",
        "                                  d_h=1, d_w=1, name='dec_deconv3_2', reuse=reuse))\n",
        "        deconv3_1 = relu(deconv2d(deconv3_2, output_shape=shapeout3, k_h=3, k_w=3,\n",
        "                                  d_h=1, d_w=1, name='dec_deconv3_1', reuse=reuse))\n",
        "\n",
        "        shapel2 = [self.batch_size, self.image_size[0]//2,\n",
        "                   self.image_size[1]//2, self.gf_dim*2]\n",
        "        shapeout3 = [self.batch_size, self.image_size[0]//2,\n",
        "                     self.image_size[1]//2, self.gf_dim]\n",
        "        depool2 = FixedUnPooling(deconv3_1, [2, 2])\n",
        "        deconv2_2 = relu(deconv2d(relu(tf.add(depool2, res_connect[1])),\n",
        "                                  output_shape=shapel2, k_h=3, k_w=3,\n",
        "                                  d_h=1, d_w=1, name='dec_deconv2_2', reuse=reuse))\n",
        "        deconv2_1 = relu(deconv2d(deconv2_2, output_shape=shapeout3, k_h=3, k_w=3,\n",
        "                                  d_h=1, d_w=1, name='dec_deconv2_1', reuse=reuse))\n",
        "\n",
        "        shapel1 = [self.batch_size, self.image_size[0],\n",
        "                   self.image_size[1], self.gf_dim]\n",
        "        shapeout1 = [self.batch_size, self.image_size[0],\n",
        "                     self.image_size[1], self.c_dim]\n",
        "        depool1 = FixedUnPooling(deconv2_1, [2, 2])\n",
        "        deconv1_2 = relu(deconv2d(relu(tf.add(depool1, res_connect[0])),\n",
        "                         output_shape=shapel1, k_h=3, k_w=3, d_h=1, d_w=1,\n",
        "                         name='dec_deconv1_2', reuse=reuse))\n",
        "        xtp1 = tanh(deconv2d(deconv1_2, output_shape=shapeout1, k_h=3, k_w=3,\n",
        "                             d_h=1, d_w=1, name='dec_deconv1_1', reuse=reuse))\n",
        "        return xtp1\n",
        "\n",
        "    def discriminator(self, image):\n",
        "        h0 = lrelu(conv2d(image, self.df_dim, name='dis_h0_conv'))\n",
        "        h1 = lrelu(batch_norm(conv2d(h0, self.df_dim*2, name='dis_h1_conv'),\n",
        "                              \"bn1\"))\n",
        "        h2 = lrelu(batch_norm(conv2d(h1, self.df_dim*4, name='dis_h2_conv'),\n",
        "                              \"bn2\"))\n",
        "        h3 = lrelu(batch_norm(conv2d(h2, self.df_dim*8, name='dis_h3_conv'),\n",
        "                              \"bn3\"))\n",
        "        h = linear(tf.reshape(h3, [self.batch_size, -1]), 1, 'dis_h3_lin')\n",
        "\n",
        "        return tf.nn.sigmoid(h), h\n",
        "\n",
        "    def save(self, sess, checkpoint_dir, step):\n",
        "        model_name = \"MCNET.model\"\n",
        "\n",
        "        if not os.path.exists(checkpoint_dir):\n",
        "            os.makedirs(checkpoint_dir)\n",
        "\n",
        "        self.saver.save(sess,\n",
        "                        os.path.join(checkpoint_dir, model_name),\n",
        "                        global_step=step)\n",
        "\n",
        "    def load(self, sess, checkpoint_dir, model_name=None):\n",
        "        print(\" [*] Reading checkpoints...\")\n",
        "\n",
        "        ckpt = tf.train.get_checkpoint_state(checkpoint_dir)\n",
        "        if ckpt and ckpt.model_checkpoint_path:\n",
        "            ckpt_name = os.path.basename(ckpt.model_checkpoint_path)\n",
        "            if model_name is None:\n",
        "                model_name = ckpt_name\n",
        "            self.saver.restore(sess, os.path.join(checkpoint_dir, model_name))\n",
        "            print(\"     Loaded model: \"+str(model_name))\n",
        "            return True, model_name\n",
        "        else:\n",
        "            return False, None\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iFmgMcNOm6Vi"
      },
      "source": [
        "# Main Function\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DqOSn70JCQYm"
      },
      "outputs": [],
      "source": [
        "def main(iters=0, lr=1e-4, batch_size=8, alpha=1.0, beta=0.001, K=4, T=1, num_iter=150000, gpu=0):\n",
        "\n",
        "    data_path = \"/content/drive/My Drive/VideoPrediction/mcnet/data/UCF101/UCF101_videos/\"\n",
        "    trainfiles = [f for f in listdir(data_path) if f.endswith(\".avi\")]\n",
        "    np.random.shuffle(trainfiles)\n",
        "    trainfiles = trainfiles[:1500]\n",
        "    margin = 0.3\n",
        "    updateD = True\n",
        "    updateG = True\n",
        "    image_size = [240, 320]\n",
        "    c_dim = 3\n",
        "    prefix = (\"UCF_MCNET\"\n",
        "              + \"_K=\"+str(K)\n",
        "              + \"_T=\"+str(T)\n",
        "              + \"_batch_size=\"+str(batch_size)\n",
        "              + \"_alpha=\"+str(alpha)\n",
        "              + \"_beta=\"+str(beta)\n",
        "              + \"_lr=\"+str(lr)+'test')\n",
        "\n",
        "    print(\"\\n\"+prefix+\"\\n\")\n",
        "    checkpoint_dir = \"models-ucf/\"+prefix+\"/\"\n",
        "    samples_dir = \"samples-ucf/\"+prefix+\"/\"\n",
        "    summary_dir = \"logs-ucf/\"+prefix+\"/\"\n",
        "\n",
        "    if not exists(checkpoint_dir):\n",
        "        makedirs(checkpoint_dir)\n",
        "    if not exists(samples_dir):\n",
        "        makedirs(samples_dir)\n",
        "    if not exists(summary_dir):\n",
        "        makedirs(summary_dir)\n",
        "\n",
        "    with tf.device(\"/gpu:%d\" % gpu[0]):\n",
        "        model = MCNET(image_size=image_size, c_dim=c_dim,\n",
        "                      K=K, batch_size=batch_size, T=T,\n",
        "                      checkpoint_dir=checkpoint_dir)\n",
        "        d_optim = tf.train.AdamOptimizer(lr, beta1=0.5).minimize(\n",
        "            model.d_loss, var_list=model.d_vars\n",
        "        )\n",
        "        g_optim = tf.train.AdamOptimizer(lr, beta1=0.5).minimize(\n",
        "            alpha*model.L_img+beta*model.L_GAN, var_list=model.g_vars\n",
        "        )\n",
        "\n",
        "    gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=1.0)\n",
        "    with tf.Session(config=tf.ConfigProto(allow_soft_placement=True,\n",
        "                    log_device_placement=False,\n",
        "                    gpu_options=gpu_options)) as sess:\n",
        "\n",
        "        tf.global_variables_initializer().run()\n",
        "\n",
        "        if model.load(sess, checkpoint_dir):\n",
        "            print(\" [*] Load SUCCESS\")\n",
        "        else:\n",
        "            print(\" [!] Load failed...\")\n",
        "            print(checkpoint_dir)\n",
        "\n",
        "        g_sum = tf.summary.merge([model.L_p_sum,\n",
        "                                  model.L_gdl_sum, model.loss_sum,\n",
        "                                  model.L_GAN_sum])\n",
        "        d_sum = tf.summary.merge([model.d_loss_real_sum, model.d_loss_sum,\n",
        "                                  model.d_loss_fake_sum])\n",
        "        writer = tf.summary.FileWriter(summary_dir, sess.graph)\n",
        "\n",
        "        counter = iters+1\n",
        "        start_time = time.time()\n",
        "\n",
        "        while iters < num_iter:\n",
        "            mini_batches = get_minibatches_idx(len(trainfiles),\n",
        "                                               batch_size, shuffle=True)\n",
        "            for _, batchidx in mini_batches:\n",
        "                for k in range(3):\n",
        "                    if len(batchidx) == batch_size:\n",
        "                        seq_batch = np.zeros((batch_size, image_size[0], image_size[1],\n",
        "                                              K+T, c_dim), dtype=\"float32\")\n",
        "                        diff_batch = np.zeros((batch_size, image_size[0], image_size[1],\n",
        "                                               K-1, 1), dtype=\"float32\")\n",
        "                        t0 = time.time()\n",
        "                        Ts = np.repeat(np.array([T]), batch_size, axis=0)\n",
        "                        Ks = np.repeat(np.array([K]), batch_size, axis=0)\n",
        "                        paths = np.repeat(data_path, batch_size, axis=0)\n",
        "                        tfiles = np.array(trainfiles)[batchidx]\n",
        "                        # print(len(tfiles))\n",
        "                        output = [load_ucf_data(f, p, trainfiles, k, t)\n",
        "                                  for f, p, k, t in zip(tfiles, paths, Ks, Ts)]\n",
        "                        # print(len(output))\n",
        "                    for i in range(batch_size):\n",
        "                        seq_batch[i] = output[i][0]\n",
        "                        diff_batch[i] = output[i][1]\n",
        "                    if updateD:\n",
        "                        _, summary_str = sess.run([d_optim, d_sum],\n",
        "                                                  feed_dict={model.diff_in: diff_batch,\n",
        "                                                             model.xt: seq_batch[:, :, :, K-1],\n",
        "                                                             model.target: seq_batch})\n",
        "                        writer.add_summary(summary_str, counter)\n",
        "\n",
        "                    if updateG:\n",
        "                        _, summary_str = sess.run([g_optim, g_sum],\n",
        "                                                  feed_dict={model.diff_in: diff_batch,\n",
        "                                                             model.xt: seq_batch[:, :, :, K-1],\n",
        "                                                             model.target: seq_batch})\n",
        "                        writer.add_summary(summary_str, counter)\n",
        "\n",
        "                    errD_fake = model.d_loss_fake.eval({model.diff_in: diff_batch,\n",
        "                                                        model.xt: seq_batch[:, :, :, K-1],\n",
        "                                                        model.target: seq_batch})\n",
        "                    errD_real = model.d_loss_real.eval({model.diff_in: diff_batch,\n",
        "                                                        model.xt: seq_batch[:, :, :, K-1],\n",
        "                                                        model.target: seq_batch})\n",
        "                    errG = model.L_GAN.eval({model.diff_in: diff_batch,\n",
        "                                             model.xt: seq_batch[:, :, :, K-1],\n",
        "                                             model.target: seq_batch})\n",
        "                    # print('.',end='')\n",
        "                    if errD_fake < margin or errD_real < margin:\n",
        "                        updateD = False\n",
        "                    if errD_fake > (1.-margin) or errD_real > (1.-margin):\n",
        "                        updateG = False\n",
        "                    if not updateD and not updateG:\n",
        "                        updateD = True\n",
        "                        updateG = True\n",
        "\n",
        "                    counter += 1\n",
        "                    if counter % 100 == 0:\n",
        "                        newlog = \"\\nCounter: [%2d] Iters: [%2d] time: %4.4f, d_loss: %.8f, L_GAN: %.8f\" % (\n",
        "                            counter, iters, time.time() - start_time, errD_fake+errD_real, errG)\n",
        "                        print(newlog)\n",
        "\n",
        "                    if np.mod(counter, 500) == 1:\n",
        "                        samples = sess.run([model.G],\n",
        "                                           feed_dict={model.diff_in: diff_batch,\n",
        "                                                      model.xt: seq_batch[:, :, :, K-1],\n",
        "                                                      model.target: seq_batch})[0]\n",
        "                        samples = np.concatenate((samples[:, :, :, 0, :],\n",
        "                                                  seq_batch[:, :, :, K, :]), axis=0)\n",
        "                        print(\"Saving sample ...\")\n",
        "                        save_images(samples[:, :, :, ::-1], [batch_size, batch_size],\n",
        "                                    samples_dir+\"train_%s.png\" % (iters))\n",
        "                    if np.mod(counter, 500) == 2:\n",
        "                        print(\"Saving Model to : \", checkpoint_dir)\n",
        "                        model.save(sess, checkpoint_dir, counter)\n",
        "\n",
        "                    iters += 1\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Thbynu7ie5Kr"
      },
      "outputs": [],
      "source": [
        "os.chdir('/content/drive/My Drive/VideoPrediction/mcnet/UCF101')\n",
        "%reload_ext tensorboard\n",
        "prefix =  'UCF_MCNET_K=10_T=3_batch_size=8_alpha=1.0_beta=0.001_lr=0.0001_test'\n",
        "summary_dir  = \"logs-ucf/\"+prefix+\"/\"\n",
        "%tensorboard --logdir {summary_dir}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "u4eHys_3mmmz"
      },
      "outputs": [],
      "source": [
        "lr=1e-4\n",
        "batch_size=8\n",
        "alpha=1.0\n",
        "beta=0.001\n",
        "K=10\n",
        "T=3\n",
        "num_iter=150000\n",
        "gpu=[0]\n",
        "iters=7200"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "background_save": true
        },
        "id": "bRPAAJFQzP4U"
      },
      "outputs": [],
      "source": [
        "tf.reset_default_graph()\n",
        "main(iters,lr,batch_size,alpha,beta,K,T,num_iter,gpu)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "88IffwYtMud1"
      },
      "outputs": [],
      "source": [
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XdzFO3tSPn2I"
      },
      "source": [
        "# Generate Videos"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7hk4BUJbRwS3",
        "outputId": "28c5b13d-a931-44f1-a1f6-55e262373a20"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Collecting pyssim\n",
            "  Downloading https://files.pythonhosted.org/packages/5f/03/65df3dde6843bcce9004e7a6a1a0657c5f84814840c4f671267c15cf1d34/pyssim-0.4.tar.gz\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from pyssim) (1.19.5)\n",
            "Requirement already satisfied: pillow in /usr/local/lib/python3.7/dist-packages (from pyssim) (7.1.2)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.7/dist-packages (from pyssim) (1.4.1)\n",
            "Building wheels for collected packages: pyssim\n",
            "  Building wheel for pyssim (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for pyssim: filename=pyssim-0.4-py2.py3-none-any.whl size=5814 sha256=9d7bd1ac240be658a85bc735f4d624137214143c3ecf1cee3c8034857af5617c\n",
            "  Stored in directory: /root/.cache/pip/wheels/cf/7c/7b/2ebeff601772f28bfd22e128a91e85bb5df4e36d33df59a26c\n",
            "Successfully built pyssim\n",
            "Installing collected packages: pyssim\n",
            "Successfully installed pyssim-0.4\n"
          ]
        }
      ],
      "source": [
        "! pip install pyssim\n",
        "\n",
        "from PIL import Image\n",
        "from PIL import ImageDraw\n",
        "import ssim\n",
        "import skimage.measure as measure"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2h8Pnvt1R5rg"
      },
      "outputs": [],
      "source": [
        "def generate_samples():\n",
        "    data_path = \"/content/drive/My Drive/VideoPrediction/mcnet/data/UCF101/UCF101_videos/\"\n",
        "    ran = np.random.randint(low=1, high=5000, size=10)\n",
        "    testfiles = [x for x in os.listdir(\n",
        "        '/content/drive/My Drive/VideoPrediction/mcnet/data/UCF101/UCF101_videos/') if 'UCF_' in x]\n",
        "    testfiles = np.random.choice(testfiles, size=10)\n",
        "    image_size = [240, 320]\n",
        "    c_dim = 3\n",
        "    iters = 0\n",
        "    prefix_mod = (\"UCF_MCNET\"\n",
        "                  + \"_K=\"+str(K)+'-'+str(K_test)\n",
        "                  + \"_T=\"+str(T)+'-'+str(T_test)\n",
        "                  + \"_batch_size=\"+str(batch_size)\n",
        "                  + \"_alpha=\"+str(alpha)\n",
        "                  + \"_beta=\"+str(beta)\n",
        "                  + \"_lr=\"+str(lr)\n",
        "                  + 'iters='+str(iters))\n",
        "\n",
        "    prefix = (\"UCF_MCNET\"\n",
        "              + \"_K=\"+str(K)\n",
        "              + \"_T=\"+str(T)\n",
        "              + \"_batch_size=\"+str(batch_size)\n",
        "              + \"_alpha=\"+str(alpha)\n",
        "              + \"_beta=\"+str(beta)\n",
        "              + \"_lr=\"+str(lr))\n",
        "    os.chdir('/content/drive/MyDrive/VideoPrediction/mcnet/UCF101')\n",
        "    checkpoint_dir = \"models-ucf/\"+prefix+\"/\"\n",
        "    best_model = None  # will pick last model\n",
        "\n",
        "    with tf.device(\"/gpu:%d\" % gpu[0]):\n",
        "        model = MCNET(image_size=image_size, batch_size=1, K=K,\n",
        "                      T=T, c_dim=c_dim, checkpoint_dir=checkpoint_dir,\n",
        "                      is_train=False)\n",
        "\n",
        "    gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.5)\n",
        "    with tf.Session(config=tf.ConfigProto(allow_soft_placement=True,\n",
        "                                          log_device_placement=False,\n",
        "                                          gpu_options=gpu_options)) as sess:\n",
        "\n",
        "        tf.global_variables_initializer().run()\n",
        "\n",
        "        loaded, model_name = model.load(sess, checkpoint_dir, best_model)\n",
        "        print(checkpoint_dir)\n",
        "        if loaded:\n",
        "            print(\" [*] Load SUCCESS\")\n",
        "        else:\n",
        "            print(\" [!] Load failed... exitting\")\n",
        "            return\n",
        "\n",
        "        quant_dir = \"results/quantitative/UCF101/\"+prefix_mod+\"/\"\n",
        "        save_path = quant_dir+\"results_model=\"+model_name+\".npz\"\n",
        "        if not exists(quant_dir):\n",
        "            makedirs(quant_dir)\n",
        "\n",
        "        vid_names = []\n",
        "        psnr_err = np.zeros((0, T_test))\n",
        "        ssim_err = np.zeros((0, T_test))\n",
        "        for i in range(0, len(testfiles)):\n",
        "            print(\" Video \"+str(i)+\"/\"+str(len(testfiles)))\n",
        "\n",
        "            vid_path = data_path + testfiles[i]\n",
        "            vid = imageio.get_reader(vid_path, \"ffmpeg\")\n",
        "            savedir = \"results/images/UCF101/\"+prefix_mod+\"/\"+str(i+1)\n",
        "\n",
        "            seq_batch = np.zeros((1, image_size[0], image_size[1], K_test+T_test, c_dim),\n",
        "                                 dtype=\"float32\")\n",
        "            diff_batch = np.zeros((1, image_size[0], image_size[1], K_test-1, 1),\n",
        "                                  dtype=\"float32\")\n",
        "            for t in range(K_test+T_test):\n",
        "                img = vid.get_data(t)[:, :, ::-1]\n",
        "                seq_batch[0, :, :, t] = transform(img)\n",
        "\n",
        "            for t in range(1, K_test):\n",
        "                prev = inverse_transform(seq_batch[0, :, :, t-1])*255\n",
        "                prev = cv2.cvtColor(prev.astype(\"uint8\"), cv2.COLOR_BGR2GRAY)\n",
        "                next = inverse_transform(seq_batch[0, :, :, t])*255\n",
        "                next = cv2.cvtColor(next.astype(\"uint8\"), cv2.COLOR_BGR2GRAY)\n",
        "                diff = next.astype(\"float32\")-prev.astype(\"float32\")\n",
        "                diff_batch[0, :, :, t-1] = diff[:, :, None]/255.\n",
        "\n",
        "            true_data = seq_batch[:, :, :, K_test:, :].copy()\n",
        "            pred_data = np.zeros(true_data.shape, dtype=\"float32\")\n",
        "            xt = seq_batch[:, :, :, K_test-1]\n",
        "            pred_data[0] = sess.run(model.G,\n",
        "                                    feed_dict={model.diff_in: diff_batch,\n",
        "                                               model.xt: xt})\n",
        "\n",
        "            if not os.path.exists(savedir):\n",
        "                os.makedirs(savedir)\n",
        "\n",
        "            cpsnr = np.zeros((K_test+T_test,))\n",
        "            cssim = np.zeros((K_test+T_test,))\n",
        "            pred_data = np.concatenate(\n",
        "                (seq_batch[:, :, :, :K_test], pred_data), axis=3)\n",
        "            true_data = np.concatenate(\n",
        "                (seq_batch[:, :, :, :K_test], true_data), axis=3)\n",
        "            for t in range(K_test+T_test):\n",
        "                pred = (inverse_transform(\n",
        "                    pred_data[0, :, :, t])*255).astype(\"uint8\")\n",
        "                target = (inverse_transform(\n",
        "                    true_data[0, :, :, t])*255).astype(\"uint8\")\n",
        "\n",
        "                cpsnr[t] = measure.compare_psnr(pred, target)\n",
        "                cssim[t] = ssim.compute_ssim(Image.fromarray(target),\n",
        "                                             Image.fromarray(pred))\n",
        "\n",
        "                pred = draw_frame(pred, t < K_test)\n",
        "                target = draw_frame(target, t < K_test)\n",
        "\n",
        "                cv2.imwrite(savedir+\"/pred_\"+\"{0:04d}\".format(t)+\".png\", pred)\n",
        "                cv2.imwrite(savedir+\"/gt_\"+\"{0:04d}\".format(t)+\".png\", target)\n",
        "\n",
        "            cmd1 = \"rm \"+savedir+\"/pred.gif\"\n",
        "            cmd2 = (\"ffmpeg -f image2 -framerate 3 -i \"+savedir +\n",
        "                    \"/pred_%04d.png \"+savedir+\"-pred.gif\")\n",
        "            cmd3 = \"rm \"+savedir+\"/pred*.png\"\n",
        "\n",
        "            system(cmd1)\n",
        "            system(cmd2)\n",
        "            system(cmd3)\n",
        "\n",
        "            cmd1 = \"rm \"+savedir+\"/gt.gif\"\n",
        "            cmd2 = (\"ffmpeg -f image2 -framerate 3 -i \"+savedir +\n",
        "                    \"/gt_%04d.png \"+savedir+\"-gt.gif\")\n",
        "            cmd3 = \"rm \"+savedir+\"/gt*.png\"\n",
        "            system(cmd1)\n",
        "            system(cmd2)\n",
        "            system(cmd3)\n",
        "\n",
        "            psnr_err = np.concatenate((psnr_err, cpsnr[None, K_test:]), axis=0)\n",
        "            ssim_err = np.concatenate((ssim_err, cssim[None, K_test:]), axis=0)\n",
        "\n",
        "            print(np.array(psnr_err).shape, np.array(\n",
        "                ssim_err).shape, np.array(pred).shape)\n",
        "\n",
        "        np.savez(save_path, psnr=psnr_err, ssim=ssim_err)\n",
        "        print(\"Results saved to \"+save_path)\n",
        "    print(\"Done.\")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zqoEdAM7TEbm"
      },
      "outputs": [],
      "source": [
        "K_test=10\n",
        "T_test=6\n",
        "tf.reset_default_graph()\n",
        "generate_samples()"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "MCNet-UCF101.ipynb",
      "toc_visible": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}