Spaces:
Running
Running
File size: 3,387 Bytes
f7ebf60 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | """Turn a streamed RerunPacket into a single displayable video frame.
The eval runtime emits `RerunPacket` protobufs (one per control tick) carrying
each camera's image plus scalar joint/box readings. We decode the packet
directly — no Rerun viewer needed — pull out the image entries, and tile the
cameras side by side into one RGB frame for the Gradio image widget.
"""
from __future__ import annotations
import io
import logging
from typing import Optional
import numpy as np
logger = logging.getLogger(__name__)
# Tile each camera to this height (px), preserving aspect ratio, then hstack.
_TILE_HEIGHT = 288
# Cameras we prefer to show left-to-right; anything else follows, sorted.
_CAMERA_ORDER = ("front", "top", "left", "right", "wrist")
def decode_frame(data: bytes) -> Optional[np.ndarray]:
"""Decode one packet's images into a single tiled RGB uint8 frame, or None."""
from armnet_runtime import _rerun_pb2 as pb
from armnet_runtime.rerun import decode_packet
try:
packet = decode_packet(data)
except Exception: # noqa: BLE001 - a bad packet must never kill the stream
logger.debug("failed to decode rerun packet", exc_info=True)
return None
images: list[tuple[str, np.ndarray]] = []
for entry in packet.entries:
if entry.WhichOneof("value") != "image":
continue
arr = _decode_image(entry.image, pb)
if arr is not None:
images.append((entry.entity_path, arr))
if not images:
return None
images.sort(key=lambda pair: _camera_rank(pair[0]))
return _tile([arr for _, arr in images])
def _decode_image(image, pb) -> Optional[np.ndarray]: # noqa: ANN001
raw = bytes(image.data)
if not raw:
return None
if image.encoding == pb.IMAGE_ENCODING_JPEG:
try:
from PIL import Image as PILImage
return np.asarray(PILImage.open(io.BytesIO(raw)).convert("RGB"))
except Exception: # noqa: BLE001
logger.debug("failed to decode JPEG image entry", exc_info=True)
return None
if image.encoding == pb.IMAGE_ENCODING_RAW_RGB:
channels = image.channels or 3
arr = np.frombuffer(raw, dtype=np.uint8)
try:
if channels > 1:
return arr.reshape(image.height, image.width, channels)[..., :3]
gray = arr.reshape(image.height, image.width)
return np.stack([gray, gray, gray], axis=-1)
except ValueError:
logger.debug("raw image entry shape mismatch", exc_info=True)
return None
return None
def _camera_rank(entity_path: str) -> tuple[int, str]:
lowered = entity_path.lower()
for i, name in enumerate(_CAMERA_ORDER):
if name in lowered:
return (i, lowered)
return (len(_CAMERA_ORDER), lowered)
def _resize_to_height(arr: np.ndarray, height: int) -> np.ndarray:
from PIL import Image as PILImage
h, w = arr.shape[0], arr.shape[1]
if h == height:
return arr
new_w = max(1, round(w * height / h))
resized = PILImage.fromarray(arr).resize((new_w, height), PILImage.BILINEAR)
return np.asarray(resized)
def _tile(arrays: list[np.ndarray]) -> np.ndarray:
if len(arrays) == 1:
return arrays[0]
resized = [_resize_to_height(a, _TILE_HEIGHT) for a in arrays]
return np.hstack(resized)
|