# Copyright 2026 Toyota Research Institute. All rights reserved. import cv2 import numpy as np import rerun as rr import torch.nn.functional as F from anydata.geometry.camera import Camera from anydata.utils.decorators import iterate1 from anydata.utils.viz import viz_depth, viz_semantic @iterate1 @iterate1 def unsqueeze0(data): return data.unsqueeze(0) class RerunDisplay: """Visualizes AnyData dataset samples using Rerun. Logs all data labels (rgb, depth, intrinsics, extrinsics, semantic, bbox3d, etc.) to a Rerun recording. Camera frustums are shown via Pinhole + Transform3D, and depth maps from all cameras are fused into a single world-space point cloud. Usage: display = RerunDisplay(dataset) display.loop() # logs all samples; open Rerun viewer to navigate """ def __init__(self, dataset, image_scale=0.25, frustum_scale=0.1): self.dataset = dataset self.euclidean = False # Downscale factor applied uniformly to depth (for reconstruction) and to # all images (for display). Camera K is scaled to match. self.image_scale = image_scale # Distance (meters) from the camera origin to the Rerun image plane. # Controls the rendered size of the camera frustum in the 3-D view. self.frustum_scale = frustum_scale self.rr = rr rr.init("AnyData", spawn=True) # ------------------------------------------------------------------ # Data processing (mirrors CamvizDisplay.process) # ------------------------------------------------------------------ def process(self, idx): """Load and pre-process a dataset sample. Returns ------- data : full sample dict (tensors) cams : {(t, i): CameraBase} in a shared world frame ego_cams : {(t, 'ego'): CameraBase} points : {(t, i): Tensor[1,3,H,W]} world-space point clouds """ data = self.dataset[idx] rgb = data.get('rgb', None) intrinsics = data.get('intrinsics', None) extrinsics = data.get('extrinsics', None) depth = data.get('depth', None) semantic = data.get('semantic', None) if rgb is not None: data['rgb'] = rgb = unsqueeze0(rgb) if intrinsics is not None: data['intrinsics'] = intrinsics = unsqueeze0(intrinsics) if extrinsics is not None: data['extrinsics'] = extrinsics = unsqueeze0(extrinsics) if depth is not None: data['depth'] = depth = unsqueeze0(depth) if semantic is not None: data['semantic'] = semantic = unsqueeze0(semantic) # Discard context timesteps — keep only the target frame (t == 0) for label in ('rgb', 'intrinsics', 'extrinsics', 'depth', 'semantic'): if data.get(label) is not None: data[label] = {k: v for k, v in data[label].items() if k[0] == 0} rgb = data.get('rgb', None) intrinsics = data.get('intrinsics', None) extrinsics = data.get('extrinsics', None) depth = data.get('depth', None) semantic = data.get('semantic', None) # Convert depth from mm (storage unit) to meters, then downsample. # Nearest-neighbour avoids blending valid depths with zeros at object edges. if depth is not None: for k in depth: depth[k] = depth[k] / 1000.0 if self.image_scale != 1.0: depth[k] = F.interpolate( depth[k], scale_factor=self.image_scale, mode='nearest' ) # Build camera objects for every (t, cam_idx) key cams = { key: Camera( K=intrinsics[key] if intrinsics is not None else None, Tcw=extrinsics[key] if extrinsics is not None else None, hw=rgb[key], ) for key in rgb.keys() } # Scale K to match the downsampled depth / display images. # cam._K is always stored as (B, 3, 3); first two rows hold fx, fy, cx, cy. if self.image_scale != 1.0: for cam in cams.values(): cam._K = cam._K.clone() cam._K[:, :2, :] *= self.image_scale # Re-express all cameras relative to camera (0, 0) base = cams[(0, 0)].Twc for key in cams.keys(): cams[key].Tcw = base @ cams[key].Tcw ego_cams = {k: v for k, v in cams.items() if k[1] == 'ego'} cams = {k: v for k, v in cams.items() if k[1] != 'ego'} # Reconstruct world-space point clouds from downsampled depth maps points = {} if depth is not None: for key, cam in cams.items(): if key in depth: points[key] = cam.reconstruct_depth_map( depth[key], euclidean=self.euclidean, to_world=True ) return data, cams, ego_cams, points # ------------------------------------------------------------------ # Rerun logging helpers # ------------------------------------------------------------------ def _cam_entity(self, t, cam_name): """Return the Rerun entity path for a camera.""" return f"world/cameras/t{t}/{cam_name}" def _resize(self, img_hw3): """Downscale an H×W×C numpy array by self.image_scale using area interpolation.""" if self.image_scale == 1.0: return img_hw3 H, W = img_hw3.shape[:2] new_W = max(1, int(W * self.image_scale)) new_H = max(1, int(H * self.image_scale)) return cv2.resize(img_hw3, (new_W, new_H), interpolation=cv2.INTER_AREA) def _log_camera(self, entity, cam): """Log camera transform and intrinsics. cam._K is already scaled by image_scale (done in process()). Transform3D pose is resolution-independent and logged as-is. Pinhole dimensions are derived from cam.hw * image_scale. """ rr = self.rr # Tcw[:3,3] = camera centre in world; Tcw[:3,:3] = cam→world rotation. # (Despite the "world2cam" label in the codebase, to_world() uses Tcw, # so Tcw is the cam-to-world transform for point/pose logging.) Tcw_np = cam.Tcw[0].cpu().numpy() rr.log( entity, rr.Transform3D( translation=Tcw_np[:3, 3], mat3x3=Tcw_np[:3, :3], ), ) K_np = cam.K[0].cpu().numpy() # 3×3, already scaled H, W = cam.hw s = self.image_scale rr.log( entity, rr.Pinhole( image_from_camera=K_np, width=max(1, int(W * s)), height=max(1, int(H * s)), image_plane_distance=self.frustum_scale, ), ) def _log_rgb(self, entity, rgb_tensor): """Log a downscaled RGB image at the camera entity path. Logging Image at the same path as Pinhole makes Rerun render the image as the textured face of the camera frustum in the 3-D view. """ rgb_np = rgb_tensor[0].permute(1, 2, 0).cpu().numpy() rgb_u8 = (rgb_np * 255).clip(0, 255).astype(np.uint8) resized = self._resize(rgb_u8) # Log at entity path (same as Pinhole) → textured frustum in 3-D view. self.rr.log(entity, self.rr.Image(resized)) # Log as child → lives in camera pixel-space, visible in 2-D Spatial View. self.rr.log(f"{entity}/rgb", self.rr.Image(resized)) def _log_depth(self, entity, depth_tensor): """Log depth as a plasma-colorized Image (already downsampled and in meters).""" # viz_depth expects a 2-D tensor (H'×W'); returns H'×W'×3 float32 in [0, 1] colored = viz_depth(depth_tensor[0, 0], filter_zeros=True) colored_u8 = (np.asarray(colored) * 255).clip(0, 255).astype(np.uint8) self.rr.log(f"{entity}/depth", self.rr.Image(colored_u8)) def _log_semantic(self, entity, semantic_tensor, ontology): """Log a downscaled semantic segmentation map as a colored image.""" colored = viz_semantic(semantic_tensor[0], ontology) # H×W×3 float [0,1] colored_u8 = (colored * 255).clip(0, 255).astype(np.uint8) self.rr.log(f"{entity}/semantic", self.rr.Image(self._resize(colored_u8))) def _log_bboxes3d(self, entity, bboxes3d_tensor, cam): """Log 3-D bounding boxes in world space.""" rr = self.rr boxes = bboxes3d_tensor boxes = boxes[boxes[:, -3] > 0] # keep valid boxes if len(boxes) == 0: return h, w, l = boxes[:, 0], boxes[:, 1], boxes[:, 2] Tcw = boxes[:, 3:-3].reshape(-1, 4, 4) corners = np.stack([ [+l, +w, +h], [+l, -w, +h], [+l, -w, -h], [+l, +w, -h], [-l, +w, +h], [-l, -w, +h], [-l, -w, -h], [-l, +w, -h], ], 2).transpose(1, 0, 2) / 2 Tcw_world = cam.Tcw.numpy() @ Tcw corners_world = ( Tcw_world[:, :3, :3] @ corners + Tcw_world[:, :3, [-1]] ).transpose(0, 2, 1) # N×8×3 # Compute center and half-sizes for each box centers = corners_world.mean(axis=1) # N×3 half_sizes = (corners_world.max(axis=1) - corners_world.min(axis=1)) / 2 # N×3 rr.log( f"{entity}/bbox3d", rr.Boxes3D( centers=centers, half_sizes=half_sizes, ), ) # ------------------------------------------------------------------ # Per-sample logging # ------------------------------------------------------------------ def log_sample(self, idx): """Process one dataset sample and log everything to Rerun.""" rr = self.rr data, cams, ego_cams, points = self.process(idx) metadata = data['metadata'] cam_names = metadata.get('cameras', []) # list of string camera names rr.set_time("sample", sequence=idx) # ---- Metadata text ----------------------------------------------- info = metadata.get('info', {}) lines = [ f"Dataset : {info.get('name', '?')}", f"Sample : {idx} / {len(self.dataset) - 1}", f"Labels : {metadata.get('labels', [])}", f"Cameras : {cam_names}", ] for key in ('station', 'tags'): if key in info: lines.append(f"{key.capitalize():8s}: {info[key]}") rr.log("info", rr.TextDocument("\n".join(lines), media_type=rr.MediaType.TEXT)) # ---- Per-camera data --------------------------------------------- for key, cam in cams.items(): t, i = key cam_name = ( str(cam_names[i]) if isinstance(cam_names, (list, tuple)) and i < len(cam_names) else f"cam_{i}" ) entity = self._cam_entity(t, cam_name) # Camera frustum (transform + pinhole) self._log_camera(entity, cam) # RGB image if 'rgb' in data and key in data['rgb']: self._log_rgb(entity, data['rgb'][key]) # Depth image if 'depth' in data and key in data['depth']: self._log_depth(entity, data['depth'][key]) # Semantic segmentation if 'semantic' in data and key in data['semantic']: ontology = metadata.get('semantic', {}).get('ontology', {}) if ontology: self._log_semantic(entity, data['semantic'][key], ontology) # 3-D bounding boxes if 'bbox3d' in data and key in data['bbox3d']: self._log_bboxes3d(entity, data['bbox3d'][key].numpy(), cam) # Per-camera point cloud if key in points: pts_np = ( points[key][0] # 3×H'×W' (downsampled) .permute(1, 2, 0) # H'×W'×3 .reshape(-1, 3) .cpu().numpy() ) valid = ( np.isfinite(pts_np).all(axis=1) & (np.abs(pts_np).sum(axis=1) > 1e-6) ) pts_valid = pts_np[valid] if len(pts_valid) > 0: kwargs = {} if 'rgb' in data and key in data['rgb']: # Downsample RGB to H'×W' so indices align with the point cloud rgb_small = F.interpolate( data['rgb'][key], # 1×3×H×W scale_factor=self.image_scale, mode='bilinear', align_corners=False, ) rgb_flat = ( rgb_small[0] # 3×H'×W' .permute(1, 2, 0) .reshape(-1, 3) .cpu().numpy() ) kwargs['colors'] = ( (rgb_flat[valid] * 255).clip(0, 255).astype(np.uint8) ) rr.log(f"world/points/{cam_name}", rr.Points3D(pts_valid, **kwargs)) # Log ego cameras (no images, just transform + pinhole) for key, cam in ego_cams.items(): t = key[0] entity = self._cam_entity(t, "ego") self._log_camera(entity, cam) # ------------------------------------------------------------------ # Main entry point # ------------------------------------------------------------------ def _get_num_frames(self): """Return the minimum frame count across all cameras from metadata.json. ``num_frames`` is removed from sample metadata during post-processing, so we temporarily disable post-processing to read the raw value. Falls back to ``len(self.dataset)`` when the key is absent (e.g. webbed datasets). """ old_pp = self.dataset.post_process self.dataset.post_process = False try: sample = self.dataset[0] num_frames = sample.get('metadata', {}).get('num_frames', None) finally: self.dataset.post_process = old_pp if num_frames is None: return len(self.dataset) # num_frames can be: int | {cam: int} | {cam: {label: int}} if isinstance(num_frames, int): total = num_frames elif isinstance(num_frames, dict): vals = [] for v in num_frames.values(): if isinstance(v, int): vals.append(v) elif isinstance(v, dict): vals.extend(int(x) for x in v.values()) total = min(vals) if vals else len(self.dataset) else: total = len(self.dataset) # Clamp to valid dataset indices (context windows shrink the usable range) return min(total, len(self.dataset)) def loop(self, stride=1): """Log every *stride*-th dataset sample to Rerun and open the viewer.""" total = self._get_num_frames() idxs = list(range(0, total, stride)) print(f"[RerunDisplay] Logging {len(idxs)} / {total} samples (stride={stride}) …") for n, idx in enumerate(idxs): self.log_sample(idx) print(f" {n + 1:>4d}/{len(idxs)} (sample {idx})") print("[RerunDisplay] Done. Rerun viewer should now be open.")