examples/vision_models_examples/Object Segmentation - Mobile SAM.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7136e08b",
   "metadata": {},
   "source": [
    "# Perform object segmentation with pollen-vision"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26ba3176-7a71-427e-be90-b4832bd9b75a",
   "metadata": {},
   "source": [
    "Learn how to perform object segmentation with the pollen-vision library, using the MobileSAM model.\n",
    "\n",
    "MobileSAM is a lighter version of SAM, a segmentation model developed by Meta AI.\n",
    "\n",
    "💡 In this notebook, we assume that you have already checked the notebook dedicated to zero shot object detection as we will also perform object detection here."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12888824-594e-4601-af85-6a6a4bb66fe9",
   "metadata": {},
   "source": [
    "![Object segmentation from Reachy's egocentric view](https://media.githubusercontent.com/media/pollen-robotics/pollen-vision/develop/examples/vision_models_examples/gif/reachy_kitchen_masks.gif)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c279456c-96c5-47a5-bff6-4ad1ede8523d",
   "metadata": {},
   "source": [
    "## A word on SAM and Mobile SAM\n",
    "\n",
    "SAM stands for Segment Anything Model. SAM is a promptable segmentation system with zero-shot generalization to unfamiliar objects and images, without the need for additional training developed by Meta AI. With SAM, you can just give a point to the model to predict the mask for a single object of interest.\n",
    "\n",
    "In 2023, researchers from Kyung Hee University developed MobileSAM, a lighter version of SAM which allows SAM to be run on mobile devices. In pollen-vison we are using the implementation of MobileSAM from its authors. Check the [MobileSAM paper](https://arxiv.org/pdf/2306.14289.pdf), its [GitHub repository](https://github.com/ChaoningZhang/MobileSAM) and the [orginal SAM paper](https://arxiv.org/pdf/2304.02643.pdf) for more information. \n",
    "\n",
    "Credits to Chaoning Zhang et al. from Kyung Hee University and to Alexander Kirillov et al. from Meta AI for developing this and making it open source!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72d45c6e-832e-4d9c-b4b5-94ec31a0607c",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## Setup environment\n",
    "\n",
    "> Note: If you are working locally on your machine and have already installed the library from source, discard the following.\n",
    "\n",
    "We need to first install the pollen-vision library. We will install the library from source, this might take a couple of minutes as there are quite heavy dependencies."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9988eae0-dea1-4540-93a3-b877fcd7bee1",
   "metadata": {},
   "source": [
    "> If you are on Colab and a warning window pops up indicating \"You must restart the runtime in order to use the newly installed versions.\" don't worry. Just press restart session and don't execute the pip install cell again, pollen-vision will already be installed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "875be5b5-3b04-4c64-a21e-60919559d642",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install \"pollen-vision[vision] @ git+https://github.com/pollen-robotics/pollen-vision.git@main\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1bfffc37-0d8f-426e-a7f4-f52e091a9792",
   "metadata": {},
   "source": [
    "## Initialize MobileSAM\n",
    "\n",
    "Let's instanciate a MobileSAM wrapper to prepare the object segmentation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d295a097",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from PIL import Image\n",
    "\n",
    "from pollen_vision.vision_models.object_segmentation import MobileSamWrapper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b741d5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "object_segmentation_wrapper = MobileSamWrapper()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4cd7d2be-7fac-4361-8def-a64a94f5ccd0",
   "metadata": {},
   "source": [
    "## Import example image\n",
    "\n",
    "Here we will import an example image to test the OwlViT wrapper. We will use an image from the [reachy-doing-things image dataset](https://huggingface.co/datasets/pollen-robotics/reachy-doing-things) available on Hugging Face. In this dataset, we captured images from an egocentric view of Reachy doing manipulation tasks while being teleoperated.\n",
    "\n",
    "Feel fry to try the object detection with your own image instead!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec28e782-0d77-45aa-b025-9c717b174aaf",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "dataset = load_dataset(\"pollen-robotics/reachy-doing-things\", split=\"train\")\n",
    "\n",
    "img = dataset[12][\"image\"]\n",
    "img"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93a85134-477f-4e66-82c6-aaaf5c0ee0fa",
   "metadata": {},
   "source": [
    "Let's perform object segmentation on objects Reachy could grasp.\n",
    "\n",
    "## First: object detection\n",
    "\n",
    "To obtain the segmentation, we first need to do object detection in the image to give inputs to MobileSAM. MobileSAM (and SAM as well) takes either a point, a list of points or a bounding box of an object to perform the segmentation. We show in this example how to use bounding boxes of objects as input. So let's get bounding boxes for objects Reachy could grasp, using the OwlViT wrapper."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "010c7f87-2259-4eb6-9417-c085aec29922",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pollen_vision.vision_models.object_detection import OwlVitWrapper\n",
    "\n",
    "object_detection_wrapper = OwlVitWrapper()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b5525cc-f8b0-4e7b-b167-8601de7c7981",
   "metadata": {},
   "source": [
    "If you chose your own image, replace the *candidate_labels* argument with your own list of objects candidates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69fd0625-dbd4-4292-a940-809606ce108e",
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions = object_detection_wrapper.infer(\n",
    "    im=np.array(img), candidate_labels=[\"blue mug\", \"paper cup\", \"kettle\", \"sponge\"], detection_threshold=0.12\n",
    ")\n",
    "predictions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59d62cd9-3573-470d-8095-ebe61dd26b20",
   "metadata": {},
   "source": [
    "We can extract the bounding boxes from the predictions, we will need them as input for the segmentation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af122b20-b969-4007-aaa0-87797f2f68fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pollen_vision.utils import get_bboxes\n",
    "\n",
    "bboxes = get_bboxes(predictions)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93ec81f0-1d19-4874-908e-b14ee19c91f0",
   "metadata": {},
   "source": [
    "N.B.: the format returned for the bounding boxes is *[xmin, ymin, xmax, ymax]*\n",
    "\n",
    "### Visualize object detections\n",
    "\n",
    "You can visualize easliy the predictions of the object detection model with the *Annotator* class from utils."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "317c4886-de01-4e0e-91a5-019b62f0517a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pollen_vision.utils import Annotator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ccaac884-49f1-4633-bdc0-ee56cb0c90c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "annotator = Annotator()\n",
    "\n",
    "img_annotated = annotator.annotate(im=img, detection_predictions=predictions)\n",
    "Image.fromarray(img_annotated)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ded606b-8a03-40fd-8baf-d09a65fed7d0",
   "metadata": {},
   "source": [
    "## At last, the segmentation!\n",
    "\n",
    "Now that we have the bounding boxes for the objects we are interested in, we can use our SAM wrapper defined earlier to obtain the segmentation of each object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59d14585-2477-4850-bd02-b9be3a17bfa0",
   "metadata": {},
   "outputs": [],
   "source": [
    "masks = object_segmentation_wrapper.infer(im=img, bboxes=bboxes)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00a873f7",
   "metadata": {},
   "source": [
    "Note: You could also call `object_segmentation_wrapper.infer(...)` with a list of list of points as input. \n",
    "Here, each list of points would correspond to points of interest for each object. An example such a list would be :\n",
    "```python\n",
    "points = [[[x1, x2], [x3, x4], ...], [[x5, x6], [x7, x8], ...], ...]\n",
    "```\n",
    "\n",
    "You could then call `object_segmentation_wrapper.infer(im=img, points_list=points)`\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1aafa45-9547-4e01-b048-e33c07fe8bc6",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_annotated = annotator.annotate(im=img, detection_predictions=predictions, masks=masks)\n",
    "Image.fromarray(img_annotated)"
   ]
  }
 ],
 "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}