notebooks/paparazzi_point_to_polygon_with_sam.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Paparazzi points to polygon with SAM\n",
    "\n",
    "The goal of this notebook is to transform points saved from paparazzi into polygons by running SAM inference on those points."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import torch\n",
    "from PIL import Image\n",
    "\n",
    "from sam2.build_sam import build_sam2\n",
    "from sam2.sam2_image_predictor import SAM2ImagePredictor\n",
    "\n",
    "from utils import (\n",
    "    download_model,\n",
    "    get_polygon,\n",
    "    show_points,\n",
    "    show_res_multi,\n",
    "    save_image_masks,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_image(annotation_name, image_dir):\n",
    "    image_name = \"_\".join(annotation_name.split(\"_\")[:-1])\n",
    "    image_file = image_dir.joinpath(image_name + \".jpg\")\n",
    "    if image_file.exists():\n",
    "        test_image = Image.open(image_file)\n",
    "        return np.array(test_image), image_file\n",
    "    else:\n",
    "        return None, None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following function will download the large SAM2 model's weights from here only if the folder has no model downloaded:\n",
    "\n",
    "https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_large.pt\n",
    "\n",
    "For all available models see here: https://github.com/facebookresearch/sam2?tab=readme-ov-file#download-checkpoints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "download_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from hydra import initialize, core\n",
    "\n",
    "core.global_hydra.GlobalHydra.instance().clear()\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# large sam2: works on gpu > 8g\n",
    "sam2_checkpoint = \"../models/sam2.1_hiera_large.pt\"\n",
    "model_cfg = \"sam2_hiera_l.yaml\"\n",
    "config_dir = \"../models/\"\n",
    "\n",
    "# base sam2: smaller version\n",
    "#sam2_checkpoint = \"../../SAM2_models/checkpoints/sam2_hiera_base_plus.pt\"\n",
    "#model_cfg = \"sam2_hiera_b+.yaml\"\n",
    "\n",
    "with initialize(version_base=None, config_path=config_dir):\n",
    "    sam2_model = build_sam2(model_cfg, sam2_checkpoint, device=device)\n",
    "\n",
    "predictor = SAM2ImagePredictor(sam2_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We should first check that the files have been downloaded and are in the proper folders."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_dir = Path(\"../data/20240813/VID_05_GP_50/\")\n",
    "print(image_dir.exists())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can loop through every image and use the points to prompt SAM2 and generate masks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results_dir = Path(\"../results/paparazzi_results\")\n",
    "results_dir.mkdir(exist_ok=True)\n",
    "\n",
    "annotation_files = image_dir.glob(\"Annotations/*.txt\")\n",
    "print(annotation_files)\n",
    "\n",
    "for file in annotation_files:\n",
    "    print(f\"\\nProcessing {file}...\")\n",
    "    image, image_file = get_image(file.stem, image_dir)\n",
    "    if image is None:\n",
    "        continue\n",
    "\n",
    "    df_annotations = pd.read_csv(\n",
    "        file,\n",
    "        delimiter=\"\\t\",\n",
    "        header=None,\n",
    "        names=[\"x\", \"y\", \"label\"]\n",
    "    )\n",
    "    # get SAM prediction for each row in point annotation file\n",
    "    print(f\"getting prediction for {len(df_annotations)} point prompts...\")\n",
    "    all_prompts = []\n",
    "    all_masks = []\n",
    "    all_scores = []\n",
    "    predictor.set_image(image)\n",
    "    for i, row in df_annotations.iterrows():\n",
    "        prompt = np.array([[row[\"x\"], row[\"y\"]]])\n",
    "\n",
    "        # predict masks using SAM2\n",
    "        masks, scores, _ = predictor.predict(\n",
    "            point_coords=prompt,\n",
    "            point_labels=np.ones(1),\n",
    "            multimask_output=False,\n",
    "        )\n",
    "        sorted_ind = np.argsort(scores)[::-1]\n",
    "        masks = masks[sorted_ind]\n",
    "        scores = scores[sorted_ind]\n",
    "\n",
    "        all_prompts.append(prompt[0])\n",
    "        all_masks.append(masks)\n",
    "        all_scores.append(scores)\n",
    "\n",
    "    # show the results\n",
    "    fig, ax = plt.subplots(1, 1, figsize=(9, 8))\n",
    "    show_res_multi(all_masks, all_scores, image=image, input_box=None, ax=ax)\n",
    "    show_points(np.vstack(all_prompts), np.ones(len(all_prompts)), ax=ax)\n",
    "    plt.show()\n",
    "\n",
    "    # save the masks\n",
    "    print(\"saving predicted masks...\")\n",
    "    save_image_masks(all_masks, image_file.stem, results_dir)\n",
    "    df_result = pd.DataFrame(columns=[\n",
    "        \"image_file\", \"prompt\", \"mask_id\", \"polygon\"\n",
    "    ])\n",
    "\n",
    "    save_dir = results_dir.joinpath(image_file.stem)\n",
    "    save_dir.mkdir(parents=True, exist_ok=True)\n",
    "    for i, mask in enumerate(all_masks):\n",
    "        # mask is 3D: 1, y, x\n",
    "        mask_np = mask[0].astype(np.uint8) * 255\n",
    "        mask_img = Image.fromarray(mask_np)\n",
    "        mask_id = i + 1\n",
    "        mask_img.save(save_dir.joinpath(f\"{mask_id:03d}.png\"))\n",
    "        #\n",
    "        df_result.loc[i, \"image_file\"] = image_file.name\n",
    "        df_result.loc[i, \"prompt\"] = str([row[\"x\"], row[\"y\"]])\n",
    "        df_result.loc[i, \"mask_id\"] = f\"{mask_id:03d}\"\n",
    "        polygon = get_polygon(mask_np)\n",
    "        df_result.loc[i, \"polygon\"] = str(polygon.ravel().tolist())\n",
    "    df_result.to_csv(save_dir.joinpath(\"polygon_masks.csv\"), index=False)\n",
    "\n",
    "print(\"\\n\\nDone!\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "sam2",
   "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.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}