Skip to content

Train a pose model

If there are no tracks and no off-the-shelf model that works on your data, you can use mosaic for the model training process: sample frames, annotate them, train a detector, and hand the weights to a tracker.

This is the most involved path in the section. Read Run a tracker first: training produces the model a tracker then uses.

1. Sample frames to annotate

mosaic run -m dataset.yaml --kind extract-frames \
    --params '{"n_frames": 20, "method": "kmeans"}'

Writes PNGs into media/frames. uniform spreads the sample evenly across each video; kmeans picks diverse frames, so the annotation budget is not spent on near duplicates of the same pose. list writes the frames you name, such as the ones where a tracker fails — see Extract the frames you choose. The method names are matched exactly. Sample from part of the dataset with --entries, --groups or --sequences rather than a key inside --params, which is refused.

2. Annotate

In CVAT, or bring annotations you already have. Three input formats are read: CVAT XML, COCO, and Lightning Pose. Converters live in mosaic/tracking/pose_training/converters/.

3. Train

Op Trains
train-pose A YOLO pose model — keypoints per detected animal
train-points A POLO point-detection model — points without boxes
train-localizer A PyTorch heatmap localizer
train-sleap SLEAP's own trainer
train-litpose Lightning Pose's own trainer
mosaic run -m dataset.yaml --kind train-pose --params @train.json

A trained model is registered as an artifact directory under models/<kind>/<run_id>/, so a later step can name a prior training run instead of carrying a weights path around.

train-pose and train-points each drive an environment mosaic does not install: ultralytics-env/ for pose, polo-env/ for points. Build whichever you need and name it with MOSAIC_ULTRALYTICS_BIN or MOSAIC_POLO_BIN — see installation. One machine can hold both.

The first pose training run fetches its base weights. The default model is the bare asset name yolo11n-pose.pt, which Ultralytics downloads from a GitHub release when the environment does not already hold it. On an air-gapped machine, or a queued job that must not write outside the dataset, pass model as a path to weights that are already there. Point training fetches nothing: polo26n.yaml is package data inside the fork.

Cancelling a training run stops it at the next epoch boundary, leaving last.pt and results.csv complete up to the epoch that finished. Ultralytics cannot be interrupted inside an epoch, so on a long one a cancel is not immediate.

SLEAP's two identity heads need labels that carry identity. multi_class_topdown and multi_class_bottomup classify which animal each detection is, so they train against the tracks in the .slp. Labels with no identity give them no classes, and training then succeeds and produces a model that learned none. Annotations reaching mosaic as COCO carry identity in each annotation's track_id; CVAT point exports carry none, so a set built from one trains the other four heads only.

4. Use it

Either run the model directly over scoped videos, which bridges into tracks/ like a tracker does:

mosaic run -m dataset.yaml --kind infer-pose --params '{"model": "<run_id>"}'

infer-points and infer-localizer are the same shape for the other two model types. infer-pose and infer-points need their environment built and named first. Their runner reads an entry's clips one after another, and an imgstore's chunk files when the chunks are video, so neither needs a join or an export for those. A store whose chunks are not video has to be exported first. Given the flags of the command above, mosaic prerequisites prints the export-store runs that it needs. An entry of several clips that includes a store is refused, because export-joined does not join stores. The refusal names the preprocess command that makes the entry one file. infer-localizer runs in mosaic's process, without an environment to build. It reads clips and stores itself, and needs a join only for clips that differ in frame rate.

Or hand the model to a tracker as its detector:

mosaic track trex -m dataset.yaml --set detect_model=<run_id>

detect_model takes the training run's run_id, exactly as infer-pose does above. A path to a weights file works too, if the weights came from somewhere else.

The difference is identity. Inference detects animals frame by frame; a tracker links those detections across frames into individuals. An inference table's id numbers the detections of each frame from 0, and one id in two frames need not be one animal. If you need to know which animal is which, you want the tracker.

5. Train one model from several datasets

A detector usually gets better with annotations from more than one experiment, and a good one is worth using everywhere. Both want the model to live somewhere that is not any single experiment's dataset. That place is a library: an ordinary mosaic dataset whose job is to hold models.

Save annotations as revisions

Annotations reach a dataset as revisions of a keypoint set, under labels_raw/keypoints/<set>/rev1, rev2, and so on. The Mosaic app saves one when you train from the set or export it. From Python:

from mosaic.core.annotations.projection import write_keypoint_set_revision

