SAM2Mice

Zero-shot multi-animal segmentation and tracking built on SAM2 + YOLOv11. State-of-the-art performance across diverse experimental settings.

GitHub Source

SAM2Mice source code, checkpoints scripts, notebooks, and documentation are available in Software/SAM2Mice/.

Installation

The demo environment uses python=3.11, torch>=2.6.0, torchvision>=0.21.0, and CUDA 12.4. Create or activate a Python environment first, then install SAM2Mice on a GPU machine:

cd Software/SAM2Mice

# Install PyTorch (CUDA 12.4)
pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 --index-url https://download.pytorch.org/whl/cu124

# Install SAM2 + SAM2Mice (notebook extras included)
pip install -e ".[notebooks]"
python setup.py build_ext --inplace

# Download the SAM2Mice checkpoint
cd checkpoints && bash download_ckpts.sh && cd ..

# For prompt-free segmentation, also download the YOLOv11 detector checkpoint
cd checkpoints_detection && bash download_ckpts.sh && cd ..

The SAM2Mice checkpoint can also be downloaded directly from Google Drive. The YOLOv11 detector (yolo11l_openfield_five_mice.pt) is trained for top-down open-field views.

Quick Start

from SAM2_Mice.segmentation import VideoSegmentationInference

# Configuration
model_cfg = "configs/sam2.1/sam2.1_hiera_b+.yaml"
checkpoint_path = "checkpoints/SAM2_Mice_base_plus.pt"
VIDEO_PATH = "your_video_path"
FRAME_DIR = "your_frame_save_dir"

# 1. video predictor
SAM2_mice_predictor = VideoSegmentationInference(
   model_cfg=model_cfg,
   checkpoint_path=checkpoint_path,
)

# 2. extract frames
SAM2_mice_predictor.extract_frames_before_seg(video_path=VIDEO_PATH, frame_dir=FRAME_DIR)

# 3. annotate the first frame with point prompts
#    use the browser-based GUI annotator (works on headless / remote servers):
from SAM2_Mice.utils import launch_annotator
launch_annotator(frames_dir=FRAME_DIR)   # saves a LabelMe-compatible .json prompt file
#    (alternatively use the LabelMe desktop GUI if you have a local display)

# 4. Run model inference
SAM2_mice_predictor.run(
    video_path=None,
    frames_dir=FRAME_DIR,              # directory of extracted frames (00000.jpg, 00001.jpg, ...)
    prompt_source="manual",             # "manual" or "detection"
    prompt_type="point",                     # "box", "point", or "mask"
    save_dir="results/base",
    fps=20,
)

Results are saved in save_dir:

frames/ — raw frames with mask overlays
segment_masks.pickle — gzip-compressed binary masks per object per frame
segment_video.mp4 — video with masks overlaid

Prompt Modes

Manual Prompts (LabelMe) — Annotate any single frame with polygon masks in LabelMe. SAM2 propagates forward and backward. Supports refinement at any keyframe.

Detection Prompts (prompt-free) — YOLOv11 detects all animals in the first frame and generates box, point, or mask prompts automatically — no manual annotation required. Zero-shot, fully automatic.

Prompt-free segmentation pipeline: detection → auto-prompt → propagation.

You could run this by setting prompt_source="detection" and specifying the path to the YOLOv11 checkpoint

# Example: prompt-free inference
SAM2_mice_predictor.run(
    video_path=None,
    frames_dir=FRAME_DIR,
    prompt_source="detection",
    detection_ckpt_path="checkpoints_detection/yolo11l_openfield_five_mice.pt",
    prompt_type="box",
    save_dir="results/detection_prompt",
    fps=20,
)

Bootstrapping for Long Videos

High-resolution recordings (1224 × 1024, 10,000+ frames) cannot be processed in a single SAM2 pass (~20 GB for 1,000 frames). SAM2Mice uses a bootstrapping strategy:

Split the video into overlapping batches of N + 1 frames.
The (N+1)-th frame is duplicated as the first frame of the next batch.
That overlapping frame carries its segmentation mask forward as a prompt, enabling seamless information transfer across batches.

Bootstrapping strategy — Bootstrapping video predictor: overlapping N+1 frames carry mask prompts forward, enabling arbitrarily long video processing with fixed GPU memory.

from SAM2_Mice.segmentation import BootstrappingVideoSegmentationInference

