i3d_speedchallenge.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "i3d_speedchallenge_reg_github.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "6VPnXvsGPeIH",
        "85mvNrD1nifc",
        "FwehewBthc7A",
        "kA9Qaaz5QUiN",
        "Kzcre83NPSp_",
        "kEBU5QtBoaPl",
        "K29GfZr4oenm",
        "1fCWCTD87lUY",
        "PYu3bsQ37uUI"
      ]
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6VPnXvsGPeIH"
      },
      "source": [
        "# I3D implementation based on\n",
        "\n",
        "https://arxiv.org/pdf/1705.07750.pdf\n",
        "\n",
        "https://github.com/piergiaj/pytorch-i3d\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yN6oZaldRwI2"
      },
      "source": [
        "# from google.colab import drive\n",
        "# drive.mount('/content/drive')\n",
        "\n",
        "# !nvidia-smi"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "85mvNrD1nifc"
      },
      "source": [
        "#Parameters"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IGpXR5kIjeZQ"
      },
      "source": [
        "class Parameters():\n",
        "  \"\"\"Part of the parameters, including:\n",
        "  - directories to folders\n",
        "  - path for the model save\n",
        "  - sequence length\n",
        "  \"\"\"\n",
        "  def __init__(self):\n",
        "    self.data_dir = 'drive/My Drive/Drive/speedchallenge/' # all the data are stored on the google drive\n",
        "\n",
        "    # Train\n",
        "    self.image_dir = self.data_dir + 'data/KITTI_COMMA_resized_10fps/' \n",
        "    self.mask_dir = self.data_dir + 'data/KITTI_COMMA_resized_masks_10fps/'  \n",
        "    self.vel_dir = self.data_dir + 'data/KITTI_COMMA_vel_GT_10fps/'\n",
        "\n",
        "    # Test\n",
        "    # self.image_dir = self.data_dir + 'data/COMMA_TEST_cropped_resized_10fps/'\n",
        "    # self.mask_dir = self.data_dir + 'data/COMMA_TEST_cropped_resized_10fps_masks/'\n",
        "\n",
        "    # Kitti folders\n",
        "    self.kitti_folder_list = [str(item) for item in list(range(1000, 1164))]\n",
        "    self.kitti_folder_list.pop(142)\n",
        "    # Comma folders\n",
        "    self.comma_folder_list = [str(item) for item in list(range(100, 185))]\n",
        "    # My data set folders\n",
        "    self.my_folders = [str(item) for item in list(range(34))]\n",
        "\n",
        "    self.video_folders = self.kitti_folder_list + self.comma_folder_list + self.my_folders\n",
        "\n",
        "    self.save_model_path = 'drive/My Drive/Drive/speedchallenge/i3d_checkpoints/comma_kitti__myData(with_car_removal)/'\n",
        "\n",
        "    self.seq_len = 10\n",
        "    self.fps = 10\n",
        "\n",
        "par = Parameters()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FwehewBthc7A"
      },
      "source": [
        "# Dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "G6NRhbh1hPDx"
      },
      "source": [
        "import torch\n",
        "import torch.utils.data as data_utl\n",
        "from torch.utils.data.dataloader import default_collate\n",
        "\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "from torch.optim import lr_scheduler\n",
        "from torch.autograd import Variable\n",
        "\n",
        "import torchvision\n",
        "from torchvision import datasets, transforms\n",
        "\n",
        "import albumentations as albu\n",
        "import numpy as np\n",
        "import json\n",
        "import csv\n",
        "import h5py\n",
        "import random\n",
        "import os\n",
        "import os.path\n",
        "import cv2\n",
        "\n",
        "def video_to_tensor(pic):\n",
        "    \"\"\"Convert a ``numpy.ndarray`` to tensor.\n",
        "    Converts a numpy.ndarray (T x H x W x C)\n",
        "    to a torch.FloatTensor of shape (C x T x H x W)\n",
        "    \n",
        "    Args:\n",
        "         pic (numpy.ndarray): Video to be converted to tensor.\n",
        "    Returns:\n",
        "         Tensor: Converted video.\n",
        "    \"\"\"\n",
        "    return torch.from_numpy(pic.transpose([3, 0, 1, 2]))\n",
        "\n",
        "def load_rgb_frames(vid, start, num):\n",
        "    frames = []\n",
        "    for i in range(start, start+num):\n",
        "        img = cv2.imread(os.path.join(par.image_dir, vid, str(i).zfill(10) + '.png'))\n",
        "        if img is None:\n",
        "            print(\"Can't load image \" + str(i) + \" from \" + vid + \", please check the img path\", os.path.join(par.image_dir, vid, str(i).zfill(10)+'.png'))\n",
        "            sys.exit(1)\n",
        "        img=img[:, :, [2, 1, 0]]\n",
        "\n",
        "        frames.append(img)\n",
        "    return np.asarray(frames)\n",
        "\n",
        "\n",
        "def make_dataset(vid_folder_list): \n",
        "    dataset = []\n",
        "    for vid in vid_folder_list: # go through all training or testing folders\n",
        "        \n",
        "        if vid in par.comma_folder_list:\n",
        "            num_frames = 120 # don't know why but os.listdir sometimes gives wrong file list from google drive, better hardcode this is possible\n",
        "        else:\n",
        "            num_frames = len(os.listdir(os.path.join(par.image_dir, vid)))\n",
        "\n",
        "        labels = np.load('{}{}.npy'.format(par.vel_dir, vid))\n",
        "        duration = num_frames/(par.fps) # seconds, do not used\n",
        "        dataset.append((vid, labels, duration, num_frames))\n",
        "    return dataset\n",
        "\n",
        "\n",
        "class Charades(data_utl.Dataset):\n",
        "\n",
        "    def __init__(self, vid_folder_list, mode, transforms=None, augmentation=False, remove_cars=False):\n",
        "        self.data = make_dataset(vid_folder_list)\n",
        "        self.vid_folder_list = vid_folder_list\n",
        "        self.transforms = transforms\n",
        "        self.augmentation = augmentation\n",
        "        self.mode = mode\n",
        "        self.remove_cars = remove_cars\n",
        "\n",
        "        self.transformations = albu.Compose(\n",
        "            [\n",
        "              albu.HorizontalFlip(p=0.5),\n",
        "              albu.OneOf([\n",
        "                  albu.IAAAdditiveGaussianNoise(),\n",
        "                  albu.GaussNoise(),\n",
        "              ], p=0.5),\n",
        "              albu.OneOf([\n",
        "                  albu.CLAHE(clip_limit=2),\n",
        "                  albu.IAASharpen(),\n",
        "                  albu.IAAEmboss(),\n",
        "                  albu.RandomGamma(p=1),\n",
        "                  albu.RGBShift(),\n",
        "                  albu.HueSaturationValue(p=0.3),\n",
        "              ], p=0.8),\n",
        "              albu.ToGray(),\n",
        "              albu.RandomBrightness(),\n",
        "              albu.RandomContrast(),\n",
        "            ], p=1\n",
        "        )\n",
        "\n",
        "    def __getitem__(self, index):\n",
        "        \"\"\"\n",
        "        Args:\n",
        "            index (int): Index\n",
        "\n",
        "        Returns:\n",
        "            tuple: (image, target)\n",
        "        \"\"\"\n",
        "\n",
        "        vid, label, dur, nf = self.data[index]\n",
        "        # print('\\nvid fold', vid)\n",
        "\n",
        "        if self.mode == 'train': # 24-.. frames\n",
        "            first_frame = 24\n",
        "        elif self.mode == 'test': # first 24 frames\n",
        "            first_frame = 1\n",
        "            nf = 24\n",
        "        \n",
        "        start_f = random.randint(first_frame, nf - (par.seq_len + 1)) # 1 .. 1990, 200 .. 1990\n",
        "        imgs = load_rgb_frames(vid, start_f, par.seq_len)\n",
        "\n",
        "        # segment cars out\n",
        "        if self.remove_cars:\n",
        "            mask_for_seq = np.zeros((224,224))\n",
        "            for i in range(start_f, start_f + par.seq_len):\n",
        "                mask = np.load(os.path.join(par.mask_dir, vid, str(i).zfill(10)+'.npy'))\n",
        "                mask_for_seq += mask # accumulate the masks\n",
        "            mask_for_seq[mask_for_seq > 1] = 1 # clip siquence mask\n",
        "\n",
        "            for i,image in enumerate(imgs):\n",
        "                imgs[i][mask_for_seq == 1] = 0\n",
        "\n",
        "        if self.augmentation:\n",
        "            additional_targets_names = {}\n",
        "            for i in range(par.seq_len-1):\n",
        "                    additional_targets_names['image' + str(i)] = 'image'\n",
        "\n",
        "            transform = albu.Compose(\n",
        "                [\n",
        "                  self.transformations\n",
        "                ],\n",
        "                additional_targets=additional_targets_names\n",
        "            )\n",
        "\n",
        "            input_dict={'image': imgs[0]}\n",
        "            for i, additional_targets_name in enumerate(additional_targets_names):\n",
        "                input_dict[additional_targets_name] = imgs[i + 1]\n",
        "\n",
        "            transformed = transform(**input_dict)\n",
        "            aug_list = []\n",
        "            for key in transformed:\n",
        "                aug_list.append(transformed[key])\n",
        "            imgs = np.asarray(aug_list)\n",
        "\n",
        "        imgs = np.array([((img/255.)*2 - 1) for img in imgs]) # Normalize\n",
        "        imgs = np.asarray(imgs, dtype=np.float32) # Convert ot float32\n",
        "\n",
        "        # print('wanted items    :', start_f, start_f+par.seq_len)\n",
        "        label = label[start_f : start_f + par.seq_len]\n",
        "\n",
        "        return video_to_tensor(imgs), torch.from_numpy(label)\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.data)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9jJTg0Nrr47Y"
      },
      "source": [
        "# i3d = InceptionI3d(400, in_channels=3)\n",
        "# i3d.eval()\n",
        "# i3d(torch.rand(1,3,64,224,224))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kA9Qaaz5QUiN"
      },
      "source": [
        "# Model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RA72C6dsQOeq"
      },
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from torch.autograd import Variable\n",
        "\n",
        "import numpy as np\n",
        "\n",
        "import os\n",
        "import sys\n",
        "from collections import OrderedDict\n",
        "\n",
        "\n",
        "class MaxPool3dSamePadding(nn.MaxPool3d):\n",
        "    \n",
        "    def compute_pad(self, dim, s):\n",
        "        if s % self.stride[dim] == 0:\n",
        "            return max(self.kernel_size[dim] - self.stride[dim], 0)\n",
        "        else:\n",
        "            return max(self.kernel_size[dim] - (s % self.stride[dim]), 0)\n",
        "\n",
        "    def forward(self, x):\n",
        "        # compute 'same' padding\n",
        "        (batch, channel, t, h, w) = x.size()\n",
        "        #print t,h,w\n",
        "        out_t = np.ceil(float(t) / float(self.stride[0]))\n",
        "        out_h = np.ceil(float(h) / float(self.stride[1]))\n",
        "        out_w = np.ceil(float(w) / float(self.stride[2]))\n",
        "        #print out_t, out_h, out_w\n",
        "        pad_t = self.compute_pad(0, t)\n",
        "        pad_h = self.compute_pad(1, h)\n",
        "        pad_w = self.compute_pad(2, w)\n",
        "        #print pad_t, pad_h, pad_w\n",
        "\n",
        "        pad_t_f = pad_t // 2\n",
        "        pad_t_b = pad_t - pad_t_f\n",
        "        pad_h_f = pad_h // 2\n",
        "        pad_h_b = pad_h - pad_h_f\n",
        "        pad_w_f = pad_w // 2\n",
        "        pad_w_b = pad_w - pad_w_f\n",
        "\n",
        "        pad = (pad_w_f, pad_w_b, pad_h_f, pad_h_b, pad_t_f, pad_t_b)\n",
        "        #print x.size()\n",
        "        #print pad\n",
        "        x = F.pad(x, pad)\n",
        "        return super(MaxPool3dSamePadding, self).forward(x)\n",
        "    \n",
        "\n",
        "class Unit3D(nn.Module):\n",
        "\n",
        "    def __init__(self, in_channels,\n",
        "                 output_channels,\n",
        "                 kernel_shape=(1, 1, 1),\n",
        "                 stride=(1, 1, 1),\n",
        "                 padding=0,\n",
        "                 activation_fn=F.relu,\n",
        "                 use_batch_norm=True,\n",
        "                 use_bias=False,\n",
        "                 name='unit_3d'):\n",
        "        \n",
        "        \"\"\"Initializes Unit3D module.\"\"\"\n",
        "        super(Unit3D, self).__init__()\n",
        "        \n",
        "        self._output_channels = output_channels\n",
        "        self._kernel_shape = kernel_shape\n",
        "        self._stride = stride\n",
        "        self._use_batch_norm = use_batch_norm\n",
        "        self._activation_fn = activation_fn\n",
        "        self._use_bias = use_bias\n",
        "        self.name = name\n",
        "        self.padding = padding\n",
        "        \n",
        "        self.conv3d = nn.Conv3d(in_channels=in_channels,\n",
        "                                out_channels=self._output_channels,\n",
        "                                kernel_size=self._kernel_shape,\n",
        "                                stride=self._stride,\n",
        "                                padding=0, # we always want padding to be 0 here. We will dynamically pad based on input size in forward function\n",
        "                                bias=self._use_bias)\n",
        "        \n",
        "        if self._use_batch_norm:\n",
        "            self.bn = nn.BatchNorm3d(self._output_channels, eps=0.001, momentum=0.01)\n",
        "\n",
        "    def compute_pad(self, dim, s):\n",
        "        if s % self._stride[dim] == 0:\n",
        "            return max(self._kernel_shape[dim] - self._stride[dim], 0)\n",
        "        else:\n",
        "            return max(self._kernel_shape[dim] - (s % self._stride[dim]), 0)\n",
        "\n",
        "            \n",
        "    def forward(self, x):\n",
        "        # compute 'same' padding\n",
        "        (batch, channel, t, h, w) = x.size()\n",
        "        #print t,h,w\n",
        "        out_t = np.ceil(float(t) / float(self._stride[0]))\n",
        "        out_h = np.ceil(float(h) / float(self._stride[1]))\n",
        "        out_w = np.ceil(float(w) / float(self._stride[2]))\n",
        "        #print out_t, out_h, out_w\n",
        "        pad_t = self.compute_pad(0, t)\n",
        "        pad_h = self.compute_pad(1, h)\n",
        "        pad_w = self.compute_pad(2, w)\n",
        "        #print pad_t, pad_h, pad_w\n",
        "\n",
        "        pad_t_f = pad_t // 2\n",
        "        pad_t_b = pad_t - pad_t_f\n",
        "        pad_h_f = pad_h // 2\n",
        "        pad_h_b = pad_h - pad_h_f\n",
        "        pad_w_f = pad_w // 2\n",
        "        pad_w_b = pad_w - pad_w_f\n",
        "\n",
        "        pad = (pad_w_f, pad_w_b, pad_h_f, pad_h_b, pad_t_f, pad_t_b)\n",
        "        #print x.size()\n",
        "        #print pad\n",
        "        x = F.pad(x, pad)\n",
        "        #print x.size()        \n",
        "\n",
        "        x = self.conv3d(x)\n",
        "        if self._use_batch_norm:\n",
        "            x = self.bn(x)\n",
        "        if self._activation_fn is not None:\n",
        "            x = self._activation_fn(x)\n",
        "        return x\n",
        "\n",
        "\n",
        "\n",
        "class InceptionModule(nn.Module):\n",
        "    def __init__(self, in_channels, out_channels, name):\n",
        "        super(InceptionModule, self).__init__()\n",
        "\n",
        "        self.b0 = Unit3D(in_channels=in_channels, output_channels=out_channels[0], kernel_shape=[1, 1, 1], padding=0,\n",
        "                         name=name+'/Branch_0/Conv3d_0a_1x1')\n",
        "        self.b1a = Unit3D(in_channels=in_channels, output_channels=out_channels[1], kernel_shape=[1, 1, 1], padding=0,\n",
        "                          name=name+'/Branch_1/Conv3d_0a_1x1')\n",
        "        self.b1b = Unit3D(in_channels=out_channels[1], output_channels=out_channels[2], kernel_shape=[3, 3, 3],\n",
        "                          name=name+'/Branch_1/Conv3d_0b_3x3')\n",
        "        self.b2a = Unit3D(in_channels=in_channels, output_channels=out_channels[3], kernel_shape=[1, 1, 1], padding=0,\n",
        "                          name=name+'/Branch_2/Conv3d_0a_1x1')\n",
        "        self.b2b = Unit3D(in_channels=out_channels[3], output_channels=out_channels[4], kernel_shape=[3, 3, 3],\n",
        "                          name=name+'/Branch_2/Conv3d_0b_3x3')\n",
        "        self.b3a = MaxPool3dSamePadding(kernel_size=[3, 3, 3],\n",
        "                                stride=(1, 1, 1), padding=0)\n",
        "        self.b3b = Unit3D(in_channels=in_channels, output_channels=out_channels[5], kernel_shape=[1, 1, 1], padding=0,\n",
        "                          name=name+'/Branch_3/Conv3d_0b_1x1')\n",
        "        self.name = name\n",
        "\n",
        "    def forward(self, x):    \n",
        "        b0 = self.b0(x)\n",
        "        b1 = self.b1b(self.b1a(x))\n",
        "        b2 = self.b2b(self.b2a(x))\n",
        "        b3 = self.b3b(self.b3a(x))\n",
        "        return torch.cat([b0,b1,b2,b3], dim=1)\n",
        "\n",
        "\n",
        "class InceptionI3d(nn.Module):\n",
        "    \"\"\"Inception-v1 I3D architecture.\n",
        "    The model is introduced in:\n",
        "        Quo Vadis, Action Recognition? A New Model and the Kinetics Dataset\n",
        "        Joao Carreira, Andrew Zisserman\n",
        "        https://arxiv.org/pdf/1705.07750v1.pdf.\n",
        "    See also the Inception architecture, introduced in:\n",
        "        Going deeper with convolutions\n",
        "        Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet, Scott Reed,\n",
        "        Dragomir Anguelov, Dumitru Erhan, Vincent Vanhoucke, Andrew Rabinovich.\n",
        "        http://arxiv.org/pdf/1409.4842v1.pdf.\n",
        "    \"\"\"\n",
        "\n",
        "    # Endpoints of the model in order. During construction, all the endpoints up\n",
        "    # to a designated `final_endpoint` are returned in a dictionary as the\n",
        "    # second return value.\n",
        "    VALID_ENDPOINTS = (\n",
        "        'Conv3d_1a_7x7',\n",
        "        'MaxPool3d_2a_3x3',\n",
        "        'Conv3d_2b_1x1',\n",
        "        'Conv3d_2c_3x3',\n",
        "        'MaxPool3d_3a_3x3',\n",
        "        'Mixed_3b',\n",
        "        'Mixed_3c',\n",
        "        'MaxPool3d_4a_3x3',\n",
        "        'Mixed_4b',\n",
        "        'Mixed_4c',\n",
        "        'Mixed_4d',\n",
        "        'Mixed_4e',\n",
        "        'Mixed_4f',\n",
        "        'MaxPool3d_5a_2x2',\n",
        "        'Mixed_5b',\n",
        "        'Mixed_5c',\n",
        "        'Logits',\n",
        "        'Predictions',\n",
        "    )\n",
        "\n",
        "    def __init__(self, num_classes=400, spatial_squeeze=True,\n",
        "                 final_endpoint='Logits', name='inception_i3d', in_channels=3, dropout_keep_prob=0.5):\n",
        "        \"\"\"Initializes I3D model instance.\n",
        "        Args:\n",
        "          num_classes: The number of outputs in the logit layer (default 400, which\n",
        "              matches the Kinetics dataset).\n",
        "          spatial_squeeze: Whether to squeeze the spatial dimensions for the logits\n",
        "              before returning (default True).\n",
        "          final_endpoint: The model contains many possible endpoints.\n",
        "              `final_endpoint` specifies the last endpoint for the model to be built\n",
        "              up to. In addition to the output at `final_endpoint`, all the outputs\n",
        "              at endpoints up to `final_endpoint` will also be returned, in a\n",
        "              dictionary. `final_endpoint` must be one of\n",
        "              InceptionI3d.VALID_ENDPOINTS (default 'Logits').\n",
        "          name: A string (optional). The name of this module.\n",
        "        Raises:\n",
        "          ValueError: if `final_endpoint` is not recognized.\n",
        "        \"\"\"\n",
        "\n",
        "        if final_endpoint not in self.VALID_ENDPOINTS:\n",
        "            raise ValueError('Unknown final endpoint %s' % final_endpoint)\n",
        "\n",
        "        super(InceptionI3d, self).__init__()\n",
        "        self._num_classes = num_classes\n",
        "        self._spatial_squeeze = spatial_squeeze\n",
        "        self._final_endpoint = final_endpoint\n",
        "        self.logits = None\n",
        "\n",
        "        if self._final_endpoint not in self.VALID_ENDPOINTS:\n",
        "            raise ValueError('Unknown final endpoint %s' % self._final_endpoint)\n",
        "\n",
        "        self.end_points = {}\n",
        "        end_point = 'Conv3d_1a_7x7'\n",
        "        self.end_points[end_point] = Unit3D(in_channels=in_channels, output_channels=64, kernel_shape=[7, 7, 7],\n",
        "                                            stride=(2, 2, 2), padding=(3,3,3),  name=name+end_point)\n",
        "        if self._final_endpoint == end_point: return\n",
        "        \n",
        "        end_point = 'MaxPool3d_2a_3x3'\n",
        "        self.end_points[end_point] = MaxPool3dSamePadding(kernel_size=[1, 3, 3], stride=(1, 2, 2),\n",
        "                                                             padding=0)\n",
        "        if self._final_endpoint == end_point: return\n",
        "        \n",
        "        end_point = 'Conv3d_2b_1x1'\n",
        "        self.end_points[end_point] = Unit3D(in_channels=64, output_channels=64, kernel_shape=[1, 1, 1], padding=0,\n",
        "                                       name=name+end_point)\n",
        "        if self._final_endpoint == end_point: return\n",
        "        \n",
        "        end_point = 'Conv3d_2c_3x3'\n",
        "        self.end_points[end_point] = Unit3D(in_channels=64, output_channels=192, kernel_shape=[3, 3, 3], padding=1,\n",
        "                                       name=name+end_point)\n",
        "        if self._final_endpoint == end_point: return\n",
        "\n",
        "        end_point = 'MaxPool3d_3a_3x3'\n",
        "        self.end_points[end_point] = MaxPool3dSamePadding(kernel_size=[1, 3, 3], stride=(1, 2, 2),\n",
        "                                                             padding=0)\n",
        "        if self._final_endpoint == end_point: return\n",
        "        \n",
        "        end_point = 'Mixed_3b'\n",
        "        self.end_points[end_point] = InceptionModule(192, [64,96,128,16,32,32], name+end_point)\n",
        "        if self._final_endpoint == end_point: return\n",
        "\n",
        "        end_point = 'Mixed_3c'\n",
        "        self.end_points[end_point] = InceptionModule(256, [128,128,192,32,96,64], name+end_point)\n",
        "        if self._final_endpoint == end_point: return\n",
        "\n",
        "        end_point = 'MaxPool3d_4a_3x3'\n",
        "        self.end_points[end_point] = MaxPool3dSamePadding(kernel_size=[3, 3, 3], stride=(2, 2, 2),\n",
        "                                                             padding=0)\n",
        "        if self._final_endpoint == end_point: return\n",
        "\n",
        "        end_point = 'Mixed_4b'\n",
        "        self.end_points[end_point] = InceptionModule(128+192+96+64, [192,96,208,16,48,64], name+end_point)\n",
        "        if self._final_endpoint == end_point: return\n",
        "\n",
        "        end_point = 'Mixed_4c'\n",
        "        self.end_points[end_point] = InceptionModule(192+208+48+64, [160,112,224,24,64,64], name+end_point)\n",
        "        if self._final_endpoint == end_point: return\n",
        "\n",
        "        end_point = 'Mixed_4d'\n",
        "        self.end_points[end_point] = InceptionModule(160+224+64+64, [128,128,256,24,64,64], name+end_point)\n",
        "        if self._final_endpoint == end_point: return\n",
        "\n",
        "        end_point = 'Mixed_4e'\n",
        "        self.end_points[end_point] = InceptionModule(128+256+64+64, [112,144,288,32,64,64], name+end_point)\n",
        "        if self._final_endpoint == end_point: return\n",
        "\n",
        "        end_point = 'Mixed_4f'\n",
        "        self.end_points[end_point] = InceptionModule(112+288+64+64, [256,160,320,32,128,128], name+end_point)\n",
        "        if self._final_endpoint == end_point: return\n",
        "\n",
        "        end_point = 'MaxPool3d_5a_2x2'\n",
        "        self.end_points[end_point] = MaxPool3dSamePadding(kernel_size=[2, 2, 2], stride=(2, 2, 2),\n",
        "                                                             padding=0)\n",
        "        if self._final_endpoint == end_point: return\n",
        "\n",
        "        end_point = 'Mixed_5b'\n",
        "        self.end_points[end_point] = InceptionModule(256+320+128+128, [256,160,320,32,128,128], name+end_point)\n",
        "        if self._final_endpoint == end_point: return\n",
        "\n",
        "        end_point = 'Mixed_5c'\n",
        "        self.end_points[end_point] = InceptionModule(256+320+128+128, [384,192,384,48,128,128], name+end_point)\n",
        "        if self._final_endpoint == end_point: return\n",
        "\n",
        "        end_point = 'Logits'\n",
        "        self.avg_pool = nn.AvgPool3d(kernel_size=[1, 7, 7], #[2, 7, 7]\n",
        "                                     stride=(1, 1, 1))\n",
        "        self.dropout = nn.Dropout(dropout_keep_prob)\n",
        "        self.logits = Unit3D(in_channels=384+384+128+128, output_channels=self._num_classes,\n",
        "                             kernel_shape=[1, 1, 1],\n",
        "                             padding=0,\n",
        "                             activation_fn=None,\n",
        "                             use_batch_norm=False,\n",
        "                             use_bias=True,\n",
        "                             name='logits')\n",
        "\n",
        "        self.build()\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "    def replace_logits(self, num_classes):\n",
        "        self._num_classes = num_classes\n",
        "        self.logits = Unit3D(in_channels=384+384+128+128, output_channels=self._num_classes,\n",
        "                             kernel_shape=[1, 1, 1],\n",
        "                             padding=0,\n",
        "                             activation_fn=None,\n",
        "                             use_batch_norm=False,\n",
        "                             use_bias=True,\n",
        "                             name='logits')\n",
        "   \n",
        "\n",
        "\n",
        "    def build(self):\n",
        "        for k in self.end_points.keys():\n",
        "            self.add_module(k, self.end_points[k])\n",
        "        \n",
        "    def forward(self, x):\n",
        "        for end_point in self.VALID_ENDPOINTS:\n",
        "            if end_point in self.end_points:\n",
        "                x = self._modules[end_point](x) # use _modules to work with dataparallel\n",
        "\n",
        "        x = self.logits(self.dropout(self.avg_pool(x)))\n",
        "        if self._spatial_squeeze:\n",
        "            logits = x.squeeze(3).squeeze(3)\n",
        "        # logits is batch X time X classes, which is what we want to work with\n",
        "\n",
        "        # add sigmoid and multiplied by the highes velocity\n",
        "        sigm = nn.Sigmoid()\n",
        "        logits = sigm(logits) * 40.0\n",
        "\n",
        "        return logits\n",
        "        \n",
        "\n",
        "    def extract_features(self, x):\n",
        "        for end_point in self.VALID_ENDPOINTS:\n",
        "            if end_point in self.end_points:\n",
        "                x = self._modules[end_point](x)\n",
        "        return self.avg_pool(x)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hhCu7rky6V0s"
      },
      "source": [
        "# i3d = InceptionI3d(num_classes=400, in_channels=3)\n",
        "# i3d.replace_logits(1)\n",
        "# i3d.eval()\n",
        "\n",
        "# input = torch.rand(2, 3, 20, 224, 224)\n",
        "\n",
        "# output = i3d(input)\n",
        "\n",
        "# t = input.size(2)\n",
        "# output_upsampled = F.upsample(output, t, mode='linear')\n",
        "\n",
        "# output = output.squeeze()\n",
        "# output.shape, output_upsampled.shape"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Kzcre83NPSp_"
      },
      "source": [
        "# Train"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kEBU5QtBoaPl"
      },
      "source": [
        "###main train"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dszF2J0DPT6v"
      },
      "source": [
        "import os\n",
        "os.environ[\"CUDA_DEVICE_ORDER\"]=\"PCI_BUS_ID\"\n",
        "import sys\n",
        "import argparse\n",
        "\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "from torch.optim import lr_scheduler\n",
        "from torch.autograd import Variable\n",
        "\n",
        "import torchvision\n",
        "from torchvision import datasets, transforms\n",
        "\n",
        "from IPython.display import clear_output\n",
        "\n",
        "import numpy as np\n",
        "\n",
        "# some params\n",
        "init_lr = 6.0e-06\n",
        "max_steps = 800\n",
        "batch_size = 30 \n",
        "num_steps_per_update = 2\n",
        "\n",
        "resume_train = False\n",
        "if resume_train:\n",
        "    checkpoint_name = '000700_val_loss_1.25_lr_6e-06'\n",
        "\n",
        "# datalaoders\n",
        "dataset = Charades(par.video_folders, 'train', augmentation=True, remove_cars=True)\n",
        "dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=30, pin_memory=True)\n",
        "\n",
        "val_dataset = Charades(par.video_folders, 'test', remove_cars=True)\n",
        "val_dataloader = torch.utils.data.DataLoader(val_dataset, batch_size=batch_size, shuffle=True, num_workers=30, pin_memory=True)    \n",
        "\n",
        "dataloaders = {'train': dataloader, 'val': val_dataloader}\n",
        "datasets = {'train': dataset, 'val': val_dataset}\n",
        "\n",
        "# model\n",
        "i3d = InceptionI3d(400, in_channels=3)\n",
        "i3d.load_state_dict(torch.load(par.data_dir + 'i3d_models/rgb_imagenet.pt'))\n",
        "i3d.replace_logits(1)\n",
        "if resume_train:\n",
        "    i3d.load_state_dict(torch.load(par.save_model_path + checkpoint_name + '.pt'))\n",
        "i3d.cuda()\n",
        "i3d = nn.DataParallel(i3d)\n",
        "\n",
        "# optimizers\n",
        "# optimizer = optim.SGD(i3d.parameters(), lr=init_lr, momentum=0.9, weight_decay=0.0000001)\n",
        "optimizer = torch.optim.Adam(i3d.parameters(), lr=init_lr)\n",
        "if resume_train:\n",
        "    optimizer.load_state_dict(torch.load(par.save_model_path + checkpoint_name + '.optimizer'))\n",
        "lr_sched = optim.lr_scheduler.MultiStepLR(optimizer, [400, 700])\n",
        "def get_lr(optimizer):\n",
        "  for param_group in optimizer.param_groups:\n",
        "    return param_group['lr']\n",
        "\n",
        "steps = 0\n",
        "if resume_train:\n",
        "    steps = int(checkpoint_name[:6]) + 1\n",
        "\n",
        "train_loss={}\n",
        "train_loss['loss'] = []\n",
        "val_loss={}\n",
        "val_loss['loss'] = []\n",
        "\n",
        "while steps < max_steps:\n",
        "\n",
        "    for phase in ['train', 'val']:\n",
        "        print('phase', phase)\n",
        "\n",
        "        if phase == 'train':\n",
        "            i3d.train(True)\n",
        "        else:\n",
        "            i3d.train(False)  # Set model to evaluate mode\n",
        "\n",
        "        loss_mean = 0\n",
        "        num_iter = 0\n",
        "        optimizer.zero_grad()\n",
        "        \n",
        "        # Iterate over data.\n",
        "        for data in dataloaders[phase]:\n",
        "            num_iter += 1\n",
        "            # get the inputs\n",
        "            inputs, labels = data\n",
        "\n",
        "            # wrap them in Variable\n",
        "            inputs = Variable(inputs.cuda())\n",
        "            t = inputs.size(2)\n",
        "            labels = Variable(labels.cuda())\n",
        "\n",
        "            per_frame_logits = i3d(inputs)\n",
        "            # upsample to input size\n",
        "            per_frame_logits = F.upsample(per_frame_logits, t, mode='linear')\n",
        "\n",
        "            loss = torch.nn.functional.mse_loss(per_frame_logits.squeeze(), labels.float())\n",
        "            loss_mean += loss\n",
        "            loss.backward()\n",
        "\n",
        "            if phase == 'train' and num_iter == num_steps_per_update:\n",
        "                print('make an optimizer.step()')\n",
        "                steps += 1\n",
        "                num_iter = 0\n",
        "                optimizer.step()\n",
        "                optimizer.zero_grad()\n",
        "                lr_sched.step()\n",
        "\n",
        "        loss_mean /= len(dataloaders[phase])\n",
        "        loss_mean_cpu = loss_mean.data.cpu().numpy()\n",
        "\n",
        "        if phase == 'train':\n",
        "            train_loss['loss'].append(loss_mean_cpu)\n",
        "        if phase == 'val':\n",
        "            if steps % 10 == 0: # save the model every 10 steps\n",
        "                torch.save(i3d.module.state_dict(), par.save_model_path+str(steps).zfill(6) + f'_val_loss_{loss_mean_cpu:2f}_lr_{get_lr(optimizer)}.pt')\n",
        "                torch.save(optimizer.state_dict(), par.save_model_path+str(steps).zfill(6) + f'_val_loss_{loss_mean_cpu:2f}_lr_{get_lr(optimizer)}.optimizer')\n",
        "            val_loss['loss'].append(loss_mean_cpu)\n",
        "\n",
        "    clear_output(True)\n",
        "    print ('Step {}/{}'.format(steps, max_steps))\n",
        "    print('LR:', get_lr(optimizer))\n",
        "    print ('-' * 10)\n",
        "\n",
        "    plt.subplot(2, 1, 1) # plt.subplot(3, 1, 1)\n",
        "    plt.title('Training loss')\n",
        "    plt.plot(train_loss['loss'], 'o', label='train loss')\n",
        "    # plt.legend(loc='lower left')\n",
        "    plt.xlabel('Steps')\n",
        "    plt.gcf().set_size_inches(15, 12)\n",
        "    plt.grid(True)\n",
        "\n",
        "    plt.subplot(2, 1, 2) # plt.subplot(3, 1, 1)\n",
        "    plt.title('Validation loss')\n",
        "    # plt.plot(train_loss['loss'], 'o', label='train loss')\n",
        "    plt.plot(val_loss['loss'], 'o', label='val loss')\n",
        "    # plt.legend(loc='lower left')\n",
        "    plt.xlabel('Steps')\n",
        "    plt.gcf().set_size_inches(15, 12)\n",
        "    plt.grid(True)\n",
        "\n",
        "    if len(val_loss['loss']) > 10:\n",
        "        print('mean val loss over last 10 steps: {:.2f}'.format(sum(val_loss['loss'][-10:]) / 10.))\n",
        "\n",
        "    plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "K29GfZr4oenm"
      },
      "source": [
        "###cross validation train"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eg8xDhYXdyiH"
      },
      "source": [
        "import os\n",
        "os.environ[\"CUDA_DEVICE_ORDER\"]=\"PCI_BUS_ID\"\n",
        "import sys\n",
        "import argparse\n",
        "\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "from torch.optim import lr_scheduler\n",
        "from torch.autograd import Variable\n",
        "\n",
        "import torchvision\n",
        "from torchvision import datasets, transforms\n",
        "\n",
        "from IPython.display import clear_output\n",
        "\n",
        "import numpy as np\n",
        "import itertools\n",
        "import time\n",
        "\n",
        "\n",
        "# params to tune:\n",
        "seq_lens = [5, 10, 20, 30]\n",
        "init_lrs = [0.006, 0.0010, 0.0002]\n",
        "optimizers = ['SGD', 'Adam']\n",
        "batch_sizes = [20, 30, 40, 50]\n",
        "imagenet_usage = [True, False]\n",
        "num_steps_per_update_l = [4, 3, 2]\n",
        "\n",
        "number_of_steps_to_try = 30\n",
        "\n",
        "video_folders = par.comma_folder_list # [str(item) for item in list(range(100,185))]\n",
        "\n",
        "results = {}\n",
        "best_val = 1000\n",
        "\n",
        "for (seq_len, optimizer, init_lr, batch_size, init_with_imagenet) in list(itertools.product(seq_lens, optimizers, init_lrs, batch_sizes, imagenet_usage)):\n",
        "# for (optimizer, init_lr, batch_size) in list(itertools.product(optimizers, init_lrs, batch_sizes)):\n",
        "\n",
        "    # Fix some of the params if needed\n",
        "    # optimizer = 'Adam'\n",
        "    # init_lr = 6.0e-04\n",
        "    # # seq_len = 20\n",
        "    # init_with_imagenet = True\n",
        "    # batch_size = 26 #30 #25\n",
        "    num_steps_per_update = 2 # accumulate gradient\n",
        "\n",
        "    par.seq_len = seq_len\n",
        "\n",
        "    print('\\nseq_len %e optimizer %e init_lr %e batch_size %e init_with_imagenet: %f' % (\n",
        "                seq_len, optimizer, init_lr, batch_size, init_with_imagenet))\n",
        "    # print('\\noptimizer %s init_lr %e batch_size %e ' % (\n",
        "    #             optimizer, init_lr, batch_size))\n",
        "\n",
        "    dataset = Charades(video_folders, 'train') #, augmentation=True, remove_cars=True\n",
        "    dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=30, pin_memory=True) #,augmentation=True\n",
        "\n",
        "    val_dataset = Charades(video_folders, 'test') # , remove_cars=True\n",
        "    val_dataloader = torch.utils.data.DataLoader(val_dataset, batch_size=batch_size, shuffle=True, num_workers=30, pin_memory=True)    \n",
        "\n",
        "    dataloaders = {'train': dataloader, 'val': val_dataloader}\n",
        "    datasets = {'train': dataset, 'val': val_dataset}\n",
        "\n",
        "    i3d = InceptionI3d(400, in_channels=3)\n",
        "    if init_with_imagenet:\n",
        "        i3d.load_state_dict(torch.load(par.data_dir + 'i3d_models/rgb_imagenet.pt'))\n",
        "    i3d.replace_logits(1)\n",
        "    i3d.cuda()\n",
        "    i3d = nn.DataParallel(i3d)\n",
        "\n",
        "    if optimizer == 'SGD':\n",
        "        optimizer = optim.SGD(i3d.parameters(), lr=init_lr, momentum=0.9, weight_decay=0.0000001)\n",
        "    elif optimizer == 'Adam':\n",
        "        optimizer = torch.optim.Adam(i3d.parameters(), lr=init_lr)\n",
        "\n",
        "    steps = 0\n",
        "    val_loss=[]\n",
        "    while steps < number_of_steps_to_try:\n",
        "\n",
        "        for phase in ['train', 'val']:\n",
        "\n",
        "            if phase == 'train':\n",
        "                i3d.train(True)\n",
        "            else:\n",
        "                i3d.train(False)  # Set model to evaluate mode\n",
        "\n",
        "            loss_mean = 0\n",
        "            optimizer.zero_grad()\n",
        "            \n",
        "            num_iter = 0\n",
        "            # Iterate over data.\n",
        "            for data in dataloaders[phase]:\n",
        "                num_iter += 1\n",
        "                inputs, labels = data\n",
        "\n",
        "                # wrap them in Variable\n",
        "                inputs = Variable(inputs.cuda())\n",
        "                t = inputs.size(2)\n",
        "                labels = Variable(labels.cuda())\n",
        "\n",
        "                per_frame_logits = i3d(inputs)\n",
        "                per_frame_logits = F.upsample(per_frame_logits, t, mode='linear')\n",
        "\n",
        "                loss = torch.nn.functional.mse_loss(per_frame_logits.squeeze(), labels.float())\n",
        "                loss_mean += loss\n",
        "                loss.backward()\n",
        "\n",
        "                if phase == 'train' and num_iter==num_steps_per_update:\n",
        "                    steps += 1\n",
        "                    optimizer.step()\n",
        "                    optimizer.zero_grad()\n",
        "\n",
        "            loss_mean /= len(dataloaders[phase])\n",
        "            \n",
        "            if phase == 'val':\n",
        "                # print('loss', loss_mean.data.cpu().numpy())\n",
        "                val_loss.append(loss_mean.data.cpu().numpy())\n",
        "\n",
        "    mean_val_acc = (sum(val_loss[-10:]) / 10)\n",
        "    print('took time', time.time() - start)\n",
        "    print('mean val loss over last 10 steps: ', mean_val_acc)\n",
        "\n",
        "\n",
        "    results[(seq_len, optimizer, init_lr, batch_size, init_with_imagenet)] = mean_val_acc\n",
        "    # # results[(optimizer, init_lr, batch_size)] = mean_val_acc\n",
        "\n",
        "    if mean_val_acc < best_val:\n",
        "        print('best val acc updated with', mean_val_acc)\n",
        "        best_val = mean_val_acc\n",
        "\n",
        "# Print out results.\n",
        "for seq_len, optimizer, init_lr, batch_size, init_with_imagenet in (results):\n",
        "    val_accuracy = results[(seq_len, optimizer, init_lr, batch_size, init_with_imagenet)]\n",
        "    print('seq_len %e optimizer %e init_lr %e batch_size %e init_with_imagenet: %f' % (\n",
        "                seq_len, optimizer, init_lr, batch_size, init_with_imagenet))\n",
        "\n",
        "print('best validation accuracy achieved during cross-validation: %f' % best_val)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1fCWCTD87lUY"
      },
      "source": [
        "# Go throught all validation set \n",
        "\n",
        "to check the global performance and look for typical errors"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "c1DZN5_Hbf6W"
      },
      "source": [
        "# with sliding window (sw), with overlap of predictions\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "i3d = InceptionI3d(1, in_channels=3)\n",
        "i3d.load_state_dict(torch.load(par.data_dir + 'i3d_checkpoints/train_final/000700_val_loss_1.25_lr_6e-06.pt'))\n",
        "i3d.cuda()\n",
        "i3d.eval()\n",
        "\n",
        "# video_folders=[str(item) for item in list(range(100,185))]\n",
        "video_folders = par.video_folders\n",
        "\n",
        "vel_pred_all_sw = []\n",
        "vel_gt_all_sw = []\n",
        "mse_loss_vel_tot = 0\n",
        "for video_folder in video_folders:\n",
        "    print('video folder: ', video_folder)\n",
        "    vel_gt = np.load('{}{}.npy'.format(par.vel_dir, video_folder))\n",
        "    vel_gt = vel_gt[:24]\n",
        "    vel_pred = np.zeros_like(vel_gt)\n",
        "    num_frames = 24\n",
        "\n",
        "    sliding_window_output=[]\n",
        "    for start_f in range(num_frames - par.seq_len+1):\n",
        "        # load\n",
        "        imgs = load_rgb_frames(video_folder, start_f, par.seq_len)\n",
        "        # mask out cars\n",
        "        mask_for_seq = np.zeros((224, 224))\n",
        "        for i in range(start_f, start_f + par.seq_len): # for i in range(start, start+num):\n",
        "            mask = np.load(os.path.join(par.mask_dir, video_folder, str(i).zfill(10)+'.npy'))\n",
        "            mask_for_seq += mask\n",
        "        mask_for_seq[mask_for_seq > 1] = 1\n",
        "        for i,image in enumerate(imgs):\n",
        "            imgs[i][mask_for_seq==1] = 0\n",
        "        # final preprocess\n",
        "        imgs = np.array([((img/255.)*2 - 1) for img in imgs]) # Normalize\n",
        "        imgs = np.asarray(imgs, dtype=np.float32) # Convert ot float32\n",
        "        # predict\n",
        "        inputs = video_to_tensor(imgs)\n",
        "        inputs = inputs.unsqueeze(0).cuda()\n",
        "        per_frame_logits = i3d(inputs)\n",
        "        t = inputs.size(2)\n",
        "        per_frame_logits = F.upsample(per_frame_logits, t, mode='linear')\n",
        "        vel_pred_seq = per_frame_logits.squeeze().data.cpu().numpy()\n",
        "        # save the prediction\n",
        "        if not sliding_window_output: # if sliding_window_output is still empty\n",
        "            sliding_window_output = [[item] for item in list(vel_pred_seq)]\n",
        "        else:\n",
        "            sliding_window_output.append([vel_pred_seq[-1]]) # add the last element (new one)\n",
        "            for i in range(len(vel_pred_seq) - 1):\n",
        "                sliding_window_output[start_f + i].append(vel_pred_seq[i])\n",
        "\n",
        "    # average over sliding window (sw) predictions\n",
        "    vel_pred = np.array([sum(item)/len(item) for item in sliding_window_output])\n",
        "    vel_pred_all_sw.append(vel_pred)\n",
        "    vel_gt_all_sw.append(vel_gt)\n",
        "\n",
        "    mse_loss_vel = np.sum((vel_pred-vel_gt) ** 2)\n",
        "    mse_loss_vel /= len(vel_gt)\n",
        "    print('mse_loss_vel', mse_loss_vel) # loss within folder\n",
        "    mse_loss_vel_tot += mse_loss_vel\n",
        "\n",
        "mse_loss_vel_tot /= len(video_folders)\n",
        "print('\\nTotal loss: ', mse_loss_vel_tot)\n",
        "\n",
        "print('\\nPlots of ', video_folder) # the last video folder, just to visualize\n",
        "plt.plot(vel_gt)\n",
        "plt.plot(vel_pred)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PYu3bsQ37uUI"
      },
      "source": [
        "# Test"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4jojbJhJKatu"
      },
      "source": [
        "# with sliding window, with overlap of predictions\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "i3d = InceptionI3d(1, in_channels=3)\n",
        "i3d.load_state_dict(torch.load(par.data_dir + 'i3d_checkpoints/train_for_test_video/000900.pt'))\n",
        "i3d.cuda()\n",
        "i3d.eval()\n",
        "\n",
        "# num_frames = 10798\n",
        "num_frames = 5399 # 10 fps\n",
        "print('num_frames',num_frames)\n",
        "\n",
        "sliding_window_output = []\n",
        "for start_f in range(num_frames - par.seq_len + 1):\n",
        "    print('start_f: ', start_f)\n",
        "\n",
        "    imgs = load_rgb_frames('000', start_f, par.seq_len) # the test folder contaains only one folder 000\n",
        "\n",
        "    mask_for_seq = np.zeros((224, 224))\n",
        "    for i in range(start_f, start_f + par.seq_len):\n",
        "        mask = np.load(os.path.join(par.mask_dir, '000', str(i).zfill(10) + '.npy'))\n",
        "        mask_for_seq += mask\n",
        "    mask_for_seq[mask_for_seq > 1] = 1\n",
        "    for i,image in enumerate(imgs):\n",
        "        imgs[i][mask_for_seq == 1] = 0\n",
        "\n",
        "    imgs = np.array([((img/255.)*2 - 1) for img in imgs]) #Normalize\n",
        "    imgs = np.asarray(imgs, dtype=np.float32)             # Convert ot float32\n",
        "  \n",
        "    inputs = video_to_tensor(imgs)\n",
        "    inputs = inputs.unsqueeze(0).cuda()\n",
        "    per_frame_logits = i3d(inputs)\n",
        "    t = inputs.size(2)\n",
        "    per_frame_logits = F.upsample(per_frame_logits, t, mode='linear')\n",
        "    vel_pred_seq = per_frame_logits.squeeze().data.cpu().numpy()\n",
        "\n",
        "\n",
        "    if not sliding_window_output: # if sliding_window_output is empty\n",
        "        sliding_window_output = [[item] for item in list(vel_pred_seq)]\n",
        "    else:\n",
        "        sliding_window_output.append([vel_pred_seq[-1]])\n",
        "        for i in range(len(vel_pred_seq) - 1):\n",
        "            sliding_window_output[start_f + i].append(vel_pred_seq[i])\n",
        "\n",
        "vel_pred_sw = np.array([sum(item)/len(item) for item in sliding_window_output])\n",
        "\n",
        "plt.plot(vel_pred_sw)\n",
        "np.save('vel_pred_test_sw.npy',vel_pred_sw)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}