{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3c3b1c46-9f5c-41c1-9101-85db8709ec0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright (c) Meta Platforms, Inc. and affiliates."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e7a0db5-7f04-4845-8b11-684fe6e9f7f2",
   "metadata": {},
   "source": [
    "# Video segmentation with SAM 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73ba7875-35e5-478b-b8ba-4b48e121dec7",
   "metadata": {},
   "source": [
    "This notebook shows how to use SAM 2 for interactive segmentation in videos. It will cover the following:\n",
    "\n",
    "- adding clicks on a frame to get and refine _masklets_ (spatio-temporal masks) \n",
    "- propagating clicks to get _masklets_ throughout the video\n",
    "- segmenting and tracking multiple objects at the same time\n",
    "\n",
    "We use the terms _segment_ or _mask_ to refer to the model prediction for an object on a single frame, and _masklet_ to refer to the spatio-temporal masks across the entire video. \n",
    "\n",
    "If running locally using jupyter, first install `segment-anything-2` in your environment using the [installation instructions](https://github.com/facebookresearch/segment-anything-2#installation) in the repository."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e5318a85-5bf7-4880-b2b3-15e4db24d796",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "08ba49d8-8c22-4eba-a2ab-46eee839287f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# use bfloat16 for the entire notebook\n",
    "torch.autocast(device_type=\"cuda\", dtype=torch.bfloat16).__enter__()\n",
    "\n",
    "if torch.cuda.get_device_properties(0).major >= 8:\n",
    "    # turn on tfloat32 for Ampere GPUs (https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices)\n",
    "    torch.backends.cuda.matmul.allow_tf32 = True\n",
    "    torch.backends.cudnn.allow_tf32 = True"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae8e0779-751f-4224-9b04-ed0f0b406500",
   "metadata": {},
   "source": [
    "### Loading the SAM 2 video predictor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f5f3245e-b4d6-418b-a42a-a67e0b3b5aec",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sam2.build_sam import build_sam2_video_predictor\n",
    "\n",
    "sam2_checkpoint = \"../checkpoints/sam2_hiera_large.pt\"\n",
    "model_cfg = \"sam2_hiera_l.yaml\"\n",
    "\n",
    "predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "1a5320fe-06d7-45b8-b888-ae00799d07fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_mask(mask, ax, obj_id=None, random_color=False):\n",
    "    if random_color:\n",
    "        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)\n",
    "    else:\n",
    "        cmap = plt.get_cmap(\"tab10\")\n",
    "        cmap_idx = 0 if obj_id is None else obj_id\n",
    "        color = np.array([*cmap(cmap_idx)[:3], 0.6])\n",
    "    h, w = mask.shape[-2:]\n",
    "    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)\n",
    "    ax.imshow(mask_image)\n",
    "\n",
    "\n",
    "def show_points(coords, labels, ax, marker_size=200):\n",
    "    pos_points = coords[labels==1]\n",
    "    neg_points = coords[labels==0]\n",
    "    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)\n",
    "    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)   "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f22aa751-b7cd-451e-9ded-fb98bf4bdfad",
   "metadata": {},
   "source": [
    "#### Select an example video"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c4c6af6-e18d-4939-beaf-2bc00f94a724",
   "metadata": {},
   "source": [
    "We assume that the video is stored as a list of JPEG frames with filenames like `<frame_index>.jpg`.\n",
    "\n",
    "For your custom videos, you can extract their JPEG frames using ffmpeg (https://ffmpeg.org/) as follows:\n",
    "```\n",
    "ffmpeg -i <your_video>.mp4 -q:v 2 -start_number 0 <output_dir>/'%05d.jpg'\n",
    "```\n",
    "where `-q:v` generates high-quality JPEG frames and `-start_number 0` asks ffmpeg to start the JPEG file from `00000.jpg`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b94c87ca-fd1a-4011-9609-e8be1cbe3230",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f9c0c03a6b0>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 588792 characters]",
      "text/plain": [
       "<Figure size 1200x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# `video_dir` a directory of JPEG frames with filenames like `<frame_index>.jpg`\n",
    "video_dir = \"./videos/bedroom\"\n",
    "\n",
    "# scan all the JPEG frame names in this directory\n",
    "frame_names = [\n",
    "    p for p in os.listdir(video_dir)\n",
    "    if os.path.splitext(p)[-1] in [\".jpg\", \".jpeg\", \".JPG\", \".JPEG\"]\n",
    "]\n",
    "frame_names.sort(key=lambda p: int(os.path.splitext(p)[0]))\n",
    "\n",
    "# take a look the first video frame\n",
    "frame_idx = 0\n",
    "plt.figure(figsize=(12, 8))\n",
    "plt.title(f\"frame {frame_idx}\")\n",
    "plt.imshow(Image.open(os.path.join(video_dir, frame_names[frame_idx])))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dff46b10-c17a-4a26-8004-8c6d80806b0a",
   "metadata": {},
   "source": [
    "#### Initialize the inference state"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f594ac71-a6b9-461d-af27-500fa1d1a420",
   "metadata": {},
   "source": [
    "SAM 2 requires stateful inference for interactive video segmentation, so we need to initialize an **inference state** on this video.\n",
    "\n",
    "During initialization, it loads all the JPEG frames in `video_path` and stores their pixels in `inference_state` (as shown in the progress bar below)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "8967aed3-eb82-4866-b8df-0f4743255c2c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "frame loading (JPEG): 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [00:09<00:00, 21.07it/s]\n"
     ]
    }
   ],
   "source": [
    "inference_state = predictor.init_state(video_path=video_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edb1f3f6-d74d-4016-934c-8d2a14d1a543",
   "metadata": {},
   "source": [
    "### Example 1: Segment & track one object"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa2d3127-67b2-45d2-9f32-8fe3e10dc5eb",
   "metadata": {},
   "source": [
    "Note: if you have run any previous tracking using this `inference_state`, please reset it first via `reset_state`.\n",
    "\n",
    "(The cell below is just for illustration; it's not needed to call `reset_state` here as this `inference_state` is just freshly initialized above.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d2646a1d-3401-438c-a653-55e0e56b7d9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "predictor.reset_state(inference_state)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26aeb04d-8cba-4f57-95da-6e5a1796003e",
   "metadata": {},
   "source": [
    "#### Step 1: Add a first click on a frame"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "695c7749-b523-4691-aad0-7558c5d1d68c",
   "metadata": {},
   "source": [
    "To get started, let's try to segment the child on the left.\n",
    "\n",
    "Here we make a **positive click** at (x, y) = (210, 350) with label `1`, by sending their coordinates and labels into the `add_new_points` API.\n",
    "\n",
    "Note: label `1` indicates a *positive click (to add a region)* while label `0` indicates a *negative click (to remove a region)*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "3e749bab-0f36-4173-bf8d-0c20cd5214b3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 590372 characters]",
      "text/plain": [
       "<Figure size 1200x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "ann_frame_idx = 0  # the frame index we interact with\n",
    "ann_obj_id = 1  # give a unique id to each object we interact with (it can be any integers)\n",
    "\n",
    "# Let's add a positive click at (x, y) = (210, 350) to get started\n",
    "points = np.array([[210, 350]], dtype=np.float32)\n",
    "# for labels, `1` means positive click and `0` means negative click\n",
    "labels = np.array([1], np.int32)\n",
    "_, out_obj_ids, out_mask_logits = predictor.add_new_points(\n",
    "    inference_state=inference_state,\n",
    "    frame_idx=ann_frame_idx,\n",
    "    obj_id=ann_obj_id,\n",
    "    points=points,\n",
    "    labels=labels,\n",
    ")\n",
    "\n",
    "# show the results on the current (interacted) frame\n",
    "plt.figure(figsize=(12, 8))\n",
    "plt.title(f\"frame {ann_frame_idx}\")\n",
    "plt.imshow(Image.open(os.path.join(video_dir, frame_names[ann_frame_idx])))\n",
    "show_points(points, labels, plt.gca())\n",
    "show_mask((out_mask_logits[0] > 0.0).cpu().numpy(), plt.gca(), obj_id=out_obj_ids[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89457875-93fa-40ed-b6dc-4e1c971a27f9",
   "metadata": {},
   "source": [
    "#### Step 2: Add a second click to refine the prediction"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a75eb21b-1413-452c-827b-a04093c30c78",
   "metadata": {},
   "source": [
    "Hmm, it seems that although we wanted to segment the child on the left, the model predicts the mask for only the shorts -- this can happen since there is ambiguity from a single click about what the target object should be. We can refine the mask on this frame via another positive click on the child's shirt.\n",
    "\n",
    "Here we make a **second positive click** at (x, y) = (250, 220) with label `1` to expand the mask.\n",
    "\n",
    "Note: we need to send **all the clicks and their labels** (i.e. not just the last click) when calling `add_new_points`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "e1ab3ec7-2537-4158-bf98-3d0977d8908d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 587620 characters]",
      "text/plain": [
       "<Figure size 1200x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "ann_frame_idx = 0  # the frame index we interact with\n",
    "ann_obj_id = 1  # give a unique id to each object we interact with (it can be any integers)\n",
    "\n",
    "# Let's add a 2nd positive click at (x, y) = (250, 220) to refine the mask\n",
    "# sending all clicks (and their labels) to `add_new_points`\n",
    "points = np.array([[210, 350], [250, 220]], dtype=np.float32)\n",
    "# for labels, `1` means positive click and `0` means negative click\n",
    "labels = np.array([1, 1], np.int32)\n",
    "_, out_obj_ids, out_mask_logits = predictor.add_new_points(\n",
    "    inference_state=inference_state,\n",
    "    frame_idx=ann_frame_idx,\n",
    "    obj_id=ann_obj_id,\n",
    "    points=points,\n",
    "    labels=labels,\n",
    ")\n",
    "\n",
    "# show the results on the current (interacted) frame\n",
    "plt.figure(figsize=(12, 8))\n",
    "plt.title(f\"frame {ann_frame_idx}\")\n",
    "plt.imshow(Image.open(os.path.join(video_dir, frame_names[ann_frame_idx])))\n",
    "show_points(points, labels, plt.gca())\n",
    "show_mask((out_mask_logits[0] > 0.0).cpu().numpy(), plt.gca(), obj_id=out_obj_ids[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df4ab457-d91d-4ac8-b350-fbcd549fd3fd",
   "metadata": {},
   "source": [
    "With this 2nd refinement click, now we get a segmentation mask of the entire child on frame 0."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f52015ac-1b7b-4c59-bca3-c2b28484cf46",
   "metadata": {},
   "source": [
    "#### Step 3: Propagate the prompts to get the masklet across the video"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30b025bd-cd58-4bfb-9572-c8d2fd0a02ef",
   "metadata": {},
   "source": [
    "To get the masklet throughout the entire video, we propagate the prompts using the `propagate_in_video` API."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "ab45e932-b0d5-4983-9718-6ee77d1ac31b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "propagate in video: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [00:08<00:00, 23.35it/s]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 213436 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 218556 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 214508 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 219452 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 218112 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 213188 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 212284 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 214576 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 219144 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 210760 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 211668 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 226680 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 228300 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 221676 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# run propagation throughout the video and collect the results in a dict\n",
    "video_segments = {}  # video_segments contains the per-frame segmentation results\n",
    "for out_frame_idx, out_obj_ids, out_mask_logits in predictor.propagate_in_video(inference_state):\n",
    "    video_segments[out_frame_idx] = {\n",
    "        out_obj_id: (out_mask_logits[i] > 0.0).cpu().numpy()\n",
    "        for i, out_obj_id in enumerate(out_obj_ids)\n",
    "    }\n",
    "\n",
    "# render the segmentation results every few frames\n",
    "vis_frame_stride = 15\n",
    "plt.close(\"all\")\n",
    "for out_frame_idx in range(0, len(frame_names), vis_frame_stride):\n",
    "    plt.figure(figsize=(6, 4))\n",
    "    plt.title(f\"frame {out_frame_idx}\")\n",
    "    plt.imshow(Image.open(os.path.join(video_dir, frame_names[out_frame_idx])))\n",
    "    for out_obj_id, out_mask in video_segments[out_frame_idx].items():\n",
    "        show_mask(out_mask, plt.gca(), obj_id=out_obj_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e801b70-72df-4a72-b3fe-84f145e5e3f6",
   "metadata": {},
   "source": [
    "#### Step 4: Add new prompts to further refine the masklet"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "478958ab-29b4-4a75-bba4-adb1b03d0a2b",
   "metadata": {},
   "source": [
    "It appears that in the output masklet above, there are some imperfections in boundary details on frame 150.\n",
    "\n",
    "With SAM 2 we can fix the model predictions interactively. We can add a **negative click** at (x, y) = (82, 415) on this frame with label `0` to refine the masklet. Here we call the `add_new_points` API with a different `frame_idx` argument to indicate the frame index we want to refine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "1a572ea9-5b7e-479c-b30c-93c38b121131",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 580760 characters]",
      "text/plain": [
       "<Figure size 1200x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 582004 characters]",
      "text/plain": [
       "<Figure size 1200x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "ann_frame_idx = 150  # further refine some details on this frame\n",
    "ann_obj_id = 1  # give a unique id to the object we interact with (it can be any integers)\n",
    "\n",
    "# show the segment before further refinement\n",
    "plt.figure(figsize=(12, 8))\n",
    "plt.title(f\"frame {ann_frame_idx} -- before refinement\")\n",
    "plt.imshow(Image.open(os.path.join(video_dir, frame_names[ann_frame_idx])))\n",
    "show_mask(video_segments[ann_frame_idx][ann_obj_id], plt.gca(), obj_id=ann_obj_id)\n",
    "\n",
    "# Let's add a negative click on this frame at (x, y) = (82, 415) to refine the segment\n",
    "points = np.array([[82, 415]], dtype=np.float32)\n",
    "# for labels, `1` means positive click and `0` means negative click\n",
    "labels = np.array([0], np.int32)\n",
    "_, _, out_mask_logits = predictor.add_new_points(\n",
    "    inference_state=inference_state,\n",
    "    frame_idx=ann_frame_idx,\n",
    "    obj_id=ann_obj_id,\n",
    "    points=points,\n",
    "    labels=labels,\n",
    ")\n",
    "\n",
    "# show the segment after the further refinement\n",
    "plt.figure(figsize=(12, 8))\n",
    "plt.title(f\"frame {ann_frame_idx} -- after refinement\")\n",
    "plt.imshow(Image.open(os.path.join(video_dir, frame_names[ann_frame_idx])))\n",
    "show_points(points, labels, plt.gca())\n",
    "show_mask((out_mask_logits > 0.0).cpu().numpy(), plt.gca(), obj_id=ann_obj_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50a3950a-acf1-435c-bd64-94297267b5e9",
   "metadata": {},
   "source": [
    "#### Step 5: Propagate the prompts (again) to get the masklet across the video"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1954ecf-c2ec-4f9c-8d10-c4f527a10cd2",
   "metadata": {},
   "source": [
    "Let's get an updated masklet for the entire video. Here we call `propagate_in_video` again to propagate all the prompts after adding the new refinement click above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "baa96690-4a38-4a24-aa17-fd2f4db0e232",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "propagate in video: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [00:08<00:00, 23.98it/s]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 213436 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 218556 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 214508 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 219452 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 218112 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 213188 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 212284 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 214576 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 219144 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 210760 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 211828 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 226676 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 228300 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 221672 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# run propagation throughout the video and collect the results in a dict\n",
    "video_segments = {}  # video_segments contains the per-frame segmentation results\n",
    "for out_frame_idx, out_obj_ids, out_mask_logits in predictor.propagate_in_video(inference_state):\n",
    "    video_segments[out_frame_idx] = {\n",
    "        out_obj_id: (out_mask_logits[i] > 0.0).cpu().numpy()\n",
    "        for i, out_obj_id in enumerate(out_obj_ids)\n",
    "    }\n",
    "\n",
    "# render the segmentation results every few frames\n",
    "vis_frame_stride = 15\n",
    "plt.close(\"all\")\n",
    "for out_frame_idx in range(0, len(frame_names), vis_frame_stride):\n",
    "    plt.figure(figsize=(6, 4))\n",
    "    plt.title(f\"frame {out_frame_idx}\")\n",
    "    plt.imshow(Image.open(os.path.join(video_dir, frame_names[out_frame_idx])))\n",
    "    for out_obj_id, out_mask in video_segments[out_frame_idx].items():\n",
    "        show_mask(out_mask, plt.gca(), obj_id=out_obj_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "607507e3-6a2b-4fd7-944c-2371bdab9d01",
   "metadata": {},
   "source": [
    "The segments now look good on all frames."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da018be8-a4ae-4943-b1ff-702c2b89cb68",
   "metadata": {},
   "source": [
    "### Example 2: Segment multiple objects simultaneously"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dea6c04c-3072-4876-b394-879321a48c4a",
   "metadata": {},
   "source": [
    "Note: if you have run any previous tracking using this `inference_state`, please reset it first via `reset_state`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "29b874c8-9f39-42d3-a667-54a0bd696410",
   "metadata": {},
   "outputs": [],
   "source": [
    "predictor.reset_state(inference_state)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48f3f7e6-4821-468c-84e4-f3a0435c9149",
   "metadata": {},
   "source": [
    "#### Step 1: Add two objects on a frame"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95158714-86d7-48a9-8365-b213f97cc9ca",
   "metadata": {},
   "source": [
    "SAM 2 can also segment and track two or more objects at the same time. One way, of course, is to do them one by one. However, it would be more efficient to batch them together (e.g. so that we can share the image features between objects to reduce computation costs).\n",
    "\n",
    "This time, let's focus on object parts and segment **the shirts of both childen** in this video. Here we add prompts for these two objects and assign each of them a unique object id."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "e22d896d-3cd5-4fa0-9230-f33e217035dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompts = {}  # hold all the clicks we add for visualization"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59d9ac57-b14a-4237-828d-927e422c518b",
   "metadata": {},
   "source": [
    "Add the first object (the left child's shirt) with a **positive click** at (x, y) = (200, 300) on frame 0.\n",
    "\n",
    "We assign it to object id `2` (it can be arbitrary integers, and only needs to be unique for each object to track), which is passed to the `add_new_points` API to distinguish the object we are clicking upon."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "d13432fc-f467-44d8-adfe-3e0c488046b7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 586580 characters]",
      "text/plain": [
       "<Figure size 1200x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "ann_frame_idx = 0  # the frame index we interact with\n",
    "ann_obj_id = 2  # give a unique id to each object we interact with (it can be any integers)\n",
    "\n",
    "# Let's add a positive click at (x, y) = (200, 300) to get started on the first object\n",
    "points = np.array([[200, 300]], dtype=np.float32)\n",
    "# for labels, `1` means positive click and `0` means negative click\n",
    "labels = np.array([1], np.int32)\n",
    "prompts[ann_obj_id] = points, labels\n",
    "_, out_obj_ids, out_mask_logits = predictor.add_new_points(\n",
    "    inference_state=inference_state,\n",
    "    frame_idx=ann_frame_idx,\n",
    "    obj_id=ann_obj_id,\n",
    "    points=points,\n",
    "    labels=labels,\n",
    ")\n",
    "\n",
    "# show the results on the current (interacted) frame\n",
    "plt.figure(figsize=(12, 8))\n",
    "plt.title(f\"frame {ann_frame_idx}\")\n",
    "plt.imshow(Image.open(os.path.join(video_dir, frame_names[ann_frame_idx])))\n",
    "show_points(points, labels, plt.gca())\n",
    "for i, out_obj_id in enumerate(out_obj_ids):\n",
    "    show_points(*prompts[out_obj_id], plt.gca())\n",
    "    show_mask((out_mask_logits[i] > 0.0).cpu().numpy(), plt.gca(), obj_id=out_obj_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1bbbd51b-e1e2-4c36-99ec-1d9a1b49b0cd",
   "metadata": {},
   "source": [
    "Hmm, this time we just want to select the child's shirt, but the model predicts the mask for the entire child. Let's refine the prediction with a **negative click** at (x, y) = (275, 175)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "95ecf61d-662b-4f98-ae62-46557b219842",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 589784 characters]",
      "text/plain": [
       "<Figure size 1200x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# add the first object\n",
    "ann_frame_idx = 0  # the frame index we interact with\n",
    "ann_obj_id = 2  # give a unique id to each object we interact with (it can be any integers)\n",
    "\n",
    "# Let's add a 2nd negative click at (x, y) = (275, 175) to refine the first object\n",
    "# sending all clicks (and their labels) to `add_new_points`\n",
    "points = np.array([[200, 300], [275, 175]], dtype=np.float32)\n",
    "# for labels, `1` means positive click and `0` means negative click\n",
    "labels = np.array([1, 0], np.int32)\n",
    "prompts[ann_obj_id] = points, labels\n",
    "_, out_obj_ids, out_mask_logits = predictor.add_new_points(\n",
    "    inference_state=inference_state,\n",
    "    frame_idx=ann_frame_idx,\n",
    "    obj_id=ann_obj_id,\n",
    "    points=points,\n",
    "    labels=labels,\n",
    ")\n",
    "\n",
    "# show the results on the current (interacted) frame\n",
    "plt.figure(figsize=(12, 8))\n",
    "plt.title(f\"frame {ann_frame_idx}\")\n",
    "plt.imshow(Image.open(os.path.join(video_dir, frame_names[ann_frame_idx])))\n",
    "show_points(points, labels, plt.gca())\n",
    "for i, out_obj_id in enumerate(out_obj_ids):\n",
    "    show_points(*prompts[out_obj_id], plt.gca())\n",
    "    show_mask((out_mask_logits[i] > 0.0).cpu().numpy(), plt.gca(), obj_id=out_obj_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "194718c1-734d-446c-a3ef-361057de2f31",
   "metadata": {},
   "source": [
    "After the 2nd negative click, now we get the left child's shirt as our first object.\n",
    "\n",
    "Let's move on to the second object (the right child's shirt) with a positive click at (x, y) = (400, 150) on frame 0. Here we assign object id `3` to this second object (it can be arbitrary integers, and only needs to be unique for each object to track).\n",
    "\n",
    "Note: when there are multiple objects, the `add_new_points` API will return a list of masks for each object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "86ca1bde-62a4-40e6-98e4-15606441e52f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 590836 characters]",
      "text/plain": [
       "<Figure size 1200x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "ann_frame_idx = 0  # the frame index we interact with\n",
    "ann_obj_id = 3  # give a unique id to each object we interact with (it can be any integers)\n",
    "\n",
    "# Let's now move on to the second object we want to track (giving it object id `3`)\n",
    "# with a positive click at (x, y) = (400, 150)\n",
    "points = np.array([[400, 150]], dtype=np.float32)\n",
    "# for labels, `1` means positive click and `0` means negative click\n",
    "labels = np.array([1], np.int32)\n",
    "prompts[ann_obj_id] = points, labels\n",
    "\n",
    "# `add_new_points` returns masks for all objects added so far on this interacted frame\n",
    "_, out_obj_ids, out_mask_logits = predictor.add_new_points(\n",
    "    inference_state=inference_state,\n",
    "    frame_idx=ann_frame_idx,\n",
    "    obj_id=ann_obj_id,\n",
    "    points=points,\n",
    "    labels=labels,\n",
    ")\n",
    "\n",
    "# show the results on the current (interacted) frame on all objects\n",
    "plt.figure(figsize=(12, 8))\n",
    "plt.title(f\"frame {ann_frame_idx}\")\n",
    "plt.imshow(Image.open(os.path.join(video_dir, frame_names[ann_frame_idx])))\n",
    "show_points(points, labels, plt.gca())\n",
    "for i, out_obj_id in enumerate(out_obj_ids):\n",
    "    show_points(*prompts[out_obj_id], plt.gca())\n",
    "    show_mask((out_mask_logits[i] > 0.0).cpu().numpy(), plt.gca(), obj_id=out_obj_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1f7add8-d577-4597-ae2f-654b8c7b05e0",
   "metadata": {},
   "source": [
    "This time the model predicts the mask of the shirt we want to track in just one click. Nice!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "448733b8-ea8b-4078-995f-b676c3b558ba",
   "metadata": {},
   "source": [
    "#### Step 2: Propagate the prompts to get masklets across the video"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60bd73de-d669-41c8-b6ba-943883f0caa2",
   "metadata": {},
   "source": [
    "Now, we propagate the prompts for both objects to get their masklets throughout the video.\n",
    "\n",
    "Note: when there are multiple objects, the `propagate_in_video` API will return a list of masks for each object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "17737191-d62b-4611-b2c6-6d0418a9ab74",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "propagate in video: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [00:10<00:00, 19.90it/s]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 214396 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 219448 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 214524 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 219696 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 218020 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 213352 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 212460 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 215108 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 220284 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 211476 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 212888 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 228888 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 229488 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 222280 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# run propagation throughout the video and collect the results in a dict\n",
    "video_segments = {}  # video_segments contains the per-frame segmentation results\n",
    "for out_frame_idx, out_obj_ids, out_mask_logits in predictor.propagate_in_video(inference_state):\n",
    "    video_segments[out_frame_idx] = {\n",
    "        out_obj_id: (out_mask_logits[i] > 0.0).cpu().numpy()\n",
    "        for i, out_obj_id in enumerate(out_obj_ids)\n",
    "    }\n",
    "\n",
    "# render the segmentation results every few frames\n",
    "vis_frame_stride = 15\n",
    "plt.close(\"all\")\n",
    "for out_frame_idx in range(0, len(frame_names), vis_frame_stride):\n",
    "    plt.figure(figsize=(6, 4))\n",
    "    plt.title(f\"frame {out_frame_idx}\")\n",
    "    plt.imshow(Image.open(os.path.join(video_dir, frame_names[out_frame_idx])))\n",
    "    for out_obj_id, out_mask in video_segments[out_frame_idx].items():\n",
    "        show_mask(out_mask, plt.gca(), obj_id=out_obj_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18a0b9d7-c78f-432b-afb0-11f2ea5b652a",
   "metadata": {},
   "source": [
    "Looks like both children's shirts are well segmented in this video.\n",
    "\n",
    "Now you can try SAM 2 on your own videos and use cases! "
   ]
  }
 ],
 "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
