feat(stereo-split): correct right-eye colour in joined H.264 conversion

jobs/stereo-split/Dockerfile

@@ -29,6 +29,7 @@ COPY jobs/stereo-split/requirements.txt /tmp/requirements.txt
 RUN --mount=type=cache,target=/root/.cache/pip \
     pip install --index-url "${PYPI_INDEX_URL}" --timeout 120 --retries 5 -r /tmp/requirements.txt
 
+COPY jobs/stereo-split/color_consistency.py /app/color_consistency.py
 COPY jobs/stereo-split/convert_mcap_stereo_h264.py /app/convert_mcap_stereo_h264.py
 COPY jobs/stereo-split/imu_decoder.py /app/imu_decoder.py
 COPY jobs/stereo-split/run_processing.py /app/run_processing.py

jobs/stereo-split/convert_mcap_stereo_h264.py

@@ -26,6 +26,11 @@
 import numpy as np
 from rosbags.typesys import Stores, get_typestore
 
+from color_consistency import (
+    ColorConsistencyCalibrator,
+    ColorConsistencyConfig,
+    ColorCorrectionPlan,
+)
 from imu_decoder import ImuSample, decode_imu_from_bgr
 
 
@@ -147,6 +152,9 @@ class ConverterConfig:
     max_bitrate: str = "16M"
     buffer_size: str = "24M"
     gop: int = 20
+    #: Fit and apply the right-to-left colour-consistency correction. Only the
+    #: joined H.264 path supports it; split inputs keep their payloads untouched.
+    apply_color_consistency: bool = True
 
 
 @dataclass
@@ -166,6 +174,12 @@ class ConvertStats:
     timestamp_log_repair_applied: bool = False
     timestamp_publish_repair_applied: bool = False
     timestamp_repaired_messages: int = 0
+    color_decision: str = ""
+    color_sampled_frames: int = 0
+    color_matches_before_filter: int = 0
+    color_matches_after_filter: int = 0
+    color_corrected_frames: int = 0
+    color_report: dict | None = None
 
 
 @dataclass(frozen=True)
@@ -863,8 +877,13 @@ def _close_streams(self) -> None:
 
 
 class StereoSplitH264Converter:
-    def __init__(self, config: ConverterConfig | None = None) -> None:
+    def __init__(
+        self,
+        config: ConverterConfig | None = None,
+        color_config: ColorConsistencyConfig | None = None,
+    ) -> None:
         self.config = config or ConverterConfig()
+        self.color_config = color_config or ColorConsistencyConfig()
         self.typestore = get_typestore(Stores.ROS2_JAZZY)
         self.msg = self.typestore.types
 
@@ -951,6 +970,10 @@ def convert(
                     f"H.264 input requires source IMU topic: {config.imu_topic}"
                 )
             generate_imu = input_codec == "jpeg" and source_imu_messages == 0
+            color_plan: ColorCorrectionPlan | None = None
+            if input_codec == "h264" and config.apply_color_consistency:
+                source.seek(0)
+                color_plan = self._calibrate_color(source, config, stats)
             source.seek(0)
             reader = make_reader(source)
 
@@ -993,6 +1016,9 @@ def process_frame(frame: np.ndarray, metadata: VideoMetadata) -> None:
                         f"frame {frame.shape[1]}x{frame.shape[0]} is smaller than "
                         f"required {required_width}x{config.eye_height}"
                     )
+                if color_plan is not None and color_plan.applied:
+                    frame = self._correct_right_eye(frame, config, color_plan)
+                    stats.color_corrected_frames += 1
                 if encoder is None:
                     # JPEG captures are upside down; H.264 captures already have final orientation.
                     encoder = DualH264Encoder(
@@ -1132,6 +1158,11 @@ def process_frame(frame: np.ndarray, metadata: VideoMetadata) -> None:
                     )
                 if video_metadata:
                     raise RuntimeError("H.264 encoder did not return every submitted frame")
+                if color_plan is not None and color_plan.applied \
+                        and stats.color_corrected_frames != stats.right_videos:
+                    raise RuntimeError(
+                        "color correction did not cover every right-eye frame: "
+                        f"{stats.color_corrected_frames}/{stats.right_videos}")
                 ordered_writer.flush_all()
                 writer.finish()
             except BaseException:
@@ -1144,6 +1175,94 @@ def process_frame(frame: np.ndarray, metadata: VideoMetadata) -> None:
         stats.copied_topics = len(copied_topic_ids)
         return stats
 
