diff --git a/src/poseforge/pose/data/synthetic/atomic_batch.py b/src/poseforge/pose/data/synthetic/atomic_batch.py index f7c0e55e..32a8f025 100644 --- a/src/poseforge/pose/data/synthetic/atomic_batch.py +++ b/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], diff --git a/src/poseforge/style_transfer/scripts/extract_dataset_bodypart.py b/src/poseforge/style_transfer/scripts/extract_dataset_bodypart.py index 59f86a38..56021cc9 100644 --- a/src/poseforge/style_transfer/scripts/extract_dataset_bodypart.py +++ b/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,12 +161,24 @@ 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) @@ -173,8 +186,8 @@ def build_endpoints( 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" diff --git a/tests/test_atomic_batch_dataset.py b/tests/test_atomic_batch_dataset.py index 4dd32d56..f4a5e0e5 100644 --- a/tests/test_atomic_batch_dataset.py +++ b/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, diff --git a/tests/test_atomic_batch_guards.py b/tests/test_atomic_batch_guards.py new file mode 100644 index 00000000..e52203bc --- /dev/null +++ b/tests/test_atomic_batch_guards.py @@ -0,0 +1,116 @@ +"""Tests for the atomic-batch loader resolution guard (audit issue #48, I1-C). + +`AtomicBatchDataset.load_atomic_batch_frames` reconstructs each per-variant +tile via a top-left crop `frame[:n_rows, start_col:end_col]`. That is only +correct when the on-disk tiles were serialized at exactly `image_size`. If an +atomic batch was extracted at a different (larger) per-variant resolution, the +crop silently grabs the wrong columns while the stored labels stay at the +extraction resolution -> corrupted training data. The loader must instead fail +loudly. These tests pin both behaviours against a REAL shipped 256-tiled batch. +""" + +from pathlib import Path + +import pytest +import torch + +from poseforge.pose.data.synthetic.atomic_batch import AtomicBatchDataset + +# The shipped atomic batches under this directory are tiled at 256x256 with +# 4 variants and a serialization spacing of 10 px. +ATOMIC_BATCHES_4VARIANTS_DIR = Path( + "bulk_data/pose_estimation/atomic_batches/4variants" +) +N_VARIANTS = 4 +SPACING = 10 +TILE_SIZE = (256, 256) + + +def _find_one_atomic_batch_frames() -> Path: + """Return the path to a single real `*_frames.mp4`, or skip the test.""" + if not ATOMIC_BATCHES_4VARIANTS_DIR.is_dir(): + pytest.skip( + f"Real atomic batches not available at {ATOMIC_BATCHES_4VARIANTS_DIR}" + ) + frames_videos = sorted( + ATOMIC_BATCHES_4VARIANTS_DIR.rglob("atomicbatch*_frames.mp4") + ) + if not frames_videos: + pytest.skip( + f"No atomicbatch *_frames.mp4 found under {ATOMIC_BATCHES_4VARIANTS_DIR}" + ) + return frames_videos[0] + + +def test_load_atomic_batch_frames_matching_resolution(): + """The loader works on a real 256-tiled batch with image_size=(256, 256).""" + frames_video = _find_one_atomic_batch_frames() + + atomic_batch = AtomicBatchDataset.load_atomic_batch_frames( + frames_video, + n_variants=N_VARIANTS, + image_size=TILE_SIZE, + n_channels=1, + spacing=SPACING, + ) + + assert isinstance(atomic_batch, torch.Tensor) + assert atomic_batch.dtype == torch.float32 + # (n_variants, n_frames, n_channels, height, width) + assert atomic_batch.ndim == 5 + assert atomic_batch.shape[0] == N_VARIANTS + assert atomic_batch.shape[2] == 1 + assert atomic_batch.shape[3:] == TILE_SIZE + assert atomic_batch.shape[1] > 0 # at least one frame + # Pixels are normalized into [0, 1] and the tile is not all-zero/all-one. + assert atomic_batch.min() >= 0.0 and atomic_batch.max() <= 1.0 + assert 0.0 < float(atomic_batch.mean()) < 1.0 + + +def test_load_atomic_batch_frames_resolution_mismatch_raises(): + """A mismatched image_size must raise a clear error, not mis-slice silently. + + The shipped tiles are 256x256; asking the loader to treat them as 512x512 + is exactly the I1-C footgun. The guard must refuse rather than top-left + crop the wrong columns. + """ + frames_video = _find_one_atomic_batch_frames() + + with pytest.raises(ValueError) as excinfo: + AtomicBatchDataset.load_atomic_batch_frames( + frames_video, + n_variants=N_VARIANTS, + image_size=(512, 512), + n_channels=1, + spacing=SPACING, + ) + + message = str(excinfo.value) + # The error must name the mismatch and how to fix it so it is actionable. + assert "inconsistent" in message + assert "image_size=(512, 512)" in message + assert "re-extract" in message + + +def test_infer_stored_variant_geometry_inversion(): + """Unit-level check of the geometry inversion that backs the guard.""" + infer = AtomicBatchDataset._infer_stored_variant_geometry + + # Real shipped layout: 4 tiles of 256 + 3*10 spacing = 1054 -> round16 1056. + assert infer(1056, 256, SPACING) == (4, 1056) + + # A 512-tiled, 4-variant video: 4*512 + 3*10 = 2078 -> round16 2080. + assert infer(2080, 512, SPACING) == (4, 2080) + + # Asking for 512-wide tiles against a 256-tiled (1056-wide) video is + # incompatible -> no integer variant count reproduces the width. + n_variants, _ = infer(1056, 512, SPACING) + assert n_variants is None + + # ...and vice versa: 256-wide tiles against a 512-tiled (2080-wide) video. + n_variants, _ = infer(2080, 256, SPACING) + assert n_variants is None + + +if __name__ == "__main__": + pytest.main([__file__, "-v"])