notebooks/Plot_masks_with_sam-batch.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "import os\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import torch\n",
    "\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 plot_mask_and_label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Input and output directory, suffix of the input files to process:\n",
    "\n",
    "input_dir = Path(\"../../data/20240813_data/VID_05_GP_50/\")\n",
    "\n",
    "print(input_dir.exists())\n",
    "\n",
    "image_suffix = \".jpg\"  # Change this to match your image suffix\n",
    "annotation_suffix = \".txt\"  # Change this to match your annotation suffix\n",
    "annotation_folder = \"Annotations\"\n",
    "\n",
    "output_dir = Path(\"../../data/20240813_data_output-sam2/VID_05_GP_50/\")\n",
    "\n",
    "#annotation_file = Path(\n",
    "#    \"../../data/training_data/1/Annotations/14717-training-images-1.csv\"\n",
    "#)\n",
    "#print(annotation_file.exists())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Download the large SAM model's weights from here:\n",
    "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth\n",
    "\n",
    "For all available models see here: https://github.com/facebookresearch/segment-anything?tab=readme-ov-file#model-checkpoints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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 = \"../../SAM2_models/checkpoints/sam2_hiera_large.pt\"\n",
    "model_cfg = \"sam2_hiera_l.yaml\"\n",
    "config_dir = \"../../SAM2_models/configs/\"\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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_file(dir_path, filename):\n",
    "    \n",
    "    in_file = os.path.join(input_dir, dir_path, filename)\n",
    "    out_folder = os.path.join(output_dir, dir_path)\n",
    "\n",
    "    # Take the name of the image without the suffix \n",
    "    basename = os.path.basename(in_file).split(image_suffix)[0]\n",
    "    annotation_file = os.path.join(input_dir, dir_path, annotation_folder, basename + \"_3840x2160\" + annotation_suffix)\n",
    "    print(annotation_file)\n",
    "    print(in_file)\n",
    "\n",
    "    # If output folder does not exist, create one\n",
    "    os.makedirs(out_folder, exist_ok=True)\n",
    " \n",
    "    # Load the annotation file related to the image\n",
    "    df_annotations = pd.read_csv(annotation_file, delimiter=\"\\t\", header=None, names=[\"x\", \"y\", \"label_name\"])\n",
    "    print(df_annotations)\n",
    "    df_points = df_annotations\n",
    "    \n",
    "    # Load the image \n",
    "    test_image = np.array(Image.open(in_file))\n",
    "    \n",
    "    out_files = []\n",
    "    point_prompts = []\n",
    "    species = []\n",
    "    for row in np.arange(0, len(df_annotations), 1):\n",
    "        out_files.append(os.path.join(output_dir, dir_path, basename + \"_\" +str(row) + image_suffix))\n",
    "        \n",
    "        # Get point coordinates; df_annotations.loc[row, \"x\"]; float(y), float(x)\n",
    "        point_prompt = [float(df_annotations.loc[row, \"x\"]),float(df_annotations.loc[row, \"y\"])] #2704 - x\n",
    "            #df_points.loc[row][\"points\"].apply(get_point_coord).to_list())\n",
    "        point_prompts.append(point_prompt)\n",
    "\n",
    "        # Get label names\n",
    "        species.append(df_points.loc[row][\"label_name\"])\n",
    "    \n",
    "\n",
    "    point_prompts = np.array(point_prompts, dtype=np.float32)\n",
    "    prompt_labels = np.ones(len(point_prompts))  # positive prompt\n",
    "    species = np.array(species)\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",
    "    all_species = []\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",
    "        all_species.append(species[i : i + 1])\n",
    "        \n",
    "    \n",
    "    for i in range(len(point_prompts)):\n",
    "        print(i)\n",
    "        plot_mask_and_label(all_masks[i], all_species[i], out_files[i], input_point=point_prompts[i], input_label=prompt_labels[i], image=test_image, saveit = True)\n",
    "\n",
    "    all_masks = None\n",
    "    all_scores = None\n",
    "    all_species = None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Recursively pass through the subfolders - start with input_folder. If you find a file ending with \".jpg\" (which is specified in \"image_suffix\") process the file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_folder_aux():\n",
    "    process_folder(\"\")\n",
    "\n",
    "def process_folder(dir_path):\n",
    "    full_dir = os.path.join(input_dir, dir_path)\n",
    "    full_out_dir = os.path.join(output_dir, dir_path)\n",
    "    \n",
    "    for item in os.listdir(full_dir):\n",
    "        item_path = os.path.join(full_dir, item)\n",
    "        out_item_path = os.path.join(full_out_dir, item)\n",
    "        \n",
    "        if os.path.isdir(item_path):\n",
    "            os.makedirs(out_item_path, exist_ok=True)\n",
    "            process_folder(os.path.join(dir_path, item))\n",
    "        elif item.endswith(image_suffix):\n",
    "            process_file(dir_path, item)\n",
    "\n",
    "\n",
    "# Run the process\n",
    "process_folder_aux()"
   ]
  }
 ],
 "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
}