diff --git a/jobs/stereo-split/Dockerfile b/jobs/stereo-split/Dockerfile index 33f8c82..cc86e80 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 fee3a71..848fd8d 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 9c9eaa0..71a211b 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 0000000..ef519e8 --- /dev/null +++ b/jobs/stereo-split/imu_retiming.py @@ -0,0 +1,527 @@ +#!/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 + #: 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. + 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_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 + if previous_us is not None and (device_us - previous_us) < repeat_limit_us: + duplicates += 1 + timestamps.append(None) + continue + previous_us = device_us + timestamps.append(clock.to_host_ns(device_us)) + + 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) + 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) > stride_tolerance: + 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] * len(messages) + + 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 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 = np.asarray([ + barcode.first_timestamp_us + + (position - position // config.packet_size) * grid.spacing_us + for position in range(packet_size) + ], 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_anchors[barcode.video_index] = barcode.first_timestamp_us + if len(verified_anchors) < config.min_sampled_frames: + return grid_timestamps, 0, "insufficient_verified_frames" + + # 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 + # 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)) + 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) + + 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 8e07c0d..7b0d415 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 0000000..459fdf6 --- /dev/null +++ b/jobs/stereo-split/tests/test_imu_retiming.py @@ -0,0 +1,252 @@ +# 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_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): + 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 0000000..695b038 --- /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 8fa1be1..735c1ec 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()