Update SAM 2 video notebook (#315)

opengeos · Sep 17, 2024 · 9ba3443 · 9ba3443
1 parent 6f798f8
commit 9ba3443
Showing 1 changed file with 131 additions and 0 deletions.
diff --git a/docs/examples/sam2_video.ipynb b/docs/examples/sam2_video.ipynb
@@ -123,6 +123,15 @@
     "The prompts can be points and boxes. The points are represented as a list of tuples, where each tuple contains the x and y coordinates of the point. The boxes are represented as a list of tuples, where each tuple contains the x, y, width, and height of the box."
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "predictor.show_images()"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -147,6 +156,13 @@
     "predictor.show_prompts(prompts, frame_idx=0)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "![image](https://github.com/user-attachments/assets/2290c685-5d08-4605-8859-ecc3d9986e8f)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -221,6 +237,121 @@
     "    \"blended\", fps=5, output_video=\"segments_blended.mp4\"\n",
     ")"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "![image](https://github.com/user-attachments/assets/111cada6-ada1-4785-8d37-385872bd7a4f)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Segment the objects from a video"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "predictor = SamGeo2(\n",
+    "    model_id=\"sam2-hiera-large\",\n",
+    "    video=True,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "url = \"https://github.com/opengeos/datasets/releases/download/videos/cars.mp4\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "video_path = url\n",
+    "predictor.set_video(video_path)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "predictor.show_images()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "prompts = {\n",
+    "    1: {\n",
+    "        \"points\": [[335, 203]],\n",
+    "        \"labels\": [1],\n",
+    "        \"frame_idx\": 0,\n",
+    "    },\n",
+    "    2: {\n",
+    "        \"points\": [[420, 201]],\n",
+    "        \"labels\": [1],\n",
+    "        \"frame_idx\": 0,\n",
+    "    },\n",
+    "}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "predictor.show_prompts(prompts, frame_idx=0)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "![image](https://github.com/user-attachments/assets/5f55c814-2a86-45a0-b75c-84e2757b88d5)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "predictor.predict_video(prompts)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "predictor.save_video_segments_blended(\"cars\", output_video=\"cars_blended.mp4\", fps=25)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "![image](https://github.com/user-attachments/assets/ef5a0310-26d1-4d64-a880-73b1e5de42c9)"
+   ]
   }
  ],
  "metadata": {