DeepLabSAM

DeepLabSAM converts raw multi-animal video into identity-resolved keypoints from a single text prompt.

What it does

DeepLabSAM is a unified, training-free behavioural analysis pipeline: it streams a video through SAM 3 for open-vocabulary detection, segmentation, and tracking, then hands each tracked instance's mask-gated crop to a pretrained DeepLabCut pose head.

Installation

Not yet published to PyPI — install from source with uv:

git clone https://github.com/juan-cobos/DeepLabSAM
cd DeepLabSAM
uv sync

facebook/sam3 weights are gated — request access on the model page, then set HF_TOKEN via .env (loaded automatically) or the environment:

echo "HF_TOKEN=hf_..." > .env

Already have the weights locally? Point --checkpoint-path/checkpoint_path= at that directory instead of the facebook/sam3 repo id to skip Hugging Face entirely. DeepLabCut Model Zoo weights need no token.

Usage

CLI

uv sync installs a deeplabsam command with three subcommands. Run deeplabsam --help or deeplabsam <command> --help for the full flag list (--help is instant — the heavy model imports are deferred until a run actually starts).

deeplabsam run — single-shot

Detect + track + pose in one pass, straight to an annotated video:

deeplabsam run path/to/video.mp4 --text mice

Every knob of the run is a flag — the SAM 3 build (--image-size, --score-threshold, --num-maskmem), the DLC pose head (--super-animal, --pose-model, --pose-input-size, --device) and the per-run options. Repeat --text/-t for several prompts; animal prompts get pose, others are mask/track only. With multiple prompts (e.g. -t mouse -t object), per-frame class-agnostic NMS (--class-agnostic, --nms-threshold) drops duplicate boxes when overlapping prompts match the same instance:

deeplabsam run path/to/video.mp4 \
    -t mouse -t object \
    --image-size 1008 \
    --max-frames 200 \
    --output-dir outputs \
    --keypoint-threshold 0.3 \
    --nms-threshold 0.5

deeplabsam segment + deeplabsam pose — curated workflow

For datasets you want to hand-curate before pose is fit (drop bad frames, fix a track ID swap, …), split the run into two stages around a standard COCO file:

# 1. Track the video, write frames/ + a COCO annotations.json for curation.
deeplabsam segment path/to/video.mp4 \
    --out-dir outputs/my_video \
    --prompt mouse \
    --max-detections 1

# 2. Curate outputs/my_video/annotations.json + frames/ by hand, then fit pose.
deeplabsam pose \
    --in outputs/my_video/annotations.json \
    --out outputs/my_video/annotations_pose.json

segment writes <out-dir>/frames/*.jpg (raw decoded frames), optionally <out-dir>/viz/*.png (masks/boxes/labels drawn, for a quick sanity check without opening an annotation tool), and <out-dir>/annotations.json (COCO detections + segmentations, the SAM 3 score, and the persistent track_id). Edit annotations.json — or drop frames — with any COCO-aware tool, then run pose, which reads the (possibly curated) file, fits pose per instance mask-gated to its segmentation, and writes the keypoints back in COCO person-keypoints form (plus a _check.mp4 QC video by default).

Python API

from deeplabsam.pipeline import Pipeline

# Build once — loads SAM 3 + the pose head onto the GPU.
# `Pipeline.default` wires a standard SAM 3 wrapper + DLC pose head for you.
pipe = Pipeline.default(super_animal="superanimal_topviewmouse")

# Run per video → writes an annotated mp4 (+ a sibling JSON) and returns its path.
out = pipe.run(
    video_path="recording.mp4",
    text="mice",  # open-vocabulary prompt; str or list[str]
    output_dir="outputs",  # <video-stem>_annotated.mp4 is written here
    max_frames=None,  # cap frames for a quick test
    keypoint_threshold=0.3,  # hide low-confidence keypoints
    export_json=True,  # also write <video-stem>_annotated.json
    save_masks=False,  # dump per-frame masks to outputs/<stem>_masks/*.npz
)
print(out)  # outputs/recording_annotated.mp4

For full control — a custom SAM 3 config, a swapped pose backend, or test fakes — build the components yourself and inject them (this is what Pipeline.default and the CLI do under the hood):

from transformers import Sam3VideoConfig
from deeplabsam.pipeline import Pipeline
from deeplabsam.pose.backends.dlc import DLCPoseHead
from deeplabsam.segment.sam3video import Sam3VideoWrapper

config = Sam3VideoConfig.from_pretrained("facebook/sam3")
config.image_size = 1008
pipe = Pipeline(
    predictor=Sam3VideoWrapper(config=config),
    pose_head=DLCPoseHead(super_animal="superanimal_topviewmouse"),
)

Pipeline.run options

Outputs

• <stem>_annotated.mp4 — masks, boxes, track IDs (#0, #1, …), and pose keypoints drawn per frame.

• <stem>_annotated.json — one row per detection via sv.JSONSink, each tagged with its frame_index and per-instance keypoints ([[x, y, conf], …], one entry per bodypart):

  {"x_min": 176.0, "y_min": 451.0, "x_max": 314.0, "y_max": 493.0,
   "class_id": 0, "confidence": 0.91, "tracker_id": 0,
   "class_name": "<prompt>", "frame_index": 0,
   "keypoints": [[245.1, 470.3, 0.94], [251.7, 462.0, 0.88]]}
  

  class_name is the text prompt that matched the detection (e.g. "mice"), and class_id is its stable index across the prompt set.

• <stem>_masks/frame_<idx>.npz (with --save-masks / save_masks=True) — each frame's masks as a compressed (N, H, W) bool array under the masks key (np.load(path)["masks"]), aligned to detection order, in a per-video folder named after the input video so runs into the same output_dir don't overwrite each other.

Pose models

Select via --super-animal on the CLI, or Pipeline.default(super_animal=...) / DLCPoseHead(super_animal=...) in Python. Weights download from the DeepLabCut Model Zoo on first use.

The pose input crop size is configurable (--pose-input-size, or DLCPoseHead(input_size=...), default 256; must be ÷ the backbone's padding divisor — 32 for HRNet). To run a community DLC model beyond the SuperAnimal zoo, build the head from a project config + snapshot with DLCPoseHead.from_dlc_project(config, snapshot_path).

Acknowledgements

DeepLabSAM stands on other projects — please cite/credit them if you use it:

• SAM 3 (Meta AI) — the open-vocabulary video detection and tracking model behind every mask and track ID.
• DeepLabCut (Mathis Lab) — the SuperAnimal PoseModel and Model Zoo weights that produce the keypoints; the pose backend vendors and runs their PyTorch model code.
• supervision (Roboflow) — the detection/annotation/serialization layer (sv.Detections, the annotators, and sv.JSONSink).
• Hugging Face — transformers serves the SAM 3 model/processor and huggingface-hub distributes the gated weights.