ucsde4e/fishsense-headtail-keypoint

FishSense head/tail keypoint model (simcc_sam, epoch 12)

A top-down snout/fork keypoint model for FishSense Lite. It takes a square crop around one fish and returns two keypoints: 0 = snout (human "head"), 1 = fork (human "tail"). Fish length is the distance between them.

Code and full write-up: UCSD-E4E/2026-10-22_headtail_keypoints (commit 26ad628; see README.md and results/SUMMARY.md).

Files

File	What it is
model.safetensors	EMA weights (state_dict of kpnet.SimCC). Metadata holds epoch and length_val.
simcc_sam_best_length.pt	The original training checkpoint: {"model", "epoch", "length_val"}. SHA-256 8918034a…0110.
kpnet.py, hk.py	Model definition, decode and crop rule, copied from the repo at 26ad628.

Model

• Backbone: timm convnext_tiny.fb_in22k_ft_in1k (last stage, stride 32).
• Input: the crop warped to 512 × 512 (longer side fills 512), ImageNet mean/std, BGR→RGB as in kpnet.to_tensor.
• Head: SimCC (RTMPose design), 2 bins per pixel per axis, decoded as the softmax-weighted mean within ±6 bins of the argmax.
• Crop rule: a square centred on the SAM 3.1 mask box of the fish under the laser dot, side = 1.4 × box diagonal + 128 px, clipped to the frame (hk.crop_from_box).
• Frame: rectified 4014 × 3016 images, decoded with fishsense-core's default RawImage(raw).

Usage

import torch, cv2
from safetensors.torch import load_file
import kpnet  # from this repo; needs timm

model = kpnet.SimCC(pretrained=False)
model.load_state_dict(load_file("model.safetensors"))
model = model.cuda().eval()

# crop = frame[y0:y1, x0:x1] with (x0, y0, x1, y1) = hk.crop_from_box(sam_box, W, H)
cv2.imwrite("crop.jpg", crop)
kps, conf = kpnet.predict_crops(model, ["crop.jpg"], torch.device("cuda"))
# kps[0] = [[snout_x, snout_y], [fork_x, fork_y]] in crop pixels; add (x0, y0) for frame pixels

predict_crops runs under CUDA bfloat16 autocast.

Training and selection

• Data: 19,654 human head/tail clicks drawn without a pre-fill, from 117 dives. Test dives (8 calibrated reef dives, 10 pool fish-model pairs) were held out by dive.
• Selection rule, fixed before any test prediction: the lowest mean |d| on validation-dive SAM-box crops, where d = predicted / human head-tail pixel length − 1. Epoch 12 of simcc_sam was chosen.

Results

Validation (22 dives, 1,739 crops; used for selection):

	Mean |d|	Median |d|	Within 10%	Swapped
SAM 3.1 mask rule	6.2%	2.7%	79%	15%
This model	2.7%	1.1%	95%	0.4%

Test, reef (7 dives, 265 frames; agreement with the human pipeline, not accuracy): mean |d| 3.0% against 5.5% for the SAM rule; paired difference −2.5 points (95% CI −3.2 to −1.8).

Test, pool (10 fish-model dives, tape truth, per-fish p90 MAE): 15.4% against 10.2% for the SAM rule, so worse by 5.2 points (CI +2.2 to +7.8). The frame medians are less biased, but a long over-length tail on two fish models moves the p90. Only 34 training frames come from the pool.

Limitations

• Do not use it in place of the SAM rule for pool fish models without more pool training data or an estimator matched to its error shape.
• The reef result rests on 7 dives.
• It learns the labellers' snout/fork convention. Inter-labeller agreement on real fish is unmeasured.