#!/usr/bin/env python3 """Fuse accepted recorded image masks into calibrated K1 LiDAR observations.""" from __future__ import annotations import argparse import hashlib import json import os import shutil from dataclasses import dataclass from datetime import UTC, datetime from pathlib import Path from typing import Any from uuid import uuid4 import numpy as np import numpy.typing as npt import rerun as rr from PIL import Image, ImageDraw from rerun.components import FillMode from k1link.device_plugins.xgrids_k1.analyze.calibrated_overlay import ( _extract_camera_frames, _load_calibration_snapshot, _select_camera_anchors, _select_lidar_samples, ) from k1link.device_plugins.xgrids_k1.analyze.calibrated_projection import ( Kb4ProjectionProfile, map_points_to_lidar, project_map_points_kb4, ) from k1link.device_plugins.xgrids_k1.mqtt.capture import read_capture_clock_origin from k1link.sessions import inspect_recorded_media_epoch SCHEMA = "missioncore.calibrated-segmentation-fusion-experiment/v1" ACCEPTED_INSTANCE_MODEL = "torchvision/maskrcnn_resnet50_fpn_v2" ACCEPTED_SEMANTIC_MODEL = "microsoft/beit-base-finetuned-ade-640-640" MIN_BOX_POINTS = 4 MAX_BOX_SPAN_M = 12.0 FloatArray = npt.NDArray[np.float64] RgbImage = npt.NDArray[np.uint8] @dataclass(frozen=True, slots=True) class ObjectFusion: instance_id: int label: str score: float candidate_points: int clustered_points: int distance_p10_m: float | None distance_median_m: float | None box_center_map: tuple[float, float, float] | None box_half_size_map: tuple[float, float, float] | None box_status: str source_indices: npt.NDArray[np.int64] @dataclass(frozen=True, slots=True) class FusedFrame: stem: str session_time_seconds: float image_rgb: RgbImage semantic_overlay_rgb: RgbImage fusion_overlay_rgb: RgbImage points_map: FloatArray projected_source_indices: npt.NDArray[np.int64] projected_pixels: FloatArray semantic_colors: RgbImage objects: tuple[ObjectFusion, ...] def _arguments() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--session", type=Path, required=True) parser.add_argument("--calibration", type=Path, required=True) parser.add_argument("--calibrated-manifest", type=Path, required=True) parser.add_argument("--segmentation", type=Path, required=True) parser.add_argument("--output-root", type=Path, required=True) parser.add_argument("--ffmpeg", type=Path, required=True) return parser.parse_args() def _read_json(path: Path) -> dict[str, Any]: value = json.loads(path.read_text(encoding="utf-8")) if not isinstance(value, dict): raise RuntimeError(f"{path.name} is not a JSON object") return value def _sha256(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as stream: for chunk in iter(lambda: stream.read(1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() def _palette(index: int) -> tuple[int, int, int]: digest = hashlib.sha256(f"mission-core-segment-{index}".encode()).digest() return (64 + digest[0] % 176, 64 + digest[1] % 176, 64 + digest[2] % 176) def _depth_cluster(indices: npt.NDArray[np.int64], depths: FloatArray) -> npt.NDArray[np.int64]: if indices.size < 2: return indices order = np.argsort(depths[indices]) ordered = indices[order] ordered_depths = depths[ordered] groups: list[npt.NDArray[np.int64]] = [] start = 0 for offset, gap in enumerate(np.diff(ordered_depths), start=1): threshold = max(0.6, 0.08 * float(ordered_depths[offset - 1])) if float(gap) > threshold: groups.append(ordered[start:offset]) start = offset groups.append(ordered[start:]) return min( groups, key=lambda group: (-int(group.size), float(np.median(depths[group]))), ) def _object_fusions( *, instance_map: npt.NDArray[np.uint16], instances: list[dict[str, Any]], pixel_x: npt.NDArray[np.int64], pixel_y: npt.NDArray[np.int64], depths: FloatArray, projected_source_indices: npt.NDArray[np.int64], points_map: FloatArray, points_lidar: FloatArray, ) -> tuple[ObjectFusion, ...]: sampled_instances = instance_map[pixel_y, pixel_x] fused: list[ObjectFusion] = [] for item in instances: instance_id = int(item["instance_id"]) candidates = np.flatnonzero(sampled_instances == instance_id).astype(np.int64) clustered = _depth_cluster(candidates, depths) source_indices = projected_source_indices[clustered] ranges = np.linalg.norm(points_lidar[source_indices], axis=1) distance_p10 = None if ranges.size == 0 else float(np.percentile(ranges, 10.0)) distance_median = None if ranges.size == 0 else float(np.median(ranges)) center: tuple[float, float, float] | None = None half_size: tuple[float, float, float] | None = None if source_indices.size < MIN_BOX_POINTS: status = f"rejected-fewer-than-{MIN_BOX_POINTS}-clustered-points" else: selected = points_map[source_indices] lower = np.percentile(selected, 5.0, axis=0) upper = np.percentile(selected, 95.0, axis=0) spans = upper - lower if not np.isfinite(spans).all() or np.any(spans > MAX_BOX_SPAN_M): status = "rejected-implausible-span" else: sizes = np.maximum(spans, 0.15) center = tuple(float(value) for value in ((lower + upper) * 0.5)) half_size = tuple(float(value) for value in (sizes * 0.5)) status = "accepted-diagnostic-axis-aligned" fused.append( ObjectFusion( instance_id=instance_id, label=str(item["label"]), score=float(item["score"]), candidate_points=int(candidates.size), clustered_points=int(clustered.size), distance_p10_m=distance_p10, distance_median_m=distance_median, box_center_map=center, box_half_size_map=half_size, box_status=status, source_indices=source_indices, ) ) return tuple(fused) def _fusion_overlay( semantic_overlay: RgbImage, projected_pixels: FloatArray, semantic_colors: RgbImage, objects: tuple[ObjectFusion, ...], instances: list[dict[str, Any]], ) -> RgbImage: canvas = Image.fromarray(semantic_overlay).convert("RGBA") draw = ImageDraw.Draw(canvas, "RGBA") for (u, v), color in zip(projected_pixels, semantic_colors, strict=True): red, green, blue = (int(value) for value in color) draw.ellipse((u - 1.5, v - 1.5, u + 1.5, v + 1.5), fill=(red, green, blue, 210)) by_id = {item.instance_id: item for item in objects} for item in instances: fused = by_id[int(item["instance_id"])] x1, y1, x2, y2 = (float(value) for value in item["box_xyxy"]) color = _palette(500 + fused.instance_id) accepted = fused.box_center_map is not None alpha = 255 if accepted else 135 draw.rectangle((x1, y1, x2, y2), outline=(*color, alpha), width=2) distance = ( "no LiDAR" if fused.distance_median_m is None else f"{fused.distance_median_m:.1f}m/{fused.clustered_points}pts" ) draw.text( (x1 + 2, max(0.0, y1 - 12)), f"{fused.label} {distance}", fill=(255, 255, 255, alpha), stroke_width=2, stroke_fill=(0, 0, 0, alpha), ) return np.asarray(canvas.convert("RGB"), dtype=np.uint8) def _object_document(item: ObjectFusion) -> dict[str, Any]: return { "instance_id": item.instance_id, "label": item.label, "score": round(item.score, 9), "candidate_projected_points": item.candidate_points, "clustered_points": item.clustered_points, "distance_p10_m": item.distance_p10_m, "distance_median_m": item.distance_median_m, "box_status": item.box_status, "box_center_map": item.box_center_map, "box_half_size_map": item.box_half_size_map, } def _write_rerun(path: Path, experiment_id: str, frames: tuple[FusedFrame, ...]) -> None: recording = rr.RecordingStream("nodedc_k1_calibrated_segmentation", recording_id=experiment_id) recording.set_sinks(rr.FileSink(path, write_footer=True)) try: recording.log( "contract", rr.TextDocument( "Recorded-only diagnostic: exact image masks sampled at factory-calibrated " "K1 LiDAR projections; 3D boxes require clustered LiDAR support." ), static=True, ) for frame in frames: recording.set_time( "session_time", duration=np.timedelta64(round(frame.session_time_seconds * 1_000_000_000), "ns"), ) recording.log("camera/raw", rr.Image(frame.image_rgb)) recording.log("camera/semantic", rr.Image(frame.semantic_overlay_rgb)) recording.log("camera/fusion", rr.Image(frame.fusion_overlay_rgb)) recording.log( "camera/projected_lidar_semantic", rr.Points2D( frame.projected_pixels.astype(np.float32), colors=frame.semantic_colors, radii=rr.Radius.ui_points(2.0), ), ) recording.log( "world/semantic_points", rr.Points3D( frame.points_map[frame.projected_source_indices].astype(np.float32), colors=frame.semantic_colors, radii=rr.Radius.ui_points(2.0), ), ) boxed = [item for item in frame.objects if item.box_center_map is not None] if boxed: recording.log( "world/boxes3d", rr.Boxes3D( centers=[item.box_center_map for item in boxed], half_sizes=[item.box_half_size_map for item in boxed], colors=[(*_palette(500 + item.instance_id), 96) for item in boxed], fill_mode=FillMode.Solid, labels=[ f"{item.label} · {item.distance_median_m:.1f} m · " f"{item.clustered_points} pts" for item in boxed ], show_labels=True, ), ) finally: recording.flush(timeout_sec=30.0) recording.disconnect() def main() -> int: args = _arguments() session = args.session.resolve(strict=True) calibration_root = args.calibration.resolve(strict=True) calibrated_manifest = _read_json(args.calibrated_manifest.resolve(strict=True)) segmentation_root = args.segmentation.resolve(strict=True) segmentation_manifest = _read_json(segmentation_root / "manifest.redacted.json") ffmpeg = args.ffmpeg.resolve(strict=True) models = segmentation_manifest.get("models") acceptance = segmentation_manifest.get("acceptance") if not isinstance(models, dict) or not isinstance(acceptance, dict): raise RuntimeError("segmentation manifest is incomplete") instance_model = models.get("instance") semantic_model = models.get("semantic") if ( not isinstance(instance_model, dict) or instance_model.get("id") != ACCEPTED_INSTANCE_MODEL or not isinstance(semantic_model, dict) or semantic_model.get("id") != ACCEPTED_SEMANTIC_MODEL or semantic_model.get("checkpoint_load") != "exact-no-missing-unexpected-or-mismatched-keys" or acceptance.get("recorded_segmentation_artifact") != "generated" ): raise RuntimeError("segmentation result is not an admitted strict-load experiment") calibration, calibration_identity = _load_calibration_snapshot(calibration_root) calibrated_input = calibrated_manifest.get("input") if ( not isinstance(calibrated_input, dict) or calibrated_input.get("session_id") != session.name or calibrated_input.get("calibration_content_identity_sha256") != calibration_identity ): raise RuntimeError("calibrated experiment does not bind this session/calibration") source_id = str(calibrated_input["source_id"]) offsets = tuple(float(value) for value in calibrated_input["video_offsets_seconds"]) profile = Kb4ProjectionProfile.from_factory_calibration(calibration, source_id) capture_root = session / "captures" / "mqtt_live" origin = read_capture_clock_origin(capture_root / "mqtt.timeline.origin.json") epoch_root = session / "media" / source_id / "epoch-1" epoch = inspect_recorded_media_epoch( epoch_root, expected_source_name=source_id, origin_epoch_ns=origin.started_at_epoch_ns, origin_monotonic_ns=origin.started_monotonic_ns, ) anchors = _select_camera_anchors( epoch_root / "index.jsonl", expected_count=len(epoch.segments), offsets_seconds=offsets, timeline_start_seconds=epoch.timeline_start_seconds, origin_monotonic_ns=origin.started_monotonic_ns, ) images = _extract_camera_frames( ffmpeg, init_path=epoch.init_path, segment_paths=tuple(item.path for item in epoch.segments), frame_sequences=tuple(item.sequence for item in anchors), width=profile.width, height=profile.height, ) lidar_samples = _select_lidar_samples( capture_root / "mqtt.raw.k1mqtt", anchors, origin_monotonic_ns=origin.started_monotonic_ns, temporal_offset_seconds=float(calibrated_input["temporal_offset_seconds"]), ) segmentation_inputs = { str(item["name"]): str(item["sha256"]) for item in segmentation_manifest["input"] if isinstance(item, dict) } segmentation_frames = segmentation_manifest["frames"] fused_frames: list[FusedFrame] = [] for image, anchor, lidar_sample in zip(images, anchors, lidar_samples, strict=True): stem = ( f"camera-{anchor.sequence:06d}-" f"session-{round(anchor.session_time_seconds * 1000):09d}ms" ) input_name = f"{stem}.png" input_artifact = args.calibrated_manifest.parent / input_name if segmentation_inputs.get(input_name) != _sha256(input_artifact): raise RuntimeError(f"segmentation input identity changed for {input_name}") if _sha256(input_artifact) != hashlib.sha256(image.tobytes()).hexdigest(): decoded = np.asarray(Image.open(input_artifact).convert("RGB"), dtype=np.uint8) if not np.array_equal(decoded, image): raise RuntimeError(f"decoded camera pixels changed for {input_name}") semantic_map = np.asarray(Image.open(segmentation_root / f"{stem}.semantic.png")) instance_map = np.asarray(Image.open(segmentation_root / f"{stem}.instances.png")) semantic_overlay = np.asarray( Image.open(segmentation_root / f"{stem}.semantic-overlay.png").convert("RGB") ) if semantic_map.shape != image.shape[:2] or instance_map.shape != image.shape[:2]: raise RuntimeError("segmentation maps do not match the camera epoch") points_map = np.asarray( [ point.scaled_xyz(lidar_sample.point_frame.header.scaler) for point in lidar_sample.point_frame.points ], dtype=np.float64, ) points_lidar = map_points_to_lidar( points_map, position_map_xyz=lidar_sample.pose_frame.position_xyz, orientation_map_from_lidar_xyzw=lidar_sample.pose_frame.orientation_xyzw, ) projection = project_map_points_kb4( points_map, position_map_xyz=lidar_sample.pose_frame.position_xyz, orientation_map_from_lidar_xyzw=lidar_sample.pose_frame.orientation_xyzw, profile=profile, ) pixel_x = np.clip( np.rint(projection.pixels_xy[:, 0]).astype(np.int64), 0, profile.width - 1 ) pixel_y = np.clip( np.rint(projection.pixels_xy[:, 1]).astype(np.int64), 0, profile.height - 1 ) semantic_ids = semantic_map[pixel_y, pixel_x].astype(np.int64) semantic_colors = np.asarray( [_palette(int(value)) for value in semantic_ids], dtype=np.uint8 ) frame_document = segmentation_frames[stem] instances = frame_document["instance"]["instances"] objects = _object_fusions( instance_map=instance_map.astype(np.uint16), instances=instances, pixel_x=pixel_x, pixel_y=pixel_y, depths=projection.depths_m, projected_source_indices=projection.source_indices, points_map=points_map, points_lidar=points_lidar, ) fusion_overlay = _fusion_overlay( semantic_overlay, projection.pixels_xy, semantic_colors, objects, instances, ) fused_frames.append( FusedFrame( stem=stem, session_time_seconds=anchor.session_time_seconds, image_rgb=image, semantic_overlay_rgb=semantic_overlay, fusion_overlay_rgb=fusion_overlay, points_map=points_map, projected_source_indices=projection.source_indices, projected_pixels=projection.pixels_xy, semantic_colors=semantic_colors, objects=objects, ) ) created_at = datetime.now(UTC) experiment_id = ( f"{created_at.strftime('%Y%m%dT%H%M%SZ')}_k1_segmentation_fusion_{uuid4().hex[:12]}" ) private_root = args.output_root.resolve() / "private" / "perception-experiments" private_root.mkdir(mode=0o700, parents=True, exist_ok=True) output = private_root / experiment_id staging = private_root / f".{experiment_id}.incomplete" staging.mkdir(mode=0o700, exist_ok=False) try: for frame in fused_frames: Image.fromarray(frame.fusion_overlay_rgb).save(staging / f"{frame.stem}.fusion.png") mosaic_rows = [ np.concatenate( (frame.image_rgb, frame.semantic_overlay_rgb, frame.fusion_overlay_rgb), axis=1, ) for frame in fused_frames ] Image.fromarray(np.concatenate(mosaic_rows, axis=0)).save(staging / "fusion-mosaic.png") _write_rerun(staging / "fusion.rrd", experiment_id, tuple(fused_frames)) identity = { "calibrated_generation_sha256": calibrated_manifest["generation_sha256"], "segmentation_generation_sha256": segmentation_manifest["generation_sha256"], "calibration_content_identity_sha256": calibration_identity, "session_id": session.name, "source_id": source_id, } outputs = sorted(staging.iterdir()) manifest = { "schema_version": SCHEMA, "experiment_id": experiment_id, "created_at_utc": created_at.isoformat(timespec="milliseconds").replace("+00:00", "Z"), "generation_sha256": hashlib.sha256( json.dumps(identity, sort_keys=True, separators=(",", ":")).encode() ).hexdigest(), "input": identity, "fusion": { "pixel_sampling": "nearest-calibrated-projection", "occlusion_policy": "largest-depth-contiguous-cluster-nearest-tiebreak", "distance": "LiDAR-origin Euclidean p10 and median", "box": { "frame": "map", "orientation": "axis-aligned-diagnostic", "minimum_clustered_points": MIN_BOX_POINTS, "robust_bounds": "p05-p95", "maximum_span_m": MAX_BOX_SPAN_M, }, }, "frames": [ { "stem": frame.stem, "session_time_seconds": frame.session_time_seconds, "projected_points": int(frame.projected_source_indices.size), "semantic_class_ids": sorted( { int(value) for value in np.asarray( Image.open(segmentation_root / f"{frame.stem}.semantic.png") ).reshape(-1) } ), "objects": [_object_document(item) for item in frame.objects], "accepted_boxes": sum( item.box_center_map is not None for item in frame.objects ), } for frame in fused_frames ], "acceptance": { "recorded_mask_to_lidar_fusion": "generated", "distance": "diagnostic-not-ground-truthed", "boxes3d": "diagnostic-not-tracked-or-ground-truthed", "live": "not-tested", "safety": "not-accepted", }, "outputs": [ {"name": path.name, "bytes": path.stat().st_size, "sha256": _sha256(path)} for path in outputs ], } (staging / "manifest.redacted.json").write_text( json.dumps(manifest, ensure_ascii=False, indent=2, sort_keys=True) + "\n", encoding="utf-8", ) os.rename(staging, output) except BaseException: shutil.rmtree(staging, ignore_errors=True) raise print(json.dumps({"experiment_id": experiment_id, "output": str(output)})) return 0 if __name__ == "__main__": raise SystemExit(main())