From 973b2876969b07f3edd26801d3a4b3ffa2eb9b8c Mon Sep 17 00:00:00 2001 From: chaoliu Date: Wed, 16 Sep 2026 19:18:30 +0800 Subject: [PATCH 1/2] feat(stereo-split): re-time IMU samples onto the device grid The recorder stamps /decxin/imu with the host time at which each message was emitted and repeats one sample at every packet boundary, so a 600 Hz sensor arrives as ~661 messages per second in bursts of eleven, and the payload header carries a per-frame interpolation whose rate is about 13% off. The image metadata column still carries the device's own microsecond timestamps for one packet per camera frame, and that packet plus the rigid link between the device grid and the frame strobe is enough to reconstruct the whole stream. - imu_retiming.py measures the grid spacing, the offset from the frame's exposure end, and the frame period per file, then checks all of them on every decoded barcode before it commits to re-timing anything. - convert_mcap_stereo_h264.py reads those barcodes in the analysis pass it already runs for colour fitting, and the conversion pass rewrites the IMU payload header and its MCAP timestamps, dropping the repeated samples. - The manifest gains a lean imu_retiming summary and the metadata keeps the full report, so a skip is always explainable. Split inputs are untouched. Measured on a real 1106-frame capture: 24794 -> 22541 messages, every step exactly 1664 us, 601.0 Hz unique rate, every surviving sample an ordered subset of the source, and header stamps equal to log times. --- jobs/stereo-split/Dockerfile | 1 + jobs/stereo-split/convert_mcap_stereo_h264.py | 191 +++++-- jobs/stereo-split/imu_decoder.py | 13 +- jobs/stereo-split/imu_retiming.py | 486 ++++++++++++++++++ jobs/stereo-split/run_processing.py | 36 +- jobs/stereo-split/tests/test_imu_retiming.py | 219 ++++++++ .../tests/test_joined_imu_retiming.py | 286 +++++++++++ .../stereo-split/tests/test_run_processing.py | 34 ++ 8 files changed, 1225 insertions(+), 41 deletions(-) create mode 100644 jobs/stereo-split/imu_retiming.py create mode 100644 jobs/stereo-split/tests/test_imu_retiming.py create mode 100644 jobs/stereo-split/tests/test_joined_imu_retiming.py diff --git a/jobs/stereo-split/Dockerfile b/jobs/stereo-split/Dockerfile index 33f8c828..cc86e807 100644 --- a/jobs/stereo-split/Dockerfile +++ b/jobs/stereo-split/Dockerfile @@ -32,6 +32,7 @@ RUN --mount=type=cache,target=/root/.cache/pip \ 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/imu_retiming.py /app/imu_retiming.py COPY jobs/stereo-split/run_processing.py /app/run_processing.py RUN python3 -m compileall -q /app diff --git a/jobs/stereo-split/convert_mcap_stereo_h264.py b/jobs/stereo-split/convert_mcap_stereo_h264.py index fee3a71f..848fd8d2 100644 --- a/jobs/stereo-split/convert_mcap_stereo_h264.py +++ b/jobs/stereo-split/convert_mcap_stereo_h264.py @@ -32,6 +32,12 @@ ColorCorrectionPlan, ) from imu_decoder import ImuSample, decode_imu_from_bgr +from imu_retiming import ( + ImuRetimingCalibrator, + ImuRetimingConfig, + ImuRetimingPlan, + barcode_from_frame, +) FOXGLOVE_SCHEMA_NAME = "foxglove.CompressedVideo" @@ -155,6 +161,9 @@ class ConverterConfig: #: 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 + #: Re-time IMU samples onto the device grid read from the image barcode. + #: Only the joined H.264 path carries that barcode. + apply_imu_retiming: bool = True @dataclass @@ -180,6 +189,10 @@ class ConvertStats: color_matches_after_filter: int = 0 color_corrected_frames: int = 0 color_report: dict | None = None + imu_retiming_decision: str = "" + imu_retimed_messages: int = 0 + imu_dropped_duplicates: int = 0 + imu_report: dict | None = None @dataclass(frozen=True) @@ -881,9 +894,11 @@ def __init__( self, config: ConverterConfig | None = None, color_config: ColorConsistencyConfig | None = None, + imu_config: ImuRetimingConfig | None = None, ) -> None: self.config = config or ConverterConfig() self.color_config = color_config or ColorConsistencyConfig() + self.imu_config = imu_config or ImuRetimingConfig() self.typestore = get_typestore(Stores.ROS2_JAZZY) self.msg = self.typestore.types @@ -971,9 +986,16 @@ def convert( ) generate_imu = input_codec == "jpeg" and source_imu_messages == 0 color_plan: ColorCorrectionPlan | None = None - if input_codec == "h264" and config.apply_color_consistency: + imu_plan: ImuRetimingPlan | None = None + if input_codec == "h264" and ( + config.apply_color_consistency or config.apply_imu_retiming + ): source.seek(0) - color_plan = self._calibrate_color(source, config, stats) + color_plan, imu_plan = self._analyse_joined_stream( + source, config, stats, + log_time_order=not timestamp_repair.requires_physical_order, + ) + imu_index = 0 source.seek(0) reader = make_reader(source) @@ -1059,6 +1081,12 @@ def process_frame(frame: np.ndarray, metadata: VideoMetadata) -> None: if channel.topic in video_output_topics: raise RuntimeError(f"input already contains output topic: {channel.topic}") copied_channel = self._copy_channel(writer, schema, channel, schema_ids, channel_ids) + copied_topic_ids.add(channel.id) + if (channel.topic == config.imu_topic + and imu_plan is not None and imu_plan.applied): + imu_index = self._write_retimed_imu( + ordered_writer, copied_channel, imu_plan, imu_index, message, stats) + continue log_time, publish_time = timestamp_cursor.message_times( message.log_time, message.publish_time ) @@ -1072,7 +1100,6 @@ def process_frame(frame: np.ndarray, metadata: VideoMetadata) -> None: stats.copied_messages += 1 if timestamp_repair.applied: stats.timestamp_repaired_messages += 1 - copied_topic_ids.add(channel.id) if channel.topic == config.imu_topic: stats.imu_messages += 1 if channel.topic == config.serial_topic and not stats.camera_serial: @@ -1163,6 +1190,11 @@ def process_frame(frame: np.ndarray, metadata: VideoMetadata) -> None: raise RuntimeError( "color correction did not cover every right-eye frame: " f"{stats.color_corrected_frames}/{stats.right_videos}") + if imu_plan is not None and imu_plan.applied \ + and imu_index != len(imu_plan.timestamps_ns): + raise RuntimeError( + "IMU retiming plan did not cover every source message: " + f"{imu_index}/{len(imu_plan.timestamps_ns)}") ordered_writer.flush_all() writer.finish() except BaseException: @@ -1175,41 +1207,57 @@ def process_frame(frame: np.ndarray, metadata: VideoMetadata) -> None: stats.copied_topics = len(copied_topic_ids) return stats - def _calibrate_color( + def _analyse_joined_stream( 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. + log_time_order: bool, + ) -> tuple[ColorCorrectionPlan | None, ImuRetimingPlan | None]: + """Decode the joined stream once, for colour fitting and IMU barcodes. + + H.264 is inter-frame coded, so every frame has to be decoded in order + even though only a few are sampled for colour matching; the metadata + barcode of each decoded frame is cheap to read in the same pass and is + what anchors the IMU device grid. """ - calibrator = ColorConsistencyCalibrator(self.color_config) + color_calibrator = ( + ColorConsistencyCalibrator(self.color_config) + if config.apply_color_consistency else None + ) + imu_calibrator = ( + ImuRetimingCalibrator(self.imu_config) + if config.apply_imu_retiming else None + ) 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 + decoded_index = 0 + video_index = 0 + pending: deque[tuple[int, int]] = deque() + + def observe(frame: np.ndarray, source_index: int, log_time: int) -> None: + nonlocal decoded_index + if color_calibrator is not None: + start = config.metadata_width + config.eye_width + color_calibrator.observe( + decoded_index, + frame[0:config.eye_height, config.metadata_width:start], + frame[0:config.eye_height, start:frame_width], + ) + if imu_calibrator is not None: + imu_calibrator.observe_frame( + source_index, + barcode_from_frame(source_index, log_time, decode_imu_from_bgr(frame)), + ) + decoded_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. + # cannot sample stops the analysis 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 @@ -1222,28 +1270,93 @@ def observe(frame: np.ndarray) -> None: 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. + # Frames before the first IDR cannot be decoded, but their + # IMU packets are still in the stream and still on the grid. + if imu_calibrator is not None: + imu_calibrator.observe_frame(video_index, None) + video_index += 1 continue decoder = H264FrameDecoder(frame_width, config.eye_height) + pending.append((video_index, message.log_time)) + video_index += 1 for frame in decoder.submit(access_unit): - observe(frame) + source_index, log_time = pending.popleft() if pending else (video_index - 1, 0) + observe(frame, source_index, log_time) 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 + source_index, log_time = pending.popleft() if pending else (video_index - 1, 0) + observe(frame, source_index, log_time) + + color_plan = None + if color_calibrator is not None: + color_plan = color_calibrator.fit() + report = color_plan.report + stats.color_decision = color_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() + + imu_plan = None + if imu_calibrator is not None and not aborted: + # The source IMU stream is read in the order the conversion pass + # emits it, so a plan position always names the same message. + source.seek(0) + for _, _, message in make_reader(source).iter_messages( + topics=[config.imu_topic], log_time_order=log_time_order + ): + imu = self.typestore.deserialize_cdr(message.data, "sensor_msgs/msg/Imu") + acceleration = imu.linear_acceleration + imu_calibrator.observe_message(( + acceleration.x * 1000.0 / 9.80665, + acceleration.y * 1000.0 / 9.80665, + acceleration.z * 1000.0 / 9.80665, + )) + imu_plan = imu_calibrator.resolve() + report = imu_plan.report + stats.imu_retiming_decision = report.decision + stats.imu_retimed_messages = report.messages_retimed + stats.imu_dropped_duplicates = report.duplicates_dropped + stats.imu_report = report.as_dict() + return color_plan, imu_plan + + def _write_retimed_imu( + self, + writer, + channel_id: int, + plan: ImuRetimingPlan, + index: int, + message, + stats: ConvertStats, + ) -> int: + """Write one IMU message with its exact device-grid timestamp. + + The payload's own header stamp is rewritten as well, so the message's + internal clock and its MCAP log time can never disagree. + """ + if index >= len(plan.timestamps_ns): + raise RuntimeError( + "IMU retiming plan has fewer entries than the source IMU stream") + timestamp_ns = plan.timestamps_ns[index] + if timestamp_ns is None: + return index + 1 + imu = self.typestore.deserialize_cdr(message.data, "sensor_msgs/msg/Imu") + seconds, nanos = divmod(timestamp_ns, 1_000_000_000) + imu.header.stamp.sec = int(seconds) + imu.header.stamp.nanosec = int(nanos) + writer.add_message( + channel_id, + timestamp_ns, + bytes(self.typestore.serialize_cdr(imu, "sensor_msgs/msg/Imu")), + timestamp_ns, + message.sequence, + ) + stats.copied_messages += 1 + stats.imu_messages += 1 + return index + 1 @staticmethod def _correct_right_eye( diff --git a/jobs/stereo-split/imu_decoder.py b/jobs/stereo-split/imu_decoder.py index 9c9eaa0a..71a211ba 100644 --- a/jobs/stereo-split/imu_decoder.py +++ b/jobs/stereo-split/imu_decoder.py @@ -32,6 +32,9 @@ class ImuFrame: device_type: list[int] device_group_count: list[int] samples: list[ImuSample] + #: V4L2 frame timestamps carried by the barcode header, in device microseconds. + exposure_start_us: int = 0 + exposure_end_us: int = 0 def decode_imu_from_bgr(bgr: np.ndarray) -> ImuFrame | None: @@ -110,7 +113,13 @@ def _decode_line(bgr: np.ndarray, row: int) -> bytes: def _decode_header(data: bytes | bytearray) -> ImuFrame: if data[:8] == data[8:16]: - return ImuFrame([DEVICE_ICM42688, 0, 0, 0, 0], [11, 0, 0, 0, 0], []) + return ImuFrame( + [DEVICE_ICM42688, 0, 0, 0, 0], + [11, 0, 0, 0, 0], + [], + exposure_start_us=int.from_bytes(data[0:4], "big"), + exposure_end_us=int.from_bytes(data[4:8], "big"), + ) return ImuFrame( [ ((data[0] & 0x0F) << 4) | ((data[1] & 0xF0) >> 4), @@ -121,6 +130,8 @@ def _decode_header(data: bytes | bytearray) -> ImuFrame: ], [data[1] & 0x0F, data[3] >> 4, data[4] & 0x0F, data[6] >> 4, data[7] & 0x0F], [], + exposure_start_us=int.from_bytes(data[8:12], "big"), + exposure_end_us=int.from_bytes(data[12:16], "big"), ) diff --git a/jobs/stereo-split/imu_retiming.py b/jobs/stereo-split/imu_retiming.py new file mode 100644 index 00000000..4e76eece --- /dev/null +++ b/jobs/stereo-split/imu_retiming.py @@ -0,0 +1,486 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: 2026 ArcheBase +# SPDX-License-Identifier: MulanPSL-2.0 + +"""Re-time DECXIN IMU samples onto the device's own rigid sampling grid. + +The recorder writes ``/decxin/imu`` with the host time at which each message was +emitted and repeats one sample at every packet boundary, so a 600 Hz sensor +shows up as ~661 messages per second in bursts of eleven, and the payload header +carries a per-frame interpolation whose rate is about 13% off. + +The image metadata column still carries the device's own microsecond timestamps +for one packet per camera frame, and that packet is enough to reconstruct the +whole stream: the device grid is rigidly locked to the camera frame strobe, so +its spacing, its offset from the frame's exposure end, and the number of grid +intervals per frame are constants that can be measured per file and then checked +frame by frame. + +The module is validation first: the caller feeds every decoded frame barcode and +every source message, and :meth:`ImuRetimingCalibrator.resolve` either returns a +plan that covers the whole stream or refuses to retime it at all. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np + +from imu_decoder import ImuFrame + +#: Decision values reported by :meth:`ImuRetimingCalibrator.resolve`. +DECISION_APPLIED = "applied" +DECISION_SKIPPED = "skipped" + +#: Algorithm identity recorded in the manifest. +ALGORITHM_VERSION = "decxin-imu-grid-v1" + + +@dataclass(frozen=True) +class ImuRetimingConfig: + """Fixed validation thresholds for the device grid.""" + + #: The device emits overlapping windows of this many samples. + packet_size: int = 11 + #: Camera frames carry this many packets, and therefore this many grid intervals. + packets_per_frame: int = 2 + #: Smallest number of decoded barcodes that can establish the grid constants. + min_sampled_frames: int = 8 + #: How far a single frame's measurement may sit from the file's median. + spacing_tolerance_us: float = 2.0 + anchor_tolerance_us: int = 200 + #: How close a source sample must be to a barcode sample to count as a match. + sample_tolerance_mg: float = 1.0 + #: Packets searched around the expected position for the first frame. + search_window_packets: int = 64 + #: Accepted device-clock rate, in nanoseconds per device microsecond. + min_clock_scale: float = 900.0 + max_clock_scale: float = 1100.0 + #: Clock-fit outlier rejection, mirroring the recorder-side fit. + clock_outlier_sigma: float = 6.0 + clock_outlier_floor_ns: float = 2_000_000.0 + + +@dataclass(frozen=True) +class FrameBarcode: + """One decoded metadata barcode: a frame's clock anchor and packet samples.""" + + video_index: int + log_time_ns: int + exposure_start_us: int + exposure_end_us: int + sample_timestamps_us: tuple[int, ...] + sample_accel_mg: tuple[tuple[float, float, float], ...] + + @property + def first_timestamp_us(self) -> int: + return self.sample_timestamps_us[0] + + def spacing_us(self) -> float: + return (self.sample_timestamps_us[-1] - self.sample_timestamps_us[0]) / ( + len(self.sample_timestamps_us) - 1 + ) + + +def barcode_from_frame( + video_index: int, log_time_ns: int, barcode: ImuFrame | None +) -> FrameBarcode | None: + """Convert one decoded barcode into a :class:`FrameBarcode`, or ``None``.""" + if barcode is None or not barcode.samples or barcode.exposure_end_us <= 0: + return None + return FrameBarcode( + video_index=video_index, + log_time_ns=log_time_ns, + exposure_start_us=barcode.exposure_start_us, + exposure_end_us=barcode.exposure_end_us, + sample_timestamps_us=tuple(sample.timestamp_us for sample in barcode.samples), + sample_accel_mg=tuple( + (sample.ax_mg, sample.ay_mg, sample.az_mg) for sample in barcode.samples + ), + ) + + +@dataclass(frozen=True) +class ImuRetimingReport: + """Serializable account of what the retimer decided and why.""" + + algorithm_version: str + decision: str + reason: str + frames_observed: int + barcodes_decoded: int + frames_without_barcode: int + grid_spacing_us: float + grid_anchor_offset_us: float + frame_period_us: float + frames_per_second: float + messages_total: int + messages_retimed: int + packets_extrapolated: int + duplicates_dropped: int + clock_scale_ns_per_us: float + clock_residual_p95_ms: float + clock_inliers: int + + def as_dict(self) -> dict[str, object]: + return { + "algorithm_version": self.algorithm_version, + "decision": self.decision, + "reason": self.reason, + "frames_observed": self.frames_observed, + "barcodes_decoded": self.barcodes_decoded, + "frames_without_barcode": self.frames_without_barcode, + "grid_spacing_us": round(self.grid_spacing_us, 3), + "grid_anchor_offset_us": round(self.grid_anchor_offset_us, 3), + "frame_period_us": round(self.frame_period_us, 3), + "frames_per_second": round(self.frames_per_second, 3), + "messages_total": self.messages_total, + "messages_retimed": self.messages_retimed, + "packets_extrapolated": self.packets_extrapolated, + "duplicates_dropped": self.duplicates_dropped, + "clock_scale_ns_per_us": round(self.clock_scale_ns_per_us, 6), + "clock_residual_p95_ms": round(self.clock_residual_p95_ms, 3), + "clock_inliers": self.clock_inliers, + } + + +@dataclass(frozen=True) +class ImuRetimingPlan: + """Per-message host timestamps, or the reason the stream was left alone.""" + + report: ImuRetimingReport + #: Host log time per source IMU message; ``None`` drops a repeated sample. + timestamps_ns: tuple[int | None, ...] = () + + @property + def applied(self) -> bool: + return self.report.decision == DECISION_APPLIED + + @property + def dropped(self) -> int: + return sum(1 for value in self.timestamps_ns if value is None) + + +@dataclass(frozen=True) +class _GridFit: + spacing_us: float = 0.0 + anchor_offset_us: float = 0.0 + frame_period_us: float = 0.0 + valid: bool = False + reason: str = "not_fitted" + + +@dataclass(frozen=True) +class _ClockFit: + scale_ns_per_us: float + origin_ns: int + origin_us: int + residual_p95_ms: float + inliers: int + + def to_host_ns(self, device_us: int) -> int: + return self.origin_ns + round((device_us - self.origin_us) * self.scale_ns_per_us) + + +class ImuRetimingCalibrator: + """Collect one recording's barcodes and IMU stream, then decide on a plan. + + Barcodes come from the decoded video frames; messages come from the source + IMU topic in the order the conversion pass will emit them. Both are indexed + by their position in that stream, which is what makes the alignment exact. + """ + + def __init__(self, config: ImuRetimingConfig | None = None) -> None: + self.config = config or ImuRetimingConfig() + self._barcodes: list[FrameBarcode] = [] + self._messages: list[tuple[float, float, float]] = [] + self._frames_observed = 0 + self._plan: ImuRetimingPlan | None = None + + @property + def barcodes(self) -> list[FrameBarcode]: + return self._barcodes + + @property + def frames_observed(self) -> int: + return self._frames_observed + + @property + def messages_observed(self) -> int: + return len(self._messages) + + def observe_frame(self, video_index: int, barcode: FrameBarcode | None) -> None: + """Record one video message: its index, and its barcode when it decoded.""" + self._frames_observed = max(self._frames_observed, video_index + 1) + if barcode is not None: + self._barcodes.append(barcode) + + def observe_message(self, accel_mg: tuple[float, float, float]) -> None: + """Record one source IMU message's acceleration, in emission order.""" + self._messages.append(accel_mg) + + def resolve(self) -> ImuRetimingPlan: + """Return the retiming plan, or a skip decision that keeps source times.""" + if self._plan is not None: + return self._plan + grid = self._fit_grid() + total = len(self._messages) + if not grid.valid: + self._plan = self._skip(grid.reason, grid, total) + return self._plan + if total < self.config.packet_size * self.config.packets_per_frame: + self._plan = self._skip("insufficient_messages", grid, total) + return self._plan + clock = self._fit_clock() + if clock is None: + self._plan = self._skip("clock_fit_failed", grid, total) + return self._plan + grid_timestamps, extrapolated, reason = self._build_grid_timestamps(grid) + if reason is not None: + self._plan = self._skip(reason, grid, total) + return self._plan + + timestamps: list[int | None] = [] + previous_ns: int | None = None + duplicates = 0 + for device_us in grid_timestamps: + if device_us is None: + timestamps.append(None) + continue + host_ns = clock.to_host_ns(device_us) + if previous_ns is not None and host_ns <= previous_ns: + # Every packet opens with a repeat of the previous packet's last + # sample; those repeats share one grid point and one timestamp. + duplicates += 1 + timestamps.append(None) + continue + previous_ns = host_ns + timestamps.append(host_ns) + + self._plan = ImuRetimingPlan( + report=ImuRetimingReport( + algorithm_version=ALGORITHM_VERSION, + decision=DECISION_APPLIED, + reason="grid_validated", + frames_observed=self._frames_observed, + barcodes_decoded=len(self._barcodes), + frames_without_barcode=self._frames_observed - len(self._barcodes), + grid_spacing_us=grid.spacing_us, + grid_anchor_offset_us=grid.anchor_offset_us, + frame_period_us=grid.frame_period_us, + frames_per_second=1e6 / grid.frame_period_us, + messages_total=total, + messages_retimed=sum(1 for value in timestamps if value is not None), + packets_extrapolated=extrapolated, + duplicates_dropped=duplicates, + clock_scale_ns_per_us=clock.scale_ns_per_us, + clock_residual_p95_ms=clock.residual_p95_ms, + clock_inliers=clock.inliers, + ), + timestamps_ns=tuple(timestamps), + ) + return self._plan + + def _skip(self, reason: str, grid: _GridFit, total: int) -> ImuRetimingPlan: + return ImuRetimingPlan( + report=ImuRetimingReport( + algorithm_version=ALGORITHM_VERSION, + decision=DECISION_SKIPPED, + reason=reason, + frames_observed=self._frames_observed, + barcodes_decoded=len(self._barcodes), + frames_without_barcode=self._frames_observed - len(self._barcodes), + grid_spacing_us=grid.spacing_us, + grid_anchor_offset_us=grid.anchor_offset_us, + frame_period_us=grid.frame_period_us, + frames_per_second=1e6 / grid.frame_period_us if grid.frame_period_us else 0.0, + messages_total=total, + messages_retimed=0, + packets_extrapolated=0, + duplicates_dropped=0, + clock_scale_ns_per_us=0.0, + clock_residual_p95_ms=0.0, + clock_inliers=0, + ), + ) + + def _fit_grid(self) -> _GridFit: + """Measure the grid constants and require them to hold on every barcode.""" + config = self.config + usable = [ + barcode for barcode in self._barcodes + if len(barcode.sample_timestamps_us) == config.packet_size + ] + if len(usable) < config.min_sampled_frames: + return _GridFit(reason="insufficient_barcodes") + + spacings = np.asarray([barcode.spacing_us() for barcode in usable]) + spacing = float(np.median(spacings)) + if spacing <= 0 or float(np.max(np.abs(spacings - spacing))) > config.spacing_tolerance_us: + return _GridFit(reason="spacing_not_constant") + anchors = np.asarray( + [barcode.first_timestamp_us - barcode.exposure_end_us for barcode in usable], + dtype=np.float64, + ) + anchor = float(np.median(anchors)) + if np.max(np.abs(anchors - anchor)) > config.anchor_tolerance_us: + return _GridFit(reason="anchor_not_constant") + + frame_period = spacing * config.packets_per_frame * (config.packet_size - 1) + ordered = sorted(usable, key=lambda barcode: barcode.video_index) + for previous, current in zip(ordered, ordered[1:]): + frame_gap = current.video_index - previous.video_index + if frame_gap <= 0: + continue + measured = (current.first_timestamp_us - previous.first_timestamp_us) / frame_gap + if abs(measured - frame_period) > frame_period * 0.05: + return _GridFit(reason="frame_stride_mismatch") + return _GridFit(spacing_us=spacing, anchor_offset_us=anchor, + frame_period_us=frame_period, valid=True) + + def _build_grid_timestamps( + self, grid: _GridFit + ) -> tuple[list[int | None], int, str | None]: + """Assign every source message a grid timestamp, verifying every barcode.""" + config = self.config + messages = self._messages + packet_size = config.packet_size + packet_count = len(messages) // packet_size + messages_per_frame = packet_size * config.packets_per_frame + grid_timestamps: list[int | None] = [None] * (packet_count * packet_size) + + barcodes = sorted(self._barcodes, key=lambda barcode: barcode.video_index) + anchor_barcode: FrameBarcode | None = None + anchor_packets: int | None = None + for barcode in barcodes: + if len(barcode.sample_timestamps_us) != packet_size: + continue + expected = barcode.video_index * config.packets_per_frame + found = None + for offset in range(0, config.search_window_packets + 1): + for candidate in ((expected + offset,) if offset == 0 + else (expected + offset, expected - offset)): + if candidate < 0 or candidate >= packet_count: + continue + if self._packet_matches(self._packet(messages, candidate), barcode): + found = candidate + break + if found is not None: + break + if found is None: + return grid_timestamps, 0, "packet_alignment_failed" + anchor_barcode = barcode + anchor_packets = found - barcode.video_index * config.packets_per_frame + break + if anchor_barcode is None or anchor_packets is None: + return grid_timestamps, 0, "no_usable_barcode" + + # Every frame is placed on the rigid grid relative to the anchor frame, + # and every barcode that decoded is checked against that placement. + def frame_start_us(video_index: int) -> float: + return (anchor_barcode.first_timestamp_us + + grid.frame_period_us * (video_index - anchor_barcode.video_index)) + + verified = 0 + for barcode in barcodes: + expected_packet = barcode.video_index * config.packets_per_frame + anchor_packets + if not 0 <= expected_packet < packet_count: + continue + if not self._packet_matches(self._packet(messages, expected_packet), barcode): + return grid_timestamps, 0, "packet_alignment_drifted" + expected_timestamps = [ + int(round(frame_start_us(barcode.video_index) + + (position - position // config.packet_size) * grid.spacing_us)) + for position in range(packet_size) + ] + if list(barcode.sample_timestamps_us) != expected_timestamps: + return grid_timestamps, 0, "barcode_grid_mismatch" + verified += 1 + if verified < config.min_sampled_frames: + return grid_timestamps, 0, "insufficient_verified_frames" + + verified_frames = {barcode.video_index for barcode in barcodes} + last_packet = packet_count - 1 + extrapolated = 0 + # The grid is rigid, so frames whose barcode could not be decoded - and + # frames that precede the first decodable one - are placed by stringing + # the same frame period out from the anchor frame. + for video_index in range(0, self._frames_observed): + base_packet = video_index * config.packets_per_frame + anchor_packets + if base_packet < 0: + continue + if base_packet > last_packet: + break + if video_index not in verified_frames: + extrapolated += 1 + start_us = frame_start_us(video_index) + for position in range(messages_per_frame): + packet_offset, in_packet = divmod(position, packet_size) + packet = base_packet + packet_offset + if packet > last_packet: + break + index = packet * packet_size + in_packet + offset = position - packet_offset + grid_timestamps[index] = int(round(start_us + offset * grid.spacing_us)) + return grid_timestamps, extrapolated, None + + def _packet(self, messages: list[tuple[float, float, float]], packet: int) -> np.ndarray: + size = self.config.packet_size + return np.asarray(messages[packet * size:(packet + 1) * size], dtype=np.float64) + + def _packet_matches(self, packet: np.ndarray, barcode: FrameBarcode) -> bool: + wanted = np.asarray(barcode.sample_accel_mg, dtype=np.float64) + if packet.shape != wanted.shape or packet.size == 0: + return False + return bool(np.max(np.abs(packet - wanted)) <= self.config.sample_tolerance_mg) + + def _fit_clock(self) -> _ClockFit | None: + """Fit device microseconds onto the recording's own host timeline.""" + config = self.config + barcodes = [barcode for barcode in self._barcodes if barcode.log_time_ns > 0] + if len(barcodes) < config.min_sampled_frames: + return None + device = np.asarray([barcode.exposure_end_us for barcode in barcodes], dtype=np.float64) + host = np.asarray([barcode.log_time_ns for barcode in barcodes], dtype=np.float64) + origin_us = int(device[0]) + origin_ns = int(host[0]) + device = device - origin_us + host = host - origin_ns + if device[-1] - device[0] < 1_000.0: + return None + + mask = np.ones(device.shape, dtype=bool) + scale = 0.0 + residuals = np.zeros_like(device) + for _ in range(5): + if np.count_nonzero(mask) < 3: + return None + selected_device = device[mask] + selected_host = host[mask] + device_center = float(selected_device.mean()) + host_center = float(selected_host.mean()) + centred = selected_device - device_center + denominator = float(np.dot(centred, centred)) + if denominator <= 0: + return None + scale = float(np.dot(centred, selected_host - host_center) / denominator) + intercept = host_center - device_center * scale + residuals = host - (intercept + device * scale) + median = float(np.median(residuals)) + mad = float(np.median(np.abs(residuals - median))) + threshold = max(config.clock_outlier_sigma * 1.4826 * mad, + config.clock_outlier_floor_ns) + updated = np.abs(residuals - median) <= threshold + if np.array_equal(updated, mask): + break + mask = updated + if not config.min_clock_scale <= scale <= config.max_clock_scale: + return None + return _ClockFit( + scale_ns_per_us=scale, + origin_ns=int(round(origin_ns + intercept)), + origin_us=origin_us, + residual_p95_ms=float(np.percentile(np.abs(residuals[mask]), 95) / 1e6), + inliers=int(np.count_nonzero(mask)), + ) diff --git a/jobs/stereo-split/run_processing.py b/jobs/stereo-split/run_processing.py index 8e07c0d0..7b0d415b 100644 --- a/jobs/stereo-split/run_processing.py +++ b/jobs/stereo-split/run_processing.py @@ -166,6 +166,34 @@ def color_consistency_summary( return summary +def imu_retiming_summary(report: dict[str, object] | None) -> dict[str, object] | None: + """Project the IMU retiming report onto the manifest's lean summary. + + The full report stays in the output metadata; the manifest only carries what + a validator needs to describe and bound the re-timing. + """ + if not report: + return None + keys = ( + "algorithm_version", + "decision", + "reason", + "barcodes_decoded", + "frames_without_barcode", + "grid_spacing_us", + "frame_period_us", + "frames_per_second", + "messages_total", + "messages_retimed", + "duplicates_dropped", + "packets_extrapolated", + "clock_scale_ns_per_us", + "clock_residual_p95_ms", + "clock_inliers", + ) + return {key: report[key] for key in keys if report.get(key) is not None} + + 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 @@ -340,8 +368,10 @@ def run(args: argparse.Namespace) -> dict[str, object]: color_consistency = color_consistency_summary( stats_payload.get("color_report"), stats_payload["color_corrected_frames"] ) + imu_retiming = imu_retiming_summary(stats_payload.get("imu_report")) manifest_stats = { - key: value for key, value in stats_payload.items() if key != "color_report" + key: value for key, value in stats_payload.items() + if key not in ("color_report", "imu_report") } manifest_calibration = None processing_mode = "timestamp_repair" if stats.input_mode == "split_h264" else "convert" @@ -367,6 +397,8 @@ def run(args: argparse.Namespace) -> dict[str, object]: metadata["camera_serial"] = camera_serial if color_consistency is not None: metadata["color_consistency"] = color_consistency + if imu_retiming is not None: + metadata["imu_retiming"] = imu_retiming write_json(local_metadata, metadata) mcap_identity = require_mcap_output(local_mcap) metadata_identity = require_output(local_metadata) @@ -399,6 +431,8 @@ def run(args: argparse.Namespace) -> dict[str, object]: manifest["camera_serial"] = camera_serial if color_consistency is not None: manifest["color_consistency"] = color_consistency + if imu_retiming is not None: + manifest["imu_retiming"] = imu_retiming if manifest_calibration is not None: manifest["calibration"] = manifest_calibration write_json(local_manifest, manifest) diff --git a/jobs/stereo-split/tests/test_imu_retiming.py b/jobs/stereo-split/tests/test_imu_retiming.py new file mode 100644 index 00000000..efc59926 --- /dev/null +++ b/jobs/stereo-split/tests/test_imu_retiming.py @@ -0,0 +1,219 @@ +# SPDX-FileCopyrightText: 2026 ArcheBase +# SPDX-License-Identifier: MulanPSL-2.0 + +from __future__ import annotations + +from pathlib import Path +import sys +import unittest + +import numpy as np + + +JOB_ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(JOB_ROOT)) +from imu_retiming import ( # noqa: E402 + DECISION_APPLIED, + DECISION_SKIPPED, + FrameBarcode, + ImuRetimingCalibrator, + ImuRetimingConfig, +) + + +SPACING_US = 1664 +PACKET_SIZE = 11 +PACKETS_PER_FRAME = 2 +FRAME_PERIOD_US = SPACING_US * (PACKET_SIZE - 1) * PACKETS_PER_FRAME +ANCHOR_OFFSET_US = -16717 +BASE_TIMESTAMP_US = 3_000_000_000 + + +def sample_values(index: int) -> tuple[float, float, float]: + return (120.0 + index * 0.35, -40.0 + index * 0.15, -900.0 + index * 0.05) + + +def device_timestamp_us(sample_index: int) -> int: + """Device timestamp of one unique sample on the rigid grid.""" + return BASE_TIMESTAMP_US + SPACING_US * (sample_index + 1) + + +def frame_barcode(video_index: int, first_sample_index: int) -> FrameBarcode: + """Barcode of one frame: packet A, which opens with the previous packet's last sample.""" + timestamps = tuple(device_timestamp_us(first_sample_index - 1 + position) + for position in range(PACKET_SIZE)) + values = tuple(sample_values(first_sample_index - 1 + position) + for position in range(PACKET_SIZE)) + exposure_end = timestamps[0] - ANCHOR_OFFSET_US + return FrameBarcode( + video_index=video_index, + log_time_ns=exposure_end * 1000, + exposure_start_us=exposure_end - 7495, + exposure_end_us=exposure_end, + sample_timestamps_us=timestamps, + sample_accel_mg=values, + ) + + +def build_stream(frames: int) -> tuple[ImuRetimingCalibrator, list[tuple[float, float, float]]]: + """Feed a calibrator with a rigid-grid barcode and a matching IMU stream. + + Each frame carries 20 new samples; packet A is the previous frame's last + sample followed by ten new ones, and packet B is packet A's last sample + followed by ten more, so every packet opens with a repeat. + """ + calibrator = ImuRetimingCalibrator(ImuRetimingConfig()) + messages: list[tuple[float, float, float]] = [] + for video_index in range(frames): + first_sample_index = 20 * video_index + calibrator.observe_frame(video_index, frame_barcode(video_index, first_sample_index)) + for position in range(PACKET_SIZE * PACKETS_PER_FRAME): + index = first_sample_index - 1 + position - (1 if position >= PACKET_SIZE else 0) + value = sample_values(index) + messages.append(value) + calibrator.observe_message(value) + return calibrator, messages + + +class ImuRetimingGridTest(unittest.TestCase): + def test_retimes_the_whole_stream_onto_the_device_grid(self) -> None: + calibrator, _ = build_stream(12) + + plan = calibrator.resolve() + + self.assertEqual(plan.report.decision, DECISION_APPLIED) + self.assertEqual(plan.report.reason, "grid_validated") + self.assertAlmostEqual(plan.report.grid_spacing_us, float(SPACING_US)) + self.assertAlmostEqual(plan.report.frame_period_us, float(FRAME_PERIOD_US)) + self.assertAlmostEqual(plan.report.frames_per_second, 1e6 / FRAME_PERIOD_US, places=3) + self.assertEqual(plan.report.messages_total, 12 * 22) + self.assertEqual(plan.report.duplicates_dropped, 12 * 2 - 1) + self.assertEqual(plan.report.messages_retimed, 12 * 20 + 1) + self.assertAlmostEqual(plan.report.clock_scale_ns_per_us, 1000.0, places=3) + + def test_retimed_timestamps_are_uniform_and_monotonic(self) -> None: + calibrator, _ = build_stream(12) + + plan = calibrator.resolve() + timestamps = [value for value in plan.timestamps_ns if value is not None] + steps = np.diff(timestamps) + + self.assertTrue(all(step == SPACING_US * 1000 for step in steps)) + self.assertTrue(all(later > earlier for earlier, later in zip(timestamps, timestamps[1:]))) + + def test_repeats_are_the_only_messages_dropped(self) -> None: + calibrator, messages = build_stream(8) + + plan = calibrator.resolve() + + kept = [message for message, value in zip(messages, plan.timestamps_ns) if value is not None] + dropped = [message for message, value in zip(messages, plan.timestamps_ns) if value is None] + self.assertEqual(len(kept) + len(dropped), len(messages)) + # A dropped message always repeats the value of its predecessor. + for index, value in enumerate(plan.timestamps_ns): + if value is None: + self.assertEqual(messages[index], messages[index - 1]) + + def test_frames_without_a_barcode_are_placed_on_the_grid(self) -> None: + calibrator = ImuRetimingCalibrator(ImuRetimingConfig()) + for video_index in range(10): + first_sample_index = 20 * video_index + barcode = (None if video_index in (3, 4) + else frame_barcode(video_index, first_sample_index)) + calibrator.observe_frame(video_index, barcode) + for position in range(PACKET_SIZE * PACKETS_PER_FRAME): + index = first_sample_index - 1 + position - (1 if position >= PACKET_SIZE else 0) + calibrator.observe_message(sample_values(index)) + + plan = calibrator.resolve() + + self.assertEqual(plan.report.decision, DECISION_APPLIED) + self.assertEqual(plan.report.frames_without_barcode, 2) + self.assertEqual(plan.report.packets_extrapolated, 2) + timestamps = [value for value in plan.timestamps_ns if value is not None] + self.assertTrue(all(step == SPACING_US * 1000 for step in np.diff(timestamps))) + + +class ImuRetimingSkipTest(unittest.TestCase): + def test_skips_without_enough_barcodes(self) -> None: + calibrator, _ = build_stream(3) + + plan = calibrator.resolve() + + self.assertEqual(plan.report.decision, DECISION_SKIPPED) + self.assertEqual(plan.report.reason, "insufficient_barcodes") + self.assertEqual(plan.timestamps_ns, ()) + + def test_skips_when_the_grid_spacing_is_not_constant(self) -> None: + calibrator, _ = build_stream(12) + barcodes = calibrator.barcodes + broken = barcodes[5] + calibrator._barcodes[5] = FrameBarcode( + video_index=broken.video_index, + log_time_ns=broken.log_time_ns, + exposure_start_us=broken.exposure_start_us, + exposure_end_us=broken.exposure_end_us, + sample_timestamps_us=tuple( + value + (900 if position == 0 else 0) + for position, value in enumerate(broken.sample_timestamps_us) + ), + sample_accel_mg=broken.sample_accel_mg, + ) + + plan = calibrator.resolve() + + self.assertEqual(plan.report.decision, DECISION_SKIPPED) + self.assertIn(plan.report.reason, ("spacing_not_constant", "frame_stride_mismatch")) + + def test_skips_when_the_exposure_anchor_moves(self) -> None: + calibrator, _ = build_stream(12) + barcodes = calibrator.barcodes + broken = barcodes[6] + calibrator._barcodes[6] = FrameBarcode( + video_index=broken.video_index, + log_time_ns=broken.log_time_ns, + exposure_start_us=broken.exposure_start_us - 5_000, + exposure_end_us=broken.exposure_end_us - 5_000, + sample_timestamps_us=broken.sample_timestamps_us, + sample_accel_mg=broken.sample_accel_mg, + ) + + plan = calibrator.resolve() + + self.assertEqual(plan.report.decision, DECISION_SKIPPED) + self.assertEqual(plan.report.reason, "anchor_not_constant") + + def test_skips_when_the_source_packets_do_not_match_the_barcode(self) -> None: + calibrator = ImuRetimingCalibrator(ImuRetimingConfig()) + for video_index in range(10): + calibrator.observe_frame(video_index, frame_barcode(video_index, 20 * video_index)) + for _ in range(PACKET_SIZE * PACKETS_PER_FRAME): + calibrator.observe_message((0.0, 0.0, 0.0)) + + plan = calibrator.resolve() + + self.assertEqual(plan.report.decision, DECISION_SKIPPED) + self.assertEqual(plan.report.reason, "packet_alignment_failed") + + def test_skips_when_the_clock_rate_is_implausible(self) -> None: + calibrator, _ = build_stream(12) + calibrator._barcodes = [ + FrameBarcode( + video_index=barcode.video_index, + log_time_ns=3_000_000_000_000 + barcode.video_index * FRAME_PERIOD_US * 2_000, + exposure_start_us=barcode.exposure_start_us, + exposure_end_us=barcode.exposure_end_us, + sample_timestamps_us=barcode.sample_timestamps_us, + sample_accel_mg=barcode.sample_accel_mg, + ) + for barcode in calibrator.barcodes + ] + + plan = calibrator.resolve() + + self.assertEqual(plan.report.decision, DECISION_SKIPPED) + self.assertEqual(plan.report.reason, "clock_fit_failed") + + +if __name__ == "__main__": + unittest.main() diff --git a/jobs/stereo-split/tests/test_joined_imu_retiming.py b/jobs/stereo-split/tests/test_joined_imu_retiming.py new file mode 100644 index 00000000..695b038c --- /dev/null +++ b/jobs/stereo-split/tests/test_joined_imu_retiming.py @@ -0,0 +1,286 @@ +# SPDX-FileCopyrightText: 2026 ArcheBase +# SPDX-License-Identifier: MulanPSL-2.0 + +"""End-to-end checks for IMU re-timing on a joined H.264 capture.""" + +from __future__ import annotations + +import re +import subprocess +import sys +import tempfile +import unittest +from pathlib import Path + +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 convert_mcap_stereo_h264 import ( # noqa: E402 + ConverterConfig, + StereoSplitH264Converter, +) +from imu_retiming import ImuRetimingConfig # noqa: E402 +sys.path.remove(str(JOB_ROOT)) + + +FRAME_WIDTH = 4000 +FRAME_HEIGHT = 1200 +EYE_WIDTH = 1920 +METADATA_WIDTH = 160 +FRAME_COUNT = 6 +PACKET_SIZE = 11 +MESSAGES_PER_FRAME = 22 +SPACING_US = 1664 +FRAME_PERIOD_US = SPACING_US * (PACKET_SIZE - 1) * 2 +ANCHOR_OFFSET_US = -16717 +BASE_TIMESTAMP_US = 3_000_000_000 + + +def encode_metadata_line(frame: np.ndarray, row: int, payload: bytes) -> None: + """Write one 16-byte barcode group into one image row.""" + if len(payload) != 16: + raise ValueError("metadata groups must contain 16 bytes") + frame[row, :, :] = 255 + frame[row, 4:8, :] = 0 + for bit_index in range(128): + bit = (payload[bit_index // 8] >> (bit_index % 8)) & 1 + start = 8 + bit_index * 8 + frame[row, start:start + 8, :] = 100 if bit else 0 + + +def icm42688_group(timestamp_us: int, accel_mg: tuple[float, float, float]) -> bytes: + def raw(value: float, scale: float) -> int: + return int(round(value / scale)) + + values = [raw(accel_mg[0], 4000.0 / 32768.0), raw(accel_mg[1], 4000.0 / 32768.0), + raw(accel_mg[2], 4000.0 / 32768.0), 100, -200, 300] + payload = timestamp_us.to_bytes(4, "big") + for value in values: + payload += int(value).to_bytes(2, "big", signed=True) + return payload + + +def barcode_frame(sample_values: list[tuple[float, float, float]], + first_timestamp_us: int) -> np.ndarray: + """Build one joined frame whose metadata column carries packet A.""" + frame = np.zeros((FRAME_HEIGHT, FRAME_WIDTH, 3), dtype=np.uint8) + frame[:, :METADATA_WIDTH] = 255 + frame[:, METADATA_WIDTH:] = 90 + exposure_end = first_timestamp_us - ANCHOR_OFFSET_US + header = (exposure_end - 7495).to_bytes(4, "big") + exposure_end.to_bytes(4, "big") + encode_metadata_line(frame, 2, header + header) + for index, values in enumerate(sample_values): + encode_metadata_line( + frame, + 2 + (index + 1) * 8, + icm42688_group(first_timestamp_us + SPACING_US * index, values), + ) + return frame + + +def sample_values(index: int) -> tuple[float, float, float]: + return (150.0 + index * 0.4, -60.0 + index * 0.2, -950.0 + index * 0.1) + + +def build_fixture() -> tuple[list[np.ndarray], list[tuple[float, float, float]]]: + """Return the joined frames and the matching source IMU message values.""" + frames = [] + messages: list[tuple[float, float, float]] = [] + for frame_index in range(FRAME_COUNT): + first_sample = 20 * frame_index + first_timestamp_us = BASE_TIMESTAMP_US + SPACING_US * (first_sample + 1) + packet_a = [sample_values(first_sample - 1 + position) for position in range(PACKET_SIZE)] + frames.append(barcode_frame(packet_a, first_timestamp_us)) + for position in range(MESSAGES_PER_FRAME): + index = first_sample - 1 + position - (1 if position >= PACKET_SIZE else 0) + messages.append(sample_values(index)) + return frames, messages + + +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{FRAME_HEIGHT}", "-framerate", "30", "-i", "pipe:0", + "-c:v", "libx264", "-preset", "ultrafast", "-tune", "zerolatency", + "-profile:v", "high", "-pix_fmt", "yuv420p", "-g", "1", "-bf", "0", + "-crf", "1", "-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 write_source(path: Path, frames: list[np.ndarray], + messages: list[tuple[float, float, float]]) -> None: + typestore = get_typestore(Stores.ROS2_JAZZY) + types = typestore.types + units = encode_access_units(frames) + image_definition, _ = typestore.generate_msgdef( + "sensor_msgs/msg/CompressedImage", ros_version=2) + imu_definition, _ = typestore.generate_msgdef("sensor_msgs/msg/Imu", 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", image_definition.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) + + for index, unit in enumerate(units): + timestamp = 1_790_000_000_000_000_000 + index * FRAME_PERIOD_US * 1000 + stamp = types["builtin_interfaces/msg/Time"](sec=index, nanosec=0) + image = types["sensor_msgs/msg/CompressedImage"]( + header=types["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, + ) + base = index * MESSAGES_PER_FRAME + for position in range(MESSAGES_PER_FRAME): + ax, ay, az = messages[base + position] + imu = types["sensor_msgs/msg/Imu"]( + header=types["std_msgs/msg/Header"](stamp=stamp, frame_id="decxin_imu"), + orientation=types["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=types["geometry_msgs/msg/Vector3"](x=0.0, y=0.0, z=0.0), + angular_velocity_covariance=np.zeros(9, dtype=np.float64), + linear_acceleration=types["geometry_msgs/msg/Vector3"]( + x=ax * 9.80665 / 1000.0, + y=ay * 9.80665 / 1000.0, + z=az * 9.80665 / 1000.0, + ), + linear_acceleration_covariance=np.zeros(9, dtype=np.float64), + ) + writer.add_message( + imu_channel, timestamp + position + 1, + bytes(typestore.serialize_cdr(imu, "sensor_msgs/msg/Imu")), + timestamp + position + 1, position, + ) + writer.finish() + + +def read_imu(path: Path) -> list[tuple[int, int, tuple[float, float, float]]]: + typestore = get_typestore(Stores.ROS2_JAZZY) + records = [] + with path.open("rb") as stream: + for _, _, message in make_reader(stream).iter_messages(topics=["/decxin/imu"]): + imu = typestore.deserialize_cdr(message.data, "sensor_msgs/msg/Imu") + records.append(( + message.log_time, + imu.header.stamp.sec * 1_000_000_000 + imu.header.stamp.nanosec, + (imu.linear_acceleration.x, imu.linear_acceleration.y, imu.linear_acceleration.z), + )) + return records + + +def test_imu_config() -> ImuRetimingConfig: + return ImuRetimingConfig(min_sampled_frames=4) + + +class JoinedImuRetimingTest(unittest.TestCase): + def convert(self, source: Path, output: Path, **kwargs): + return StereoSplitH264Converter( + ConverterConfig(apply_color_consistency=False), + imu_config=kwargs.pop("imu_config", test_imu_config()), + ).convert(source, output) + + def test_retimes_and_drops_repeats_on_the_grid(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + frames, messages = build_fixture() + source = root / "source.mcap" + output = root / "output.mcap" + write_source(source, frames, messages) + + stats = self.convert(source, output) + + self.assertEqual(stats.imu_retiming_decision, "applied") + self.assertEqual(stats.imu_retimed_messages, FRAME_COUNT * 20 + 1) + self.assertEqual(stats.imu_dropped_duplicates, FRAME_COUNT * 2 - 1) + self.assertEqual(stats.imu_messages, stats.imu_retimed_messages) + self.assertEqual(stats.imu_report["grid_spacing_us"], float(SPACING_US)) + self.assertEqual(stats.imu_report["frame_period_us"], float(FRAME_PERIOD_US)) + + records = read_imu(output) + self.assertEqual(len(records), stats.imu_retimed_messages) + for log_time, stamp, _ in records: + self.assertEqual(log_time, stamp) + steps = [later[0] - earlier[0] for earlier, later in zip(records, records[1:])] + self.assertTrue(all(step == SPACING_US * 1000 for step in steps)) + kept = [record[2] for record in records] + cursor = 0 + for value in kept: + while cursor < len(messages) and not all( + abs(actual - expected * 9.80665 / 1000.0) < 1e-6 + for actual, expected in zip(value, messages[cursor]) + ): + cursor += 1 + self.assertLess(cursor, len(messages), "output value is not a source sample") + cursor += 1 + + def test_leaves_the_source_stream_untouched_when_the_grid_is_broken(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + frames, messages = build_fixture() + for frame in frames: + for row in range(2, 2 + 12 * 8, 8): + frame[row, :, :] = 90 # erase the barcode rows entirely + source = root / "source.mcap" + output = root / "output.mcap" + write_source(source, frames, messages) + + stats = self.convert(source, output) + + self.assertNotEqual(stats.imu_retiming_decision, "applied") + self.assertEqual(stats.imu_retiming_decision, "skipped") + self.assertTrue(stats.imu_report["reason"]) + self.assertEqual(stats.imu_messages, len(messages)) + + def test_can_be_disabled(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + frames, messages = build_fixture() + source = root / "source.mcap" + output = root / "output.mcap" + write_source(source, frames, messages) + + stats = StereoSplitH264Converter( + ConverterConfig(apply_color_consistency=False, apply_imu_retiming=False) + ).convert(source, output) + + self.assertEqual(stats.imu_retiming_decision, "") + self.assertEqual(stats.imu_messages, len(messages)) + self.assertIsNone(stats.imu_report) + + +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 8fa1be1f..735c1ec7 100644 --- a/jobs/stereo-split/tests/test_run_processing.py +++ b/jobs/stereo-split/tests/test_run_processing.py @@ -716,5 +716,39 @@ def test_returns_none_without_a_report(self) -> None: self.assertIsNone(processing_runner.color_consistency_summary({}, 0)) +class ImuRetimingSummaryTest(unittest.TestCase): + def test_projects_the_retiming_report_onto_manifest_fields(self) -> None: + report = { + "algorithm_version": "decxin-imu-grid-v1", + "decision": "applied", + "reason": "grid_validated", + "barcodes_decoded": 1106, + "frames_without_barcode": 21, + "grid_spacing_us": 1664.0, + "frame_period_us": 33280.0, + "frames_per_second": 30.048, + "messages_total": 24794, + "messages_retimed": 22541, + "packets_extrapolated": 21, + "duplicates_dropped": 2253, + "clock_scale_ns_per_us": 999.964617, + "clock_residual_p95_ms": 5.237, + "clock_inliers": 1104, + } + + summary = processing_runner.imu_retiming_summary(report) + + self.assertEqual(summary["decision"], "applied") + self.assertEqual(summary["messages_retimed"], 22541) + self.assertEqual(summary["duplicates_dropped"], 2253) + self.assertEqual(summary["grid_spacing_us"], 1664.0) + self.assertEqual(summary["clock_inliers"], 1104) + self.assertNotIn("messages", summary) + + def test_returns_none_without_a_report(self) -> None: + self.assertIsNone(processing_runner.imu_retiming_summary(None)) + self.assertIsNone(processing_runner.imu_retiming_summary({})) + + if __name__ == "__main__": unittest.main() From 77cf414eab236e85da7e8a44df3a4e3de7f99a28 Mon Sep 17 00:00:00 2001 From: chaoliu Date: Thu, 17 Sep 2026 14:08:01 +0800 Subject: [PATCH 2/2] fix(stereo-split): tolerate microsecond rounding in the IMU grid Replaying five fresh Ego Portal Stereo captures surfaced two ways the grid check was stricter than the device actually is: - Device timestamps are whole microseconds, so a frame's samples can sit one microsecond off the ideal grid and a repeated sample can land one microsecond after its predecessor. The grid check now allows a two microsecond tolerance and treats anything closer than half a grid step as a repeat, instead of demanding exact equality. - A capture does not always contain a whole number of 11-sample packets (84606 = 11 * 7691 + 5 on the newest one), so the tail is now continued on the same rigid grid instead of being left uncovered, which used to abort the conversion. --- jobs/stereo-split/imu_retiming.py | 89 ++++++++++++++------ jobs/stereo-split/tests/test_imu_retiming.py | 33 ++++++++ 2 files changed, 98 insertions(+), 24 deletions(-) diff --git a/jobs/stereo-split/imu_retiming.py b/jobs/stereo-split/imu_retiming.py index 4e76eece..ef519e89 100644 --- a/jobs/stereo-split/imu_retiming.py +++ b/jobs/stereo-split/imu_retiming.py @@ -50,6 +50,9 @@ class ImuRetimingConfig: #: How far a single frame's measurement may sit from the file's median. spacing_tolerance_us: float = 2.0 anchor_tolerance_us: int = 200 + #: Device timestamps are microseconds with rounding jitter, so the grid is + #: checked against a tolerance rather than for exact equality. + timestamp_tolerance_us: int = 2 #: How close a source sample must be to a barcode sample to count as a match. sample_tolerance_mg: float = 1.0 #: Packets searched around the expected position for the first frame. @@ -242,21 +245,23 @@ def resolve(self) -> ImuRetimingPlan: return self._plan timestamps: list[int | None] = [] - previous_ns: int | None = None + previous_us: int | None = None duplicates = 0 + # Every packet opens with a repeat of the previous packet's last sample. + # The device timestamps carry microsecond rounding, so a repeat can land + # one microsecond off its predecessor; anything closer than half a grid + # step is a repeat, because the real grid never gets that close. + repeat_limit_us = grid.spacing_us / 2 for device_us in grid_timestamps: if device_us is None: timestamps.append(None) continue - host_ns = clock.to_host_ns(device_us) - if previous_ns is not None and host_ns <= previous_ns: - # Every packet opens with a repeat of the previous packet's last - # sample; those repeats share one grid point and one timestamp. + if previous_us is not None and (device_us - previous_us) < repeat_limit_us: duplicates += 1 timestamps.append(None) continue - previous_ns = host_ns - timestamps.append(host_ns) + previous_us = device_us + timestamps.append(clock.to_host_ns(device_us)) self._plan = ImuRetimingPlan( report=ImuRetimingReport( @@ -329,12 +334,13 @@ def _fit_grid(self) -> _GridFit: frame_period = spacing * config.packets_per_frame * (config.packet_size - 1) ordered = sorted(usable, key=lambda barcode: barcode.video_index) + stride_tolerance = max(4.0, spacing * 0.05) for previous, current in zip(ordered, ordered[1:]): frame_gap = current.video_index - previous.video_index if frame_gap <= 0: continue measured = (current.first_timestamp_us - previous.first_timestamp_us) / frame_gap - if abs(measured - frame_period) > frame_period * 0.05: + if abs(measured - frame_period) > stride_tolerance: return _GridFit(reason="frame_stride_mismatch") return _GridFit(spacing_us=spacing, anchor_offset_us=anchor, frame_period_us=frame_period, valid=True) @@ -348,7 +354,7 @@ def _build_grid_timestamps( packet_size = config.packet_size packet_count = len(messages) // packet_size messages_per_frame = packet_size * config.packets_per_frame - grid_timestamps: list[int | None] = [None] * (packet_count * packet_size) + grid_timestamps: list[int | None] = [None] * len(messages) barcodes = sorted(self._barcodes, key=lambda barcode: barcode.video_index) anchor_barcode: FrameBarcode | None = None @@ -376,31 +382,40 @@ def _build_grid_timestamps( if anchor_barcode is None or anchor_packets is None: return grid_timestamps, 0, "no_usable_barcode" - # Every frame is placed on the rigid grid relative to the anchor frame, - # and every barcode that decoded is checked against that placement. - def frame_start_us(video_index: int) -> float: - return (anchor_barcode.first_timestamp_us - + grid.frame_period_us * (video_index - anchor_barcode.video_index)) - - verified = 0 + # Every decoded barcode is checked against its own anchor: the packet + # values must match the source stream and the sample timestamps must sit + # on the grid that the anchor implies. + verified_anchors: dict[int, int] = {} for barcode in barcodes: expected_packet = barcode.video_index * config.packets_per_frame + anchor_packets if not 0 <= expected_packet < packet_count: continue if not self._packet_matches(self._packet(messages, expected_packet), barcode): return grid_timestamps, 0, "packet_alignment_drifted" - expected_timestamps = [ - int(round(frame_start_us(barcode.video_index) - + (position - position // config.packet_size) * grid.spacing_us)) + expected_timestamps = np.asarray([ + barcode.first_timestamp_us + + (position - position // config.packet_size) * grid.spacing_us for position in range(packet_size) - ] - if list(barcode.sample_timestamps_us) != expected_timestamps: + ], dtype=np.int64) + measured = np.asarray(barcode.sample_timestamps_us, dtype=np.int64) + if int(np.max(np.abs(measured - expected_timestamps))) > config.timestamp_tolerance_us: return grid_timestamps, 0, "barcode_grid_mismatch" - verified += 1 - if verified < config.min_sampled_frames: + verified_anchors[barcode.video_index] = barcode.first_timestamp_us + if len(verified_anchors) < config.min_sampled_frames: return grid_timestamps, 0, "insufficient_verified_frames" - verified_frames = {barcode.video_index for barcode in barcodes} + # Frames whose barcode decoded anchor themselves; the rest are placed on + # the same rigid grid relative to the nearest verified frame. + reference_video = min(verified_anchors) + reference_us = verified_anchors[reference_video] + + def frame_start_us(video_index: int) -> float: + anchor = verified_anchors.get(video_index) + if anchor is not None: + return float(anchor) + return reference_us + grid.frame_period_us * (video_index - reference_video) + + verified_frames = set(verified_anchors) last_packet = packet_count - 1 extrapolated = 0 # The grid is rigid, so frames whose barcode could not be decoded - and @@ -423,8 +438,34 @@ def frame_start_us(video_index: int) -> float: index = packet * packet_size + in_packet offset = position - packet_offset grid_timestamps[index] = int(round(start_us + offset * grid.spacing_us)) + self._fill_grid_gaps(grid_timestamps, grid.spacing_us) return grid_timestamps, extrapolated, None + @staticmethod + def _fill_grid_gaps(timestamps: list[int | None], spacing_us: float) -> None: + """Continue the grid over any message a whole packet could not cover. + + A recording does not always contain a whole number of packets, and the + tail still sits on the same rigid grid, so it is filled forwards (or + backwards when the stream opens mid-packet) rather than dropped. + """ + step = int(round(spacing_us)) + last: int | None = None + for index, value in enumerate(timestamps): + if value is not None: + last = value + elif last is not None: + last = last + step + timestamps[index] = last + following: int | None = None + for index in range(len(timestamps) - 1, -1, -1): + value = timestamps[index] + if value is not None: + following = value + elif following is not None: + following = following - step + timestamps[index] = following + def _packet(self, messages: list[tuple[float, float, float]], packet: int) -> np.ndarray: size = self.config.packet_size return np.asarray(messages[packet * size:(packet + 1) * size], dtype=np.float64) diff --git a/jobs/stereo-split/tests/test_imu_retiming.py b/jobs/stereo-split/tests/test_imu_retiming.py index efc59926..459fdf61 100644 --- a/jobs/stereo-split/tests/test_imu_retiming.py +++ b/jobs/stereo-split/tests/test_imu_retiming.py @@ -114,6 +114,39 @@ def test_repeats_are_the_only_messages_dropped(self) -> None: if value is None: self.assertEqual(messages[index], messages[index - 1]) + def test_a_repeat_that_rounds_one_microsecond_late_is_still_dropped(self) -> None: + """Device timestamps round to whole microseconds, so repeats can drift.""" + calibrator, _ = build_stream(12) + barcodes = calibrator.barcodes + broken = barcodes[4] + calibrator._barcodes[4] = FrameBarcode( + video_index=broken.video_index, + log_time_ns=broken.log_time_ns, + exposure_start_us=broken.exposure_start_us, + exposure_end_us=broken.exposure_end_us, + sample_timestamps_us=tuple(value + 1 for value in broken.sample_timestamps_us), + sample_accel_mg=broken.sample_accel_mg, + ) + + plan = calibrator.resolve() + + self.assertEqual(plan.report.decision, DECISION_APPLIED) + self.assertEqual(plan.report.duplicates_dropped, 12 * 2 - 1) + + def test_trailing_messages_are_placed_on_the_grid(self) -> None: + """A capture may not contain a whole number of packets.""" + calibrator, _ = build_stream(12) + tail = [sample_values(500 + index) for index in range(5)] + for value in tail: + calibrator.observe_message(value) + + plan = calibrator.resolve() + + self.assertEqual(plan.report.decision, DECISION_APPLIED) + self.assertEqual(len(plan.timestamps_ns), 12 * 22 + 5) + timestamps = [value for value in plan.timestamps_ns if value is not None] + self.assertTrue(all(step == SPACING_US * 1000 for step in np.diff(timestamps))) + def test_frames_without_a_barcode_are_placed_on_the_grid(self) -> None: calibrator = ImuRetimingCalibrator(ImuRetimingConfig()) for video_index in range(10):