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

[sibocw]

This work was implemented by Claude Code (Opus 4.8) upon user instruction.

Addresses audit issue #48: a pvio import break plus two data-pipeline footguns (I1-C and I2-C).

pvio private-symbol import break (the package-breaking bug)

src/poseforge/pose/data/synthetic/atomic_batch.py imported pvio's private _default_ffmpeg_params_for_video_writing, then mutated its -level entry in place to allow higher-resolution (900x900) videos:

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"

The installed pvio no longer exports that internal default, so this raised ImportError: cannot import name '_default_ffmpeg_params_for_video_writing' at import time — which broke the entire poseforge.pose.data.synthetic and poseforge.pose.contrast packages (anything importing them failed at collection/startup).

Fix (public API): pvio's write_frames_to_video(video_path, frames, fps, codec='libx264', ffmpeg_params: list[str] | None = None, ...) accepts ffmpeg_params. We now define the params as a module-level constant mirroring 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, drop the private import, and pass the constant through write_frames_to_video(..., ffmpeg_params=...). Behavior is preserved; no private symbol is touched.

This unblocks importing the pose.data.synthetic and pose.contrast packages (verified below).

I1-C: loader resolution guard (load_atomic_batch_frames)

The loader reconstructs each per-variant tile via a top-left crop frame[:n_rows, start_col:end_col], where (n_rows, n_cols) = image_size and the column stride is n_cols + spacing. This is correct only if each on-disk tile was serialized at exactly image_size. If an atomic batch were extracted at a larger per-variant resolution, the crop would silently grab the wrong columns (slicing across tile boundaries) and clip the fly, while the stored labels (keypoints / segmentation maps) remained at the extraction resolution — silently corrupted training data. The previous assertion only caught tiles that were too narrow; it happily accepted tiles that were too wide. On the currently shipped data the tiles are 256, so this was latent, but it is an unguarded footgun that blocks higher-resolution work.

Fix: add a loud guard that reconstructs the geometry save_atomic_batch_frames would have produced — width round_up_16(V * n_cols + (V-1) * spacing) for some V >= n_variants, and height round_up_16(n_rows) — and raises a clear, actionable ValueError on any mismatch, naming the stored-vs-expected sizes and how to fix it. The check correctly accounts for the <16px encoder padding that save_atomic_batch_frames applies (so it does not false-positive on the real 256 data, whose 4-variant width is 1056, not the naive 1054). It also still recovers the true stored variant count to drive the existing "extra variants" warning. We intentionally do not silently resize, since that would not rescale the stored labels; a clear error is the safe fix.

I2-C: broken extraction script (extract_dataset_bodypart.py)

• TypeError on launch: main() called list_nmf_simulations_and_num_frames(nmf_rendering_dir, video_filename), but that function (in extract_dataset.py) requires a third num_workers: int argument. We add a num_workers config in main() (defaulting to os.cpu_count()) and thread it through.
• Hardcoded SCALING: build_endpoints hardcoded SCALING = 900.0/256.0 for spotlight keypoint coords, silently wrong if either resolution changes. We replace the is_spotlight flag with an explicit scaling parameter; the synthetic path passes 1.0, and the spotlight path derives scaling = saved_image_width / pred_coord_resolution from the actual saved-frame size and a configurable pred_coord_resolution (default 256), with a guard that the spotlight frame is square.

Findings from #48 (Refs #48)

• The pvio break is package-level: it was an ImportError at module import, not a runtime error, so it took down all of pose.data.synthetic and pose.contrast — including the existing test suite, which couldn't even be collected.
• The reorganized data now lives under bulk_data/pose_estimation/atomic_batches/4variants/...; the pre-existing tests/test_atomic_batch_dataset.py still points at the old atomic_batches/BO_Gal4_fly1_trial001/ path (see test results).
• extract_dataset.py's own main() already passed num_workers correctly — only the bodypart script had the missing-arg bug.

Test results

Import verification (worktree source):

module + AtomicBatchDataset import OK
poseforge.pose.contrast import OK

(Note: load_atomic_batch_frames is a @staticmethod of AtomicBatchDataset, not a module-level function, so importing the bare name is expected to fail — importing the module and the class succeeds.)

Before the pvio fix, the existing suite could not even be collected:

ImportError: cannot import name '_default_ffmpeg_params_for_video_writing' from 'pvio.io'
!!!!!!!!!!!!!!!!!!!! Interrupted: 1 error during collection !!!!!!!!!!!!!!!!!!!!

After the fix, pytest tests/test_atomic_batch_dataset.py tests/test_synthetic_data_sampler.py tests/test_atomic_batch_guards.py:

5 passed, 1 failed

• New tests/test_atomic_batch_guards.py: 3 passed — real 256-tiled batch loads correctly at image_size=(256,256); the guard raises a clear ValueError at image_size=(512,512); geometry-inversion unit checks pass.
• tests/test_synthetic_data_sampler.py: 2 passed.
• tests/test_atomic_batch_dataset.py::test_atomic_batch_dataset_loading: 1 failed, but only due to a stale hardcoded data path predating the 4variants/ reorg (ValueError: Provided path bulk_data/pose_estimation/atomic_batches/BO_Gal4_fly1_trial001 is not a directory). This is unrelated to the changes here — the pvio fix made the module importable, so the failure is now a clean missing-directory error rather than an ImportError at collection. Updating that test's data path is out of scope for the files this PR is allowed to touch.

🤖 Generated with Claude Code

---

Update (Claude Code, upon user instruction)

Also repointed the stale data path in tests/test_atomic_batch_dataset.py (atomic_batches/BO_Gal4_fly1_trial001/ → atomic_batches/4variants/BO_Gal4_fly1_trial001/). Combined with the pvio import fix in this PR, that test now passes (it previously errored at collection). Human verification: pytest tests/test_atomic_batch_dataset.py with the data present.