# Configuration
model_cfg = "configs/sam2.1/sam2.1_hiera_b+.yaml"
checkpoint_path = "checkpoints/SAM2_Mice_base_plus.pt"
VIDEO_PATH = "your_video_path"
FRAME_DIR = "your_frame_save_dir"

# 1. bootstrapping video predictor
SAM2_mice_boots_predictor = BootstrappingVideoSegmentationInference(
   model_cfg=model_cfg,
   checkpoint_path=checkpoint_path,
)

# 2. extract frames
SAM2_mice_boots_predictor.extract_bootstrapping_frames(video_path=VIDEO_PATH, batch_size=batch_size, batch_save_dir=FRAME_DIR )

# 3. label the masks in the first frame with LabelMe (manual prompts)

# 4. Run model inference
SAM2_mice_boots_predictor.run_bootstrapping(
   video_path=None,
   frames_dir=FRAME_DIR,
   frame_interval=batch_size,
   extract_frame=False,                        # extract frames if not already done
   prompt_source="manual",                 # "manual" or "detection"
   prompt_type="point",                        # "box", "point", or "mask"
   save_dir="results/bootstrapping",
   fps=20,
)

Benchmark Performance

SAM2Mice was trained on >4,500 frames with 13,354 expert-labeled masklets spanning diverse experimental conditions: home cages, enriched habitats, varying group sizes (1–5 mice), light/dark, multiple camera viewpoints.

Diverse environments — Generalization across conditions: open field, enriched habitat, social interactions, and varying illumination.

Method	MOTA ↑ (3-mouse)	MOTA ↑ (5-mouse)	Recall ↑	IDF1 ↑
SuperAnimal	0.478 ± 0.040	0.179 ± 0.016	0.724	0.395
DeepLabCut	0.881 ± 0.055	0.703 ± 0.042	0.940	0.940
SAM2Mice (ours)	0.979 ± 0.029	0.958 ± 0.026	0.990	0.990

Google Drive Data Structure

SAM2Mice demo data, benchmark data, checkpoints, and YOLO training resources are hosted in the public SAM2Mice Google Drive folder. The detailed folder listing is also maintained in docs/google_drive_data.md.

SAM2_Mice/
├── benchmark_data/
│   ├── DLC_trained_output/
│   ├── ground_truth_bb_labels/        # 00000.txt ... 01000.txt
│   ├── openfield_five_mouse/          # 00000.jpg ... 00999.jpg + prompt JSON
│   ├── openfield_five_mouse_seg/      # frames, prompts, SAM2Mice labels, masks, video
│   └── SuperAnimal_output/
├── ckpt/
│   ├── SAM2_Mice_base_plus.pt
│   └── yolo11l_openfield_five_mice.pt
├── demo_data/
│   ├── 1_habitat_ultra/               # Notebook 01
│   ├── 2_auto_prompt/                 # Notebook 01
│   ├── 3_bootstrapping/               # Notebook 02
│   ├── 4_dynamic_id/                  # Notebook 03
│   ├── 5_zero_shot_web_dataset/       # Notebook 04
│   └── 6_example_dataset.zip          # Notebook 05
└── YOLO_data/
    ├── images/
    │   ├── train/
    │   ├── train_vis/
    │   └── val/
    └── labels/
        ├── train/
        ├── val/
        ├── train.cache
        └── val.cache

Demo data are organized by notebook: basic segmentation and auto-prompt examples use demo_data/1_habitat_ultra/ and demo_data/2_auto_prompt/; bootstrapping uses demo_data/3_bootstrapping/; dynamic-ID tracking uses demo_data/4_dynamic_id/; public-dataset examples use demo_data/5_zero_shot_web_dataset/; and advanced VOS examples use demo_data/6_example_dataset.zip.

Notebooks

Interactive notebooks are in Software/SAM2Mice/notebooks_SAM2-MICE/ in the GitHub repository:

01_basic_video_segmentation.ipynb — Basic video segmentation with manual prompts and YOLO-generated prompts.
02_bootstrapping_long_video.ipynb — Long-video segmentation via bootstrapped frame batches.
03_auto_tracking_dynamic_ids.ipynb — Automatic multi-animal tracking with dynamic IDs.
04_public_datasets_demo.ipynb — Public mouse-behaviour dataset examples.
05_vos_inference_advanced.ipynb — Advanced video object segmentation utilities.