examples/vision_models_examples/Object Detection OWL-ViT.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9052f047-0ed1-435f-9ab8-b37b5c22558f",
   "metadata": {},
   "source": [
    "# Perform object detection with pollen-vision\n",
    "\n",
    "Learn how to perform zero shot object detection with the pollen-vision library, using the OWL-ViT model.\n",
    "\n",
    "This notebook will show you how to use our wrapper for the OWL-ViT object detection model developed by the Google Research lab. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "915566f0-7009-4f03-899c-92491831aad8",
   "metadata": {},
   "source": [
    "![Gif Object detection from Reachy's egocentric view](https://media.githubusercontent.com/media/pollen-robotics/pollen-vision/develop/examples/vision_models_examples/gif/reachy_kitchen_detection.gif)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d11993d6-d4d2-4991-b9aa-1210d6c5141d",
   "metadata": {},
   "source": [
    "## A word on OWL-ViT\n",
    "OWL-ViT stands for Vision Transformer for Open-World Localization. It is a zero shot object detection model, meaning the model is able to perform object detection based on text queries, without needing to retrain the model on any labeled data, as it is the case with traditional Deep Learning object detection models.\n",
    "\n",
    "You can find more information on the model on the dedicated page of the [Hugging Face documentation](https://huggingface.co/docs/transformers/model_doc/owlvit). The implementation of the wrapper actually uses Hugging Face's [transformers library](https://huggingface.co/docs/transformers/index)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16d58e8a",
   "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": "00c3fe9f-6c79-4eaa-91a0-92abbaabd178",
   "metadata": {},
   "source": [
    "> If you are on Colab 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": "0562c4c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install \"pollen-vision[vision] @ git+https://github.com/pollen-robotics/pollen-vision.git@main\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce010bf4",
   "metadata": {},
   "source": [
    "## Use OWL-ViT\n",
    "\n",
    "Let's use the OwlViT wrapper to perform zero shot object detection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d007292-0b1f-41db-899a-6bcbf1a65449",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from PIL import Image\n",
    "\n",
    "from pollen_vision.vision_models.object_detection import OwlVitWrapper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f8e4188-581f-479d-8594-136226a75d83",
   "metadata": {},
   "outputs": [],
   "source": [
    "object_detection_wrapper = OwlVitWrapper()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e449d38-8cf3-42c6-a76d-4c6b83bedce5",
   "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": "7e0551d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "dataset = load_dataset(\"pollen-robotics/reachy-doing-things\", split=\"train\")\n",
    "\n",
    "img = dataset[11][\"image\"]\n",
    "img"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5afac224-c67c-4258-b0b1-a5a02332249e",
   "metadata": {},
   "source": [
    "## Run inference with the model\n",
    "\n",
    "As explained, the OWL-ViT model is a zero shot object detection model and takes text queries as input. The inference is performed with the *infer* method. Just pass as argument a list of the candidate for the object detection that you want to detect. OWL-ViT will only try to detect classes that are in the list.\n",
    "\n",
    "NB: Please note that the image passed as argument for the *infer* method must be a **numpy array object**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "766be514-3453-4b0e-8baa-f2315e39eca1",
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions = object_detection_wrapper.infer(\n",
    "    im=np.array(img),\n",
    "    candidate_labels=[\"kettle\", \"black mug\", \"sink\", \"blue mug\", \"sponge\", \"bag of chips\"],\n",
    "    detection_threshold=0.15,\n",
    ")\n",
    "\n",
    "predictions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3238f284-6683-4d33-9bd2-efb4278197cd",
   "metadata": {},
   "source": [
    "Change the candidates list and check what you can detect!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c02a3b13-fd0c-47cd-96c0-0b07c1301cde",
   "metadata": {},
   "source": [
    "### Visualize detection results\n",
    "\n",
    "You can visualize easliy the predictions of the model with the *Annotator* class from utils."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81e164bd-1aa5-4764-bbf1-dcd7fa995e19",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pollen_vision.utils import Annotator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79cf2daf-b161-41f1-a279-21eab2c5b5a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "annotator = Annotator()\n",
    "\n",
    "img_annotated = annotator.annotate(im=np.array(img), detection_predictions=predictions)\n",
    "Image.fromarray(img_annotated)  # annotator returns a numpy array object"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e3eec07-521f-4301-9757-9f3a89710e8b",
   "metadata": {},
   "source": [
    "## Final notes\n",
    "\n",
    "That's all folks! You can use [this script](https://github.com/pollen-robotics/pollen-vision/blob/99-make-the-notebooks-runnable-on-google-colab/scripts/annotate_video.py) if you want to perform zero shot object detection on a recorded video. The scripts gathers every commands that you saw here in the notebook.\n",
    "\n",
    "Check out the [other notebooks](https://drive.google.com/drive/folders/1Xx42Pk4exkS95iyD-5arHIYQLXyRWTXw?usp=drive_link) if you want to learn how to use other vision models like RAM for image tagging or SAM to perform object segmentation."
   ]
  }
 ],
 "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
}