examples/notebooks/samv2_automatic_segmentation_with_wandb_tables.ipynb

{
 "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": []
  }
 ]
}