thermokourt.track

Stage 2 of the pipeline: identity tracking via idtracker.ai, plus post-tracking utilities for crossing review and identity QC.

thermokourt.track.posttrack

.. py:module:: thermokourt.track.posttrack

╔═══════════════════════════════════════════════════════════════════════╗
║   ThermoKourt — Post-Tracking Pipeline                    v1.0      ║
║   Assign identities, crop individuals, render overlay videos        ║
╚═══════════════════════════════════════════════════════════════════════╝

After idtracker.ai has produced trajectories for a 3-fly arena, this tool:

 1. IDENTITY ASSIGNMENT — Opens an interactive GUI where the user clicks
    on each fly to assign: male_teal, male_orange, female.
    Saves the mapping as a JSON sidecar.

 2. INDIVIDUAL CROPS — For each identified animal, produces a centred
    crop video following that animal. These are suitable as input for
    SLEAP or DeepLabCut single-animal pose estimation.

 3. OVERLAY VIDEO — Renders teal and orange semi-transparent auras
    around the two males on the full arena video, for manual ethogram
    scoring in GameThogram.

Usage: thermokourt-posttrack arena_video.mp4 –session session_folder/ thermokourt-posttrack arena_video.mp4 –trajectories without_gaps.npy thermokourt-posttrack arena_video.mp4 –session session_folder/ –skip_gui (reuses previously saved identity assignment)

Requirements: pip install numpy opencv-python (idtracker.ai output files)

License: MIT · Author: Bart R.H. Geurten

.. py:function:: load_trajectories(session_dir: str = ‘’, traj_path: str = ‘’) -> numpy.ndarray :module: thermokourt.track.posttrack

Load trajectories from idtracker.ai output.

Returns array of shape (N_frames, N_animals, 2) — xy centroids. NaN values indicate frames where the animal was not detected.

.. py:function:: find_clear_frame(traj: numpy.ndarray, video_path: str, min_distance_frac: float = 0.15) -> ~typing.Tuple[int, numpy.ndarray] :module: thermokourt.track.posttrack

Find a frame where all animals are well-separated and detected.

Searches for a frame where:

  • All animals have valid (non-NaN) positions

  • The minimum pairwise distance is > min_distance_frac * frame_diagonal

Returns (frame_index, frame_image_rgb).

.. py:class:: IdentityAssigner(frame_rgb: numpy.ndarray, positions: numpy.ndarray, animal_indices: ~typing.List[int]) :module: thermokourt.track.posttrack

Bases: :py:class:object

OpenCV GUI for assigning identities to tracked animals.

Shows a frame with animal positions marked. User clicks near each animal to assign identities in order: male_teal first, then male_orange. The remaining animal is auto-assigned as female.

.. py:attribute:: IdentityAssigner.WIN :module: thermokourt.track.posttrack :value: ‘Identity Assignment’

.. py:method:: IdentityAssigner.init(frame_rgb: numpy.ndarray, positions: numpy.ndarray, animal_indices: ~typing.List[int]) :module: thermokourt.track.posttrack

  :param frame_rgb: The frame image (RGB).
  :param positions: (N_animals, 2) array of centroid positions.
  :param animal_indices: list of idtracker animal indices (0-based).

.. py:method:: IdentityAssigner.get_assignments() -> ~typing.Dict[str, int] | None :module: thermokourt.track.posttrack

  Returns {identity_name: animal_index} or None if aborted.

.. py:function:: create_individual_crops(video_path: str, traj: numpy.ndarray, assignments: ~typing.Dict[str, int], output_dir: str, crop_size: int = 256, basename: str = ‘’) :module: thermokourt.track.posttrack

Create centred crop videos for each identified animal.

For each frame, extracts a crop_size x crop_size window centred on the animal’s centroid. Pads with white if the crop extends beyond frame boundaries. Writes one video per identity.

:param video_path: Path to the arena video. :param traj: Trajectory array (N_frames, N_animals, 2). :param assignments: {identity_name: animal_index} mapping. :param output_dir: Directory for output videos. :param crop_size: Side length of the square crop. :param basename: Prefix for output filenames.

.. py:function:: create_overlay_video(video_path: str, traj: numpy.ndarray, assignments: ~typing.Dict[str, int], output_path: str, aura_alpha: float = 0.35, aura_radius_mult: float = 2.5) :module: thermokourt.track.posttrack

Render an overlay video with coloured auras around identified animals.

Males get teal/orange semi-transparent circles. The female gets a subtle grey marker. Output is suitable for GameThogram annotation.

:param video_path: Path to the arena video. :param traj: Trajectory array (N_frames, N_animals, 2). :param assignments: {identity_name: animal_index} mapping. :param output_path: Path for the output video. :param aura_alpha: Transparency of the coloured overlay. :param aura_radius_mult: Aura radius as multiple of estimated body length.

.. py:function:: parse_args() :module: thermokourt.track.posttrack

.. py:function:: main() :module: thermokourt.track.posttrack

thermokourt.track.cli

.. py:module:: thermokourt.track.cli

CLI entry point – placeholder for future implementation.

.. py:function:: main() :module: thermokourt.track.cli