fix(lab): сжать и адресовать RAV004 overlay

src/k1link/laboratory/canonical_rerun_overlay.py

@@ -14,6 +14,7 @@ import io
 import json
 import os
 import subprocess
+import sys
 import tempfile
 import threading
 import zipfile
@@ -30,9 +31,9 @@ from PIL import Image
 
 APPLICATION_ID: Final = "nodedc_mission_core_recorded"
 SESSION_TIMELINE: Final = "session_time"
-RENDERER_VERSION: Final = "upstream-rerun-0.36.3-v2"
+RENDERER_VERSION: Final = "upstream-rerun-0.36.3-encoded-optimized-v4"
 MAX_SOURCE_BYTES: Final = 768 * 1024 * 1024
-MAX_OVERLAY_BYTES: Final = 512 * 1024 * 1024
+MAX_OVERLAY_BYTES: Final = 256 * 1024 * 1024
 
 
 class CanonicalLabOverlayError(RuntimeError):
@@ -136,6 +137,7 @@ def canonical_lab_overlay(
             source = _verified_camera_source(root, route, jobs)
             proxy = _camera_proxy(source, int(route["frame_count"]), ffmpeg, cache)
             _render_overlay(temporary, root, route, recording_id, proxy)
+            _optimize_overlay(temporary)
             stat = temporary.stat()
             if stat.st_size < 4 or stat.st_size > MAX_OVERLAY_BYTES:
                 raise CanonicalLabOverlayError("canonical LAB overlay size is invalid")
@@ -324,6 +326,7 @@ def _render_overlay(
     frame_times = _frame_times(root, route)
     layers = route["layers"]
     archives: dict[str, zipfile.ZipFile] = {}
+    palettes: dict[str, tuple[int, ...]] = {}
     recording = rr.RecordingStream(APPLICATION_ID, recording_id=recording_id)
     try:
         recording.set_sinks(rr.FileSink(output, write_footer=True))
@@ -346,6 +349,7 @@ def _render_overlay(
             classes = taxonomy.get("classes") if isinstance(taxonomy, dict) else None
             if not isinstance(classes, list):
                 raise CanonicalLabOverlayError("semantic taxonomy is invalid")
+            palettes[layer_id] = _semantic_palette(classes)
             context = rr.AnnotationContext(
                 [
                     rr.ClassDescription(
@@ -386,7 +390,15 @@ def _render_overlay(
             for layer_id, mask in masks.items():
                 recording.log(
                     f"/perception/camera/segmentation/{layer_id}",
-                    rr.SegmentationImage(mask, opacity=0.72, draw_order=1.0),
+                    rr.EncodedImage(
+                        contents=_encoded_semantic_png(
+                            mask,
+                            palettes[layer_id],
+                        ),
+                        media_type="image/png",
+                        opacity=0.72,
+                        draw_order=1.0,
+                    ),
                 )
             boxes, labels = semantic_component_boxes(masks["city"], index)
             if boxes:
@@ -414,6 +426,55 @@ def _render_overlay(
             recording.disconnect()
 
 
+def _optimize_overlay(path: Path) -> None:
+    """Compact thousands of per-frame chunks before the browser admits the store."""
+
+    optimized = path.with_name(f".{path.stem}.{uuid4().hex}.optimized.rrd")
+    try:
+        completed = subprocess.run(
+            [
+                sys.executable,
+                "-m",
+                "rerun",
+                "rrd",
+                "optimize",
+                "--profile",
+                "object-store",
+                "--max-size",
+                "4MiB",
+                "--max-rows",
+                "512",
+                "--num-pass",
+                "20",
+                str(path),
+                "-o",
+                str(optimized),
+            ],
+            check=False,
+            capture_output=True,
+            timeout=120,
+        )
+        if (
+            completed.returncode != 0
+            or not optimized.is_file()
+            or optimized.is_symlink()
+            or optimized.stat().st_size < 4
+            or optimized.stat().st_size > MAX_OVERLAY_BYTES
+        ):
+            raise CanonicalLabOverlayError(
+                f"canonical LAB overlay optimization failed: {completed.stderr[-1000:]!r}"
+            )
+        with optimized.open("rb") as stream:
+            if stream.read(4) != b"RRF2":
+                raise CanonicalLabOverlayError("canonical LAB overlay optimization is invalid")
+        os.chmod(optimized, 0o600)
+        os.replace(optimized, path)
+    except subprocess.TimeoutExpired as exc:
+        raise CanonicalLabOverlayError("canonical LAB overlay optimization timed out") from exc
+    finally:
+        optimized.unlink(missing_ok=True)
+
+
 def _frame_times(root: Path, route: dict[str, Any]) -> np.ndarray:
     descriptor = route.get("timeline")
     relative = descriptor.get("path") if isinstance(descriptor, dict) else None
@@ -465,6 +526,54 @@ def _read_mask(archive: zipfile.ZipFile, sequence: int) -> np.ndarray:
     return mask
 
 
+def _semantic_palette(classes: list[object]) -> tuple[int, ...]:
+    """Return one complete indexed-PNG palette from a sealed taxonomy."""
+
+    palette = [0] * (256 * 3)
+    seen: set[int] = set()
+    for item in classes:
+        if not isinstance(item, dict):
+            raise CanonicalLabOverlayError("semantic taxonomy class is invalid")
+        class_id = item.get("class_id")
+        color = item.get("color_rgb")
+        if (
+            not isinstance(class_id, int)
+            or isinstance(class_id, bool)
+            or not 0 <= class_id <= 255
+            or class_id in seen
+            or not isinstance(color, list)
+            or len(color) != 3
+            or any(
+                not isinstance(channel, int)
+                or isinstance(channel, bool)
+                or not 0 <= channel <= 255
+                for channel in color
+            )
+        ):
+            raise CanonicalLabOverlayError("semantic taxonomy palette is invalid")
+        seen.add(class_id)
+        offset = class_id * 3
+        palette[offset : offset + 3] = color
+    if not seen:
+        raise CanonicalLabOverlayError("semantic taxonomy palette is empty")
+    return tuple(palette)
+
+
+def _encoded_semantic_png(mask: np.ndarray, palette: tuple[int, ...]) -> bytes:
+    """Keep semantic frames compressed in the Rerun store instead of expanding 7 GiB."""
+
+    if mask.shape != (600, 800) or mask.dtype != np.uint8 or len(palette) != 256 * 3:
+        raise CanonicalLabOverlayError("semantic frame encoding input is invalid")
+    image = Image.fromarray(mask)
+    image.putpalette(palette)
+    output = io.BytesIO()
+    image.save(output, format="PNG", compress_level=6)
+    payload = output.getvalue()
+    if not payload.startswith(b"\x89PNG\r\n\x1a\n"):
+        raise CanonicalLabOverlayError("semantic frame encoding failed")
+    return payload
+
+
 def semantic_component_boxes(
     mask: np.ndarray,
     _sequence: int,
@@ -589,8 +698,11 @@ def _restore_cached(
 
 def _artifact_is_regular(artifact: CanonicalLabOverlayArtifact) -> bool:
     try:
+        with artifact.path.open("rb") as stream:
+            magic = stream.read(4)
         return (
             not artifact.path.is_symlink()
+            and magic == b"RRF2"
             and artifact.path.stat().st_size == artifact.byte_length
             and _sha256(artifact.path) == artifact.sha256
         )