Data pipeline: fix pvio import break + loader resolution guard (I1-C) + extraction script (I2-C)

src/poseforge/pose/data/synthetic/atomic_batch.py

@@ -5,12 +5,34 @@
 import torch.nn.functional as F
 from torch.utils.data import Dataset, DataLoader
 from pathlib import Path
-from pvio.io import read_frames_from_video, write_frames_to_video, _default_ffmpeg_params_for_video_writing
-level_idx = _default_ffmpeg_params_for_video_writing.index("-level")
-_default_ffmpeg_params_for_video_writing[level_idx + 1] = "5.0"  # allow higher resolution videos for 900x900
+from pvio.io import read_frames_from_video, write_frames_to_video
 
 from poseforge.util.sys import get_hardware_availability
 
+# ffmpeg parameters used when serializing atomic batches to video.
+#
+# Previously we imported pvio's private `_default_ffmpeg_params_for_video_writing`
+# and patched its `-level` entry in place. That symbol is no longer exported by
+# pvio (it is an internal default), so importing it raised ImportError and broke
+# the whole `pose.data.synthetic` / `pose.contrast` packages on import. We now
+# define the params locally and pass them through the public
+# `write_frames_to_video(..., ffmpeg_params=...)` API instead.
+#
+# These mirror pvio's documented high-quality H.264 defaults
+# (["-crf", "15", "-preset", "slow", "-profile:v", "high", "-level", "4.0"]),
+# but with `-level` bumped to 5.0 so that larger frames (e.g. the 900x900
+# renders) stay within the H.264 level limits.
+_DEFAULT_FFMPEG_PARAMS_FOR_VIDEO_WRITING = [
+    "-crf",
+    "15",  # Lower CRF = higher quality (15 is very high quality)
+    "-preset",
+    "slow",  # Slower preset = better compression efficiency
+    "-profile:v",
+    "high",  # Use high profile for better compression
+    "-level",
+    "5.0",  # H.264 level (bumped from pvio default 4.0 to allow larger frames)
+]
+
 # Emit the BODY SEGMENTATION RESIZE WARNING only once per process to avoid
 # flooding the logs when many atomic batches need on-the-fly mask resizing.
 _BODY_SEGMENTATION_RESIZE_WARNING_EMITTED = False
