efish4-track weights
These are the model weights of efish4-track, a pipeline that tracks four fish in one top-down tank setup. For each frame and each fish, it gives the identity (A, B, C, D), the point, the visible mask, the full-body (amodal) mask and a four-point midline. The code downloads this repository at a pinned revision and checks the sha256 of every file (MANIFEST.sha256).
Intended use
Use these weights only with the efish4-track code and only for videos of the setup they were trained on:
- top-down camera, grey background, the same tank;
- 912 x 912 pixels at 50 frames per second;
- the same four fish, with A > B > C > D by body size.
The identity network learned these four individuals. On other fish, another tank or another camera, the letters and masks are not reliable. Such use is out of scope.
Files
| File | Model | Used by stage |
|---|---|---|
identity/namer.pt |
CNN that classifies a fish crop as A, B, C or D | identity |
contact/prompt.pt, contact/prompt2.pt |
two prompted U-Nets that find one fish's centre inside a blob of several fish | contact |
body/body.pt |
U-Net that predicts a fish's visible and full-body masks and four keypoint heatmaps in contact | body_masks |
keypoints/bu_model/best.ckpt, training_config.yaml |
sleap-nn bottom-up model (mouth, head, middle, tail) | keypoints |
keypoints/likelihood.json |
Gaussian-mixture likelihood model for the association of detections to fish | keypoints, recover |
The PyTorch files hold only tensors ({"state_dict": ...}) and load with torch.load(..., weights_only=True). The
sleap-nn checkpoint holds the model state and the training epoch counters only. sleap-nn loads it with pickle, so load
only the pinned revision after the sha256 check, as the code does.
Training data
All models were trained on labelled frames of two training recordings of this tank: one February session and one October session. Only the frames of the labelled "search" blocks were used; held-out blocks of the same recordings were kept for evaluation. idtracker.ai, the first stage, is not part of these weights: it trains its own small network on each new video. No video and no label is published.
Quality
With these weights and the CPU path of the pipeline:
- Labelled February training frames (in-sample): about 99.86 % of fish-frames carry the correct letter.
- Held-out labelled February frames: every fish-frame carried the correct letter, also in contacts and overlaps.
- Nine unlabelled February recordings: 29 flagged or random contact exits were checked by eye; none had a letter swap.
- October recordings: the project's acceptance rules preferred another configuration because of two October contact measures. This release targets February-type recordings.
Limitations
- In 0.1-0.6 % of fish-frames, a visible fish has no visible mask.
- Short C/D letter flashes can occur after some crossings.
- When a non-fish object (for example a hand) is in the tank, the affected frames get
needs_reviewand the object's masks are removed. - idtracker.ai is not deterministic, so two runs on one video can differ slightly.
- The code runs on CPU only. A GPU path is not validated.
Resources
About 6 hours per 20-minute video on 24 CPU cores (measured: 6.2 h on Intel Sapphire Rapids). The peak memory was 93 GB (request 160 GB), and the temporary disk is about 2.5 GB.
License and citation
CC-BY-4.0 (LICENSE). The code is GPL-3.0. Please cite efish4-track (CITATION.cff in the code repository), idtracker.ai (Romero-Ferrero et al., Nature Methods 16, 179-182, 2019) and SLEAP (Pereira et al., Nature Methods 19, 486-495, 2022).