saved = write_keypoint_set_revision(
    ds, set_key="openfield", annotations=annotation_set, origin={"note": "after review"}
)
saved.revision   # 3
saved.written    # False if nothing had changed since revision 2

annotation_set is a PoseAnnotationSet from mosaic.core.annotations.pose_annotations: every frame with whether it is finished, every pose with its aliases and left-right keypoint pairs, and every object. Save the whole set, not just what one model should see; choosing that is the next step's job.

Each revision holds that state as annotations.mosaic.json, and beside it a COCO Keypoints export, annotations.coco.json, that CVAT and other COCO tools open. The export holds the finished frames only, one category per pose, and each object's alias under attributes.

A revision is never rewritten, and saving a state that did not change writes nothing. That is what lets a model say exactly which annotations it was trained on.

Build the library and claim what to train on

mosaic init libraries/lab

mosaic sources add -m libraries/lab/dataset.yaml --kind labels --series keypoints \
    --id mice-openfield --path /data/mice/labels_raw/keypoints/openfield \
    --file rev3/annotations.mosaic.json
mosaic sources add -m libraries/lab/dataset.yaml --kind labels --series keypoints \
    --id rats-arena --path /data/rats/labels_raw/keypoints/arena \
    --file rev1/annotations.mosaic.json

mosaic scan -m libraries/lab/dataset.yaml --kind labels

Name the revisions you want with --file. The list is then also the record of what this library has trained on. To train on a later revision, add it with mosaic sources add-files and scan again.

Prepare, then train

mosaic run -m libraries/lab/dataset.yaml --kind prepare-training-data \
    --params '{"sets": [{"set_key": "openfield"}, {"set_key": "arena"}], "target": "yolo-pose"}'

This merges the sets into one training dataset and copies the images in, so the library can be moved, and the datasets it drew from can be archived, without breaking the model. target is yolo-pose, polo, sleap or litpose. Frames from one recording are kept together in one split, which is what makes a validation score honest; "split_by": "frame" turns that off, and its scores are optimistic.

Only frames marked finished are trained on. A set whose objects use more than one pose needs the one to train named, as "pose": "mouse"; objects of other poses are then background. "class_by": "alias" trains each alias as its own class, and refuses a set with an object of that pose that has none, since leaving it unlabelled would teach the model it is background. A box the annotator did not draw is the padded one the annotator saw; "bbox" replaces that padding for this run. A pose with left-right keypoint pairs trains yolo-pose with horizontal flips, swapping each pair; one without pairs trains without flips.

Leave revision out to take the latest one claimed, or pin it with {"set_key": "openfield", "revision": 3}. Either way the run is named by what the revision contains, so the same annotations always give the same prepared dataset.

Then hand the trainer the run id the preparation returned:

mosaic run -m libraries/lab/dataset.yaml --kind train-pose \
    --params '{"data": "prepare-training-data.0.2-<digest>", "epochs": 100}'

Pass the run id rather than a path to its data.yaml. A path is a location, so the same data at two paths would count as two different models.

Use the model from any dataset

mosaic libraries add -m /data/mice/dataset.yaml --id lab --path ../../libraries/lab
mosaic run -m /data/mice/dataset.yaml --kind infer-pose --params '{"model": "<run_id>"}'

Once a dataset links a library, a model there is named by its run id exactly as one trained locally is. Give the link as a relative path when the two datasets move together.

Ask a model what it saw

mosaic models provenance -m /data/mice/dataset.yaml train-pose.0.2-<digest>

Prints the prepared dataset behind the model, the pose that dataset was narrowed to, and each annotation revision behind it, with whatever was recorded when the revision was saved. If a dataset it drew from has since been archived, the answer says which revisions are no longer on disk.

Augmentation is opt-in

Build the training environment with uv sync --python 3.12 --extra augment to add albumentations. Ultralytics then applies Blur, MedianBlur, ToGray and CLAHE at p=0.01 during YOLO and POLO training. Nothing records which way a run went, so decide before you build: the choice belongs to the environment, not to a mosaic extra.

Worked examples

Two notebooks train a model from tracks nobody labelled by hand, on data they download themselves:

  • calms21-pose-training-and-tracking.ipynb turns CalMS21's MARS keypoints into YOLO pose annotations through mosaic's AnnotationSet, trains with train-pose, and hands the result to TRex as a detector.
  • shiners-polo-tracking.ipynb does the same for POLO point labels, where a label row is <class> <radius> <x> <y> and the radius is derived from the data rather than typed in.

Both are bootstrapping rather than annotation -- the labels are as good as the tracker that produced them -- and both say so.