{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "60a7e08e-93c6-4370-9778-3bb102dce78b",
   "metadata": {},
   "source": [
    "Copyright (c) Meta Platforms, Inc. and affiliates."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3081cd8f-f6f9-4a1a-8c36-8a857b0c3b03",
   "metadata": {},
   "source": [
    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/facebookresearch/co-tracker/blob/main/notebooks/demo.ipynb\">\n",
    "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
    "</a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9f3240f-0354-4802-b8b5-9070930fc957",
   "metadata": {},
   "source": [
    "# CoTracker: It is Better to Track Together\n",
    "This is a demo for <a href=\"https://co-tracker.github.io/\">CoTracker</a>, a model that can track any point in a video."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36ff1fd0-572e-47fb-8221-1e73ac17cfd1",
   "metadata": {},
   "source": [
    "<img src=\"https://www.robots.ox.ac.uk/~nikita/storage/cotracker/bmx-bumps.gif\" alt=\"Logo\" width=\"50%\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88c6db31",
   "metadata": {},
   "source": [
    "Don't forget to turn on GPU support if you're running this demo in Colab. \n",
    "\n",
    "**Runtime** -> **Change runtime type** -> **Hardware accelerator** -> **GPU**\n",
    "\n",
    "Let's install dependencies for Colab:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "278876a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "!git clone https://github.com/facebookresearch/co-tracker\n",
    "%cd co-tracker\n",
    "!pip install -e .\n",
    "!pip install opencv-python einops timm matplotlib moviepy flow_vis\n",
    "!mkdir checkpoints\n",
    "%cd checkpoints\n",
    "!wget https://huggingface.co/facebook/cotracker/resolve/main/cotracker2.pth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1745a859-71d4-4ec3-8ef3-027cabe786d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "%cd ..\n",
    "import os\n",
    "import torch\n",
    "\n",
    "from base64 import b64encode\n",
    "from cotracker.utils.visualizer import Visualizer, read_video_from_path\n",
    "from IPython.display import HTML"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7894bd2d-2099-46fa-8286-f0c56298ecd1",
   "metadata": {},
   "source": [
    "Read a video from CO3D:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f1f9ca4d-951e-49d2-8844-91f7bcadfecd",
   "metadata": {},
   "outputs": [],
   "source": [
    "video = read_video_from_path('./assets/apple.mp4')\n",
    "video = torch.from_numpy(video).permute(0, 3, 1, 2)[None].float()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "fb4c2e9d-0e85-4c10-81a2-827d0759bf87",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<video width=\"640\" height=\"480\" autoplay loop controls><source src=\"data:video/mp4;base64,[encoded data, 1626496 characters]\"></video>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def show_video(video_path):\n",
    "    video_file = open(video_path, \"r+b\").read()\n",
    "    video_url = f\"data:video/mp4;base64,{b64encode(video_file).decode()}\"\n",
    "    return HTML(f\"\"\"<video width=\"640\" height=\"480\" autoplay loop controls><source src=\"{video_url}\"></video>\"\"\")\n",
    " \n",
    "show_video(\"./assets/apple.mp4\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f89ae18-54d0-4384-8a79-ca9247f5f31a",
   "metadata": {},
   "source": [
    "Import CoTrackerPredictor and create an instance of it. We'll use this object to estimate tracks:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d59ac40b-bde8-46d4-bd57-4ead939f22ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "from cotracker.predictor import CoTrackerPredictor\n",
    "\n",
    "model = CoTrackerPredictor(\n",
    "    checkpoint=os.path.join(\n",
    "        './checkpoints/cotracker2.pth'\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3f2a4485",
   "metadata": {},
   "outputs": [],
   "source": [
    "if torch.cuda.is_available():\n",
    "    model = model.cuda()\n",
    "    video = video.cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8398155-6dae-4ff0-95f3-dbb52ac70d20",
   "metadata": {},
   "source": [
    "Track points sampled on a regular grid of size 30\\*30 on the first frame:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "17fcaae9-7b3c-474c-977a-cce08a09d580",
   "metadata": {},
   "outputs": [],
   "source": [
    "pred_tracks, pred_visibility = model(video, grid_size=30)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50a58521-a9ba-4f8b-be02-cfdaf79613a2",
   "metadata": {},
   "source": [
    "Visualize and save the result: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "7e793ce0-7b77-46ca-a629-155a6a146000",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IMAGEIO FFMPEG_WRITER WARNING: input image is not divisible by macro_block_size=16, resizing from (1496, 920) to (1504, 928) to ensure video compatibility with most codecs and players. To prevent resizing, make your input image divisible by the macro_block_size or set the macro_block_size to 1 (risking incompatibility).\n",
      "[swscaler @ 0x5e9d040] Warning: data is not aligned! This can lead to a speed loss\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Video saved to ./videos/teaser.mp4\n"
     ]
    }
   ],
   "source": [
    "vis = Visualizer(save_dir='./videos', pad_value=100)\n",
    "vis.visualize(video=video, tracks=pred_tracks, visibility=pred_visibility, filename='teaser');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "2d0733ba-8fe1-4cd4-b963-2085202fba13",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<video width=\"640\" height=\"480\" autoplay loop controls><source src=\"data:video/mp4;base64,[encoded data, 2650272 characters]\"></video>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_video(\"./videos/teaser.mp4\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73d88a5f-057c-4b9f-828d-ee0b97d1e72f",
   "metadata": {},
   "source": [
    "## Tracking manually selected points"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a75bca85-b872-4f4e-be19-ff16f0984037",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "We will start by tracking points queried manually.\n",
    "We define a queried point as: [time, x coord, y coord] \n",
    "\n",
    "So, the code below defines points with different x and y coordinates sampled on frames 0, 10, 20, and 30:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "c6422e7c-8c6f-4269-92c3-245344afe35b",
   "metadata": {},
   "outputs": [],
   "source": [
    "queries = torch.tensor([\n",
    "    [0., 400., 350.],  # point tracked from the first frame\n",
    "    [10., 600., 500.], # frame number 10\n",
    "    [20., 750., 600.], # ...\n",
    "    [30., 900., 200.]\n",
    "])\n",
    "if torch.cuda.is_available():\n",
    "    queries = queries.cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13697a2a-7304-4d18-93be-bfbebf3dc12a",
   "metadata": {},
   "source": [
    "That's what our queried points look like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "d7141079-d7e0-40b3-b031-a28879c4bd6d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 24748 characters]",
      "text/plain": [
       "<Figure size 640x480 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "# Create a list of frame numbers corresponding to each point\n",
    "frame_numbers = queries[:,0].int().tolist()\n",
    "\n",
    "fig, axs = plt.subplots(2, 2)\n",
    "axs = axs.flatten()\n",
    "\n",
    "for i, (query, frame_number) in enumerate(zip(queries, frame_numbers)):\n",
    "    ax = axs[i]\n",
    "    ax.plot(query[1].item(), query[2].item(), 'ro') \n",
    "    \n",
    "    ax.set_title(\"Frame {}\".format(frame_number))\n",
    "    ax.set_xlim(0, video.shape[4])\n",
    "    ax.set_ylim(0, video.shape[3])\n",
    "    ax.invert_yaxis()\n",
    "    \n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aec7693b-9d74-48b3-b612-360290ff1e7a",
   "metadata": {},
   "source": [
    "We pass these points as input to the model and track them:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "09008ca9-6a87-494f-8b05-6370cae6a600",
   "metadata": {},
   "outputs": [],
   "source": [
    "pred_tracks, pred_visibility = model(video, queries=queries[None])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b00d2a35-3daf-482d-b40b-b6d4f548ca40",
   "metadata": {},
   "source": [
    "Finally, we visualize the results with tracks leaving traces from the frame where the tracking starts.\n",
    "Color encodes time:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "01467f8d-667c-4f41-b418-93132584c659",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Video saved to ./videos/queries.mp4\n"
     ]
    }
   ],
   "source": [
    "vis = Visualizer(\n",
    "    save_dir='./videos',\n",
    "    linewidth=6,\n",
    "    mode='cool',\n",
    "    tracks_leave_trace=-1\n",
    ")\n",
    "vis.visualize(\n",
    "    video=video,\n",
    "    tracks=pred_tracks,\n",
    "    visibility=pred_visibility,\n",
    "    filename='queries');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "fe23d210-ed90-49f1-8311-b7e354c7a9f6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<video width=\"640\" height=\"480\" autoplay loop controls><source src=\"data:video/mp4;base64,[encoded data, 1460268 characters]\"></video>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_video(\"./videos/queries.mp4\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "199ecb50-f707-406a-bc64-94bfded827a8",
   "metadata": {},
   "source": [
    "Notice that points queried at frames 10, 20, and 30 are tracked **incorrectly** before the query frame. This is because CoTracker is an online algorithm and only tracks points in one direction. However, we can also run it backward from the queried point to track in both directions. Let's correct this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "b40775f2-6ab0-4bc6-9099-f903935657c9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Video saved to ./videos/queries_backward.mp4\n"
     ]
    }
   ],
   "source": [
    "pred_tracks, pred_visibility = model(video, queries=queries[None], backward_tracking=True)\n",
    "vis.visualize(\n",
    "    video=video,\n",
    "    tracks=pred_tracks,\n",
    "    visibility=pred_visibility,\n",
    "    filename='queries_backward');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "d3120f31-9365-4867-8c85-5638b0708edc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<video width=\"640\" height=\"480\" autoplay loop controls><source src=\"data:video/mp4;base64,[encoded data, 1474588 characters]\"></video>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_video(\"./videos/queries_backward.mp4\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87f2a3b4-a8b3-4aeb-87d2-28f056c624ba",
   "metadata": {},
   "source": [
    "## Points on a regular grid"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9aac679-19f8-4b78-9cc9-d934c6f83b01",
   "metadata": {},
   "source": [
    "### Tracking forward from the frame number x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0aeabca9-cc34-4d0f-8b2d-e6a6f797cb20",
   "metadata": {},
   "source": [
    "Let's now sample points on a regular grid and start tracking from the frame number 20 with a grid of 30\\*30. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "c880f3ca-cf42-4f64-9df6-a0e8de6561dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "grid_size = 30\n",
    "grid_query_frame = 20"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "3cd58820-7b23-469e-9b6d-5fa81257981f",
   "metadata": {},
   "outputs": [],
   "source": [
    "pred_tracks, pred_visibility = model(video, grid_size=grid_size, grid_query_frame=grid_query_frame)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "25a85a1d-dce0-4e6b-9f7a-aaf31ade0600",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IMAGEIO FFMPEG_WRITER WARNING: input image is not divisible by macro_block_size=16, resizing from (1496, 920) to (1504, 928) to ensure video compatibility with most codecs and players. To prevent resizing, make your input image divisible by the macro_block_size or set the macro_block_size to 1 (risking incompatibility).\n",
      "[swscaler @ 0x6363e40] Warning: data is not aligned! This can lead to a speed loss\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Video saved to ./videos/grid_query_20.mp4\n"
     ]
    }
   ],
   "source": [
    "vis = Visualizer(save_dir='./videos', pad_value=100)\n",
    "vis.visualize(\n",
    "    video=video,\n",
    "    tracks=pred_tracks,\n",
    "    visibility=pred_visibility,\n",
    "    filename='grid_query_20',\n",
    "    query_frame=grid_query_frame);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce0fb5b8-d249-4f4e-b59a-51b4f03972c4",
   "metadata": {},
   "source": [
    "Note that tracking starts only from points sampled on a frame in the middle of the video. This is different from the grid in the first example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "f0b01d51-9222-472b-a714-188c38d83ad9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<video width=\"640\" height=\"480\" autoplay loop controls><source src=\"data:video/mp4;base64,[encoded data, 2204176 characters]\"></video>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_video(\"./videos/grid_query_20.mp4\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10baad8f-0cb8-4118-9e69-3fb24575715c",
   "metadata": {},
   "source": [
    "### Tracking forward **and backward** from the frame number x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "506233dc-1fb3-4a3c-b9eb-5cbd5df49128",
   "metadata": {},
   "outputs": [],
   "source": [
    "grid_size = 30\n",
    "grid_query_frame = 20"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "495b5fb4-9050-41fe-be98-d757916d0812",
   "metadata": {},
   "source": [
    "Let's activate backward tracking:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "677cf34e-6c6a-49e3-a21b-f8a4f718f916",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IMAGEIO FFMPEG_WRITER WARNING: input image is not divisible by macro_block_size=16, resizing from (1496, 920) to (1504, 928) to ensure video compatibility with most codecs and players. To prevent resizing, make your input image divisible by the macro_block_size or set the macro_block_size to 1 (risking incompatibility).\n",
      "[swscaler @ 0x5734e40] Warning: data is not aligned! This can lead to a speed loss\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Video saved to ./videos/grid_query_20_backward.mp4\n"
     ]
    }
   ],
   "source": [
    "pred_tracks, pred_visibility = model(video, grid_size=grid_size, grid_query_frame=grid_query_frame, backward_tracking=True)\n",
    "vis.visualize(\n",
    "    video=video,\n",
    "    tracks=pred_tracks,\n",
    "    visibility=pred_visibility,\n",
    "    filename='grid_query_20_backward');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "585a0afa-2cfc-4a07-a6f0-f65924b9ebce",
   "metadata": {},
   "source": [
    "As you can see, we are now tracking points queried in the middle from the first frame:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "c8d64ab0-7e92-4238-8e7d-178652fc409c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<video width=\"640\" height=\"480\" autoplay loop controls><source src=\"data:video/mp4;base64,[encoded data, 2616040 characters]\"></video>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_video(\"./videos/grid_query_20_backward.mp4\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb55fb01-6d8e-4e06-9346-8b2e9ef489c2",
   "metadata": {},
   "source": [
    "## Regular grid + Segmentation mask"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e93a6b0a-b173-46fa-a6d2-1661ae6e6779",
   "metadata": {},
   "source": [
    "Let's now sample points on a grid and filter them with a segmentation mask.\n",
    "This allows us to track points sampled densely on an object because we consume less GPU memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "b759548d-1eda-473e-9c90-99e5d3197e20",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from PIL import Image\n",
    "grid_size = 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "14ae8a8b-fec7-40d1-b6f2-10e333b75db4",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_mask = './assets/apple_mask.png'\n",
    "segm_mask = np.array(Image.open(input_mask))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e3a1520-64bf-4a0d-b6e9-639430e31940",
   "metadata": {},
   "source": [
    "That's a segmentation mask for the first frame:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "4d2efd4e-22df-4833-b9a0-a0763d59ee22",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f4ad3d79e70>"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 48784 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.imshow((segm_mask[...,None]/255.*video[0,0].permute(1,2,0).cpu().numpy()/255.))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "b42dce24-7952-4660-8298-4c362d6913cf",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IMAGEIO FFMPEG_WRITER WARNING: input image is not divisible by macro_block_size=16, resizing from (1496, 920) to (1504, 928) to ensure video compatibility with most codecs and players. To prevent resizing, make your input image divisible by the macro_block_size or set the macro_block_size to 1 (risking incompatibility).\n",
      "[swscaler @ 0x7227e40] Warning: data is not aligned! This can lead to a speed loss\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Video saved to ./videos/segm_grid.mp4\n"
     ]
    }
   ],
   "source": [
    "pred_tracks, pred_visibility = model(video, grid_size=grid_size, segm_mask=torch.from_numpy(segm_mask)[None, None])\n",
    "vis = Visualizer(\n",
    "    save_dir='./videos',\n",
    "    pad_value=100,\n",
    "    linewidth=2,\n",
    ")\n",
    "vis.visualize(\n",
    "    video=video,\n",
    "    tracks=pred_tracks,\n",
    "    visibility=pred_visibility,\n",
    "    filename='segm_grid');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a386308-0d20-4ba3-bbb9-98ea79823a47",
   "metadata": {},
   "source": [
    "We are now only tracking points on the object (and around):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "1810440f-00f4-488a-a174-36be05949e42",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<video width=\"640\" height=\"480\" autoplay loop controls><source src=\"data:video/mp4;base64,[encoded data, 2250072 characters]\"></video>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_video(\"./videos/segm_grid.mp4\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a63e89e4-8890-4e1b-91ec-d5dfa3f93309",
   "metadata": {},
   "source": [
    "## Dense Tracks"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ae764d8-db7c-41c2-a712-1876e7b4372d",
   "metadata": {},
   "source": [
    "### Tracking forward **and backward** from the frame number x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0dde3237-ecad-4c9b-b100-28b1f1b3cbe6",
   "metadata": {},
   "source": [
    "CoTracker also has a mode to track **every pixel** in a video in a **dense** manner but it is much slower than in previous examples. Let's downsample the video in order to make it faster: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "379557d9-80ea-4316-91df-4da215193b41",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 50, 3, 720, 1296])"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "video.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "c6db5cc7-351d-4d9e-9b9d-3a40f05b077a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn.functional as F\n",
    "video_interp = F.interpolate(video[0], [200, 360], mode=\"bilinear\")[None]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ba32cb3-97dc-46f5-b2bd-b93a094dc819",
   "metadata": {},
   "source": [
    "The video now has a much lower resolution:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "0918f246-5556-43b8-9f6d-88013d5a487e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 50, 3, 200, 360])"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "video_interp.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "e4451ae5-c132-4ef2-9329-133b31305db7",
   "metadata": {},
   "outputs": [],
   "source": [
    "grid_query_frame=20"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc7d3a2c-5e87-4c8d-ad10-1f9c6d2ffbed",
   "metadata": {},
   "source": [
    "Again, let's track points in both directions. This will only take a couple of minutes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "3b852606-5229-4abd-b166-496d35da1009",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step 0 / 4\n",
      "step 1 / 4\n",
      "step 2 / 4\n",
      "step 3 / 4\n"
     ]
    }
   ],
   "source": [
    "pred_tracks, pred_visibility = model(video_interp, grid_query_frame=grid_query_frame, backward_tracking=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4143ab14-810e-4e65-93f1-5775957cf4da",
   "metadata": {},
   "source": [
    "Visualization with an optical flow color encoding:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "5394b0ba-1fc7-4843-91d5-6113a6e86bdf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Video saved to ./videos/dense.mp4\n"
     ]
    }
   ],
   "source": [
    "vis = Visualizer(\n",
    "    save_dir='./videos',\n",
    "    pad_value=20,\n",
    "    linewidth=1,\n",
    "    mode='optical_flow'\n",
    ")\n",
    "vis.visualize(\n",
    "    video=video_interp,\n",
    "    tracks=pred_tracks,\n",
    "    visibility=pred_visibility,\n",
    "    filename='dense');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "9113c2ac-4d25-4ef2-8951-71a1c1be74dd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<video width=\"640\" height=\"480\" autoplay loop controls><source src=\"data:video/mp4;base64,[encoded data, 119020 characters]\"></video>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_video(\"./videos/dense.mp4\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95e9bce0-382b-4d18-9316-7f92093ada1d",
   "metadata": {},
   "source": [
    "That's all, now you can use CoTracker in your projects!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "5c3af07c9b62e9e5a46520c50c603ac08cca123627f098bd5ae6d6af0e7d15f9"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
