MouseLite

Real-time mouse detection, segmentation and pose estimation.

MouseLite wraps fine-tuned RF-DETR models in a video pipeline: it predicts per frame, links predictions across frames with a multi-object tracker, writes an annotated video, and exports the predictions as a COCO dataset you can re-track or analyse later.

Three model kinds are available:

Installation

Requires Python ≥ 3.12.

pip install mouselite

With uv:

uv add mouselite

The Gradio demo is an optional extra:

pip install "mouselite[app]"

From a checkout:

git clone https://github.com/juan-cobos/mouselite
cd mouselite
uv sync --extra app

Model weights are downloaded from the Hugging Face Hub on first use and cached; pass --checkpoint to use your own instead.

CLI

mouselite --help

run — inference on a video

mouselite run video.mp4 --kind keypoints

Writes two things under --output-dir (default output/):

output/
├── video_annotated.mp4          # annotated video
└── video_coco/
    ├── annotations.json         # COCO export (boxes, masks, keypoints, track ids)
    └── images/                  # the frames inference ran on

Common options:

Two animals, pose, half the frames, with a preview window:

mouselite run video.mp4 --kind keypoints --top-k 2 --every 2 --tracker ocsort --show

retrack — re-run tracking without re-running inference

Tracking is usually what you end up tuning, and it is far cheaper than inference. retrack replays an existing COCO export through a different tracker:

mouselite retrack output/video_coco/annotations.json --tracker ocsort

Writes output/video_retracked.mp4 and updates each annotation's track_id in annotations.json in place, so the export always reflects the last tracking pass (-1 for detections the tracker did not confirm). The frame rate is read from the export (run records it, divided by --every), falling back to 30; pass --fps to override.

The two knobs that matter most for mice are how long a track survives an occlusion and how loosely a detection may match it:

Both are forwarded as-is to the tracker class.

mouselite retrack output/video_coco/annotations.json --tracker ocsort --lost-track-buffer 90 --minimum-iou-threshold 0.15

app — Gradio demo

mouselite app                       # needs the [app] extra
mouselite app --no-share --port 7860

Upload a video, pick a model and tracker, run, and retrack the same predictions with a different tracker without paying for inference again.

list-models / list-trackers

$ mouselite list-models
detection: nano, small, medium, large
segmentation: nano, small, medium, large
keypoints

$ mouselite list-trackers
botsort
ocsort
bytetrack
sort
cbiou
mcbyte

Python API

The CLI is a thin wrapper over three pieces: a model, a tracker, and a Pipeline that joins them.

import supervision as sv

from mouselite.models import get_model
from mouselite.pipeline import Pipeline
from mouselite.tracker import get_tracker

model = get_model("keypoints")
fps = sv.VideoInfo.from_video_path("video.mp4").fps
tracker = get_tracker("ocsort", frame_rate=fps)

pipeline = Pipeline(model, tracker, threshold=0.5, top_k=2)
annotated_path = pipeline.run("video.mp4", output_dir="output")

get_model

model = get_model(
    "segmentation",       # "detection", "segmentation" or "keypoints"
    size="large",         # ignored for "keypoints"
    checkpoint=None,      # path to your own weights; otherwise pulled from the Hub
    dtype="float32",
    batch_size=1,
    compile=False,
)

Returns an RF-DETR model already put in inference mode. Any object with a predict(frame, threshold) -> sv.Detections | sv.KeyPoints method and a class_names attribute works in its place — that is the whole MLModel protocol the pipeline depends on.

get_tracker

tracker = get_tracker("bytetrack", frame_rate=30)

Any name from mouselite.tracker.TRACKERS; keyword arguments go straight to the underlying trackers class.

Pipeline

pipeline = Pipeline(
    model,
    tracker,
    threshold=0.5,        # confidence floor
    nms_threshold=0.5,    # NMS IoU threshold
    top_k=None,           # cap on predictions per frame, by confidence
    every=1,              # run inference on 1 of every N frames
)

annotated_path = pipeline.run(
    "video.mp4",
    output_dir="output",
    save_path=None,       # override the annotations.json location
    show=False,           # live preview window
    hud=False,            # FPS overlay
    show_progress=True,
)

run returns the path of the annotated video and writes the COCO export beside it. Keypoint predictions are converted to sv.Detections for tracking — keeping the model's own box rather than a box fitted to the keypoints — and carried through to the export as COCO keypoints/num_keypoints fields.

retrack

from mouselite.tracker import retrack

retracked_path = retrack(
    "output/video_coco/annotations.json",
    "ocsort",
    output_dir="output",
    fps=None,               # recorded by `run`, else 30
    lost_track_buffer=90,   # any further kwargs go to the tracker class
)

Training

The code behind the released models — fine-tuning RF-DETR and the DeepLabCut SuperAnimal baseline, plus the scripts that scored them — lives in training/. It is for reproducing the paper; to just run the models, use the package above.

Acknowledgements

MouseLite is built on work by others:

• RF-DETR — the real-time detection transformer behind every MouseLite model. The detection, segmentation and keypoints-preview architectures are RF-DETR's; MouseLite fine-tunes them on mice.
• supervision — detection and keypoint containers, NMS, annotators, video I/O and the COCO format helpers. It is the vocabulary the whole pipeline is written in.
• trackers — every multi-object tracker MouseLite offers. ByteTrack, BoT-SORT, OC-SORT, SORT, C-BIoU and McByte all come from it unchanged; MouseLite only picks one and hands it detections.
• DeepLabCut — the reference point for markerless animal pose estimation, and the SuperAnimal baseline MouseLite is evaluated against. This project exists because of the problem DeepLabCut defined and the community it built around it.