|
| 1 | +"""A small 3-D project on disk: real idtracker.ai v5 session folders with known depths. |
| 2 | +
|
| 3 | +The folders are read by a built-in reader, so a spawned worker process (which sees none of the |
| 4 | +tests' monkeypatches) runs exactly what the calling process runs. |
| 5 | +""" |
| 6 | + |
| 7 | +from __future__ import annotations |
| 8 | + |
| 9 | +import json |
| 10 | +from datetime import UTC, datetime |
| 11 | +from pathlib import Path |
| 12 | + |
| 13 | +import numpy as np |
| 14 | + |
| 15 | +from track2data.core.models import ( |
| 16 | + CalibrationConfig, |
| 17 | + FusionSettings, |
| 18 | + MetricSelection, |
| 19 | + ProjectManifest, |
| 20 | + ProjectMode, |
| 21 | + SessionRef, |
| 22 | + ViewPair, |
| 23 | +) |
| 24 | + |
| 25 | +N_FRAMES = 100 |
| 26 | +FPS = 25.0 |
| 27 | +#: Side-view rows of fish 0, 1, 2: depth 0.25 / 0.5 / 0.75 between surface row 100 and floor 300. |
| 28 | +SIDE_ROWS = (150.0, 200.0, 250.0) |
| 29 | +DEPTHS = (0.25, 0.5, 0.75) |
| 30 | +TANK_CM = 20.0 |
| 31 | +LABELS = ("0", "1", "2") # the v5 reader has no labels, so fish are named by column |
| 32 | + |
| 33 | + |
| 34 | +def write_v5_session(folder: Path, xy: np.ndarray, *, fps: float = FPS) -> Path: |
| 35 | + """Write ``xy`` (n_frames, n_animals, 2) as an idtracker.ai v5 session folder.""" |
| 36 | + (folder / "trajectories").mkdir(parents=True) |
| 37 | + np.save(folder / "trajectories" / "trajectories.npy", xy) |
| 38 | + n_frames, n_animals = xy.shape[:2] |
| 39 | + np.save( |
| 40 | + folder / "video_object.npy", |
| 41 | + { |
| 42 | + "fps": fps, |
| 43 | + "frame_number": n_frames, |
| 44 | + "height": 400, |
| 45 | + "width": 400, |
| 46 | + "paths_to_video_files": ["video.mp4"], |
| 47 | + }, |
| 48 | + ) |
| 49 | + (folder / "session_progress.json").write_text( |
| 50 | + json.dumps( |
| 51 | + { |
| 52 | + "session_name": folder.name, |
| 53 | + "number_of_animals": n_animals, |
| 54 | + "frames_per_second": fps, |
| 55 | + "video_height": 400, |
| 56 | + "video_width": 400, |
| 57 | + "status": "finished", |
| 58 | + } |
| 59 | + ), |
| 60 | + encoding="utf-8", |
| 61 | + ) |
| 62 | + return folder |
| 63 | + |
| 64 | + |
| 65 | +def top_xy(shift: float = 0.0) -> np.ndarray: |
| 66 | + t = np.arange(N_FRAMES, dtype=float) |
| 67 | + out = np.empty((N_FRAMES, 3, 2)) |
| 68 | + for k in range(3): |
| 69 | + out[:, k, 0] = 20 + 10 * k + t + shift |
| 70 | + out[:, k, 1] = 50 + 20 * k |
| 71 | + return out |
| 72 | + |
| 73 | + |
| 74 | +def side_xy(shift: float = 0.0) -> np.ndarray: |
| 75 | + out = top_xy(shift) |
| 76 | + for k in range(3): |
| 77 | + out[:, k, 1] = SIDE_ROWS[k] |
| 78 | + return out |
| 79 | + |
| 80 | + |
| 81 | +def fusion_settings(**kw: object) -> FusionSettings: |
| 82 | + base: dict[str, object] = {"surface_row": 100.0, "floor_row": 300.0, "tank_height_cm": TANK_CM} |
| 83 | + base.update(kw) |
| 84 | + return FusionSettings(**base) # type: ignore[arg-type] |
| 85 | + |
| 86 | + |
| 87 | +def pair(top: str, side: str, *, fusion: FusionSettings | None = None) -> ViewPair: |
| 88 | + return ViewPair( |
| 89 | + top_session_id=top, |
| 90 | + side_session_id=side, |
| 91 | + fish_map={lab: lab for lab in LABELS}, |
| 92 | + fusion=fusion if fusion is not None else fusion_settings(), |
| 93 | + ) |
| 94 | + |
| 95 | + |
| 96 | +def build_scene( |
| 97 | + base: Path, |
| 98 | + *, |
| 99 | + side2_fps: float = FPS, |
| 100 | + pairs: list[ViewPair] | None = None, |
| 101 | + with_lone: bool = True, |
| 102 | +) -> ProjectManifest: |
| 103 | + """Two top/side pairs (t1+s1, t2+s2) and, with ``with_lone``, a top session in no pair. |
| 104 | +
|
| 105 | + ``side2_fps`` other than ``FPS`` makes the second pair fail to fuse (frame rates differ). |
| 106 | + """ |
| 107 | + write_v5_session(base / "t1", top_xy()) |
| 108 | + write_v5_session(base / "s1", side_xy()) |
| 109 | + write_v5_session(base / "t2", top_xy(5.0)) |
| 110 | + write_v5_session(base / "s2", side_xy(5.0), fps=side2_fps) |
| 111 | + ids = ["t1", "s1", "t2", "s2"] |
| 112 | + if with_lone: |
| 113 | + write_v5_session(base / "lone", top_xy(9.0)) |
| 114 | + ids.append("lone") |
| 115 | + now = datetime.now(tz=UTC) |
| 116 | + return ProjectManifest( |
| 117 | + project_name="scene3d", |
| 118 | + created_at=now, |
| 119 | + updated_at=now, |
| 120 | + sessions=[SessionRef(session_id=i, folder=base / i, sha256="") for i in ids], |
| 121 | + mode=ProjectMode(dimension="3d", layout="two_videos"), |
| 122 | + calibration=CalibrationConfig(mode="scalar", px_per_cm=10.0), |
| 123 | + metrics=MetricSelection(individual=["IL-1", "IL-15"], group=[], zone=[], diagnostic=[]), |
| 124 | + view_pairs=pairs if pairs is not None else [pair("t1", "s1"), pair("t2", "s2")], |
| 125 | + ) |
0 commit comments