run.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import division\n",
    "import argparse\n",
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.pyplot import imsave, imread\n",
    "import getopt\n",
    "import matplotlib\n",
    "matplotlib.use('Agg')\n",
    "import numpy as np\n",
    "import math\n",
    "import os\n",
    "import PIL\n",
    "import PIL.Image\n",
    "import sys\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import math, copy, time\n",
    "from torch.autograd import Variable\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 52
    },
    "colab_type": "code",
    "id": "4xPiu3vI_H52",
    "outputId": "e0227caf-4ce1-4b4d-8e2f-bd0b9bc8fbca"
   },
   "source": [
    "#%tensorflow_version 2.x\n",
    "import tensorflow as tf\n",
    "device_name = tf.test.gpu_device_name()\n",
    "if device_name != '/device:GPU:0':\n",
    "  raise SystemError('GPU device not found')\n",
    "print('Found GPU at: {}'.format(device_name))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "rS_7CRfIUP9y"
   },
   "outputs": [],
   "source": [
    "#!/usr/bin/env python\n",
    "\n",
    "#from video_processing import frame_capture\n",
    "\n",
    "\n",
    "##########################################################\n",
    "\n",
    "assert(int(str('').join(torch.__version__.split('.')[0:3])) >= 41) # requires at least pytorch version 0.4.1\n",
    "\n",
    "torch.set_grad_enabled(False) # make sure to not compute gradients for computational performance\n",
    "\n",
    "torch.backends.cudnn.enabled = True # make sure to use cudnn for computational performance\n",
    "\n",
    "##########################################################\n",
    "\n",
    "arguments_strModel = 'sintel-final'\n",
    "arguments_strFirst = './images/first.png'\n",
    "arguments_strSecond = './images/second.png'\n",
    "arguments_strOut = './out.flo'\n",
    "\n",
    "parser = argparse.ArgumentParser()\n",
    "parser.add_argument('-gpuid', nargs=1, type=str, default='0')  # python3 main.py -gpuid=1,2,3\n",
    "#args = parser.parse_args()\n",
    "#os.environ['_VISIBLE_DEVICES'] = args.gpuid[0]\n",
    "#print(os.environ['CUDA_VISIBLE_DEVICES'])\n",
    "#\n",
    "# for strOption, strArgument in getopt.getopt(sys.argv[1:], '', [ strParameter[2:] + '=' for strParameter in sys.argv[1::2] ])[0]:\n",
    "# \tif strOption == '--model' and strArgument != '': arguments_strModel = strArgument # which model to use, see below\n",
    "# \tif strOption == '--first' and strArgument != '': arguments_strFirst = strArgument # path to the first frame\n",
    "# \tif strOption == '--second' and strArgument != '': arguments_strSecond = strArgument # path to the second frame\n",
    "# \tif strOption == '--out' and strArgument != '': arguments_strOut = strArgument # path to where the output should be stored\n",
    "# end\n",
    "\n",
    "##########################################################\n",
    "\n",
    "Backward_tensorGrid = {}\n",
    "\n",
    "def Backward(tensorInput, tensorFlow):\n",
    "\tif str(tensorFlow.size()) not in Backward_tensorGrid:\n",
    "\t\ttensorHorizontal = torch.linspace(-1.0, 1.0, tensorFlow.size(3)).view(1, 1, 1, tensorFlow.size(3)).expand(tensorFlow.size(0), -1, tensorFlow.size(2), -1)\n",
    "\t\ttensorVertical = torch.linspace(-1.0, 1.0, tensorFlow.size(2)).view(1, 1, tensorFlow.size(2), 1).expand(tensorFlow.size(0), -1, -1, tensorFlow.size(3))\n",
    "\n",
    "\t\tBackward_tensorGrid[str(tensorFlow.size())] = torch.cat([ tensorHorizontal, tensorVertical ], 1).cuda()#\n",
    "\t# end\n",
    "\n",
    "\ttensorFlow = torch.cat([ tensorFlow[:, 0:1, :, :] / ((tensorInput.size(3) - 1.0) / 2.0), tensorFlow[:, 1:2, :, :] / ((tensorInput.size(2) - 1.0) / 2.0) ], 1)\n",
    "\n",
    "\treturn torch.nn.functional.grid_sample(input=tensorInput, grid=(Backward_tensorGrid[str(tensorFlow.size())] + tensorFlow).permute(0, 2, 3, 1), mode='bilinear', padding_mode='border')\n",
    "# end\n",
    "\n",
    "##########################################################\n",
    "\n",
    "class Network(torch.nn.Module):\n",
    "\tdef __init__(self):\n",
    "\t\tsuper(Network, self).__init__()\n",
    "\n",
    "\t\tclass Preprocess(torch.nn.Module):\n",
    "\t\t\tdef __init__(self):\n",
    "\t\t\t\tsuper(Preprocess, self).__init__()\n",
    "\t\t\t# end\n",
    "\n",
    "\t\t\tdef forward(self, tensorInput):\n",
    "\t\t\t\ttensorBlue = (tensorInput[:, 0:1, :, :] - 0.406) / 0.225\n",
    "\t\t\t\ttensorGreen = (tensorInput[:, 1:2, :, :] - 0.456) / 0.224\n",
    "\t\t\t\ttensorRed = (tensorInput[:, 2:3, :, :] - 0.485) / 0.229\n",
    "\n",
    "\t\t\t\treturn torch.cat([ tensorRed, tensorGreen, tensorBlue ], 1)\n",
    "\t\t\t# end\n",
    "\t\t# end\n",
    "\n",
    "\t\tclass Basic(torch.nn.Module):\n",
    "\t\t\tdef __init__(self, intLevel):\n",
    "\t\t\t\tsuper(Basic, self).__init__()\n",
    "\n",
    "\t\t\t\tself.moduleBasic = torch.nn.Sequential(\n",
    "\t\t\t\t\ttorch.nn.Conv2d(in_channels=8, out_channels=32, kernel_size=7, stride=1, padding=3),\n",
    "\t\t\t\t\ttorch.nn.ReLU(inplace=False),\n",
    "\t\t\t\t\ttorch.nn.Conv2d(in_channels=32, out_channels=64, kernel_size=7, stride=1, padding=3),\n",
    "\t\t\t\t\ttorch.nn.ReLU(inplace=False),\n",
    "\t\t\t\t\ttorch.nn.Conv2d(in_channels=64, out_channels=32, kernel_size=7, stride=1, padding=3),\n",
    "\t\t\t\t\ttorch.nn.ReLU(inplace=False),\n",
    "\t\t\t\t\ttorch.nn.Conv2d(in_channels=32, out_channels=16, kernel_size=7, stride=1, padding=3),\n",
    "\t\t\t\t\ttorch.nn.ReLU(inplace=False),\n",
    "\t\t\t\t\ttorch.nn.Conv2d(in_channels=16, out_channels=2, kernel_size=7, stride=1, padding=3)\n",
    "\t\t\t\t)\n",
    "\t\t\t# end\n",
    "\n",
    "\t\t\tdef forward(self, tensorInput):\n",
    "\t\t\t\treturn self.moduleBasic(tensorInput)\n",
    "\t\t\t# end\n",
    "\t\t# end\n",
    "\n",
    "\t\tself.modulePreprocess = Preprocess()\n",
    "\n",
    "\t\tself.moduleBasic = torch.nn.ModuleList([ Basic(intLevel) for intLevel in range(6) ])\n",
    "\n",
    "\t\tself.load_state_dict(torch.load('./network-' + arguments_strModel + '.pytorch'))\n",
    "\t# end\n",
    "\n",
    "\tdef forward(self, tensorFirst, tensorSecond):\n",
    "\t\ttensorFirst = [ self.modulePreprocess(tensorFirst) ]\n",
    "\t\ttensorSecond = [ self.modulePreprocess(tensorSecond) ]\n",
    "\n",
    "\t\tfor intLevel in range(5):\n",
    "\t\t\tif tensorFirst[0].size(2) > 32 or tensorFirst[0].size(3) > 32:\n",
    "\t\t\t\ttensorFirst.insert(0, torch.nn.functional.avg_pool2d(input=tensorFirst[0], kernel_size=2, stride=2, count_include_pad=False))\n",
    "\t\t\t\ttensorSecond.insert(0, torch.nn.functional.avg_pool2d(input=tensorSecond[0], kernel_size=2, stride=2, count_include_pad=False))\n",
    "\t\t\t# end\n",
    "\t\t# end\n",
    "\n",
    "\t\ttensorFlow = tensorFirst[0].new_zeros([ tensorFirst[0].size(0), 2, int(math.floor(tensorFirst[0].size(2) / 2.0)), int(math.floor(tensorFirst[0].size(3) / 2.0)) ])\n",
    "\n",
    "\t\tfor intLevel in range(len(tensorFirst)):\n",
    "\t\t\ttensorUpsampled = torch.nn.functional.interpolate(input=tensorFlow, scale_factor=2, mode='bilinear', align_corners=True) * 2.0\n",
    "\n",
    "\t\t\tif tensorUpsampled.size(2) != tensorFirst[intLevel].size(2): tensorUpsampled = torch.nn.functional.pad(input=tensorUpsampled, pad=[ 0, 0, 0, 1 ], mode='replicate')\n",
    "\t\t\tif tensorUpsampled.size(3) != tensorFirst[intLevel].size(3): tensorUpsampled = torch.nn.functional.pad(input=tensorUpsampled, pad=[ 0, 1, 0, 0 ], mode='replicate')\n",
    "\n",
    "\t\t\ttensorFlow = self.moduleBasic[intLevel](torch.cat([ tensorFirst[intLevel], Backward(tensorInput=tensorSecond[intLevel], tensorFlow=tensorUpsampled), tensorUpsampled ], 1)) + tensorUpsampled\n",
    "\t\t# end\n",
    "\n",
    "\t\treturn tensorFlow\n",
    "\t# end\n",
    "# end\n",
    "\n",
    "moduleNetwork = Network().cuda().eval()\n",
    "\n",
    "##########################################################\n",
    "def estimate(tensorFirst, tensorSecond):\n",
    "\tassert(tensorFirst.size(1) == tensorSecond.size(1))\n",
    "\tassert(tensorFirst.size(2) == tensorSecond.size(2))\n",
    "\n",
    "\tintWidth = tensorFirst.size(2)\n",
    "\tintHeight = tensorFirst.size(1)\n",
    "\n",
    "#\tassert(intWidth == 1024) # remember that there is no guarantee for correctness, comment this line out if you acknowledge this and want to continue\n",
    "#\tassert(intHeight == 416) # remember that there is no guarantee for correctness, comment this line out if you acknowledge this and want to continue\n",
    "\n",
    "\ttensorPreprocessedFirst = tensorFirst.cuda().view(1, 3, intHeight, intWidth)\n",
    "\ttensorPreprocessedSecond = tensorSecond.cuda().view(1, 3, intHeight, intWidth)\n",
    "\n",
    "\tintPreprocessedWidth = int(math.floor(math.ceil(intWidth / 32.0) * 32.0))\n",
    "\tintPreprocessedHeight = int(math.floor(math.ceil(intHeight / 32.0) * 32.0))\n",
    "\n",
    "\ttensorPreprocessedFirst = torch.nn.functional.interpolate(input=tensorPreprocessedFirst, size=(intPreprocessedHeight, intPreprocessedWidth), mode='bilinear', align_corners=False)\n",
    "\ttensorPreprocessedSecond = torch.nn.functional.interpolate(input=tensorPreprocessedSecond, size=(intPreprocessedHeight, intPreprocessedWidth), mode='bilinear', align_corners=False)\n",
    "\n",
    "\ttensorFlow = torch.nn.functional.interpolate(input=moduleNetwork(tensorPreprocessedFirst, tensorPreprocessedSecond), size=(intHeight, intWidth), mode='bilinear', align_corners=False)\n",
    "\n",
    "\ttensorFlow[:, 0, :, :] *= float(intWidth) / float(intPreprocessedWidth)\n",
    "\ttensorFlow[:, 1, :, :] *= float(intHeight) / float(intPreprocessedHeight)\n",
    "\treturn tensorFlow[0, :, :, :].cpu()\n",
    "# end\n",
    "\n",
    "##########################################################\n",
    "def make_color_wheel():\n",
    "    \"\"\"\n",
    "    Generate color wheel according Middlebury color code\n",
    "    :return: Color wheel\n",
    "    \"\"\"\n",
    "    RY = 15\n",
    "    YG = 6\n",
    "    GC = 4\n",
    "    CB = 11\n",
    "    BM = 13\n",
    "    MR = 6\n",
    "\n",
    "    ncols = RY + YG + GC + CB + BM + MR\n",
    "\n",
    "    colorwheel = np.zeros([ncols, 3])\n",
    "\n",
    "    col = 0\n",
    "\n",
    "    # RY\n",
    "    colorwheel[0:RY, 0] = 255\n",
    "    colorwheel[0:RY, 1] = np.transpose(np.floor(255*np.arange(0, RY) / RY))\n",
    "    col += RY\n",
    "\n",
    "    # YG\n",
    "    colorwheel[col:col+YG, 0] = 255 - np.transpose(np.floor(255*np.arange(0, YG) / YG))\n",
    "    colorwheel[col:col+YG, 1] = 255\n",
    "    col += YG\n",
    "\n",
    "    # GC\n",
    "    colorwheel[col:col+GC, 1] = 255\n",
    "    colorwheel[col:col+GC, 2] = np.transpose(np.floor(255*np.arange(0, GC) / GC))\n",
    "    col += GC\n",
    "\n",
    "    # CB\n",
    "    colorwheel[col:col+CB, 1] = 255 - np.transpose(np.floor(255*np.arange(0, CB) / CB))\n",
    "    colorwheel[col:col+CB, 2] = 255\n",
    "    col += CB\n",
    "\n",
    "    # BM\n",
    "    colorwheel[col:col+BM, 2] = 255\n",
    "    colorwheel[col:col+BM, 0] = np.transpose(np.floor(255*np.arange(0, BM) / BM))\n",
    "    col += + BM\n",
    "\n",
    "    # MR\n",
    "    colorwheel[col:col+MR, 2] = 255 - np.transpose(np.floor(255 * np.arange(0, MR) / MR))\n",
    "    colorwheel[col:col+MR, 0] = 255\n",
    "\n",
    "    return colorwheel\n",
    "\n",
    "\n",
    "\n",
    "def compute_color(u, v):\n",
    "\t\"\"\"\n",
    "    compute optical flow color map\n",
    "    :param u: optical flow horizontal map\n",
    "    :param v: optical flow vertical map\n",
    "    :return: optical flow in color code\n",
    "    \"\"\"\n",
    "\t[h, w] = u.shape\n",
    "\timg = np.zeros([h, w, 3])\n",
    "\tnanIdx = np.isnan(u) | np.isnan(v)\n",
    "\tu[nanIdx] = 0\n",
    "\tv[nanIdx] = 0\n",
    "\n",
    "\tcolorwheel = make_color_wheel()\n",
    "\tncols = np.size(colorwheel, 0)\n",
    "\n",
    "\trad = np.sqrt(u**2+v**2)\n",
    "\n",
    "\ta = np.arctan2(-v, -u) / np.pi\n",
    "\n",
    "\tfk = (a+1) / 2 * (ncols - 1) + 1\n",
    "\n",
    "\tk0 = np.floor(fk).astype(int)\n",
    "\n",
    "\tk1 = k0 + 1\n",
    "\tk1[k1 == ncols+1] = 1\n",
    "\tf = fk - k0\n",
    "\n",
    "\tfor i in range(0, np.size(colorwheel,1)):\n",
    "\t\ttmp = colorwheel[:, i]\n",
    "\t\tcol0 = tmp[k0-1] / 255\n",
    "\t\tcol1 = tmp[k1-1] / 255\n",
    "\t\tcol = (1-f) * col0 + f * col1\n",
    "\n",
    "\t\tidx = rad <= 1\n",
    "\t\tcol[idx] = 1-rad[idx]*(1-col[idx])\n",
    "\t\tnotidx = np.logical_not(idx)\n",
    "\n",
    "\t\tcol[notidx] *= 0.75\n",
    "\t\timg[:, :, i] = np.uint8(np.floor(255 * col*(1-nanIdx)))\n",
    "\n",
    "\treturn img\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 285
    },
    "colab_type": "code",
    "id": "2UmHNJ8-FCxF",
    "outputId": "b303ef91-b918-47b4-df59-d688abb6e6d7"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing video 'ship'\n",
      "output folder already exists, not running analysis for file\n",
      "Processing video 'video_4'\n",
      "output folder already exists, not running analysis for file\n",
      "Processing video 'gardens'\n",
      "output folder already exists, not running analysis for file\n",
      "Processing video 'sharks'\n",
      "output folder already exists, not running analysis for file\n",
      "Processing video 'video_8'\n",
      "output folder already exists, not running analysis for file\n",
      "Processing video 'skyhouse'\n",
      "output folder already exists, not running analysis for file\n",
      "Processing video 'video_7'\n",
      "output folder already exists, not running analysis for file\n",
      "Processing video 'minecraft'\n",
      "output folder already exists, not running analysis for file\n",
      "Processing video 'snowplanet'\n",
      "Processing video 'sculptures'\n",
      "Processing video 'video_2'\n",
      "finished comparing frame 100 video video_2\n",
      "finished comparing frame 200 video video_2\n",
      "finished comparing frame 300 video video_2\n",
      "finished comparing frame 400 video video_2\n",
      "finished comparing frame 500 video video_2\n",
      "finished comparing frame 600 video video_2\n",
      "finished comparing frame 700 video video_2\n"
     ]
    }
   ],
   "source": [
    "def calculate_and_output_optical_flows(input_dir, flow_output_dir, video_name):\n",
    "    if not os.path.exists(flow_output_dir):\n",
    "        os.mkdir(flow_output_dir)\n",
    "\n",
    "    input_video_dir = f\"{input_dir}/{video_name}\"\n",
    "    flow_output_video_dir = f\"{flow_output_dir}/{video_name}\"\n",
    "    \n",
    "    if not os.path.exists(flow_output_video_dir):\n",
    "        os.mkdir(flow_output_video_dir)\n",
    "    else:\n",
    "        print(\"output folder already exists, not running analysis for file\")\n",
    "        return None\n",
    "    \n",
    "    num_frames = len(os.listdir(input_video_dir))\n",
    "    if num_frames == 0:\n",
    "        print(f\"Unable to find any input files for video {video_name} in folder {input_video_dir}\")\n",
    "        \n",
    "    for i in range(1, num_frames - 1):\n",
    "        first_file = f\"{input_video_dir}/{video_name}{i:05d}.jpeg\"\n",
    "        second_file = f\"{input_video_dir}/{video_name}{(i+1):05d}.jpeg\"\n",
    "\n",
    "        first_image_np = np.array(PIL.Image.open(first_file)) \n",
    "        second_image_np = np.array(PIL.Image.open(second_file)) \n",
    "\n",
    "        tensorFirst = torch.FloatTensor(\n",
    "        first_image_np[:, :, ::-1].transpose(2, 0, 1).astype(np.float32) * (1.0 / 255.0))\n",
    "        tensorSecond = torch.FloatTensor(\n",
    "        second_image_np[:, :, ::-1].transpose(2, 0, 1).astype(np.float32) * (1.0 / 255.0))\n",
    "\n",
    "        tensorOutput = estimate(tensorFirst, tensorSecond)\n",
    "        tensorOutput = tensorOutput.numpy()\n",
    "        np.save(\n",
    "            f\"{flow_output_video_dir}/{video_name}{i:05d}\",\n",
    "            tensorOutput)\n",
    "\n",
    "        if i % 100==0:\n",
    "            print(\"finished comparing frame \" + str(i) + \" video \" + video_name)\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    input_dir = \"videos/4-scaled-down-jpegs\"\n",
    "    flow_output_dir=\"videos/6-optical-flows-scaled-down\"\n",
    "    videos = os.listdir(input_dir)\n",
    "    for video_name in videos:\n",
    "        print(f\"Processing video '{video_name}'\")\n",
    "        calculate_and_output_optical_flows(\n",
    "            input_dir=input_dir,\n",
    "            flow_output_dir=flow_output_dir,\n",
    "            video_name=video_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calculate the sum of norms of the flows for each frame\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calc_sum_of_norms(flow_array):\n",
    "    if shape(flow_array)[0] != 2:\n",
    "        raise ValueError(\"Wrong input shape. Has to contain flow in two dimensions\")\n",
    "    squared = np.power(flow_array, 2)\n",
    "    squared_sum = squared[0] + squared[1]\n",
    "    norms = np.sqrt(squared_sum)\n",
    "    return norms.sum()\n",
    "\n",
    "def calc_flow_and_save(input_dir, output_dir, video_name):\n",
    "    flows = []\n",
    "    video_dir = f\"{input_dir}/{video_name}\"\n",
    "    video_files = os.listdir(video_dir)\n",
    "    video_files.sort()\n",
    "\n",
    "    for video_file in video_files:\n",
    "        x = np.load(f\"{video_dir}/{video_file}\")\n",
    "        flows.append(calc_sum_of_norms(x))\n",
    "    np.save(f\"{output_dir}/{video_name}_flow\", np.asarray(flows))\n",
    "\n",
    "input_dir=\"videos/6-optical-flows-scaled-down\"\n",
    "output_dir=\"videos/7-final-aggregated-optical-flows-scaled-down\"\n",
    "videos = os.listdir(input_dir)\n",
    "for video_name in videos:\n",
    "    print(f\"Processing video '{video_name}'\")\n",
    "    calc_flow_and_save(input_dir, output_dir, video_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "machine_shape": "hm",
   "name": "run.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}