Skip to content

Latest commit

 

History

2 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Zebrafish Re-ID + 4D Tracking Pipeline

A local, human-in-the-loop desktop pipeline for non-invasive individual zebrafish tracking: record a tank of fish, and week-to-week re-identify the same individuals from their natural body/fin patterning — no tagging, no clipping. Built out from 2D per-frame outlines toward full 4D (3D + time) reconstruction, so downstream stages and any machine learning built on this data get real 3D shape/trajectory information, not just a flat 2D silhouette.

The system is a sequence of sub-networks/stages. Each is trained and validated before its output feeds the next:

  1. Whole-body outline + orientation (Stage 1) — dense, user-correctable outline per frame (sortable by low confidence) plus head/tail keypoints, giving which side the fish shows (left/right) and its up/down swim angle.
  2. Multi-view 3D lift (new) — once a session is recorded with a calibrated multi-camera rig, Stage 1's per-camera 2D outlines/keypoints are triangulated into 3D, and a dynamic Gaussian-Splatting reconstruction turns the whole session into a scrubbable 4D volume (see 4D reconstruction below).
  3. Best-frame body-part capture (planned, not yet wired) — per side, the sharpest "straight-on" snapshots of whole body, head, body, dorsal fin, caudal fin, anal fin.
  4. Pattern characterization (planned) — learnable colour/contrast/LUT tooling → a mathematical representation (embedding) of each part's pattern. Colours are learned, never assumed.
  5. Re-identification (planned) — match this session's fish to a prior session by those representations, with confidence + a visual "why", user approval, and drag-and-drop rematch (which also becomes training signal).

Plus a data layer: projects → tanks → fish → recording sessions, arbitrary per-fish values, move fish between tanks, per-fish PDF export, cross-session value export.

Camera setup: 1 camera vs. 2–3 cameras

Both are supported end-to-end; you don't choose a code path, the pipeline picks one per session automatically based on whether that session has a calibrated rig attached:

1 fixed camera 2–3 synchronized cameras (recommended)
Setup cost none — works with existing footage ChArUco calibration once per rig arrangement
Stage 1 (outline/orientation) full 2D result full 2D result, per camera
3D keypoints (triangulate) not available real triangulated 3D
4D reconstruction Deformable 3D Gaussians (monocular; depth is appearance-inferred, geometrically approximate) Dynamic 3D Gaussians (multi-view; real triangulated geometry + persistent per-point 3D tracks)
Best for getting started now / footage you already have anything downstream that needs real 3D — ML features, accurate morphology over time

export-4dgs reads whether the session's sessions.rig_id is set and prints which path it picked; see 4D reconstruction below for the multi-camera walkthrough, or just skip straight to stage1-propose on a single-camera session for the lower-setup-cost path.

Design principles (from the project brief)

  • Reuse open source; cite it. See Open-source components.
  • No predetermined conclusions. Models learn from your labelled data; nothing about colour, shape, or pattern is hardcoded.
  • Human in the loop everywhere. Every stage surfaces its lowest-confidence outputs for you to correct, and corrections are used for training (active learning).

Status

Phased build.

  • Phase 1a — foundation (done): project/data store + video ingestion & frame sampling. Standard library + ffmpeg only.
  • Phase 1b — Stage 1 outline + orientation (implemented, needs the GPU env to run): motion localization → SAM 2 outline (+ confidence) → head/tail → left/right side & up/down angle, stored as reviewable annotations, plus a napari review window.
  • Multi-camera calibration + 3D triangulation (implemented, pure OpenCV/numpy): ChArUco rig calibration, DLT multi-view triangulation of Stage-1 keypoints. Unit-tested against synthetic cameras (no real rig needed to verify the math).
  • 4D Gaussian-Splatting reconstruction (dataset export implemented; training/viewing depend on external repos + the GPU env): see below.
  • Stage 2–4 (body-part capture, pattern embeddings, re-ID matching): not yet built — next up once Stage 1 + the 4D lift have real reviewed data to work from.

Using Stage 1 (single- or multi-camera)

pip install -e ".[stage1]"   # after the base install below

python -m zebrafish_pipeline.cli --root ./_projects import-video --project P --video "<clip>" --session Week1
python -m zebrafish_pipeline.cli --root ./_projects sample        --project P --video-id 1 --count 60
python -m zebrafish_pipeline.cli --root ./_projects stage1-propose --project P --video-id 1 \
    --checkpoint models/sam2/sam2.1_hiera_tiny.pt --model-cfg sam2.1_hiera_t.yaml
python -m zebrafish_pipeline.cli --root ./_projects stage1-review  --project P --video-id 1

