Automate Surgical Tool Segmentation using SAM2

Installation

SAM 2 needs to be installed first before use. The code requires python>=3.10, as well as torch>=2.3.1 and torchvision>=0.18.1. Please follow the instructions here to install both PyTorch and TorchVision dependencies. You can install SAM 2 on a GPU machine using:

git clone https://github.com/facebookresearch/sam2.git && cd sam2

pip install -e .

If you are installing on Windows, it's strongly recommended to use Windows Subsystem for Linux (WSL) with Ubuntu.

To use the SAM 2 predictor and run the example notebooks, jupyter and matplotlib are required and can be installed by:

pip install -e ".[notebooks]"

Note:

1. It's recommended to create a new Python environment via Anaconda for this installation and install PyTorch 2.3.1 (or higher) via pip following https://pytorch.org/. If you have a PyTorch version lower than 2.3.1 in your current environment, the installation command above will try to upgrade it to the latest PyTorch version using pip.
2. The step above requires compiling a custom CUDA kernel with the nvcc compiler. If it isn't already available on your machine, please install the CUDA toolkits with a version that matches your PyTorch CUDA version.
3. If you see a message like Failed to build the SAM 2 CUDA extension during installation, you can ignore it and still use SAM 2 (some post-processing functionality may be limited, but it doesn't affect the results in most cases).

Please see INSTALL.md for FAQs on potential issues and solutions.

Getting Started

Download Checkpoints (we use sam2 instead of sam2.1 checkpoints here)

First, we need to download a model checkpoint. All the model checkpoints can be downloaded by running:

cd checkpoints && \
./download_ckpts.sh && \
cd ..

or individually from:

• sam2.1_hiera_tiny.pt
• sam2.1_hiera_small.pt
• sam2.1_hiera_base_plus.pt
• sam2.1_hiera_large.pt

(note that these are the improved checkpoints denoted as SAM 2.1; see Model Description for details.)

Then SAM 2 can be used in a few lines as follows for image and video prediction.

Automate annotation using ChatGPT

export OPENAI_API_KEY="your_openai_api_key"
python automate_GPT/generate_sam2_prompts.py --input_dir /mnt/disk0/haoding/cholec80/annotated_data/video01/ws_0/images --input_json /mnt/disk0/haoding/cholec80/annotated_data/video01/ws_0/prompts.json

Prediction with surgical tool

SAM2/sam2/video_infer.py
python SAM2/sam2/video_infer.py --input_dir cholec80/annotated_data/video43/ws_0 --output_dir cholec80/annotated_data/video43/ws_0