colabs/trajan_demo.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fetSbKRrczhD"
      },
      "source": [
        "Copyright 2020 DeepMind Technologies Limited.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "you may not use this file except in compliance with the License.\n",
        "You may obtain a copy of the License at\n",
        "\n",
        "https://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "Unless required by applicable law or agreed to in writing, software\n",
        "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "See the License for the specific language governing permissions and\n",
        "limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zuYtT7GpuVgQ"
      },
      "source": [
        "\u003cp align=\"center\"\u003e\n",
        "  \u003ch1 align=\"center\"\u003eTRAJAN: Direct Motion Models for Assessing Generated Videos\u003c/h1\u003e\n",
        "  \u003cp align=\"center\"\u003e\n",
        "    \u003ca href=\"https://k-r-allen.github.io/\"\u003eKelsey Allen*\u003c/a\u003e\n",
        "    ·\n",
        "    \u003ca href=\"https://www.carldoersch.com/\"\u003eCarl Doersch\u003c/a\u003e\n",
        "    ·\n",
        "    \u003ca href=\"https://stanniszhou.github.io/\"\u003eGuangyao Zhou\u003c/a\u003e\n",
        "    ·\n",
        "    \u003ca href=\"https://mohammedsuhail.net/\"\u003eMohammed Suhail\u003c/a\u003e\n",
        "    ·\n",
        "    \u003ca href=\"https://dannydriess.github.io/\"\u003eDanny Driess\u003c/a\u003e\n",
        "    ·\n",
        "    \u003ca href=\"https://www.irocco.info/\"\u003eIgnacio Rocco\u003c/a\u003e\n",
        "    ·\n",
        "    \u003ca href=\"https://yuliarubanova.github.io/\"\u003eYulia Rubanova\u003c/a\u003e\n",
        "    ·\n",
        "    \u003ca href=\"https://tkipf.github.io/\"\u003eThomas Kipf\u003c/a\u003e\n",
        "    ·\n",
        "    \u003ca href=\"https://msajjadi.com/\"\u003eMehdi S. M. Sajjadi\u003c/a\u003e\n",
        "    ·\n",
        "    \u003ca href=\"https://scholar.google.com/citations?user=MxxZkEcAAAAJ\u0026hl=en\"\u003eKevin Murphy\u003c/a\u003e\n",
        "    ·\n",
        "    \u003ca href=\"https://scholar.google.co.uk/citations?user=IUZ-7_cAAAAJ\"\u003eJoao Carreira\u003c/a\u003e\n",
        "    ·\n",
        "    \u003ca href=\"https://www.sjoerdvansteenkiste.com/\"\u003eSjoerd van Steenkiste*\u003c/a\u003e\n",
        "  \u003c/p\u003e\n",
        "  \u003ch3 align=\"center\"\u003e\u003ca href=\"\"\u003ePaper\u003c/a\u003e | \u003ca href=\"https://trajan-paper.github.io\"\u003eProject Page\u003c/a\u003e | \u003ca href=\"https://github.com/deepmind/tapnet\"\u003eGitHub\u003c/a\u003e \u003c/h3\u003e\n",
        "  \u003cdiv align=\"center\"\u003e\u003c/div\u003e\n",
        "\u003c/p\u003e\n",
        "\n",
        "\u003cp align=\"center\"\u003e\n",
        "  \u003ca href=\"\"\u003e\n",
        "    \u003cimg src=\"https://storage.googleapis.com/dm-tapnet/swaying_gif.gif\" alt=\"Logo\" width=\"50%\"\u003e\n",
        "  \u003c/a\u003e\n",
        "\u003c/p\u003e\n",
        "\n",
        ""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "mCmDvfFvxnGB"
      },
      "outputs": [],
      "source": [
        "# @title Install Code and Dependencies {form-width: \"25%\"}\n",
        "!pip install git+https://github.com/google-deepmind/tapnet.git"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zyEo9-Kv78S7"
      },
      "outputs": [],
      "source": [
        "MODEL_TYPE = 'bootstapir'  # 'tapir' or 'bootstapir'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HaswJZMq9B3c"
      },
      "outputs": [],
      "source": [
        "# @title Download Model {form-width: \"25%\"}\n",
        "\n",
        "%mkdir tapnet/checkpoints\n",
        "\n",
        "if MODEL_TYPE == \"tapir\":\n",
        "  !wget -P tapnet/checkpoints https://storage.googleapis.com/dm-tapnet/tapir_checkpoint_panning.npy\n",
        "else:\n",
        "  !wget -P tapnet/checkpoints https://storage.googleapis.com/dm-tapnet/bootstap/bootstapir_checkpoint_v2.npy\n",
        "\n",
        "!wget -P tapnet/checkpoints https://storage.googleapis.com/dm-tapnet/trajan/track_autoencoder_ckpt.npz\n",
        "\n",
        "%ls tapnet/checkpoints"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FxlHY242m-6Q"
      },
      "outputs": [],
      "source": [
        "# @title Imports {form-width: \"25%\"}\n",
        "from google.colab import output\n",
        "import jax\n",
        "import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "import mediapy as media\n",
        "import numpy as np\n",
        "from tapnet.models import tapir_model\n",
        "from tapnet.utils import model_utils\n",
        "from tapnet.utils import transforms\n",
        "from tapnet.utils import viz_utils\n",
        "from tapnet.tapvid import evaluation_datasets\n",
        "from tapnet.trajan import track_autoencoder\n",
        "\n",
        "output.enable_custom_widget_manager()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7rfy2yobnHqw"
      },
      "outputs": [],
      "source": [
        "# @title Load Checkpoint {form-width: \"25%\"}\n",
        "\n",
        "if MODEL_TYPE == 'tapir':\n",
        "  checkpoint_path = 'tapnet/checkpoints/tapir_checkpoint_panning.npy'\n",
        "else:\n",
        "  checkpoint_path = 'tapnet/checkpoints/bootstapir_checkpoint_v2.npy'\n",
        "ckpt_state = np.load(checkpoint_path, allow_pickle=True).item()\n",
        "params, state = ckpt_state['params'], ckpt_state['state']\n",
        "\n",
        "kwargs = dict(bilinear_interp_with_depthwise_conv=False, pyramid_level=0)\n",
        "if MODEL_TYPE == 'bootstapir':\n",
        "  kwargs.update(\n",
        "      dict(pyramid_level=1, extra_convs=True, softmax_temperature=10.0)\n",
        "  )\n",
        "tapir = tapir_model.ParameterizedTAPIR(params, state, tapir_kwargs=kwargs)\n",
        "\n",
        "trajan_checkpoint_path = 'tapnet/checkpoints/track_autoencoder_ckpt.npz'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Y1HkuvyT6SJF"
      },
      "outputs": [],
      "source": [
        "# @title Load an Exemplar Video {form-width: \"25%\"}\n",
        "\n",
        "%mkdir tapnet/examplar_videos\n",
        "\n",
        "!wget -P tapnet/examplar_videos https://storage.googleapis.com/dm-tapnet/horsejump-high.mp4\n",
        "\n",
        "video = media.read_video(\"tapnet/examplar_videos/horsejump-high.mp4\")\n",
        "media.show_video(video, fps=10)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ogRTRVgfSq0W"
      },
      "outputs": [],
      "source": [
        "# @title Utility Functions {form-width: \"25%\"}\n",
        "\n",
        "\n",
        "def inference(frames, query_points):\n",
        "  \"\"\"Inference on one video.\n",
        "\n",
        "  Args:\n",
        "    frames: [num_frames, height, width, 3], [0, 255], np.uint8\n",
        "    query_points: [num_points, 3], [0, num_frames/height/width], [t, y, x]\n",
        "\n",
        "  Returns:\n",
        "    tracks: [num_points, 3], [-1, 1], [t, y, x]\n",
        "    visibles: [num_points, num_frames], bool\n",
        "  \"\"\"\n",
        "  # Preprocess video to match model inputs format\n",
        "  frames = model_utils.preprocess_frames(frames)\n",
        "  query_points = query_points.astype(np.float32)\n",
        "  frames, query_points = frames[None], query_points[None]  # Add batch dimension\n",
        "\n",
        "  outputs = tapir(\n",
        "      video=frames,\n",
        "      query_points=query_points,\n",
        "      is_training=False,\n",
        "      query_chunk_size=32,\n",
        "  )\n",
        "  tracks, occlusions, expected_dist = (\n",
        "      outputs['tracks'],\n",
        "      outputs['occlusion'],\n",
        "      outputs['expected_dist'],\n",
        "  )\n",
        "\n",
        "  # Binarize occlusions\n",
        "  visibles = model_utils.postprocess_occlusions(occlusions, expected_dist)\n",
        "  return tracks[0], visibles[0]\n",
        "\n",
        "\n",
        "inference = jax.jit(inference)\n",
        "\n",
        "\n",
        "def sample_random_points(frame_max_idx, height, width, num_points):\n",
        "  \"\"\"Sample random points with (time, height, width) order.\"\"\"\n",
        "  y = np.random.randint(0, height, (num_points, 1))\n",
        "  x = np.random.randint(0, width, (num_points, 1))\n",
        "  t = np.random.randint(0, frame_max_idx + 1, (num_points, 1))\n",
        "  points = np.concatenate((t, y, x), axis=-1).astype(\n",
        "      np.int32\n",
        "  )  # [num_points, 3]\n",
        "  return points"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "QpTDg7oQJDC3"
      },
      "outputs": [],
      "source": [
        "# @title Efficient Chunked Point Track Prediction {form-width: \"25%\"}\n",
        "\n",
        "resize_height = 256  # @param {type: \"integer\"}\n",
        "resize_width = 256  # @param {type: \"integer\"}\n",
        "num_points = 4096  # @param {type: \"integer\"}\n",
        "\n",
        "frames = media.resize_video(video, (resize_height, resize_width))\n",
        "frames = model_utils.preprocess_frames(frames[None])\n",
        "feature_grids = tapir.get_feature_grids(frames, is_training=False)\n",
        "query_points = sample_random_points(\n",
        "    frames.shape[1], frames.shape[2], frames.shape[3], num_points\n",
        ")\n",
        "chunk_size = 32\n",
        "\n",
        "\n",
        "def chunk_inference(query_points):\n",
        "  query_points = query_points.astype(np.float32)[None]\n",
        "\n",
        "  outputs = tapir(\n",
        "      video=frames,\n",
        "      query_points=query_points,\n",
        "      is_training=False,\n",
        "      query_chunk_size=chunk_size,\n",
        "      feature_grids=feature_grids,\n",
        "  )\n",
        "  tracks, occlusions, expected_dist = (\n",
        "      outputs[\"tracks\"],\n",
        "      outputs[\"occlusion\"],\n",
        "      outputs[\"expected_dist\"],\n",
        "  )\n",
        "\n",
        "  # Binarize occlusions\n",
        "  visibles = model_utils.postprocess_occlusions(occlusions, expected_dist)\n",
        "  return tracks[0], visibles[0]\n",
        "\n",
        "\n",
        "chunk_inference = jax.jit(chunk_inference)\n",
        "\n",
        "all_tracks = []\n",
        "all_visibles = []\n",
        "for chunk in range(0, query_points.shape[0], chunk_size):\n",
        "  tracks, visibles = chunk_inference(query_points[chunk : chunk + chunk_size])\n",
        "  all_tracks.append(np.array(tracks))\n",
        "  all_visibles.append(np.array(visibles))\n",
        "\n",
        "tracks = np.concatenate(all_tracks, axis=0)\n",
        "visibles = np.concatenate(all_visibles, axis=0)\n",
        "\n",
        "# Visualize sparse point tracks\n",
        "height, width = video.shape[1:3]\n",
        "tracks = transforms.convert_grid_coordinates(\n",
        "    tracks, (resize_width, resize_height), (width, height)\n",
        ")\n",
        "video_viz = viz_utils.paint_point_track(video, tracks, visibles)\n",
        "media.show_video(video_viz, fps=10)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lOWN6_baASZ5"
      },
      "source": [
        "## Apply TRAJAN"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "jU_rIbnDooNu"
      },
      "outputs": [],
      "source": [
        "# @title Imports\n",
        "\n",
        "from __future__ import annotations\n",
        "\n",
        "import dataclasses\n",
        "import einops\n",
        "import os\n",
        "import io"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "KkbHj_Oo_8dt"
      },
      "outputs": [],
      "source": [
        "# @title Load parameters {form-width: \"25%\"}\n",
        "\n",
        "\n",
        "def npload(fname):\n",
        "  loaded = np.load(fname, allow_pickle=False)\n",
        "  if isinstance(loaded, np.ndarray):\n",
        "    return loaded\n",
        "  else:\n",
        "    return dict(loaded)\n",
        "\n",
        "\n",
        "def recover_tree(flat_dict):\n",
        "  tree = (\n",
        "      {}\n",
        "  )  # Initialize an empty dictionary to store the resulting tree structure\n",
        "  for (\n",
        "      k,\n",
        "      v,\n",
        "  ) in (\n",
        "      flat_dict.items()\n",
        "  ):  # Iterate over each key-value pair in the flat dictionary\n",
        "    parts = k.split(\n",
        "        '/'\n",
        "    )  # Split the key into parts using \"/\" as a delimiter to build the tree structure\n",
        "    node = tree  # Start at the root of the tree\n",
        "    for part in parts[\n",
        "        :-1\n",
        "    ]:  # Loop through each part of the key, except the last one\n",
        "      if (\n",
        "          part not in node\n",
        "      ):  # If the current part doesn't exist as a key in the node, create an empty dictionary for it\n",
        "        node[part] = {}\n",
        "      node = node[part]  # Move down the tree to the next level\n",
        "    node[parts[-1]] = v  # Set the value at the final part of the key\n",
        "  return tree  # Return the reconstructed tree\n",
        "\n",
        "params = recover_tree(npload(trajan_checkpoint_path))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "0T-63_sJoSps"
      },
      "outputs": [],
      "source": [
        "# @title Preprocessor for Tracks\n",
        "\n",
        "@dataclasses.dataclass(kw_only=True, frozen=True, eq=True)\n",
        "class ProcessTracksForTrackAutoencoder:\n",
        "  \"\"\"Samples tracks and fills out support_tracks, query_points etc.\n",
        "\n",
        "  TrackAutoencoder format which will be output from this transform:\n",
        "   video: float[\"*B T H W 3\"]\n",
        "   support_tracks: float[\"*B QS T 2\"]\n",
        "   support_tracks_visible: float[\"*B QS T 1\"]\n",
        "   query_points: float[\"*B Q 3\"]\n",
        "  \"\"\"\n",
        "\n",
        "  # note that we do not use query points in the encoding, so it is expected\n",
        "  # that num_support_tracks \u003e\u003e num_target_tracks\n",
        "\n",
        "  num_support_tracks: int\n",
        "  num_target_tracks: int\n",
        "\n",
        "  # If true, assume that everything after the boundary_frame is padding,\n",
        "  # so don't sample any query points after the boundary_frame, and only sample\n",
        "  # target tracks that have at least one visible frame before the boundary.\n",
        "  before_boundary: bool = True\n",
        "  episode_length: int = 150\n",
        "\n",
        "  # Keys.\n",
        "  video_key: str = \"video\"\n",
        "  tracks: str = \"tracks\"  # [time, num_points, 2]\n",
        "  visible_key: str = \"visible\"  # [time, num_points, 1]\n",
        "\n",
        "  def random_map(self, features):\n",
        "\n",
        "    # set tracks xy and compute visibility\n",
        "    tracks_xy = features[self.tracks][..., :2]\n",
        "    tracks_xy = np.asarray(tracks_xy, np.float32)\n",
        "    boundary_frame = features[\"video\"].shape[0]\n",
        "\n",
        "    # visibles already post-processed by compute_point_tracks.py\n",
        "    visibles = np.asarray(features[self.visible_key], np.float32)\n",
        "\n",
        "    # pad to 'episode_length' frames\n",
        "    if self.before_boundary:\n",
        "      # if input video is longer than episode_length, crop to episode_length\n",
        "      if self.episode_length - visibles.shape[0] \u003c 0:\n",
        "        visibles = visibles[: self.episode_length]\n",
        "        tracks_xy = tracks_xy[: self.episode_length]\n",
        "\n",
        "      visibles = np.pad(\n",
        "          visibles,\n",
        "          [[0, self.episode_length - visibles.shape[0]], [0, 0]],\n",
        "          constant_values=0.0,\n",
        "      )\n",
        "      tracks_xy = np.pad(\n",
        "          tracks_xy,\n",
        "          [[0, self.episode_length - tracks_xy.shape[0]], [0, 0], [0, 0]],\n",
        "          constant_values=0.0,\n",
        "      )\n",
        "\n",
        "    # Samples indices for support tracks and target tracks.\n",
        "    num_input_tracks = tracks_xy.shape[1]\n",
        "    idx = np.arange(num_input_tracks)\n",
        "    np.random.shuffle(idx)\n",
        "\n",
        "    assert (\n",
        "        num_input_tracks \u003e= self.num_support_tracks + self.num_target_tracks\n",
        "    ), (\n",
        "        (\n",
        "            f\"num_input_tracks {num_input_tracks} must be greater than\"\n",
        "            f\" num_support_tracks {self.num_support_tracks} + num_target_tracks\"\n",
        "            f\" {self.num_target_tracks}\"\n",
        "        ),\n",
        "    )\n",
        "\n",
        "    idx_support = idx[-self.num_support_tracks :]\n",
        "    idx_target = idx[: self.num_target_tracks]\n",
        "\n",
        "    # Gathers support tracks from `features`.  Features are of shape\n",
        "    # [time, num_points, 2]\n",
        "    support_tracks = tracks_xy[..., idx_support, :]\n",
        "    support_tracks_visible = visibles[..., idx_support]\n",
        "\n",
        "    # Gathers target tracks from `features`.\n",
        "    target_tracks = tracks_xy[..., idx_target, :]\n",
        "    target_tracks_visible = visibles[..., idx_target]\n",
        "\n",
        "    # transpose to [num_points, time, ...]\n",
        "    support_tracks = np.transpose(support_tracks, [1, 0, 2])\n",
        "    support_tracks_visible = np.expand_dims(\n",
        "        np.transpose(support_tracks_visible, [1, 0]), axis=-1\n",
        "    )\n",
        "\n",
        "    # [time, point_id, x/y] -\u003e [point_id, time, x/y]\n",
        "    target_tracks = np.transpose(target_tracks, [1, 0, 2])\n",
        "    target_tracks_visible = np.transpose(target_tracks_visible, [1, 0])\n",
        "\n",
        "    # Sample query points as random visible points\n",
        "    num_target_tracks = target_tracks_visible.shape[0]\n",
        "    num_frames = target_tracks_visible.shape[1]\n",
        "    random_frame = np.zeros(num_target_tracks, dtype=np.int64)\n",
        "\n",
        "    for i in range(num_target_tracks):\n",
        "      visible_indices = np.where(target_tracks_visible[i] \u003e 0)[0]\n",
        "      if len(visible_indices) \u003e 0:\n",
        "        # Choose a random frame index from the visible ones\n",
        "        random_frame[i] = np.random.choice(visible_indices)\n",
        "      else:\n",
        "        # If no frame is visible for a track, default to frame 0\n",
        "        # (or handle as appropriate for your use case)\n",
        "        random_frame[i] = 0\n",
        "\n",
        "    # Create one-hot encoding based on the randomly selected frame for each track\n",
        "    idx = np.eye(num_frames, dtype=np.float32)[\n",
        "        random_frame\n",
        "    ]  # [num_target_tracks, num_frames]\n",
        "\n",
        "    # Use the one-hot index to select the coordinates at the chosen frame\n",
        "    target_queries_xy = np.sum(\n",
        "        target_tracks * idx[..., np.newaxis], axis=1\n",
        "    )  # [num_target_tracks, 2]\n",
        "\n",
        "    # Stack frame index and coordinates: [t, x, y]\n",
        "    target_queries = np.stack(\n",
        "        [\n",
        "            random_frame.astype(np.float32),\n",
        "            target_queries_xy[..., 0],\n",
        "            target_queries_xy[..., 1],\n",
        "        ],\n",
        "        axis=-1,\n",
        "    )  # [num_target_tracks, 3]\n",
        "\n",
        "    # Add channel dimension to target_tracks_visible\n",
        "    target_tracks_visible = np.expand_dims(target_tracks_visible, axis=-1)\n",
        "\n",
        "    # Updates `features` to contain these *new* features and add batch dim.\n",
        "    features_new = {\n",
        "        \"support_tracks\": support_tracks[None, :],\n",
        "        \"support_tracks_visible\": support_tracks_visible[None, :],\n",
        "        \"query_points\": target_queries[None, :],\n",
        "        \"target_points\": target_tracks[None, :],\n",
        "        \"boundary_frame\": np.array([boundary_frame]),\n",
        "        \"target_tracks_visible\": target_tracks_visible[None, :],\n",
        "    }\n",
        "    features.update(features_new)\n",
        "    return features"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "b9rL_bEFbBHk"
      },
      "outputs": [],
      "source": [
        "# @title Run Model\n",
        "\n",
        "# Create model and define forward pass.\n",
        "model = track_autoencoder.TrackAutoEncoder(decoder_scan_chunk_size=32)\n",
        "\n",
        "@jax.jit\n",
        "def forward(params, inputs):\n",
        "  outputs = model.apply({'params': params}, inputs)\n",
        "  return outputs\n",
        "\n",
        "# Create preprocessor\n",
        "preprocessor = ProcessTracksForTrackAutoencoder(\n",
        "    num_support_tracks=2048,\n",
        "    num_target_tracks=2048,\n",
        "    video_key=\"video\",\n",
        "    before_boundary=True,\n",
        ")\n",
        "\n",
        "# Preprocess Batch\n",
        "batch = {\n",
        "    \"video\": video,\n",
        "    \"tracks\": einops.rearrange(\n",
        "        transforms.convert_grid_coordinates(\n",
        "            tracks + 0.5, (width, height), (1, 1)\n",
        "        ),\n",
        "        \"q t c -\u003e t q c\",\n",
        "    ),\n",
        "    \"visible\": einops.rearrange(visibles, \"q t -\u003e t q\"),\n",
        "}\n",
        "\n",
        "batch = preprocessor.random_map(batch)\n",
        "batch.pop(\"tracks\", None)\n",
        "\n",
        "# Run forward pass\n",
        "outputs = forward(params, batch)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vbhZvQE2eOAY"
      },
      "outputs": [],
      "source": [
        "# @title Visualize reconstructed point tracks\n",
        "height, width = video.shape[1:3]\n",
        "\n",
        "reconstructed_tracks = transforms.convert_grid_coordinates(\n",
        "    outputs.tracks[0], (1, 1), (width, height)\n",
        ")\n",
        "\n",
        "support_tracks_vis = transforms.convert_grid_coordinates(\n",
        "    batch['support_tracks'][0], (1, 1), (width, height)\n",
        ")\n",
        "\n",
        "target_tracks_vis = transforms.convert_grid_coordinates(\n",
        "    batch['target_points'][0], (1, 1), (width, height)\n",
        ")\n",
        "\n",
        "reconstructed_visibles = model_utils.postprocess_occlusions(\n",
        "    outputs.visible_logits, outputs.certain_logits\n",
        ")\n",
        "\n",
        "# NOTE: uncomment the lines below to also visualize the support \u0026 target tracks.\n",
        "video_length = video.shape[0]\n",
        "\n",
        "# video_viz = viz_utils.paint_point_track(\n",
        "#     video,\n",
        "#     support_tracks_vis[:, : video.shape[1]],\n",
        "#     batch['support_tracks_visible'][0, :, :video_length],\n",
        "# )\n",
        "# media.show_video(video_viz, fps=10)\n",
        "\n",
        "# video_viz = viz_utils.paint_point_track(\n",
        "#     video,\n",
        "#     target_tracks_vis[:, :video_length],\n",
        "#     batch['target_tracks_visible'][0, :, :video_length],\n",
        "# )\n",
        "# media.show_video(video_viz, fps=10)\n",
        "\n",
        "video_viz = viz_utils.paint_point_track(\n",
        "    video,\n",
        "    reconstructed_tracks[:, :video_length],\n",
        "    batch['target_tracks_visible'][0, :, :video_length],\n",
        ")\n",
        "media.show_video(video_viz, fps=10)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4KXLEMBJkGOL"
      },
      "outputs": [],
      "source": [
        "#@title Compute Metrics\n",
        "\n",
        "# Query from the first frame onward.\n",
        "query_points = np.zeros((\n",
        "    reconstructed_visibles.shape[0],\n",
        "    batch['target_tracks_visible'].shape[1],\n",
        "    1,\n",
        "))\n",
        "\n",
        "# Compute TapVid metrics\n",
        "metrics = evaluation_datasets.compute_tapvid_metrics(\n",
        "    query_points=query_points,\n",
        "    gt_occluded=1 - batch['target_tracks_visible'][..., :video_length, 0],\n",
        "    gt_tracks=target_tracks_vis[None, ..., :video_length, :],\n",
        "    pred_occluded=reconstructed_visibles[..., :video_length, 0],\n",
        "    pred_tracks=reconstructed_tracks[..., :video_length, :],\n",
        "    query_mode='strided',\n",
        "    get_trackwise_metrics=False,\n",
        ")\n",
        "\n",
        "jaccard = np.mean([metrics[f'jaccard_{d}'] for d in [1, 2, 4, 8, 16]])\n",
        "print('Average Jaccard:', jaccard)\n",
        "print('Occlusion Accuracy:', metrics['occlusion_accuracy'].mean())"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "last_runtime": {
        "build_target": "//learning/grp/tools/ml_python/gpu:ml_notebook",
        "kind": "private"
      },
      "private_outputs": true,
      "provenance": [],
      "toc_visible": true
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}