@@ -217,7 +239,12 @@ def save_atomic_batch_frames(
                 selection = (selection * 255).astype(np.uint8)
                 image[:n_rows, start_col:end_col, :] = selection
             output_frames.append(image.squeeze())
-        write_frames_to_video(output_path, output_frames, fps=fps, ffmpeg_params=_default_ffmpeg_params_for_video_writing)
+        write_frames_to_video(
+            output_path,
+            output_frames,
+            fps=fps,
+            ffmpeg_params=_DEFAULT_FFMPEG_PARAMS_FOR_VIDEO_WRITING,
+        )
 
     @staticmethod
     def load_atomic_batch_frames(
@@ -247,15 +274,61 @@ def load_atomic_batch_frames(
             raise ValueError(f"Error: no frames found in video {video_path}")
         n_rows_total, n_cols_total = frames[0].shape[:2]
         n_rows, n_cols = image_size
-        assert (
-            n_variants * n_cols + (n_variants - 1) * spacing <= n_cols_total
-        ), "Error: width of concatenated video does not match expectation"
-
-        # Detect over-supply of variants in the on-disk video. The width is
-        # n_actual_variants * n_cols + (n_actual_variants - 1) * spacing, so
-        # invert that to recover the stored count and warn once if some are
-        # being silently discarded.
-        actual_n_variants = (n_cols_total + spacing) // (n_cols + spacing)
+
+        # --- I1-C: loud guard on stored-vs-requested tile resolution ---------
+        #
+        # This loader reconstructs each per-variant tile by a TOP-LEFT CROP
+        # `frame[:n_rows, start_col:end_col]`, where the column stride is
+        # `n_cols + spacing` and `(n_rows, n_cols) == image_size`. This is only
+        # correct if each on-disk tile was serialized at exactly `image_size`.
+        # If the atomic batch was extracted at a LARGER per-variant resolution
+        # than `image_size`, the crop silently grabs the wrong columns (slicing
+        # across tile boundaries) and clips the fly, while the stored labels
+        # (keypoints / segmentation maps) remain in the extraction resolution
+        # -> silently corrupted training data. The old assertion only caught
+        # tiles that were too NARROW; it happily accepted tiles that were too
+        # wide. We instead reconstruct the geometry that `save_atomic_batch_frames`
+        # would have produced and require the stored frame size to match it.
+        #
+        # `save_atomic_batch_frames` rounds both the total width and the height
+        # UP to a multiple of 16 for the video encoder, so the stored size is
+        # NOT exactly the logical size -- we must allow that <16px padding.
+        n_variants_in_video, expected_n_cols_total = (
+            AtomicBatchDataset._infer_stored_variant_geometry(
+                n_cols_total, n_cols, spacing
+            )
+        )
+        expected_n_rows_total = _round_up_to_multiple(n_rows, 16)
+        if (
+            n_variants_in_video is None
+            or n_variants_in_video < n_variants
+            or n_rows_total != expected_n_rows_total
+            or n_cols_total != expected_n_cols_total
+        ):
+            raise ValueError(
+                f"Atomic batch {video_path} has stored frame size "
+                f"(height={n_rows_total}, width={n_cols_total}) that is "
+                f"inconsistent with the requested image_size={image_size} and "
+                f"n_variants={n_variants} (spacing={spacing}).\n"
+                f"  Expected, for tiles of size {image_size} laid out side by "
+                f"side and rounded up to a multiple of 16 for encoding: "
+                f"height={expected_n_rows_total}, and a width matching "
+                f"round_up_16(V * {n_cols} + (V - 1) * {spacing}) for some "
+                f"V >= {n_variants}.\n"
+                f"  This usually means the atomic batch was extracted at a "
+                f"DIFFERENT per-variant resolution than image_size. Loading it "
+                f"with this image_size would top-left-crop each tile to the "
+                f"wrong columns while leaving the stored labels at the original "
+                f"resolution, silently corrupting the data.\n"
+                f"  Fix: either pass the image_size the batch was actually "
+                f"extracted at, or re-extract the atomic batches at "
+                f"image_size={image_size}. (Resizing here is intentionally NOT "
+                f"done, because it would not rescale the stored labels.)"
+            )
+
+        # Number of variants physically present in the stored video (>= the
+        # requested `n_variants`, guaranteed by the guard above).
+        actual_n_variants = n_variants_in_video
         if actual_n_variants > n_variants:
             global _EXTRA_VARIANTS_WARNING_EMITTED
             if not _EXTRA_VARIANTS_WARNING_EMITTED:
@@ -287,6 +360,40 @@ def load_atomic_batch_frames(
 
         return torch.from_numpy(atomic_batch).to(torch.float32)
 
+    @staticmethod
+    def _infer_stored_variant_geometry(
+        n_cols_total: int, n_cols: int, spacing: int
+    ) -> tuple[int | None, int]:
+        """Invert the stored-video column geometry.
+
+        ``save_atomic_batch_frames`` lays ``V`` tiles of width ``n_cols`` side
+        by side with ``spacing`` pixels between them and then rounds the total
+        width UP to a multiple of 16 for the encoder, i.e. the stored width is::
+
+            round_up_16(V * n_cols + (V - 1) * spacing)
+
+        Given a stored ``n_cols_total`` and a candidate per-tile width
+        ``n_cols``, find the integer ``V >= 1`` that reproduces it exactly
+        (within the <16px rounding). If no such ``V`` exists, the stored tiles
+        were NOT serialized at width ``n_cols``.
+
+        Returns:
+            (n_variants_in_video, expected_n_cols_total):
+                ``n_variants_in_video`` is the recovered variant count, or
+                ``None`` if the stored width is incompatible with ``n_cols``.
+                ``expected_n_cols_total`` is the width that a single variant
+                (``V = 1``) would have produced -- only used to build a helpful
+                error message when the recovered count is ``None``.
+        """
+        # Lower bound on V (ignoring rounding); the true V can be at most this
+        # because rounding only ever increases the stored width.
+        v_upper = (n_cols_total + spacing) // (n_cols + spacing)
+        for v in range(v_upper, 0, -1):
+            logical_width = v * n_cols + (v - 1) * spacing
+            if _round_up_to_multiple(logical_width, 16) == n_cols_total:
+                return v, n_cols_total
+        return None, _round_up_to_multiple(n_cols, 16)
+
     @staticmethod
     def save_atomic_batch_sim_data(
         sim_data: dict[str, np.ndarray],

src/poseforge/style_transfer/scripts/extract_dataset_bodypart.py

@@ -1,4 +1,5 @@
 import logging
+import os
 import re
 from collections import defaultdict
 from dataclasses import dataclass
@@ -160,21 +161,33 @@ def list_spotlight_recordings_and_usable_frames(
 def build_endpoints(
     target_segment_names: list[str],
     xy_lookup: dict[str, np.ndarray],
-    is_spotlight: bool = False,
+    scaling: float = 1.0,
 ) -> np.ndarray:
-    if is_spotlight:
-        SCALING = 900.0/256.0
-    else:
-        SCALING = 1.0
+    """Assemble (proximal, distal) endpoint coordinates for each segment.
+
+    Args:
+        target_segment_names: Segments to build endpoints for, in order.
+        xy_lookup: Maps segment name -> (x, y) keypoint in the SOURCE pixel
+            space of the coordinates (e.g. the resolution the keypoint model
+            ran at).
+        scaling: Multiplicative factor converting the source coordinate space
+            to the saved-image pixel space, i.e.
+            ``target_resolution / source_resolution``. Use ``1.0`` when the
+            coordinates are already in the saved image's pixel space (as for
+            the synthetic renders). Previously this was hardcoded to
+            ``900.0 / 256.0`` for the spotlight branch, which silently produced
+            wrong labels whenever either resolution changed; it is now derived
+            from the actual frame/prediction resolutions by the caller.
+    """
     k = len(target_segment_names)
     xy = np.full((k, 2, 2), np.nan, dtype=np.float32)
 
     for seg_idx, seg_name in enumerate(target_segment_names):
         prox_xy = xy_lookup[seg_name]
         distal_seg = DISTAL_CHILD_BY_SEGMENT[seg_name]
         dist_xy = xy_lookup[distal_seg]
-        xy[seg_idx, 0, :] = prox_xy * SCALING
-        xy[seg_idx, 1, :] = dist_xy * SCALING
+        xy[seg_idx, 0, :] = prox_xy * scaling
+        xy[seg_idx, 1, :] = dist_xy * scaling
 
     return xy
 
@@ -308,7 +321,9 @@ def extract_synthetic_with_annotations(
 
                 frame_coords = camera_coords_ds[frame_idx]  # (K, 3)
                 xy_lookup = frame_xy_vis_from_synthetic(frame_coords, seg_names)
-                xy = build_endpoints(TARGET_SEGMENT_NAMES, xy_lookup, False)
+                # Synthetic keypoints are rendered in the same pixel space as
+                # the saved image, so no rescaling is needed.
+                xy = build_endpoints(TARGET_SEGMENT_NAMES, xy_lookup, scaling=1.0)
 
                 rel_path = str(output_image_path.relative_to(output_dir))
                 sample = AnnotationSample(
@@ -325,8 +340,22 @@ def extract_spotlight_with_annotations(
     output_dir: Path,
     trial_info_by_dir: dict[Path, SpotlightTrialInfo],
     accumulators: dict[str, AnnotationAccumulator],
+    pred_coord_resolution: int,
 ) -> None:
+    """Extract spotlight frames and rescale their predicted keypoints.
+
+    Args:
+        pred_coord_resolution: The (square) pixel resolution that the spotlight
+            keypoint predictions in ``pred_xy`` are expressed in. Predicted
+            coordinates are rescaled to the saved-image pixel space by
+            ``saved_image_width / pred_coord_resolution`` per frame, replacing
+            the old hardcoded ``900.0 / 256.0`` factor.
+    """
     print("Extracting spotlight images + annotations...")
+    if pred_coord_resolution <= 0:
+        raise ValueError(
+            f"pred_coord_resolution must be positive, got {pred_coord_resolution}"
+        )
 
     specs_by_trial = defaultdict(list)
     for trial_dir, usable_row_idx, split_dir_name in spotlight_specs:
@@ -385,7 +414,17 @@ def extract_spotlight_with_annotations(
                     pred_xy_frame=pred_xy_frame,
                     pred_names=pred_names,
                 )
-                xy = build_endpoints(TARGET_SEGMENT_NAMES, xy_lookup, True)
+                # Spotlight predictions live in a square `pred_coord_resolution`
+                # space; rescale them to the saved-image pixel space. The
+                # aligned behavior frames are square, so assert that and derive
+                # the factor from the actual sizes instead of hardcoding it.
+                if img_h != img_w:
+                    raise ValueError(
+                        f"Expected a square spotlight frame for {trial_dir.name} "
+                        f"(predictions are isotropic), got (h={img_h}, w={img_w})"
+                    )
+                scaling = img_w / pred_coord_resolution
+                xy = build_endpoints(TARGET_SEGMENT_NAMES, xy_lookup, scaling=scaling)
 
                 rel_path = str(output_image_path.relative_to(output_dir))
                 sample = AnnotationSample(
@@ -409,6 +448,19 @@ def main() -> None:
 
     video_filename = "processed_nmf_sim_render_grayscale.mp4"
 
+    # Number of worker threads used to index/validate the NMF videos. This is
+    # required by `list_nmf_simulations_and_num_frames`; previously it was
+    # omitted, raising a TypeError on launch.
+    num_workers = os.cpu_count() or 1
+
+    # Pixel resolution that the spotlight keypoint predictions (`pred_xy`) are
+    # expressed in. Predicted coords are rescaled to the saved-image pixel
+    # space by `saved_image_width / pred_coord_resolution`. Previously this was
+    # baked into a hardcoded `900.0 / 256.0` factor; pulling it out here keeps
+    # the extraction correct if either the prediction or the saved-frame
+    # resolution changes.
+    pred_coord_resolution = 256
+
     # Sampling config.
     random_seed = 42
     np.random.seed(random_seed)
@@ -420,7 +472,9 @@ def main() -> None:
     split_to_out = {"train": "train", "test": "test", "val": "val"}
 
     # Index available frames.
-    nmf_num_frames = list_nmf_simulations_and_num_frames(nmf_rendering_dir, video_filename)
+    nmf_num_frames = list_nmf_simulations_and_num_frames(
+        nmf_rendering_dir, video_filename, num_workers
+    )
     spotlight_info = list_spotlight_recordings_and_usable_frames(
         spotlight_data_dir,
     )
@@ -501,6 +555,7 @@ def main() -> None:
         output_dir=output_dir,
         trial_info_by_dir=spotlight_info,
         accumulators=accumulators,
+        pred_coord_resolution=pred_coord_resolution,
     )
 
     annotations_dir = output_dir / "annotations"

tests/test_atomic_batch_dataset.py

@@ -6,7 +6,7 @@
 
 def test_atomic_batch_dataset_loading():
     data_dirs = [
-        Path("bulk_data/pose_estimation/atomic_batches/BO_Gal4_fly1_trial001/")
+        Path("bulk_data/pose_estimation/atomic_batches/4variants/BO_Gal4_fly1_trial001/")
     ]
     dataset = AtomicBatchDataset(
         data_dirs=data_dirs,