+    def _calibrate_color(
+        self,
+        source: BinaryIO,
+        config: ConverterConfig,
+        stats: ConvertStats,
+    ) -> ColorCorrectionPlan:
+        """Decode the joined H.264 stream once and fit one fixed colour correction.
+
+        H.264 is inter-frame coded, so the sampled frames come from a full
+        sequential decode; only feature matching on the sampled frames is
+        expensive. The fitted model is frozen for the whole recording, so the
+        correction cannot flicker over time.
+        """
+        calibrator = ColorConsistencyCalibrator(self.color_config)
+        frame_width = config.metadata_width + config.eye_width * 2
+        decoder: H264FrameDecoder | None = None
+        aborted = False
+        index = 0
+
+        def observe(frame: np.ndarray) -> None:
+            nonlocal index
+            start = config.metadata_width + config.eye_width
+            calibrator.observe(
+                index,
+                frame[0:config.eye_height, config.metadata_width:start],
+                frame[0:config.eye_height, start:frame_width],
+            )
+            index += 1
+
+        for schema, _, message in make_reader(source).iter_messages(
+            topics=[config.input_topic], log_time_order=False
+        ):
+            # The main pass owns input validation. Anything this advisory pass
+            # cannot sample stops the calibration and lets the main pass report
+            # the canonical error, including mixed JPEG and H.264 frames.
+            if schema is None or schema.name != "sensor_msgs/msg/CompressedImage":
+                aborted = True
+                break
+            compressed = self.typestore.deserialize_cdr(
+                message.data, "sensor_msgs/msg/CompressedImage"
+            )
+            if _compressed_image_codec(compressed.format) != "h264":
+                aborted = True
+                break
+            access_unit = bytes(compressed.data)
+            if decoder is None:
+                if 5 not in _h264_nal_types(access_unit):
+                    # Frames before the first IDR cannot be decoded. The main
+                    # pass skips the same messages, so skipping them here keeps
+                    # both passes counting the same decoded frames.
+                    continue
+                decoder = H264FrameDecoder(frame_width, config.eye_height)
+            for frame in decoder.submit(access_unit):
+                observe(frame)
+        if decoder is not None:
+            if aborted:
+                decoder.abort()
+            else:
+                for frame in decoder.finish():
+                    observe(frame)
+
+        plan = calibrator.fit()
+        report = plan.report
+        stats.color_decision = plan.decision
+        stats.color_sampled_frames = report.sampled_frames
+        stats.color_matches_before_filter = report.matches_before_filter
+        stats.color_matches_after_filter = report.matches_after_filter
+        stats.color_report = report.as_dict()
+        return plan
+
+    @staticmethod
+    def _correct_right_eye(
+        frame: np.ndarray,
+        config: ConverterConfig,
+        plan: ColorCorrectionPlan,
+    ) -> np.ndarray:
+        """Return a writable frame whose right-eye crop carries the correction.
+
+        The crop offsets match the encoder filter graph exactly: joined H.264
+        input is never rotated, so this crop and the FFmpeg crop are the same.
+        """
+        start = config.metadata_width + config.eye_width
+        corrected = np.ascontiguousarray(frame).copy()
+        corrected[0:config.eye_height, start:start + config.eye_width] = plan.apply(
+            frame[0:config.eye_height, start:start + config.eye_width]
+        )
+        return corrected
+
     def _capture_source_camera_serial(
         self, stats: ConvertStats, schema, message
     ) -> None:

jobs/stereo-split/run_processing.py

@@ -125,6 +125,47 @@ def write_json(path: Path, value: dict[str, object]) -> None:
         os.fsync(stream.fileno())
 
 
+def color_consistency_summary(
+    report: dict[str, object] | None, corrected_frames: int
+) -> dict[str, object] | None:
+    """Project the fitted colour report onto the manifest's lean summary.
+
+    The full report, including the fitted gain curves and spatial field, stays
+    in the output metadata; the manifest only carries what a validator needs to
+    describe and bound the correction.
+    """
+    if not report:
+        return None
+    summary: dict[str, object] = {"corrected_frames": corrected_frames}
+    for key in (
+        "algorithm_version",
+        "reference_eye",
+        "decision",
+        "reason",
+        "sampled_frames",
+        "matches_before_filter",
+        "matches_after_filter",
+        "train_samples",
+        "validation_samples",
+        "baseline_ciede2000_median",
+        "corrected_ciede2000_median",
+    ):
+        if report.get(key) is not None:
+            summary[key] = report[key]
+    model = report.get("model")
+    if isinstance(model, dict):
+        for source, target in (
+            ("sha256", "model_sha256"),
+            ("gain_bins", "gain_bins"),
+            ("spatial_grid", "spatial_grid"),
+            ("strength", "strength"),
+            ("highlight_protect", "highlight_protect"),
+        ):
+            if model.get(source) is not None:
+                summary[target] = model[source]
+    return summary
+
+
 def require_scratch_capacity(scratch: Path, source_size_bytes: int) -> None:
     scratch.mkdir(parents=True, exist_ok=True)
     required = source_size_bytes * SCRATCH_SPACE_MULTIPLIER
@@ -296,6 +337,12 @@ def run(args: argparse.Namespace) -> dict[str, object]:
     )
     stats_payload = asdict(stats)
     camera_serial = str(stats_payload.pop("camera_serial", "")).strip()
+    color_consistency = color_consistency_summary(
+        stats_payload.get("color_report"), stats_payload["color_corrected_frames"]
+    )
+    manifest_stats = {
+        key: value for key, value in stats_payload.items() if key != "color_report"
+    }
     manifest_calibration = None
     processing_mode = "timestamp_repair" if stats.input_mode == "split_h264" else "convert"
     metadata: dict[str, object] = {
@@ -318,6 +365,8 @@ def run(args: argparse.Namespace) -> dict[str, object]:
     }
     if camera_serial:
         metadata["camera_serial"] = camera_serial
+    if color_consistency is not None:
+        metadata["color_consistency"] = color_consistency
     write_json(local_metadata, metadata)
     mcap_identity = require_mcap_output(local_mcap)
     metadata_identity = require_output(local_metadata)
@@ -342,12 +391,14 @@ def run(args: argparse.Namespace) -> dict[str, object]:
             "mcap": {"name": OUTPUT_MCAP_NAME, **asdict(mcap_identity)},
             "metadata": {"name": OUTPUT_METADATA_NAME, **asdict(metadata_identity)},
         },
-        "stats": stats_payload,
+        "stats": manifest_stats,
         "started_at": started_at,
         "finished_at": utc_now(),
     }
     if camera_serial:
         manifest["camera_serial"] = camera_serial
+    if color_consistency is not None:
+        manifest["color_consistency"] = color_consistency
     if manifest_calibration is not None:
         manifest["calibration"] = manifest_calibration
     write_json(local_manifest, manifest)