In the review window: drag outline vertices to fix the shape, move the green (head) / red (tail) dots, delete any non-fish outline (e.g. a reflection, press d), then Approve & Next (press n). Approved frames become Stage-1 training data.

4D (3D+time) reconstruction

Requires a calibrated multi-camera rig (2–3 synchronized cameras around the tank — one fixed camera alone cannot give real depth). Print a ChArUco board, then:

pip install -e ".[calibration,reconstruction4d]"

# 1. capture: <intrinsics-dir>/<camera>/*.jpg (board waved around each camera alone)
#             <synced-dir>/<camera>/*.jpg     (board held still, seen by every camera at once)
python -m zebrafish_pipeline.cli --root ./_projects calibrate-rig --project P \
    --rig front-left-top --cameras front,left,top \
    --intrinsics-dir calib/intrinsics --synced-dir calib/synced

# 2. record + ingest a session tagged with that rig (one --import-video call per camera)
python -m zebrafish_pipeline.cli --root ./_projects import-video --project P \
    --video front.mp4 --session Week1 --camera front --rig front-left-top --cameras front,left,top
python -m zebrafish_pipeline.cli --root ./_projects import-video --project P \
    --video left.mp4  --session Week1 --camera left
python -m zebrafish_pipeline.cli --root ./_projects import-video --project P \
    --video top.mp4   --session Week1 --camera top

# 3. sample synced frame groups, run Stage 1 per camera, approve in review, then:
python -m zebrafish_pipeline.cli --root ./_projects sample-synced --project P --session-id 1 --count 200
python -m zebrafish_pipeline.cli --root ./_projects triangulate   --project P --session-id 1
python -m zebrafish_pipeline.cli --root ./_projects export-4dgs   --project P --session-id 1

# 4. train (needs the external repo cloned -- see reconstruction3d/gaussian4d.py) and view:
python -m zebrafish_pipeline.cli --root ./_projects train-4dgs --project P --session-id 1 \
    --external-repo-dir path/to/Dynamic3DGaussians

export-4dgs writes a dataset in the format the external Gaussian-splatting repo expects into --scratch-root (see Setup below for why); train-4dgs then shells out to that repo's own train.py (see reconstruction3d/gaussian4d.py for exactly which repo, why, and the citation). Only have one camera so far? export-4dgs/train-4dgs automatically fall back to the monocular method for sessions with no rig — see Camera setup above.

Both of these are compute-heavy and are meant to run on a cluster, not your desktop — see Running at scale on Compute Canada below.

Setup

python -m venv .venv
.\.venv\Scripts\python.exe -m pip install --upgrade pip
.\.venv\Scripts\python.exe -m pip install -e .
# Stage 1 (GPU): CUDA torch + SAM 2 + napari review UI
.\.venv\Scripts\python.exe -m pip install torch torchvision --index-url https://download.pytorch.org/whl/cu124
.\.venv\Scripts\python.exe -m pip install -e ".[stage1]"
.\.venv\Scripts\python.exe -m pip install "git+https://github.com/facebookresearch/sam2.git"
# SAM 2 tiny checkpoint -> models/sam2/sam2.1_hiera_tiny.pt   (see models/ note)

# Multi-camera calibration + 3D triangulation (CPU only)
.\.venv\Scripts\python.exe -m pip install -e ".[calibration]"

# 4D reconstruction + viewer
.\.venv\Scripts\python.exe -m pip install -e ".[reconstruction4d]"
# then clone whichever Gaussian-splat repo you need (see reconstruction3d/gaussian4d.py)

Run the test suite (no GPU, no ffmpeg, no camera rig required — the calibration/ triangulation tests use synthetic cameras):

pip install -e ".[dev]"
pytest

