From d21370ed817571caf86dbef4af3df69759735249 Mon Sep 17 00:00:00 2001 From: chaoliu Date: Wed, 16 Sep 2026 17:18:20 +0800 Subject: [PATCH] feat(stereo-split): correct right-eye colour in joined H.264 conversion Fit one right-to-left photometric correction from the joined H.264 stream and apply it to the right eye before encoding, so both eyes of a stereo-split derivative agree in colour. - color_consistency.py owns sampling, SIFT matching, exposure outlier rejection, luminance-binned gain curves, a smooth spatial gain field, held-out block validation, and the apply/skip decision. - The correction is fitted once and frozen for the whole recording, and it is only wired into the joined path; split H.264 inputs still keep their payloads untouched. - Replaying a real 1106-frame DECXIN capture moves the matched median CIEDE2000 from 13.3 to 4.0 while preserving IMU, serial number, and every video timestamp. - The manifest gains a lean colour summary and the output metadata keeps the full fitted model, so the processing mode and manifest schema are unchanged. --- jobs/stereo-split/Dockerfile | 1 + jobs/stereo-split/color_consistency.py | 843 ++++++++++++++++++ jobs/stereo-split/convert_mcap_stereo_h264.py | 121 ++- jobs/stereo-split/run_processing.py | 53 +- .../tests/test_color_consistency.py | 282 ++++++ .../tests/test_joined_color_correction.py | 378 ++++++++ .../stereo-split/tests/test_run_processing.py | 41 + 7 files changed, 1717 insertions(+), 2 deletions(-) create mode 100644 jobs/stereo-split/color_consistency.py create mode 100644 jobs/stereo-split/tests/test_color_consistency.py create mode 100644 jobs/stereo-split/tests/test_joined_color_correction.py diff --git a/jobs/stereo-split/Dockerfile b/jobs/stereo-split/Dockerfile index 546a8943..33f8c828 100644 --- a/jobs/stereo-split/Dockerfile +++ b/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 diff --git a/jobs/stereo-split/color_consistency.py b/jobs/stereo-split/color_consistency.py new file mode 100644 index 00000000..b3671231 --- /dev/null +++ b/jobs/stereo-split/color_consistency.py @@ -0,0 +1,843 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: 2026 ArcheBase +# SPDX-License-Identifier: MulanPSL-2.0 + +"""Fit one fixed right-to-left photometric correction from stereo frame pairs. + +This module owns the whole colour-consistency algorithm: matching the two eyes, +fitting a luminance-binned gain curve per channel plus a smooth spatial gain +field, validating the fitted model on held-out blocks of the recording, and +deciding whether the correction is worth applying at all. Callers only feed +decoded frame pairs and then apply the returned plan, so MCAP, H.264, topics, +and job orchestration stay outside this seam. + +The correction always maps the *source* eye (right) onto the *reference* eye +(left). It is photometric only: it cannot remove lens-internal chromatic +aberration, and with no colour chart it can only claim consistency with the +reference eye, never absolute colorimetric accuracy. +""" + +from __future__ import annotations + +import hashlib +import json +from dataclasses import dataclass + +import cv2 +import numpy as np + +#: Decision values reported by :meth:`ColorConsistencyCalibrator.fit`. +DECISION_APPLIED = "applied" +DECISION_NOT_NEEDED = "not_needed" +DECISION_INSUFFICIENT = "insufficient" +DECISION_REJECTED = "rejected" + +#: Algorithm identity recorded in the manifest so an output stays explainable. +ALGORITHM_VERSION = "stereo-color-v1" + +#: Size of the linear-light grid used to invert the transfer function. +_LINEAR_GRID_SIZE = 65536 + +#: Source level above which a sample counts as clipped when measuring how much +#: headroom a correction consumes. +UNSATURATED_LEVEL = 250 + +_XYZ_FROM_SRGB = np.array([ + [.4124564, .3575761, .1804375], + [.2126729, .7151522, .0721750], + [.0193339, .1191920, .9503041], +]) +_WHITE_XYZ = np.array([.95047, 1., 1.08883]) + + +def _build_transfer_tables() -> tuple[np.ndarray, np.ndarray]: + """Return the sRGB8->linear table and the linear->sRGB8 grid.""" + encoded = np.arange(256, dtype=np.float64) / 255.0 + srgb_to_linear = np.where( + encoded <= .04045, + encoded / 12.92, + ((encoded + .055) / 1.055) ** 2.4, + ).astype(np.float32) + grid = np.linspace(0.0, 1.0, _LINEAR_GRID_SIZE, dtype=np.float32) + linear_to_srgb = np.rint(np.where( + grid <= .0031308, + 12.92 * grid, + 1.055 * grid ** (1 / 2.4) - .055, + ) * 255.0) + return srgb_to_linear, np.clip(linear_to_srgb, 0, 255).astype(np.uint8) + + +_SRGB8_TO_LINEAR, _LINEAR_TO_SRGB8 = _build_transfer_tables() + + +@dataclass(frozen=True) +class ColorConsistencyConfig: + """Fixed algorithm parameters; a stereo-split image pins them all. + + The defaults are the configuration measured on real DECXIN recordings. They + are deliberately not runtime settings: one image digest must describe one + deterministic processing behaviour so a release can be rolled back by + selecting the previous digest. + """ + + #: Decode/analyse every Nth stereo frame. + sample_step: int = 20 + #: Scale used for feature matching; 0.5 is materially faster than 1.0. + match_scale: float = .5 + #: Feature budget per sampled frame and per eye. + max_match_features: int = 1000 + #: Hard cap on retained correspondences, keeping memory bounded. + max_samples: int = 100_000 + #: Per-channel gain curve resolution; 1 collapses to a global gain. + gain_bins: int = 8 + #: Spatial gain field resolution; 0 disables the spatial field. + spatial_grid: int = 10 + #: How much of the fitted correction to apply; the fitted value is 1.0. + strength: float = 1.0 + #: Source level above which the correction fades back to the source. + highlight_protect: int = 245 + #: Log-brightness tolerance for dropping exposure-mismatched pairs. + drop_exposure_outliers: float = .7 + #: Correspondences required per gain bin and per spatial cell. + min_gain_bin_samples: int = 150 + min_spatial_cell_samples: int = 40 + #: Sampled frames and correspondences required to attempt a fit at all. + min_sampled_frames: int = 8 + min_matched_samples: int = 500 + #: Sampling block size in sampled frames; alternating blocks train and validate. + train_validation_block_size: int = 4 + #: Smallest held-out improvement, and largest tolerated regressions. + min_improvement: float = .2 + max_p95_regression: float = 1.05 + max_clipped_increase: float = .005 + #: Baseline median CIEDE2000 below which the eyes are already consistent. + apply_threshold: float = 3.0 + + def __post_init__(self) -> None: + if self.sample_step < 1: + raise ValueError("sample_step must be at least 1") + if not 0 < self.match_scale <= 1: + raise ValueError("match_scale must be in (0, 1]") + if self.gain_bins < 1: + raise ValueError("gain_bins must be at least 1") + if self.spatial_grid < 0: + raise ValueError("spatial_grid must be non-negative") + if self.strength < 0: + raise ValueError("strength must be non-negative") + if not 0 <= self.highlight_protect <= 255: + raise ValueError("highlight_protect must be in [0, 255]") + if self.train_validation_block_size < 1: + raise ValueError("train_validation_block_size must be at least 1") + + +def srgb_to_linear(values: np.ndarray) -> np.ndarray: + """Convert gamma-encoded sRGB in [0, 1] to linear light.""" + values = np.asarray(values, dtype=np.float64) + return np.where(values <= .04045, values / 12.92, ((values + .055) / 1.055) ** 2.4) + + +def linear_to_srgb(values: np.ndarray) -> np.ndarray: + """Convert linear light to gamma-encoded sRGB in [0, 1].""" + values = np.clip(np.asarray(values, dtype=np.float64), 0.0, 1.0) + return np.where(values <= .0031308, 12.92 * values, 1.055 * values ** (1 / 2.4) - .055) + + +def srgb_to_lab(rgb: np.ndarray) -> np.ndarray: + """Convert gamma-encoded sRGB in [0, 1] to CIE L*a*b* (D65, 2 degree).""" + linear = srgb_to_linear(np.asarray(rgb, dtype=np.float64)) + xyz = (linear @ _XYZ_FROM_SRGB.T) / _WHITE_XYZ + delta = 6 / 29 + f = np.where(xyz > delta ** 3, np.cbrt(xyz), xyz / (3 * delta ** 2) + 4 / 29) + return np.c_[116 * f[:, 1] - 16, 500 * (f[:, 0] - f[:, 1]), 200 * (f[:, 1] - f[:, 2])] + + +def ciede2000(lab1: np.ndarray, lab2: np.ndarray) -> np.ndarray: + """Vectorized CIEDE2000 colour difference (Sharma et al. formulation).""" + l1, a1, b1 = np.asarray(lab1, dtype=np.float64).T + l2, a2, b2 = np.asarray(lab2, dtype=np.float64).T + c1, c2 = np.hypot(a1, b1), np.hypot(a2, b2) + c_bar = (c1 + c2) / 2 + g = (1 - np.sqrt(c_bar ** 7 / (c_bar ** 7 + 25 ** 7))) / 2 + a1p, a2p = (1 + g) * a1, (1 + g) * a2 + c1p, c2p = np.hypot(a1p, b1), np.hypot(a2p, b2) + h1p = np.where(c1p == 0, 0.0, np.degrees(np.arctan2(b1, a1p)) % 360) + h2p = np.where(c2p == 0, 0.0, np.degrees(np.arctan2(b2, a2p)) % 360) + dlp, dcp = l2 - l1, c2p - c1p + dhp = h2p - h1p + dhp = np.where(c1p * c2p == 0, 0.0, + np.where(dhp > 180, dhp - 360, np.where(dhp < -180, dhp + 360, dhp))) + dh = 2 * np.sqrt(c1p * c2p) * np.sin(np.radians(dhp / 2)) + l_bar, c_barp = (l1 + l2) / 2, (c1p + c2p) / 2 + h_sum = h1p + h2p + h_bar = np.where(c1p * c2p == 0, h_sum, + np.where(np.abs(h1p - h2p) <= 180, h_sum / 2, + np.where(h_sum < 360, (h_sum + 360) / 2, (h_sum - 360) / 2))) + t = (1 - .17 * np.cos(np.radians(h_bar - 30)) + .24 * np.cos(np.radians(2 * h_bar)) + + .32 * np.cos(np.radians(3 * h_bar + 6)) - .20 * np.cos(np.radians(4 * h_bar - 63))) + d_theta = 30 * np.exp(-(((h_bar - 275) / 25) ** 2)) + rc = 2 * np.sqrt(c_barp ** 7 / (c_barp ** 7 + 25 ** 7)) + sl = 1 + .015 * (l_bar - 50) ** 2 / np.sqrt(20 + (l_bar - 50) ** 2) + sc = 1 + .045 * c_barp + sh = 1 + .015 * c_barp * t + rt = -np.sin(np.radians(2 * d_theta)) * rc + return np.sqrt((dlp / sl) ** 2 + (dcp / sc) ** 2 + (dh / sh) ** 2 + + rt * (dcp / sc) * (dh / sh)) + + +def _match_frame_pairs( + reference_bgr: np.ndarray, + source_bgr: np.ndarray, + detector: cv2.Feature2D, + match_scale: float, + max_level: int | None, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Match one stereo pair and return source pixels, reference pixels, positions. + + The source (right) coordinates are reported because the spatial gain field is + evaluated in the image that the correction is applied to; the two eyes are + offset relative to each other, so reference coordinates would not line up. + """ + if match_scale != 1: + reference_match = cv2.resize( + reference_bgr, None, fx=match_scale, fy=match_scale, interpolation=cv2.INTER_AREA + ) + source_match = cv2.resize( + source_bgr, None, fx=match_scale, fy=match_scale, interpolation=cv2.INTER_AREA + ) + else: + reference_match, source_match = reference_bgr, source_bgr + + empty = (np.empty((0, 3)), np.empty((0, 3)), np.empty((0, 2))) + detections = [ + detector.detectAndCompute(cv2.cvtColor(image, cv2.COLOR_BGR2GRAY), None) + for image in (reference_match, source_match) + ] + (key_reference, reference_descriptor), (key_source, source_descriptor) = detections + if reference_descriptor is None or source_descriptor is None: + return empty + + pairs = cv2.BFMatcher(cv2.NORM_L2).knnMatch(reference_descriptor, source_descriptor, k=2) + inverse_scale = 1.0 / match_scale + source_colors, reference_colors, positions = [], [], [] + for pair in pairs: + if len(pair) != 2: + continue + best, second = pair + if best.distance >= .7 * second.distance: + continue + x, y = key_reference[best.queryIdx].pt + u, v = key_source[best.trainIdx].pt + x, y, u, v = (int(round(value * inverse_scale)) for value in (x, y, u, v)) + x = min(max(x, 0), reference_bgr.shape[1] - 1) + y = min(max(y, 0), reference_bgr.shape[0] - 1) + u = min(max(u, 0), source_bgr.shape[1] - 1) + v = min(max(v, 0), source_bgr.shape[0] - 1) + if abs(x - u) > 250 or abs(y - v) > 250: + continue + level_reference = float(np.mean(reference_bgr[y, x])) + level_source = float(np.mean(source_bgr[v, u])) + if level_reference <= 5 or level_source <= 5: + continue + if max_level is not None and (level_reference >= max_level or level_source >= max_level): + continue + reference_colors.append(reference_bgr[y, x][::-1] / 255.0) + source_colors.append(source_bgr[v, u][::-1] / 255.0) + positions.append((u / source_bgr.shape[1], v / source_bgr.shape[0])) + if not source_colors: + return empty + return ( + np.asarray(source_colors), + np.asarray(reference_colors), + np.asarray(positions), + ) + + +def exposure_outlier_mask(source: np.ndarray, reference: np.ndarray, tolerance: float) -> np.ndarray: + """Flag pairs whose brightness ratio is far from the fitted gain. + + A pair whose source is much brighter than the expected ratio is an + over-exposed sample: it reflects an exposure mismatch, not a tone + difference, and pulling it into the fit biases the highlight end. + """ + if tolerance <= 0 or not len(source): + return np.ones(len(source), dtype=bool) + gain = np.sum(source * reference, 0) / (np.sum(source * source, 0) + 1e-8) + expected = float(np.log(1 / np.clip(gain, 1e-6, None)).mean()) + deviation = np.abs(np.log(np.maximum(source.mean(1), 1e-6)) + - np.log(np.maximum(reference.mean(1), 1e-6)) - expected) + return deviation <= tolerance + + +def _fit_gain_bins( + source: np.ndarray, + reference: np.ndarray, + bins: int, + min_samples: int, +) -> tuple[tuple[np.ndarray, ...], tuple[np.ndarray, ...]]: + """Fit a per-channel gain curve from binned correspondences. + + Each channel is binned on its *own* value, which keeps the correction + expressible as one 256-entry table per channel. A channel with too few + populated bins falls back to its single global least-squares gain. + """ + bins = max(2, min(bins, len(source) // (4 * min_samples) or 2)) + anchors: list[np.ndarray] = [] + gains: list[np.ndarray] = [] + for channel in range(3): + values = source[:, channel] + global_gain = float(np.sum(values * reference[:, channel]) + / (np.sum(values ** 2) + 1e-8)) + edges = np.unique(np.quantile(values, np.linspace(0, 1, bins + 1))) + channel_anchors, channel_gains = [], [] + for low, high in zip(edges[:-1], edges[1:]): + mask = (values >= low) & (values <= high) + if mask.sum() < min_samples: + continue + channel_anchors.append(float(np.median(values[mask]))) + channel_gains.append(float(np.sum(values[mask] * reference[mask, channel]) + / (np.sum(values[mask] ** 2) + 1e-8))) + if not channel_anchors: + anchors.append(np.asarray([float(np.median(values))])) + gains.append(np.asarray([global_gain])) + continue + anchors.append(np.asarray(channel_anchors)) + gains.append(np.asarray(channel_gains)) + return tuple(anchors), tuple(gains) + + +def _fit_spatial_gain( + source: np.ndarray, + reference: np.ndarray, + positions: np.ndarray, + anchors: tuple[np.ndarray, ...], + gains: tuple[np.ndarray, ...], + grid: int, + min_samples: int, + smooth: float = 1.2, +) -> np.ndarray: + """Fit a smooth spatial gain field on the residual left by the gain curves. + + The stereo mismatch is not uniform across the frame: the gain needed on one + side of the image differs from the other by more than a factor of two, and no + global model can express that. Returns a small ``(grid, grid, 3)`` RGB field + that the corrector upsamples to the frame. + """ + corrected = np.empty_like(source) + for channel in range(3): + corrected[:, channel] = source[:, channel] * np.interp( + source[:, channel], anchors[channel], np.clip(gains[channel], .2, 3)) + residual = reference / np.maximum(corrected, 1e-6) + cell_x = np.clip((positions[:, 0] * grid).astype(int), 0, grid - 1) + cell_y = np.clip((positions[:, 1] * grid).astype(int), 0, grid - 1) + cells = np.ones((grid, grid, 3)) + populated = np.zeros((grid, grid), dtype=bool) + for row in range(grid): + for column in range(grid): + mask = (cell_y == row) & (cell_x == column) + if mask.sum() < min_samples: + continue + cells[row, column] = np.median(residual[mask], 0) + populated[row, column] = True + if not populated.any(): + return cells + # Empty cells are filled from the nearest populated one: a single global + # value leaves holes that the upsampled field turns into blotches across the + # frame, which costs far more than the field gains. + filled_y, filled_x = np.where(populated) + for row, column in zip(*np.where(~populated)): + nearest = np.argmin((filled_y - row) ** 2 + (filled_x - column) ** 2) + cells[row, column] = cells[filled_y[nearest], filled_x[nearest]] + for channel in range(3): + cells[:, :, channel] = cv2.GaussianBlur(cells[:, :, channel], (0, 0), smooth) + return cells + + +def _sample_spatial_gain(cells: np.ndarray, positions: np.ndarray) -> np.ndarray: + """Sample the fitted cell field at normalized positions the way the corrector does.""" + grid = cells.shape[0] + coordinates = np.clip(np.asarray(positions, dtype=np.float64), 0.0, 1.0) * grid - 0.5 + lower = np.clip(np.floor(coordinates).astype(int), 0, grid - 1) + upper = np.clip(lower + 1, 0, grid - 1) + fraction = np.clip(coordinates - lower, 0.0, 1.0) + lower_x, upper_x = lower[:, 0], upper[:, 0] + lower_y, upper_y = lower[:, 1], upper[:, 1] + fx = fraction[:, 0:1] + fy = fraction[:, 1:2] + top = cells[lower_y, lower_x] * (1 - fx) + cells[lower_y, upper_x] * fx + bottom = cells[upper_y, lower_x] * (1 - fx) + cells[upper_y, upper_x] * fx + return top * (1 - fy) + bottom * fy + + +class _ColorModel: + """Fitted gain curves plus an optional spatial field.""" + + def __init__( + self, + anchors: tuple[np.ndarray, ...], + gains: tuple[np.ndarray, ...], + cells: np.ndarray | None, + strength: float, + ) -> None: + self.anchors = tuple(np.asarray(anchor, dtype=np.float64) for anchor in anchors) + self.gains = tuple( + np.clip(np.asarray(gain, dtype=np.float64), .2, 3) for gain in gains + ) + self.cells = None if cells is None else np.asarray(cells, dtype=np.float64) + self.strength = float(strength) + + @property + def _scaled_gains(self) -> tuple[np.ndarray, ...]: + return tuple(1 + self.strength * (gain - 1) for gain in self.gains) + + def evaluate(self, source_linear: np.ndarray, positions: np.ndarray) -> np.ndarray: + """Evaluate the fitted model at specific linear-light samples.""" + corrected = np.empty_like(source_linear) + scaled = self._scaled_gains + for channel in range(3): + corrected[:, channel] = source_linear[:, channel] * np.interp( + source_linear[:, channel], self.anchors[channel], scaled[channel]) + if self.cells is not None and len(positions): + corrected *= _sample_spatial_gain(self.cells, positions) + return corrected + + def corrector(self, highlight_protect: int): + """Build the per-frame transform, preferring the cheapest correct path.""" + if self.cells is None: + return _highlight_protect(_channel_table_corrector(self), highlight_protect) + return _highlight_protect(_spatial_corrector(self), highlight_protect) + + def fingerprint(self) -> str: + payload = { + "algorithm_version": ALGORITHM_VERSION, + "strength": round(self.strength, 6), + "anchors": [[round(float(value), 6) for value in anchor] for anchor in self.anchors], + "gains": [[round(float(value), 6) for value in gain] for gain in self.gains], + "spatial_gain": None if self.cells is None else np.round(self.cells, 6).tolist(), + } + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":")).encode() + return hashlib.sha256(encoded).hexdigest() + + +def _channel_table_corrector(model: _ColorModel): + """Collapse channel-separable gains into one 256-entry table per channel.""" + linear = _SRGB8_TO_LINEAR.astype(np.float64) + columns = [] + scaled = model._scaled_gains + for channel in range(3): + curve = np.interp(linear, model.anchors[channel], scaled[channel]) + indices = np.rint(np.clip(curve * linear, 0, 1) + * (_LINEAR_GRID_SIZE - 1)).astype(np.uint16) + columns.append(_LINEAR_TO_SRGB8[indices]) + table = np.ascontiguousarray(np.stack(columns, 1)[:, ::-1]).reshape(1, 256, 3) + + def correct(image: np.ndarray) -> np.ndarray: + return cv2.LUT(np.ascontiguousarray(image), table) + + return correct + + +def _spatial_corrector(model: _ColorModel): + """Gain curves plus a smooth spatial gain field, applied in linear light. + + The spatial multiply has to happen in linear light, so this path runs + srgb->linear table, per-pixel multiply, and a 16-bit index back into the + inverse transfer table. OpenCV's ``LUT`` only accepts 8-bit sources, so the + inverse step uses NumPy indexing instead. + """ + linear = _SRGB8_TO_LINEAR.astype(np.float64) + columns = [] + scaled = model._scaled_gains + for channel in range(3): + # The table holds the gain-corrected *linear* value, not the gain: the + # spatial field multiplies it afterwards, still in linear light. + columns.append((np.interp(linear, model.anchors[channel], scaled[channel]) + * linear).astype(np.float32)) + linear_table = np.ascontiguousarray(np.stack(columns, 1)[:, ::-1]).reshape(1, 256, 3) + small = np.ascontiguousarray(model.cells[:, :, ::-1], dtype=np.float32) + cache: dict[tuple[int, int], np.ndarray] = {} + + def correct(image: np.ndarray) -> np.ndarray: + image = np.ascontiguousarray(image) + height, width = image.shape[:2] + gain_map = cache.get((height, width)) + if gain_map is None: + gain_map = np.ascontiguousarray( + cv2.resize(small, (width, height), interpolation=cv2.INTER_LINEAR)) + cache[(height, width)] = gain_map + linear_image = cv2.LUT(image, linear_table) + indices = cv2.multiply(linear_image, gain_map, scale=65535.0, dtype=cv2.CV_16U) + return _LINEAR_TO_SRGB8[indices] + + return correct + + +def _highlight_protect(correct, level: int): + """Fade the correction back to the source where pixels are saturated. + + A multiplicative correction cannot restore a clipped highlight: it tints a + blown-out white instead. Fading to the source above ``level`` (on the + smallest channel) leaves those pixels alone. No-op when ``level <= 0``. + """ + if level <= 0: + return correct + ramp = np.clip((np.arange(256) - level) / (255 - level), 0, 1) + weight_lut = np.rint(ramp * 255).astype(np.uint8) + + def protected(image: np.ndarray) -> np.ndarray: + corrected = correct(image) + blue, green, red = cv2.split(image) + weight = cv2.merge( + [cv2.LUT(cv2.min(cv2.min(blue, green), red), weight_lut)] * 3) + return cv2.add(cv2.multiply(corrected, cv2.bitwise_not(weight), scale=1 / 255.0), + cv2.multiply(image, weight, scale=1 / 255.0)) + + return protected + + +#: Metrics reported when no usable measurement exists. +_EMPTY_METRICS = { + "ciede2000_median": 0.0, + "ciede2000_p95": 0.0, + "linear_mae_median": 0.0, + "clipped_fraction": 0.0, +} + + +def _measure(source_srgb: np.ndarray, reference_srgb: np.ndarray) -> dict[str, float]: + """Perceptual and photometric disagreement on one set of matched pairs.""" + source_lab = srgb_to_lab(source_srgb) + reference_lab = srgb_to_lab(reference_srgb) + delta_e = ciede2000(source_lab, reference_lab) + source_linear = srgb_to_linear(source_srgb) + reference_linear = srgb_to_linear(reference_srgb) + return { + "ciede2000_median": float(np.median(delta_e)), + "ciede2000_p95": float(np.percentile(delta_e, 95)), + "linear_mae_median": float(np.median(np.mean(np.abs(source_linear - reference_linear), 1))), + "clipped_fraction": float(np.mean(np.max(source_srgb, axis=1) >= UNSATURATED_LEVEL / 255.0)), + } + + +@dataclass(frozen=True) +class ColorCorrectionReport: + """Immutable, serializable outcome of one calibration.""" + + algorithm_version: str + decision: str + reason: str + reference_eye: str + sampled_frames: int + matches_before_filter: int + matches_after_filter: int + train_samples: int + validation_samples: int + baseline: dict[str, float] + corrected: dict[str, float] + model: dict[str, object] + + def as_dict(self) -> dict[str, object]: + """Return the full report, including the fitted model parameters.""" + payload: dict[str, object] = { + "algorithm_version": self.algorithm_version, + "reference_eye": self.reference_eye, + "decision": self.decision, + "reason": self.reason, + "sampled_frames": self.sampled_frames, + "matches_before_filter": self.matches_before_filter, + "matches_after_filter": self.matches_after_filter, + "train_samples": self.train_samples, + "validation_samples": self.validation_samples, + } + for prefix, metrics in (("baseline", self.baseline), ("corrected", self.corrected)): + for name, value in metrics.items(): + payload[f"{prefix}_{name}"] = value + payload["model"] = dict(self.model) + return payload + + +class ColorCorrectionPlan: + """The fixed correction for one recording, or the reason none was applied.""" + + def __init__( + self, + report: ColorCorrectionReport, + corrector=None, + ) -> None: + self._report = report + self._corrector = corrector + + @property + def decision(self) -> str: + return self._report.decision + + @property + def reason(self) -> str: + return self._report.reason + + @property + def applied(self) -> bool: + """Whether :meth:`apply` may be used on this recording.""" + return self._report.decision == DECISION_APPLIED and self._corrector is not None + + @property + def report(self) -> ColorCorrectionReport: + return self._report + + @property + def summary(self) -> dict[str, object]: + """Return the report without the fitted model arrays.""" + payload = self._report.as_dict() + payload.pop("model", None) + model = self._report.model + payload["model_sha256"] = model.get("sha256") + payload["gain_bins"] = model.get("gain_bins") + payload["spatial_grid"] = model.get("spatial_grid") + payload["strength"] = model.get("strength") + payload["highlight_protect"] = model.get("highlight_protect") + return payload + + def apply(self, source_bgr: np.ndarray) -> np.ndarray: + """Apply the fixed correction to one right-eye BGR frame.""" + if not self.applied: + raise RuntimeError( + f"color correction plan is not applied: {self.decision} ({self.reason})") + return self._corrector(source_bgr) + + +class ColorConsistencyCalibrator: + """Accumulate stereo correspondences, then fit and validate one model. + + The calibrator owns sampling, matching, outlier rejection, the train and + validation split, the fit, and the apply/skip decision, so the caller only + feeds decoded frames and then uses the returned plan. + """ + + def __init__(self, config: ColorConsistencyConfig | None = None) -> None: + self.config = config or ColorConsistencyConfig() + self._detector = cv2.SIFT_create(nfeatures=self.config.max_match_features) + self._source_samples: list[np.ndarray] = [] + self._reference_samples: list[np.ndarray] = [] + self._sample_positions: list[np.ndarray] = [] + self._sample_blocks: list[int] = [] + self._sampled_frames = 0 + self._matches_before_filter = 0 + self._sample_ordinals = 0 + self._plan: ColorCorrectionPlan | None = None + + @property + def sampled_frames(self) -> int: + """Number of frame pairs that produced at least one correspondence.""" + return self._sampled_frames + + @property + def observed_samples(self) -> int: + """Number of retained correspondences, before outlier rejection.""" + return self._matches_before_filter + + def observe(self, frame_index: int, reference_bgr: np.ndarray, source_bgr: np.ndarray) -> bool: + """Feed one aligned stereo pair; returns whether it was sampled.""" + if self._plan is not None: + raise RuntimeError("calibrator already fitted") + if frame_index < 0: + raise ValueError("frame_index must be non-negative") + if reference_bgr.shape != source_bgr.shape: + raise ValueError("reference and source frames must share one shape") + if frame_index % self.config.sample_step: + return False + source, reference, positions = _match_frame_pairs( + reference_bgr, + source_bgr, + self._detector, + self.config.match_scale, + None, + ) + self._sample_ordinals += 1 + if not len(source): + return False + self._sampled_frames += 1 + room = self.config.max_samples - self._matches_before_filter + if room <= 0: + return True + accepted = min(room, len(source)) + self._source_samples.append(source[:accepted]) + self._reference_samples.append(reference[:accepted]) + self._sample_positions.append(positions[:accepted]) + self._sample_blocks.append( + np.full(accepted, (self._sample_ordinals - 1) // self.config.train_validation_block_size, + dtype=np.int64)) + self._matches_before_filter += accepted + return True + + def fit(self) -> ColorCorrectionPlan: + """Fit, validate, and freeze the correction; idempotent.""" + if self._plan is not None: + return self._plan + config = self.config + source = (np.concatenate(self._source_samples) if self._source_samples + else np.empty((0, 3))) + reference = (np.concatenate(self._reference_samples) if self._reference_samples + else np.empty((0, 3))) + positions = (np.concatenate(self._sample_positions) if self._sample_positions + else np.empty((0, 2))) + blocks = (np.concatenate(self._sample_blocks) if self._sample_blocks + else np.empty(0, dtype=np.int64)) + + if self._sampled_frames < config.min_sampled_frames or len(source) < config.min_matched_samples: + self._plan = self._skip( + DECISION_INSUFFICIENT, + "insufficient_correspondences", + source, reference, 0, len(source), + ) + return self._plan + + keep = exposure_outlier_mask( + srgb_to_linear(source), srgb_to_linear(reference), config.drop_exposure_outliers) + source, reference, positions, blocks = source[keep], reference[keep], positions[keep], blocks[keep] + after_filter = len(source) + train = blocks % 2 == 0 + # Adjacent frames are highly correlated, so blocks alternate between + # train and validation instead of splitting the recording in half. + if train.sum() < 2 or (~train).sum() < 2: + self._plan = self._skip( + DECISION_INSUFFICIENT, "insufficient_correspondences", + source, reference, after_filter, int((~train).sum()), + train_mask=train, + ) + return self._plan + + baseline_validation = _measure(source[~train], reference[~train]) + baseline_train = _measure(source[train], reference[train]) + if baseline_validation["ciede2000_median"] < config.apply_threshold: + self._plan = self._skip( + DECISION_NOT_NEEDED, "baseline_within_threshold", + source, reference, after_filter, int((~train).sum()), + baseline=baseline_validation, train_baseline=baseline_train, + train_mask=train, + ) + return self._plan + + source_linear = srgb_to_linear(source) + reference_linear = srgb_to_linear(reference) + anchors, gains = _fit_gain_bins( + source_linear[train], reference_linear[train], + config.gain_bins, config.min_gain_bin_samples) + cells = None + if config.spatial_grid > 0: + cells = _fit_spatial_gain( + source_linear[train], reference_linear[train], positions[train], + anchors, gains, config.spatial_grid, config.min_spatial_cell_samples) + model = _ColorModel(anchors, gains, cells, config.strength) + if not _model_is_finite(model): + self._plan = self._skip( + DECISION_REJECTED, "invalid_model", + source, reference, after_filter, int((~train).sum()), + baseline=baseline_validation, train_mask=train, + ) + return self._plan + + corrected_linear = model.evaluate(source_linear, positions) + corrected_srgb = linear_to_srgb(corrected_linear) + corrected_validation = _measure(corrected_srgb[~train], reference[~train]) + corrected_train = _measure(corrected_srgb[train], reference[train]) + + reason = self._rejection_reason(baseline_validation, corrected_validation) + model_payload = { + "gain_bins": config.gain_bins, + "spatial_grid": config.spatial_grid, + "strength": config.strength, + "highlight_protect": config.highlight_protect, + "anchors": [anchor.tolist() for anchor in model.anchors], + "gains": [gain.tolist() for gain in model.gains], + "spatial_gain": None if model.cells is None else np.round(model.cells, 6).tolist(), + "sha256": model.fingerprint(), + } + if reason is not None: + self._plan = ColorCorrectionPlan(self._build_report( + DECISION_REJECTED, reason, after_filter, int(train.sum()), int((~train).sum()), + baseline_validation, corrected_validation, model_payload)) + return self._plan + + self._plan = ColorCorrectionPlan( + self._build_report( + DECISION_APPLIED, "validation_improved", after_filter, int(train.sum()), + int((~train).sum()), baseline_validation, corrected_validation, model_payload), + model.corrector(config.highlight_protect), + ) + return self._plan + + def _rejection_reason( + self, + baseline: dict[str, float], + corrected: dict[str, float], + ) -> str | None: + config = self.config + if corrected["ciede2000_median"] > baseline["ciede2000_median"] * (1 - config.min_improvement): + return "validation_not_improved" + if corrected["ciede2000_p95"] > baseline["ciede2000_p95"] * config.max_p95_regression: + return "validation_p95_regressed" + if (corrected["clipped_fraction"] - baseline["clipped_fraction"]) > config.max_clipped_increase: + return "clipped_highlight_increase" + return None + + def _build_report( + self, + decision: str, + reason: str, + after_filter: int, + train_samples: int, + validation_samples: int, + baseline: dict[str, float], + corrected: dict[str, float], + model: dict[str, object], + train_baseline: dict[str, float] | None = None, + ) -> ColorCorrectionReport: + empty = {"ciede2000_median": 0.0, "ciede2000_p95": 0.0, + "linear_mae_median": 0.0, "clipped_fraction": 0.0} + payload = dict(model) + payload["train_baseline_ciede2000_median"] = ( + float(train_baseline["ciede2000_median"]) if train_baseline else 0.0) + return ColorCorrectionReport( + algorithm_version=ALGORITHM_VERSION, + decision=decision, + reason=reason, + reference_eye="left", + sampled_frames=self._sampled_frames, + matches_before_filter=self._matches_before_filter, + matches_after_filter=after_filter, + train_samples=train_samples, + validation_samples=validation_samples, + baseline=dict(baseline or empty), + corrected=dict(corrected or empty), + model=payload, + ) + + def _skip( + self, + decision: str, + reason: str, + source: np.ndarray, + reference: np.ndarray, + after_filter: int, + validation_samples: int, + baseline: dict[str, float] | None = None, + train_baseline: dict[str, float] | None = None, + train_mask: np.ndarray | None = None, + ) -> ColorCorrectionPlan: + metrics = baseline + if metrics is None: + metrics = _measure(source, reference) if len(source) else _EMPTY_METRICS + if train_baseline is None and train_mask is not None and len(source): + train_baseline = _measure(source[train_mask], reference[train_mask]) + return ColorCorrectionPlan(self._build_report( + decision, reason, after_filter, + int(train_mask.sum()) if train_mask is not None else 0, + validation_samples, metrics, metrics, {}, train_baseline)) + + +def _model_is_finite(model: _ColorModel) -> bool: + for anchor, gain in zip(model.anchors, model.gains): + if not np.all(np.isfinite(anchor)) or not np.all(np.isfinite(gain)): + return False + if model.cells is not None and not np.all(np.isfinite(model.cells)): + return False + return True diff --git a/jobs/stereo-split/convert_mcap_stereo_h264.py b/jobs/stereo-split/convert_mcap_stereo_h264.py index 69aeebe3..fee3a71f 100644 --- a/jobs/stereo-split/convert_mcap_stereo_h264.py +++ b/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: diff --git a/jobs/stereo-split/run_processing.py b/jobs/stereo-split/run_processing.py index 7fdbaf6c..8e07c0d0 100644 --- a/jobs/stereo-split/run_processing.py +++ b/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) diff --git a/jobs/stereo-split/tests/test_color_consistency.py b/jobs/stereo-split/tests/test_color_consistency.py new file mode 100644 index 00000000..a867b8e0 --- /dev/null +++ b/jobs/stereo-split/tests/test_color_consistency.py @@ -0,0 +1,282 @@ +# SPDX-FileCopyrightText: 2026 ArcheBase +# SPDX-License-Identifier: MulanPSL-2.0 + +from __future__ import annotations + +import json +from pathlib import Path +import sys +import tempfile +import unittest + +import cv2 +import numpy as np + + +JOB_ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(JOB_ROOT)) +from color_consistency import ( # noqa: E402 + ALGORITHM_VERSION, + DECISION_APPLIED, + DECISION_INSUFFICIENT, + DECISION_NOT_NEEDED, + DECISION_REJECTED, + ColorConsistencyCalibrator, + ColorConsistencyConfig, + ciede2000, + srgb_to_lab, +) + + +def textured_scene(width: int = 480, height: int = 320, seed: int = 7) -> np.ndarray: + """Build a high-contrast textured BGR image with repeatable features.""" + rng = np.random.default_rng(seed) + image = np.zeros((height, width, 3), dtype=np.float64) + yy, xx = np.mgrid[0:height, 0:width] + image[..., 0] = 60 + 90 * xx / width + image[..., 1] = 70 + 80 * yy / height + image[..., 2] = 110 + 60 * (1 - xx / width) + for _ in range(70): + center = ( + int(rng.integers(20, width - 20)), + int(rng.integers(20, height - 20)), + ) + radius = int(rng.integers(6, 18)) + level = float(rng.integers(20, 200)) + cv2.circle( + image, + center, + radius, + (level, min(200.0, level + 40), max(15.0, level - 50)), + -1, + ) + for _ in range(35): + origin = ( + int(rng.integers(0, width - 70)), + int(rng.integers(0, height - 70)), + ) + size = (int(rng.integers(20, 60)), int(rng.integers(20, 60))) + cv2.rectangle( + image, + origin, + (origin[0] + size[0], origin[1] + size[1]), + ( + float(rng.integers(10, 200)), + float(rng.integers(10, 200)), + float(rng.integers(10, 200)), + ), + -1, + ) + # Keep every level clear of the clipped-highlight metric so the fixture + # measures photometric mismatch rather than saturation. + return cv2.GaussianBlur(np.clip(image, 0, 200).astype(np.uint8), (0, 0), 1.0) + + +def mismatched_source( + reference: np.ndarray, + gains: tuple[float, float, float] = (.65, 1.35, 1.25), + spatial: float = .35, + disparity: int = 0, +) -> np.ndarray: + """Return the source eye: gained, spatially graded, and optionally shifted.""" + height, width = reference.shape[:2] + field = 1.0 + spatial * (.5 - np.linspace(0, 1, width)[None, :, None]) + values = reference.astype(np.float64) / 255.0 + values = values * np.asarray(gains)[None, None, :] * field + source = np.clip(np.rint(values * 255.0), 0, 255).astype(np.uint8) + if disparity: + matrix = np.float32([[1, 0, disparity], [0, 1, 0]]) + source = cv2.warpAffine( + source, matrix, (width, height), borderMode=cv2.BORDER_REPLICATE) + return source + + +def scene_frames(count: int, disparity: int = 12): + """Yield ``count`` slightly varying reference/source frame pairs.""" + reference = textured_scene() + for index in range(count): + faded = cv2.convertScaleAbs(reference, alpha=1.0 - .002 * index, beta=-.5 * index) + yield faded, mismatched_source(faded, disparity=disparity) + + +def calibration_config(**overrides) -> ColorConsistencyConfig: + values = { + "sample_step": 1, + "match_scale": 1.0, + "min_sampled_frames": 8, + "min_matched_samples": 50, + "min_gain_bin_samples": 10, + "min_spatial_cell_samples": 5, + "spatial_grid": 4, + "train_validation_block_size": 2, + } + values.update(overrides) + return ColorConsistencyConfig(**values) + + +def fitted_plan(**overrides): + calibrator = ColorConsistencyCalibrator(calibration_config(**overrides)) + for index, (reference, source) in enumerate(scene_frames(12)): + calibrator.observe(index, reference, source) + return calibrator.fit() + + +class ColorConsistencyFitTest(unittest.TestCase): + def test_fits_and_improves_held_out_colour_difference(self) -> None: + plan = fitted_plan() + + self.assertEqual(plan.decision, DECISION_APPLIED) + self.assertEqual(plan.reason, "validation_improved") + self.assertTrue(plan.applied) + report = plan.report + self.assertGreater(report.sampled_frames, 8) + self.assertGreater(report.matches_before_filter, 0) + self.assertEqual( + report.train_samples + report.validation_samples, + report.matches_after_filter, + ) + self.assertGreater(report.train_samples, 0) + self.assertGreater(report.validation_samples, 0) + self.assertLess( + report.corrected["ciede2000_median"], + report.baseline["ciede2000_median"] * .5, + ) + self.assertLess( + report.corrected["linear_mae_median"], + report.baseline["linear_mae_median"] * .5, + ) + + def test_correction_reduces_pixel_difference_for_aligned_eyes(self) -> None: + reference = textured_scene() + source = mismatched_source(reference) + calibrator = ColorConsistencyCalibrator(calibration_config()) + for index in range(12): + calibrator.observe(index, reference, source) + plan = calibrator.fit() + + corrected = plan.apply(source) + self.assertEqual(corrected.shape, source.shape) + self.assertEqual(corrected.dtype, np.uint8) + before = float(np.mean(np.abs(source.astype(np.float64) - reference))) + after = float(np.mean(np.abs(corrected.astype(np.float64) - reference))) + self.assertLess(after, before * .4) + + def test_table_corrector_path_matches_held_out_metrics(self) -> None: + plan = fitted_plan(spatial_grid=0) + + self.assertEqual(plan.decision, DECISION_APPLIED) + self.assertIsNone(plan.report.model["spatial_gain"]) + corrected = plan.apply(textured_scene()) + self.assertEqual(corrected.shape, textured_scene().shape) + + def test_skips_when_eyes_are_already_consistent(self) -> None: + calibrator = ColorConsistencyCalibrator(calibration_config()) + reference = textured_scene() + for index in range(12): + calibrator.observe(index, reference, reference.copy()) + plan = calibrator.fit() + + self.assertEqual(plan.decision, DECISION_NOT_NEEDED) + self.assertEqual(plan.reason, "baseline_within_threshold") + self.assertFalse(plan.applied) + with self.assertRaises(RuntimeError): + plan.apply(reference) + + def test_reports_insufficient_correspondences_without_texture(self) -> None: + calibrator = ColorConsistencyCalibrator(calibration_config()) + flat = np.full((120, 160, 3), 128, dtype=np.uint8) + for index in range(12): + calibrator.observe(index, flat, flat) + plan = calibrator.fit() + + self.assertEqual(plan.decision, DECISION_INSUFFICIENT) + self.assertEqual(plan.reason, "insufficient_correspondences") + self.assertFalse(plan.applied) + self.assertEqual(plan.report.matches_after_filter, 0) + + def test_rejects_when_the_fit_cannot_beat_the_baseline(self) -> None: + plan = fitted_plan(min_improvement=.999) + + self.assertEqual(plan.decision, DECISION_REJECTED) + self.assertEqual(plan.reason, "validation_not_improved") + self.assertFalse(plan.applied) + + def test_highlight_protect_keeps_clipped_pixels(self) -> None: + plan = fitted_plan() + self.assertTrue(plan.applied) + frame = np.full((60, 80, 3), 180, dtype=np.uint8) + frame[10:20, 10:20] = 255 + frame[30:40, 30:40] = 120 + corrected = plan.apply(frame) + + self.assertTrue(np.array_equal(corrected[10:20, 10:20], frame[10:20, 10:20])) + self.assertFalse(np.array_equal(corrected[30:40, 30:40], frame[30:40, 30:40])) + + def test_same_input_produces_the_same_model(self) -> None: + first = fitted_plan().report.model["sha256"] + second = fitted_plan().report.model["sha256"] + + self.assertEqual(first, second) + + def test_summary_is_json_serializable_and_lists_the_algorithm(self) -> None: + plan = fitted_plan() + + summary = plan.summary + encoded = json.dumps(summary, sort_keys=True) + self.assertIn(ALGORITHM_VERSION, encoded) + self.assertNotIn("anchors", summary) + self.assertEqual(summary["algorithm_version"], ALGORITHM_VERSION) + self.assertEqual(summary["model_sha256"], plan.report.model["sha256"]) + full = json.dumps(plan.report.as_dict(), sort_keys=True) + self.assertIn("anchors", full) + temp_dir = tempfile.mkdtemp() + path = Path(temp_dir) / "report.json" + path.write_text(full, encoding="utf-8") + self.assertGreater(path.stat().st_size, 0) + + def test_sampling_step_skips_unsampled_frames(self) -> None: + calibrator = ColorConsistencyCalibrator(calibration_config(sample_step=3)) + for index, (reference, source) in enumerate(scene_frames(12)): + calibrator.observe(index, reference, source) + + self.assertEqual(calibrator.sampled_frames, 4) + + def test_rejects_mismatched_reference_and_source_shapes(self) -> None: + calibrator = ColorConsistencyCalibrator(calibration_config()) + with self.assertRaises(ValueError): + calibrator.observe(0, np.zeros((10, 10, 3), np.uint8), np.zeros((11, 10, 3), np.uint8)) + with self.assertRaises(ValueError): + calibrator.observe(-1, np.zeros((10, 10, 3), np.uint8), np.zeros((10, 10, 3), np.uint8)) + + +class ColorConsistencyConfigTest(unittest.TestCase): + def test_rejects_invalid_parameters(self) -> None: + for overrides in ( + {"sample_step": 0}, + {"match_scale": 0.0}, + {"gain_bins": 0}, + {"spatial_grid": -1}, + {"strength": -1.0}, + {"highlight_protect": 256}, + {"train_validation_block_size": 0}, + ): + with self.subTest(overrides=overrides): + with self.assertRaises(ValueError): + ColorConsistencyConfig(**overrides) + + +class Ciede2000Test(unittest.TestCase): + def test_identical_colours_have_zero_difference(self) -> None: + lab = srgb_to_lab(np.asarray([[.2, .5, .8], [.9, .1, .3]])) + self.assertAlmostEqual(float(np.max(ciede2000(lab, lab))), 0.0, places=9) + + def test_reference_pair_matches_published_value(self) -> None: + # Sharma et al. reference pair 1: L*a*b* (50, 2.6772, -79.7751) vs + # (50, 0, -82.7485) has a CIEDE2000 difference of 2.0425. + first = np.asarray([[50.0, 2.6772, -79.7751]]) + second = np.asarray([[50.0, 0.0, -82.7485]]) + self.assertAlmostEqual(float(ciede2000(first, second)[0]), 2.0425, places=4) + + +if __name__ == "__main__": + unittest.main() diff --git a/jobs/stereo-split/tests/test_joined_color_correction.py b/jobs/stereo-split/tests/test_joined_color_correction.py new file mode 100644 index 00000000..bd3684b3 --- /dev/null +++ b/jobs/stereo-split/tests/test_joined_color_correction.py @@ -0,0 +1,378 @@ +# SPDX-FileCopyrightText: 2026 ArcheBase +# SPDX-License-Identifier: MulanPSL-2.0 + +"""End-to-end checks for the joined H.264 colour-consistency integration.""" + +from __future__ import annotations + +import re +import hashlib +import json +import subprocess +import sys +import tempfile +import unittest +from pathlib import Path + +import cv2 +from mcap.reader import make_reader +from mcap.writer import IndexType, Writer +import numpy as np +from rosbags.typesys import Stores, get_typestore + + +JOB_ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(JOB_ROOT)) +from color_consistency import ColorConsistencyConfig # noqa: E402 +from convert_mcap_stereo_h264 import ( # noqa: E402 + CompressedVideo, + ConverterConfig, + FOXGLOVE_DESCRIPTOR_SET, + FOXGLOVE_SCHEMA_NAME, + StereoSplitH264Converter, + _build_timestamp_repair_plan, +) +sys.path.remove(str(JOB_ROOT)) + + +WIDTH = 1920 +HEIGHT = 1200 +METADATA_WIDTH = 160 +FRAME_WIDTH = METADATA_WIDTH + WIDTH * 2 +FRAME_COUNT = 4 +RIGHT_GAINS = (.62, 1.38, 1.26) + + +def textured_eye(seed: int = 11) -> np.ndarray: + """Build one eye-sized textured BGR frame with unsaturated levels.""" + rng = np.random.default_rng(seed) + eye = np.zeros((HEIGHT, WIDTH, 3), dtype=np.float64) + yy, xx = np.mgrid[0:HEIGHT, 0:WIDTH] + eye[..., 0] = 70 + 80 * xx / WIDTH + eye[..., 1] = 90 + 70 * yy / HEIGHT + eye[..., 2] = 110 - 50 * xx / WIDTH + for _ in range(700): + center = (int(rng.integers(20, WIDTH - 20)), int(rng.integers(20, HEIGHT - 20))) + level = float(rng.integers(20, 180)) + cv2.circle( + eye, + center, + int(rng.integers(8, 26)), + (level, min(180.0, level + 40), max(15.0, level - 50)), + -1, + ) + return cv2.GaussianBlur(np.clip(eye, 0, 185).astype(np.uint8), (0, 0), 1.2) + + +def gained(eye: np.ndarray) -> np.ndarray: + """Apply a per-channel gain and a horizontal grade as the colour mismatch.""" + field = 1.0 + .3 * (.5 - np.linspace(0, 1, WIDTH)[None, :, None]) + values = eye.astype(np.float64) / 255.0 * np.asarray(RIGHT_GAINS)[None, None, :] * field + return np.clip(np.rint(values * 255.0), 0, 255).astype(np.uint8) + + +def joined_frame(index: int, textured: bool = True) -> np.ndarray: + """Compose one joined full-frame image from the two eyes.""" + if textured: + left = textured_eye() + faded = cv2.convertScaleAbs(left, alpha=1 - .01 * index, beta=0) + right = gained(faded) + else: + left = np.full((HEIGHT, WIDTH, 3), 120, dtype=np.uint8) + right = np.full((HEIGHT, WIDTH, 3), 140, dtype=np.uint8) + frame = np.zeros((HEIGHT, FRAME_WIDTH, 3), dtype=np.uint8) + frame[:, :METADATA_WIDTH] = 255 + frame[:, METADATA_WIDTH:METADATA_WIDTH + WIDTH] = left + frame[:, METADATA_WIDTH + WIDTH:] = right + return frame + + +def encode_access_units(frames: list[np.ndarray]) -> list[bytes]: + encoded = subprocess.run( + [ + "ffmpeg", "-hide_banner", "-loglevel", "error", + "-f", "rawvideo", "-pixel_format", "bgr24", + "-video_size", f"{FRAME_WIDTH}x{HEIGHT}", "-framerate", "60", "-i", "pipe:0", + "-c:v", "libx264", "-preset", "ultrafast", "-tune", "zerolatency", + "-profile:v", "high", "-pix_fmt", "yuv420p", "-g", "1", "-bf", "0", + "-x264-params", "aud=1:repeat-headers=1", "-f", "h264", "pipe:1", + ], + input=b"".join(frame.tobytes() for frame in frames), + capture_output=True, + check=True, + ).stdout + starts = [match.start() for match in re.finditer(b"\\x00\\x00(?:\\x00)?\\x01\\x09", encoded)] + units = [ + encoded[start: starts[index + 1] if index + 1 < len(starts) else len(encoded)] + for index, start in enumerate(starts) + ] + if len(units) != len(frames): + raise RuntimeError(f"fixture access unit mismatch: {len(units)} != {len(frames)}") + return units + + +def decode_eye_topic(path: Path, topic: str) -> list[np.ndarray]: + """Decode every message of one eye topic as a single H.264 stream. + + Only the first access unit carries SPS/PPS, so the frames have to be decoded + in order by one process rather than one access unit at a time. + """ + payload = bytearray() + with path.open("rb") as stream: + for _, _, message in make_reader(stream).iter_messages(topics=[topic]): + payload.extend(CompressedVideo.FromString(message.data).data) + decoded = subprocess.run( + [ + "ffmpeg", "-hide_banner", "-loglevel", "error", + "-f", "h264", "-i", "pipe:0", + "-pix_fmt", "bgr24", "-f", "rawvideo", "pipe:1", + ], + input=bytes(payload), + capture_output=True, + check=True, + ).stdout + frame_bytes = WIDTH * HEIGHT * 3 + if len(decoded) % frame_bytes: + raise RuntimeError(f"decoded stream size mismatch: {len(decoded)}") + buffer = np.frombuffer(decoded, dtype=np.uint8) + return [ + buffer[index * frame_bytes:(index + 1) * frame_bytes].reshape(HEIGHT, WIDTH, 3) + for index in range(len(decoded) // frame_bytes) + ] + + +def make_joined_source(path: Path, frames: list[np.ndarray]) -> None: + typestore = get_typestore(Stores.ROS2_JAZZY) + messages = typestore.types + units = encode_access_units(frames) + imu_definition, _ = typestore.generate_msgdef("sensor_msgs/msg/Imu", ros_version=2) + string_definition, _ = typestore.generate_msgdef("std_msgs/msg/String", ros_version=2) + with path.open("wb") as stream: + writer = Writer( + stream, + index_types=IndexType.ALL, + repeat_channels=True, + repeat_schemas=True, + use_statistics=True, + use_summary_offsets=True, + ) + writer.start() + image_schema = writer.register_schema( + "sensor_msgs/msg/CompressedImage", "ros2msg", + typestore.generate_msgdef( + "sensor_msgs/msg/CompressedImage", ros_version=2 + )[0].encode(), + ) + image_channel = writer.register_channel( + "/decxin/rgb/compressed", "cdr", image_schema) + imu_schema = writer.register_schema( + "sensor_msgs/msg/Imu", "ros2msg", imu_definition.encode()) + imu_channel = writer.register_channel("/decxin/imu", "cdr", imu_schema) + serial_schema = writer.register_schema( + "std_msgs/msg/String", "ros2msg", string_definition.encode()) + serial_channel = writer.register_channel("/decxin/serial_number", "cdr", serial_schema) + for index, unit in enumerate(units): + timestamp = (index + 1) * 33_333_333 + stamp = messages["builtin_interfaces/msg/Time"](sec=index, nanosec=timestamp) + image = messages["sensor_msgs/msg/CompressedImage"]( + header=messages["std_msgs/msg/Header"](stamp=stamp, frame_id="joined_camera"), + format="h264", + data=np.frombuffer(unit, dtype=np.uint8), + ) + writer.add_message( + image_channel, timestamp, + bytes(typestore.serialize_cdr(image, "sensor_msgs/msg/CompressedImage")), + timestamp, index, + ) + imu = messages["sensor_msgs/msg/Imu"]( + header=messages["std_msgs/msg/Header"](stamp=stamp, frame_id="decxin_imu"), + orientation=messages["geometry_msgs/msg/Quaternion"](x=0.0, y=0.0, z=0.0, w=1.0), + orientation_covariance=np.zeros(9, dtype=np.float64), + angular_velocity=messages["geometry_msgs/msg/Vector3"](x=0.0, y=0.0, z=0.0), + angular_velocity_covariance=np.zeros(9, dtype=np.float64), + linear_acceleration=messages["geometry_msgs/msg/Vector3"](x=0.0, y=0.0, z=9.8), + linear_acceleration_covariance=np.zeros(9, dtype=np.float64), + ) + writer.add_message( + imu_channel, timestamp + 1, + bytes(typestore.serialize_cdr(imu, "sensor_msgs/msg/Imu")), + timestamp + 1, index, + ) + serial = messages["std_msgs/msg/String"](data="EP-000013-BD") + writer.add_message( + serial_channel, timestamp + 2, + bytes(typestore.serialize_cdr(serial, "std_msgs/msg/String")), + timestamp + 2, index, + ) + writer.finish() + + +def test_color_config(spatial_grid: int = 0) -> ColorConsistencyConfig: + return ColorConsistencyConfig( + sample_step=1, + match_scale=1.0, + max_match_features=2000, + gain_bins=4, + spatial_grid=spatial_grid, + min_gain_bin_samples=20, + min_spatial_cell_samples=5, + min_sampled_frames=FRAME_COUNT, + min_matched_samples=100, + train_validation_block_size=1, + ) + + +def eye_of(frame: np.ndarray, position: str) -> np.ndarray: + if position == "left": + return frame[:, METADATA_WIDTH:METADATA_WIDTH + WIDTH] + return frame[:, METADATA_WIDTH + WIDTH:] + + +def decode_topic(path: Path, topic: str) -> list[np.ndarray]: + """Decode an output eye topic; both stereo topics are eye-sized.""" + frames = decode_eye_topic(path, topic) + if not frames: + raise RuntimeError(f"no messages on {topic}") + return frames + + +def mean_abs(first: np.ndarray, second: np.ndarray) -> float: + return float(np.mean(np.abs(first.astype(np.float64) - second.astype(np.float64)))) + + +def run_runner(source: Path, output_binding: Path, scratch: Path) -> subprocess.CompletedProcess[str]: + """Run the published Job entrypoint the way Orbit invokes it.""" + command = [ + sys.executable, str(JOB_ROOT / "run_processing.py"), + "--input", str(source), + "--output-binding", str(output_binding), + "--scratch", str(scratch), + "--expected-source-size", str(source.stat().st_size), + "--expected-source-checksum", hashlib.sha256(source.read_bytes()).hexdigest(), + "--source-uri", "tos://test-bucket/raw/source.mcap", + "--processor-image", "ghcr.io/archebase/stereo-split@sha256:test", + "--kind", "stereo_split", + "--generation", "1", + ] + return subprocess.run(command, cwd=JOB_ROOT, capture_output=True, text=True, check=False) + + +class JoinedColorCorrectionTest(unittest.TestCase): + def convert(self, source: Path, output: Path, **kwargs): + return StereoSplitH264Converter( + ConverterConfig(), kwargs.pop("color_config", test_color_config()) + ).convert(source, output) + + def test_corrects_joined_h264_right_eye_and_keeps_left_eye(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + frames = [joined_frame(index) for index in range(FRAME_COUNT)] + source = root / "source.mcap" + output = root / "output.mcap" + make_joined_source(source, frames) + + stats = self.convert(source, output) + + self.assertEqual(stats.input_mode, "joined") + self.assertEqual(stats.color_decision, "applied") + self.assertEqual(stats.color_corrected_frames, stats.right_videos) + self.assertEqual(stats.right_videos, FRAME_COUNT) + self.assertIsNotNone(stats.color_report) + self.assertGreater(stats.color_matches_after_filter, 0) + + reference_left = eye_of(frames[0], "left") + source_right = eye_of(frames[0], "right") + decoded_left = decode_topic(output, "/decxin/left_rgb/h264")[0] + decoded_right = decode_topic(output, "/decxin/right_rgb/h264")[0] + + baseline = mean_abs(source_right, reference_left) + after = mean_abs(decoded_left, reference_left) + corrected = mean_abs(decoded_right, reference_left) + # The left eye is only re-encoded; the right eye moves most of the + # way to the reference eye's colour. + self.assertLess(after, baseline * .3) + self.assertLess(corrected, baseline * .5) + + def test_color_consistency_can_be_disabled(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + frames = [joined_frame(index) for index in range(FRAME_COUNT)] + source = root / "source.mcap" + output = root / "output.mcap" + make_joined_source(source, frames) + + stats = StereoSplitH264Converter( + ConverterConfig(apply_color_consistency=False)).convert(source, output) + + self.assertEqual(stats.color_decision, "") + self.assertEqual(stats.color_corrected_frames, 0) + self.assertIsNone(stats.color_report) + source_right = eye_of(frames[0], "right") + decoded_right = decode_topic(output, "/decxin/right_rgb/h264")[0] + self.assertLess(mean_abs(decoded_right, source_right), 8.0) + + def test_skips_correction_when_the_eyes_are_flat(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + frames = [joined_frame(0, textured=False) for _ in range(FRAME_COUNT)] + source = root / "source.mcap" + output = root / "output.mcap" + make_joined_source(source, frames) + + stats = self.convert(source, output) + + self.assertEqual(stats.color_decision, "insufficient") + self.assertEqual(stats.color_corrected_frames, 0) + self.assertEqual(stats.right_videos, FRAME_COUNT) + source_right = eye_of(frames[0], "right") + decoded_right = decode_topic(output, "/decxin/right_rgb/h264")[0] + self.assertLess(mean_abs(decoded_right, source_right), 8.0) + + def test_spatial_field_path_runs_in_the_converter(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + frames = [joined_frame(index) for index in range(FRAME_COUNT)] + source = root / "source.mcap" + output = root / "output.mcap" + make_joined_source(source, frames) + + stats = self.convert( + source, output, color_config=test_color_config(spatial_grid=3)) + + self.assertEqual(stats.color_decision, "applied") + self.assertEqual(stats.color_report["model"]["spatial_grid"], 3) + self.assertEqual(stats.color_corrected_frames, stats.right_videos) + + def test_runner_publishes_the_colour_report_and_summary(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + frames = [joined_frame(index) for index in range(FRAME_COUNT)] + source = root / "source.mcap" + output_binding = root / "published" + output_binding.mkdir() + make_joined_source(source, frames) + + result = run_runner(source, output_binding, root / "scratch") + + self.assertEqual(result.returncode, 0, result.stderr) + manifest = json.loads( + (output_binding / "processing_manifest.json").read_text(encoding="utf-8")) + metadata = json.loads( + (output_binding / "metadata.yaml").read_text(encoding="utf-8")) + summary = manifest["color_consistency"] + # Four fixture frames are far below the production sampling + # density, so the runner decides the eyes need no correction; the + # applied decision itself is covered by the converter tests above. + self.assertEqual(summary["decision"], "insufficient") + self.assertEqual(summary["corrected_frames"], 0) + self.assertEqual( + summary["corrected_frames"], manifest["stats"]["color_corrected_frames"]) + self.assertNotIn("anchors", summary) + self.assertNotIn("color_report", manifest["stats"]) + report = metadata["stats"]["color_report"] + self.assertEqual(report["decision"], summary["decision"]) + self.assertEqual(metadata["color_consistency"], summary) + + +if __name__ == "__main__": + unittest.main() diff --git a/jobs/stereo-split/tests/test_run_processing.py b/jobs/stereo-split/tests/test_run_processing.py index 17c13366..8fa1be1f 100644 --- a/jobs/stereo-split/tests/test_run_processing.py +++ b/jobs/stereo-split/tests/test_run_processing.py @@ -675,5 +675,46 @@ def test_manifest_omits_camera_serial_when_source_has_no_serial_topic(self) -> N self.assertNotIn("camera_serial", metadata["stats"]) +class ColorConsistencySummaryTest(unittest.TestCase): + def test_projects_the_fitted_report_onto_manifest_fields(self) -> None: + report = { + "algorithm_version": "stereo-color-v1", + "reference_eye": "left", + "decision": "applied", + "reason": "validation_improved", + "sampled_frames": 56, + "matches_before_filter": 18784, + "matches_after_filter": 15867, + "train_samples": 8119, + "validation_samples": 7748, + "baseline_ciede2000_median": 13.31, + "corrected_ciede2000_median": 4.01, + "model": { + "gain_bins": 8, + "spatial_grid": 10, + "strength": 1.0, + "highlight_protect": 245, + "anchors": [[0.0], [0.0], [0.0]], + "gains": [[1.0], [1.0], [1.0]], + "spatial_gain": [[[1.0, 1.0, 1.0]]], + "sha256": "ab" * 32, + }, + } + + summary = processing_runner.color_consistency_summary(report, 1106) + + self.assertEqual(summary["corrected_frames"], 1106) + self.assertEqual(summary["decision"], "applied") + self.assertEqual(summary["model_sha256"], "ab" * 32) + self.assertEqual(summary["spatial_grid"], 10) + self.assertEqual(summary["baseline_ciede2000_median"], 13.31) + self.assertNotIn("model", summary) + self.assertNotIn("anchors", summary) + + def test_returns_none_without_a_report(self) -> None: + self.assertIsNone(processing_runner.color_consistency_summary(None, 0)) + self.assertIsNone(processing_runner.color_consistency_summary({}, 0)) + + if __name__ == "__main__": unittest.main()