feat: add automatic segmentation notebook (#9)

* feat: add automatic segmentation nb

* feat: add automatic mask generation nb

* docs: add notebooks in README

.github/README.md

@@ -38,6 +38,9 @@ model = load_model(
 )
 ```
 
+* [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/SauravMaheshkar/samv2/blob/main/examples/notebooks/samv2_prompted_segmentation_with_wandb_tables.ipynb) Example Notebook to run prompted segmentation on images logging predictions as W&B Tables.
+* [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/SauravMaheshkar/samv2/blob/main/examples/notebooks/samv2_automatic_segmentation_with_wandb_tables.ipynb) Example Notebook to run automatic segmentation on images logging predictions as W&B Tables.
+
 ## Citation
 
 ```bibtex

examples/notebooks/samv2_automatic_segmentation_with_wandb_tables.ipynb

@@ -0,0 +1,192 @@
+{
+ "nbformat": 4,
+ "nbformat_minor": 0,
+ "metadata": {
+  "colab": {
+   "private_outputs": true,
+   "provenance": [],
+   "authorship_tag": "ABX9TyN/CnVgbjl/x6mP3u4DTjra",
+   "include_colab_link": true
+  },
+  "kernelspec": {
+   "name": "python3",
+   "display_name": "Python 3"
+  },
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "view-in-github",
+    "colab_type": "text"
+   },
+   "source": "<a href=\"https://colab.research.google.com/github/SauravMaheshkar/blob/main/examples/notebooks/samv2_automatic_segmentation_with_wandb_tables.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "Please add a `W&B` named secret containing your API Key to Colab."
+   ],
+   "metadata": {
+    "id": "E6Ld-l5-xvkr"
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "## 📦 Packages and Basic Setup\n",
+    "---"
+   ],
+   "metadata": {
+    "id": "v5L5aWquYb7d"
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "pMEw72heX9xv"
+   },
+   "outputs": [],
+   "source": [
+    "%%capture\n",
+    "!pip install git+https://github.com/SauravMaheshkar/samv2.git wandb\n",
+    "!wget https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_tiny.pt\n",
+    "\n",
+    "url = \"https://github.com/SauravMaheshkar/SauravMaheshkar/blob/main/assets/text2img/llama_spiderman_coffee.png?raw=true\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "source": [
+    "import os\n",
+    "\n",
+    "import wandb\n",
+    "from google.colab import userdata\n",
+    "\n",
+    "os.environ[\"WANDB_API_KEY\"] = userdata.get(\"W&B\")\n",
+    "\n",
+    "run = wandb.init(project=\"samv2\", entity=\"sauravmaheshkar\")\n",
+    "\n",
+    "columns = [\"image\", \"mask\"]\n",
+    "wandb_table = wandb.Table(columns=columns)"
+   ],
+   "metadata": {
+    "id": "twQBVqZ_yCWp"
+   },
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
+    "import requests\n",
+    "from PIL import Image\n",
+    "\n",
+    "image = Image.open(requests.get(url, stream=True).raw)\n",
+    "image = np.array(image.convert(\"RGB\"))"
+   ],
+   "metadata": {
+    "id": "kQv5hOJYYho3"
+   },
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "source": [
+    "plt.imshow(image)"
+   ],
+   "metadata": {
+    "id": "hpkQZJ2iYjrt"
+   },
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "source": [
+    "from sam2.automatic_mask_generator import SAM2AutomaticMaskGenerator\n",
+    "from sam2.build_sam import build_sam2\n",
+    "from sam2.utils.misc import variant_to_config_mapping\n",
+    "from sam2.utils.visualization import show_masks\n",
+    "\n",
+    "model = build_sam2(\n",
+    "    variant_to_config_mapping[\"tiny\"],\n",
+    "    \"/content/sam2_hiera_tiny.pt\",\n",
+    ")\n",
+    "\n",
+    "mask_generator = SAM2AutomaticMaskGenerator(model)"
+   ],
+   "metadata": {
+    "id": "-aZvjlutYl3C"
+   },
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "source": [
+    "masks = mask_generator.generate(image)"
+   ],
+   "metadata": {
+    "id": "EJV3nmvpY3m0"
+   },
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "source": [
+    "output_mask = show_masks(\n",
+    "    image=image, masks=masks, scores=None, only_best=False, autogenerated_mask=True\n",
+    ")"
+   ],
+   "metadata": {
+    "id": "gVHN3tXZY51B"
+   },
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "source": [
+    "output_mask"
+   ],
+   "metadata": {
+    "id": "LLdov3xtyQtM"
+   },
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "source": [
+    "wandb_table.add_data(wandb.Image(image), wandb.Image(output_mask))"
+   ],
+   "metadata": {
+    "id": "Fv-5saRqyRYC"
+   },
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "source": [
+    "run.log({\"samv2_automatic_mask_generation\": wandb_table})\n",
+    "\n",
+    "wandb.finish()"
+   ],
+   "metadata": {
+    "id": "FtKdl40TyYlg"
+   },
+   "execution_count": null,
+   "outputs": []
+  }
+ ]
+}

examples/notebooks/samv2_prompted_segmentation_with_wandb_tables.ipynb

@@ -9,6 +9,12 @@
    "source": "<a href=\"https://colab.research.google.com/github/SauravMaheshkar/samv2/blob/main/examples/notebooks/samv2_prompted_segmentation_with_wandb_tables.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>",
    "id": "c5d558e73f170b5b"
   },
+  {
+   "metadata": {},
+   "cell_type": "markdown",
+   "source": "Please add a `W&B` named secret containing your API Key to Colab.",
+   "id": "fe8f4fdc77367125"
+  },
   {
    "cell_type": "markdown",
    "source": [