Every command above also takes --scratch-root alongside --root: --root is persistent project storage (videos, project.db, calibrations — defaults to ./_projects), --scratch-root is disposable working storage for heavy per-frame intermediates (defaults to the same place as --root if you don't set it separately). On a laptop that distinction usually doesn't matter; on a shared cluster filesystem it does — see the next section.

Running at scale on Compute Canada

This pipeline's actual data-analysis/ML runs happen on the Digital Research Alliance of Canada (Compute Canada) rather than locally. hpc/ has everything for that:

  • hpc/setup_env.sh — one-time environment bootstrap (venv, wheelhouse installs, cloning + building the SAM 2 / Gaussian-splat external repos).
  • hpc/data_layout.mdwhere to put data: $PROJECT (durable — videos, project.db, calibrations, finished reconstructions) vs. $SCRATCH (disposable — sampled frames, in-progress training) vs. $SLURM_TMPDIR (per-job node-local staging for the many-small-files parts, which matters a lot on a shared parallel filesystem).
  • hpc/jobs/ — one sbatch script per stage (calibrate_rig.sh, sample_synced.sh, stage1_propose.sh, triangulate.sh, export_4dgs.sh, train_4dgs.sh), each a thin wrapper around the same CLI commands shown above.
  • hpc/README.md — how to submit them.

Interactive/human steps (stage1-review, view4d) still run on your own desktop; the GPU/CPU-heavy, unattended steps run on the cluster, with the two sides synced via the project's project.db + --root.

Open-source components

Stage Tool License Purpose
1 SAM 2 (Meta) Apache-2.0 Interactive video segmentation of the fish outline, with per-frame confidence and click-correction
1 SLEAP / DeepLabCut BSD / LGPL Head/tail (and fin) keypoints for orientation (a geometric outline-taper heuristic bootstraps this before enough labels exist — see stage1/orientation.py)
1 OpenCV cv2.aruco (ChArUco) Apache-2.0 Multi-camera rig calibration (Zhang, A Flexible New Technique for Camera Calibration, TPAMI 2000; Garrido-Jurado et al., Automatic generation and detection of highly reliable fiducial markers under occlusion, Pattern Recognition 2014)
1→4D linear-eigen DLT (Hartley & Zisserman, Multiple View Geometry, 2nd ed.) Multi-view triangulation of Stage-1 keypoints into 3D, in the style of Anipose (Karashchuk et al., Anipose: a toolkit for robust markerless 3D pose estimation, Cell Reports 2021) and DANNCE (Dunn et al., Geometric deep learning enables 3D kinematic profiling across species and environments, Nature Methods 2021)
4D Dynamic 3D Gaussians (Luiten et al., 3DV 2024) MIT Multi-view 4D reconstruction with persistent per-point tracking — the primary target once a rig is set up, chosen because it yields per-point 3D trajectories directly usable as ML features
4D Deformable 3D Gaussians (Yang et al., CVPR 2024) MIT Monocular fallback 4D reconstruction, usable with today's single-camera footage while a rig is being built
4D 3D Gaussian Splatting (Kerbl et al., SIGGRAPH 2023) Inria/MPII non-commercial research license Base static method both of the above extend
4D viser (Nerfstudio) MIT Interactive browser-based 4D point-cloud/camera-frustum viewer
2 Ultralytics YOLO11 AGPL-3.0 Body-part segmentation/keypoints for snapshot capture (planned)
2 OpenCV (var-of-Laplacian), CPBD Apache-2.0 / MIT No-reference sharpness / blur scoring for best-frame selection (planned)
3–4 WildlifeTools / MegaDescriptor MIT Metric-learning embeddings for individual re-ID (planned)
UI napari + Qt (PySide6) BSD Stage-1 annotation / review desktop UI
infra ffmpeg, SQLite (stdlib) LGPL / public domain Video decode, project storage

Also informing the overall re-ID method:

  • Haurum et al., Re-Identification of Zebrafish using Metric Learning, WACV 2020.
  • Stripe pattern differences can be used to distinguish individual adult zebrafish, PLOS ONE, 2024.
  • Čermák et al., WildlifeDatasets: An open-source toolkit for animal re-identification, 2024.

Layout

pyproject.toml              # package metadata, optional-dependency groups
requirements.txt            # manual/offline install reference, GPU torch notes
src/zebrafish_pipeline/
  config.py                 # settings, default paths
  cli.py                    # command-line entry point
  data/
    schema.sql              # SQLite schema (tanks/fish/sessions/rigs/... )
    store.py                # data access layer
  video/
    ingest.py                # ffmpeg probe + single/multi-camera frame sampling
  calibration/
    charuco.py               # ChArUco intrinsic + multi-camera extrinsic calibration
    rig.py                    # RigCalibration container, save/load
  stage1/
    motion.py                 # background-subtraction motion localization
    sam2_outline.py            # SAM 2 wrapper -> dense outline + confidence
    orientation.py              # head/tail -> side + vertical swim angle
    review.py                    # napari review window
  reconstruction3d/
    triangulate.py               # multi-view DLT triangulation
    gaussian4d.py                 # dataset export + training/loading for 4D-GS
  viz/
    viewer4d.py                    # viser-based 4D scrubber
tests/                          # pytest; calibration/triangulation use synthetic cameras
_projects/                      # runtime data: one folder + project.db per project (gitignored reconstruction output)
hpc/                            # Compute Canada: env setup + SLURM job scripts, see hpc/README.md
  setup_env.sh
  data_layout.md                # $PROJECT vs $SCRATCH vs $SLURM_TMPDIR mapping
  jobs/                         # one sbatch script per stage
docs/
  brief.docx                     # original project brief
  assets/, logs/                 # demo images, install logs

About

No description, website, or topics provided.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages