notebooks/BIIGLE_point_to_polygon_with_sam.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# BIIGLE points to polygon with SAM\n",
    "\n",
    "The goal of this notebook is to transform points saved from BIIGLE 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",
    "\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_image,\n",
    "    get_point_coord,\n",
    "    get_polygon,\n",
    "    show_points,\n",
    "    show_res_multi,\n",
    "    save_image_masks,\n",
    ")"
   ]
  },
  {
   "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(\"../results/BIIGLE_results\")\n",
    "print(image_dir.exists())\n",
    "\n",
    "annotation_file = Path(\"../results/BIIGLE_results/2520_csv_image_annotation_report/14717-training-images-1.csv\")\n",
    "print(annotation_file.exists())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's load the annotation file and get the points."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_annotations = pd.read_csv(annotation_file)\n",
    "df_annotations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_points = df_annotations[df_annotations[\"shape_name\"] == \"Point\"]\n",
    "df_points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_points.groupby(\"filename\").count()"
   ]
  },
  {
   "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": 7,
   "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": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "user_id = 2813\n",
    "first_name = \"Nils\"\n",
    "last_name = \"Jacobsen\"\n",
    "last_ann_label_id = int(df_annotations[\"annotation_label_id\"].max())\n",
    "last_label_id = int(df_annotations[\"label_id\"].max())\n",
    "last_annotation_id = int(df_annotations[\"annotation_id\"].max())\n",
    "polygon_id = 3\n",
    "\n",
    "df_new_annotations = df_annotations.copy()"
   ]
  },
  {
   "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/BIIGLE_results\")\n",
    "results_dir.mkdir(exist_ok=True)\n",
    "\n",
    "for image_name, row_locations in df_points.groupby(\"filename\").groups.items():\n",
    "    print(f\"\\nProcessing {image_name}\")\n",
    "    if len(row_locations) < 3:\n",
    "        # the image has only the Laser Points\n",
    "        print(\"the image has only the laser points\")\n",
    "        continue\n",
    "\n",
    "    # load the image\n",
    "    test_image = get_image(image_name, image_dir)\n",
    "    # get point coordinates\n",
    "    point_prompts = df_points.loc[row_locations][\"points\"].apply(get_point_coord).to_list()\n",
    "    point_prompts = np.array(point_prompts, dtype=np.float32)\n",
    "    prompt_labels = np.ones(len(point_prompts))  # positive prompt\n",
    "    # plot image + points\n",
    "    fig, ax = plt.subplots(1, 1, figsize=(9, 8))\n",
    "    ax.imshow(test_image)\n",
    "    show_points(point_prompts, prompt_labels, ax, marker_size=15)\n",
    "    plt.show()\n",
    "\n",
    "    # get the SAM predictions for each point\n",
    "    print(f\"getting predictions for {len(point_prompts)} point prompts...\")\n",
    "    all_masks = []\n",
    "    all_scores = []\n",
    "    predictor.set_image(test_image)\n",
    "    for i in range(len(point_prompts)):\n",
    "        masks, scores, logits = predictor.predict(\n",
    "            point_coords=point_prompts[i: i + 1],\n",
    "            point_labels=prompt_labels[i: i + 1],\n",
    "            multimask_output=False,\n",
    "        )\n",
    "        all_masks.append(masks)\n",
    "        all_scores.append(scores)\n",
    "    # show the results\n",
    "    show_res_multi(all_masks, all_scores, image=test_image, input_box=None)\n",
    "\n",
    "    # save all masks\n",
    "    image_id = int(df_points.loc[row_locations[0], \"image_id\"])\n",
    "    filename = df_points.loc[row_locations[0], \"filename\"]\n",
    "    image_file = image_dir.joinpath(filename)\n",
    "    save_image_masks(all_masks, image_file.stem, results_dir)\n",
    "\n",
    "    # mask to polygons\n",
    "    print(\"getting polygons from masks...\")\n",
    "    for i, masks in enumerate(all_masks):\n",
    "        mask = masks[0]\n",
    "        polygon = get_polygon(mask)\n",
    "        # add an annotation row into the annotation csv table\n",
    "        last_ann_label_id += 1\n",
    "        last_label_id += 1\n",
    "        last_annotation_id += 1\n",
    "        row_idx = len(df_new_annotations.index)\n",
    "        df_new_annotations.loc[row_idx, \"image_id\"] = image_id\n",
    "        df_new_annotations.loc[row_idx, \"filename\"] = filename\n",
    "        df_new_annotations.loc[row_idx, \"user_id\"] = user_id\n",
    "        df_new_annotations.loc[row_idx, \"firstname\"] = first_name\n",
    "        df_new_annotations.loc[row_idx, \"lastname\"] = last_name\n",
    "        df_new_annotations.loc[row_idx, \"annotation_label_id\"] = last_ann_label_id\n",
    "        df_new_annotations.loc[row_idx, \"label_id\"] = last_label_id\n",
    "        df_new_annotations.loc[row_idx, \"annotation_id\"] = last_annotation_id\n",
    "        df_new_annotations.loc[row_idx, \"label_hierarchy\"] = df_points.loc[row_locations[i], \"label_hierarchy\"]\n",
    "        df_new_annotations.loc[row_idx, \"label_name\"] = df_points.loc[row_locations[i], \"label_name\"]\n",
    "        df_new_annotations.loc[row_idx, \"shape_id\"] = polygon_id\n",
    "        df_new_annotations.loc[row_idx, \"shape_name\"] = \"Polygon\"\n",
    "        df_new_annotations.loc[row_idx, \"points\"] = str(polygon.ravel().tolist())\n",
    "\n",
    "\n",
    "print(\"\\n\\nDone!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_new_annotations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_new_annotations.to_csv(results_dir.joinpath(\"new_annotation.csv\"), index=False)"
   ]
  }
 ],
 "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
}