{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5fa21d44",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright (c) Meta Platforms, Inc. and affiliates.\n",
    "# Lightly adapted from https://github.com/facebookresearch/segment-anything/blob/main/notebooks/automatic_mask_generator_example.ipynb"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7c0041e",
   "metadata": {},
   "source": [
    "# Automatically generating object masks with SAM"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "289bb0b4",
   "metadata": {},
   "source": [
    "Since SAM 2 can efficiently process prompts, masks for the entire image can be generated by sampling a large number of prompts over an image.\n",
    "\n",
    "The class `SAM2AutomaticMaskGenerator` implements this capability. It works by sampling single-point input prompts in a grid over the image, from each of which SAM can predict multiple masks. Then, masks are filtered for quality and deduplicated using non-maximal suppression. Additional options allow for further improvement of mask quality and quantity, such as running prediction on multiple crops of the image or postprocessing masks to remove small disconnected regions and holes."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4290fb06-a63f-4624-a70c-f7c9aae4b5d5",
   "metadata": {},
   "source": [
    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/facebookresearch/sam2/blob/main/notebooks/automatic_mask_generator_example.ipynb\">\n",
    "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
    "</a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0b71431",
   "metadata": {},
   "source": [
    "## Environment Set-up"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47e5a78f",
   "metadata": {},
   "source": [
    "If running locally using jupyter, first install `SAM 2` in your environment using the installation instructions in the repository.\n",
    "\n",
    "If running from Google Colab, set `using_colab=True` below and run the cell. In Colab, be sure to select 'GPU' under 'Edit'->'Notebook Settings'->'Hardware accelerator'. Note that it's recommended to use **A100 or L4 GPUs when running in Colab** (T4 GPUs might also work, but could be slow and might run out of memory in some cases)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a941fd2f-a960-4e5f-916b-a5a385bf3a37",
   "metadata": {},
   "outputs": [],
   "source": [
    "using_colab = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c3e9a446-0d12-4b6f-ba93-03fe7453ace7",
   "metadata": {},
   "outputs": [],
   "source": [
    "if using_colab:\n",
    "    import torch\n",
    "    import torchvision\n",
    "    print(\"PyTorch version:\", torch.__version__)\n",
    "    print(\"Torchvision version:\", torchvision.__version__)\n",
    "    print(\"CUDA is available:\", torch.cuda.is_available())\n",
    "    import sys\n",
    "    !{sys.executable} -m pip install opencv-python matplotlib\n",
    "    !{sys.executable} -m pip install 'git+https://github.com/facebookresearch/sam2.git'\n",
    "\n",
    "    !mkdir -p images\n",
    "    !wget -P images https://raw.githubusercontent.com/facebookresearch/sam2/main/notebooks/images/cars.jpg\n",
    "\n",
    "    !mkdir -p ../checkpoints/\n",
    "    !wget -P ../checkpoints/ https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_large.pt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd2bc687",
   "metadata": {},
   "source": [
    "## Set-up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "effad654-436d-400e-97cc-8bd36141370f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "# if using Apple MPS, fall back to CPU for unsupported ops\n",
    "os.environ[\"PYTORCH_ENABLE_MPS_FALLBACK\"] = \"1\"\n",
    "import numpy as np\n",
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "560725a2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "using device: cuda\n"
     ]
    }
   ],
   "source": [
    "# select the device for computation\n",
    "if torch.cuda.is_available():\n",
    "    device = torch.device(\"cuda\")\n",
    "elif torch.backends.mps.is_available():\n",
    "    device = torch.device(\"mps\")\n",
    "else:\n",
    "    device = torch.device(\"cpu\")\n",
    "print(f\"using device: {device}\")\n",
    "\n",
    "if device.type == \"cuda\":\n",
    "    # use bfloat16 for the entire notebook\n",
    "    torch.autocast(\"cuda\", dtype=torch.bfloat16).__enter__()\n",
    "    # turn on tfloat32 for Ampere GPUs (https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices)\n",
    "    if torch.cuda.get_device_properties(0).major >= 8:\n",
    "        torch.backends.cuda.matmul.allow_tf32 = True\n",
    "        torch.backends.cudnn.allow_tf32 = True\n",
    "elif device.type == \"mps\":\n",
    "    print(\n",
    "        \"\\nSupport for MPS devices is preliminary. SAM 2 is trained with CUDA and might \"\n",
    "        \"give numerically different outputs and sometimes degraded performance on MPS. \"\n",
    "        \"See e.g. https://github.com/pytorch/pytorch/issues/84936 for a discussion.\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "74b6e5f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(3)\n",
    "\n",
    "def show_anns(anns, borders=True):\n",
    "    if len(anns) == 0:\n",
    "        return\n",
    "    sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True)\n",
    "    ax = plt.gca()\n",
    "    ax.set_autoscale_on(False)\n",
    "\n",
    "    img = np.ones((sorted_anns[0]['segmentation'].shape[0], sorted_anns[0]['segmentation'].shape[1], 4))\n",
    "    img[:, :, 3] = 0\n",
    "    for ann in sorted_anns:\n",
    "        m = ann['segmentation']\n",
    "        color_mask = np.concatenate([np.random.random(3), [0.5]])\n",
    "        img[m] = color_mask \n",
    "        if borders:\n",
    "            import cv2\n",
    "            contours, _ = cv2.findContours(m.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE) \n",
    "            # Try to smooth contours\n",
    "            contours = [cv2.approxPolyDP(contour, epsilon=0.01, closed=True) for contour in contours]\n",
    "            cv2.drawContours(img, contours, -1, (0, 0, 1, 0.4), thickness=1) \n",
    "\n",
    "    ax.imshow(img)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27c41445",
   "metadata": {},
   "source": [
    "## Example image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ad354922",
   "metadata": {},
   "outputs": [],
   "source": [
    "image = Image.open('images/cars.jpg')\n",
    "image = np.array(image.convert(\"RGB\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e0ac8c67",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 2813648 characters]",
      "text/plain": [
       "<Figure size 2000x2000 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure(figsize=(20, 20))\n",
    "plt.imshow(image)\n",
    "plt.axis('off')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8c2824a",
   "metadata": {},
   "source": [
    "## Automatic mask generation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9ef74c5",
   "metadata": {},
   "source": [
    "To run automatic mask generation, provide a version of SAM 2 to the `SAM2AutomaticMaskGenerator` class. Set the path below to the SAM 2 checkpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "1848a108",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sam2.build_sam import build_sam2\n",
    "from sam2.automatic_mask_generator import SAM2AutomaticMaskGenerator\n",
    "\n",
    "sam2_checkpoint = \"../checkpoints/sam2.1_hiera_large.pt\"\n",
    "model_cfg = \"configs/sam2.1/sam2.1_hiera_l.yaml\"\n",
    "\n",
    "sam2 = build_sam2(model_cfg, sam2_checkpoint, device=device, apply_postprocessing=False)\n",
    "\n",
    "mask_generator = SAM2AutomaticMaskGenerator(sam2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6b1ea21",
   "metadata": {},
   "source": [
    "To generate masks, just run `generate` on an image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "391771c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "masks = mask_generator.generate(image)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e36a1a39",
   "metadata": {},
   "source": [
    "Mask generation returns a list over masks, where each mask is a dictionary containing various data about the mask. These keys are:\n",
    "* `segmentation` : the mask\n",
    "* `area` : the area of the mask in pixels\n",
    "* `bbox` : the boundary box of the mask in XYWH format\n",
    "* `predicted_iou` : the model's own prediction for the quality of the mask\n",
    "* `point_coords` : the sampled input point that generated this mask\n",
    "* `stability_score` : an additional measure of mask quality\n",
    "* `crop_box` : the crop of the image used to generate this mask in XYWH format"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "4fae8d66",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "59\n",
      "dict_keys(['segmentation', 'area', 'bbox', 'predicted_iou', 'point_coords', 'stability_score', 'crop_box'])\n"
     ]
    }
   ],
   "source": [
    "print(len(masks))\n",
    "print(masks[0].keys())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53009a1f",
   "metadata": {},
   "source": [
    "Show all the masks overlayed on the image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "77ac29c5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 2507256 characters]",
      "text/plain": [
       "<Figure size 2000x2000 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure(figsize=(20, 20))\n",
    "plt.imshow(image)\n",
    "show_anns(masks)\n",
    "plt.axis('off')\n",
    "plt.show() "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00b3d6b2",
   "metadata": {},
   "source": [
    "## Automatic mask generation options"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "183de84e",
   "metadata": {},
   "source": [
    "There are several tunable parameters in automatic mask generation that control how densely points are sampled and what the thresholds are for removing low quality or duplicate masks. Additionally, generation can be automatically run on crops of the image to get improved performance on smaller objects, and post-processing can remove stray pixels and holes. Here is an example configuration that samples more masks:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "68364513",
   "metadata": {},
   "outputs": [],
   "source": [
    "mask_generator_2 = SAM2AutomaticMaskGenerator(\n",
    "    model=sam2,\n",
    "    points_per_side=64,\n",
    "    points_per_batch=128,\n",
    "    pred_iou_thresh=0.7,\n",
    "    stability_score_thresh=0.92,\n",
    "    stability_score_offset=0.7,\n",
    "    crop_n_layers=1,\n",
    "    box_nms_thresh=0.7,\n",
    "    crop_n_points_downscale_factor=2,\n",
    "    min_mask_region_area=25.0,\n",
    "    use_m2m=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "bebcdaf1",
   "metadata": {},
   "outputs": [],
   "source": [
    "masks2 = mask_generator_2.generate(image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "fb702ae3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 2524468 characters]",
      "text/plain": [
       "<Figure size 2000x2000 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure(figsize=(20, 20))\n",
    "plt.imshow(image)\n",
    "show_anns(masks2)\n",
    "plt.axis('off')\n",
    "plt.show() "
   ]
  }
 ],
 "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
