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

2026-08-30 14:05:32 +03:00
parent 07453142c2
commit 35daf73d5c
5 changed files with 316 additions and 58 deletions
@@ -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
        )
@@ -22,6 +22,7 @@ from PIL import Image
 from pydantic import BaseModel, ConfigDict, Field

 from k1link.laboratory.canonical_rerun_overlay import (
+    CanonicalLabOverlayArtifact,
    CanonicalLabOverlayError,
    _mask_component_boxes,
    canonical_lab_overlay,
@@ -298,12 +299,12 @@ def _build_vegetation_lab_router(
            },
        )

-    async def canonical_rerun_overlay_response(
+    async def canonical_rerun_overlay_artifact(
        result_id: str,
        request: CanonicalLabRerunRequest,
        *,
        expected_base_generation_sha256: str | None = None,
-    ) -> FileResponse:
+    ) -> CanonicalLabOverlayArtifact:
        """Project LAB-only evidence into the base recording's native clock."""

        if (
@@ -349,6 +350,11 @@ def _build_vegetation_lab_router(
                status_code=503,
                detail="Canonical LAB Rerun overlay failed verification",
            ) from exc
+        return artifact
+
+    def canonical_rerun_overlay_file_response(
+        artifact: CanonicalLabOverlayArtifact,
+    ) -> FileResponse:
        return FileResponse(
            artifact.path,
            media_type="application/vnd.rerun.rrd",
@@ -366,7 +372,47 @@ def _build_vegetation_lab_router(
    ) -> FileResponse:
        """Resolve the sealed sidecar for bounded non-viewer consumers."""

-        return await canonical_rerun_overlay_response(result_id, request)
+        artifact = await canonical_rerun_overlay_artifact(result_id, request)
+        return canonical_rerun_overlay_file_response(artifact)
+
+    @router.head("/{result_id}/canonical-overlay.rrd")
+    async def describe_canonical_rerun_overlay(
+        result_id: str,
+        application_id: Annotated[
+            Literal["nodedc_mission_core_recorded"],
+            Query(),
+        ],
+        recording_id: Annotated[
+            str,
+            Query(
+                min_length=1,
+                max_length=128,
+                pattern=r"^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$",
+            ),
+        ],
+        generation: Annotated[str, Query(pattern=r"^[a-f0-9]{64}$")],
+    ) -> Response:
+        """Describe the exact AI generation before Rerun opens its immutable URL."""
+
+        artifact = await canonical_rerun_overlay_artifact(
+            result_id,
+            CanonicalLabRerunRequest(
+                application_id=application_id,
+                recording_id=recording_id,
+            ),
+            expected_base_generation_sha256=generation,
+        )
+        return Response(
+            status_code=200,
+            media_type="application/vnd.rerun.rrd",
+            headers={
+                "Cache-Control": "private, no-store",
+                "Content-Length": str(artifact.byte_length),
+                "ETag": f'"{artifact.sha256}"',
+                "X-Content-Type-Options": "nosniff",
+                "X-Rerun-Format": "RRF2",
+            },
+        )

    @router.get("/{result_id}/canonical-overlay.rrd")
    async def stream_canonical_rerun_overlay(
@@ -384,10 +430,11 @@ def _build_vegetation_lab_router(
            ),
        ],
        generation: Annotated[str, Query(pattern=r"^[a-f0-9]{64}$")],
+        overlay_generation: Annotated[str, Query(pattern=r"^[a-f0-9]{64}$")],
    ) -> FileResponse:
        """Stream one immutable LAB sidecar through Rerun's native HTTP receiver."""

-        return await canonical_rerun_overlay_response(
+        artifact = await canonical_rerun_overlay_artifact(
            result_id,
            CanonicalLabRerunRequest(
                application_id=application_id,
@@ -395,6 +442,9 @@ def _build_vegetation_lab_router(
            ),
            expected_base_generation_sha256=generation,
        )
+        if overlay_generation != artifact.sha256:
+            raise HTTPException(status_code=412, detail="Canonical overlay generation changed")
+        return canonical_rerun_overlay_file_response(artifact)

    @router.get("/{result_id}/timeline")
    def get_canonical_route_timeline(result_id: str) -> dict[str, object]: