import torch import numpy as np from typing import Dict from copy import deepcopy from camviz import Draw from camviz import Camera as CameraCV from camviz import BBox2D, BBox3D from camviz.utils.cmaps import jet from anydata.dataloaders.Base import BaseDataset from anydata.geometry.camera import Camera from anydata.geometry.camera_utils import invert_extrinsics from anydata.geometry.depth import calculate_normals from anydata.utils.data import interleave_dict, modrem from anydata.utils.viz import viz_depth, viz_optflow, viz_semantic, viz_normals from anydata.utils.decorators import iterate1 from anydata.utils.colmap import qvec2rotmat, qvec2rotmat_batched from scipy.spatial.transform import Rotation as R def change_key(data, c, t, n): steps = sorted(set([key[0] for key in data.keys()])) return (c + n) % len(steps), steps[(steps.index(t) + n) % len(steps)] @iterate1 @iterate1 @iterate1 @iterate1 def unsqueeze0(data: dict | torch.Tensor) -> dict | torch.Tensor: return data if isinstance(data, (str, np.ndarray)) else data.unsqueeze(0) def reduce_dict(data): if all([data[0] == data[i] for i in range(len(data))]): data = data[0] return data class CamvizDisplay: def __init__(self, dataset: BaseDataset, virtual_pose=None, zero_origin=False, fov=90.0): self.idx = 0 self.tgt = (0,0) self.dataset = dataset self.fov = fov self.virtual_pose = virtual_pose self.scale = 0.15 self.zero_origin = zero_origin self.relative_action = False self.euclidean = False self.viewer_pose = None # (3.50314, -4.16307, 3.07801, 0.47160, 0.77606, 0.41221, -0.07346) self.filter_depth_outliers = False self.depth_percentile_hi = 90 # filter points above this percentile self.depth_percentile_lo = 5 # filter points below this percentile self.tasks = ['rgb', 'depth', 'normals', 'semantic', 'optflow', 'optflow'] self.cam_colors = ['red', 'blu', 'gre', 'yel', 'mag', 'cya'] * 100 self.offset = [None, None, None, None, 'fwd', 'bwd'] data, wh, keys, offsets, cams, ego_cams, cams_idx, points, points_normals, timestep = self.process() # Compute initial viewer pose: look at median camera position from a distance if self.viewer_pose =='median': cam_positions = np.stack([cams[k].Tcw[0, :3, 3].numpy() for k in cams.keys()], axis=0) median_pos = np.median(cam_positions, axis=0) # Place viewer behind and above, looking at the median offset = np.array([0.0, -1.0, 0.5]) viewer_pos = median_pos + offset self.viewer_pose = tuple(viewer_pos) + (0.47160, 0.77606, 0.41221, -0.07346) self.create_draw(data, wh, cams, virtual_pose) def create_draw(self, data, wh, cams, virtual_pose): if self.dataset.cameras_core_sample is not None: num_cams = self.dataset.cameras_core_sample if self.dataset.cameras_aux_sample is not None: num_cams += self.dataset.cameras_aux_sample else: if len(self.dataset.cameras_core) > 0 and isinstance(self.dataset.cameras_core[0], list): self.dataset.cameras_core = self.dataset.cameras_core[0] # Account for list of lists if len(self.dataset.cameras_aux) > 0 and isinstance(self.dataset.cameras_aux[0], list): self.dataset.cameras_aux = self.dataset.cameras_aux[0] # Account for list of lists num_cams = len(set(self.dataset.cameras_core + self.dataset.cameras_aux)) self.wh = wh import subprocess, shlex try: out = subprocess.check_output(shlex.split( "system_profiler SPDisplaysDataType -json"), text=True, timeout=5) import json displays = json.loads(out)['SPDisplaysDataType'] res = displays[0]['spdisplays_ndrvs'][0]['_spdisplays_resolution'] # e.g. "3440 x 1440 @ 100.00Hz" or "2560 x 1600" res = res.split('@')[0].strip() parts = res.split(' x ') screen_w, screen_h = int(parts[0]), int(parts[1]) except Exception: screen_w, screen_h = 2560, 1600 # Cap window to 90% of screen in both dimensions aspect = (wh[1] * 3) / (wh[0] * 5) # internal height/width ratio max_w = int(screen_w * 0.9) max_h = int(screen_h * 0.9) # Try fitting by width, check if height exceeds limit draw_width = max_w draw_height = int(draw_width * aspect) if draw_height > max_h: draw_height = max_h draw_width = int(draw_height / aspect) self.draw = Draw((wh[0] * 5, wh[1] * 3), width=draw_width) self.draw.add2DimageGrid('img', (0.0, 0.0, 0.33, 1.0), n=(max(3, (num_cams + 1) // 2), 2), res=wh) self.draw.add3Dworld('wld', (0.33, 0.0, 1.0, 1.0), pose=self.viewer_pose, enable_blending=False) self.draw.addTexture('img', n=num_cams) self.draw.addBuffer3f('pts', 1000000, n=num_cams) self.draw.addBuffer3f('clr', 1000000, n=num_cams) self.draw.addBuffer3f('pts_nrm', 1000000, n=num_cams) self.draw.addBuffer3f('clr_nrm', 1000000, n=num_cams) self.draw.addBuffer3f('rays', 1000000, n=num_cams) def process(self): data: Dict[str, Dict[str, torch.Tensor | dict]] = self.dataset[self.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) optflow = data.get('optflow', None) timestep = data.get('timestep', None) action = data.get('action', None) bbox2d = data.get('bbox2d', None) bbox3d = data.get('bbox3d', 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) if optflow is not None: data['optflow'] = optflow = unsqueeze0(optflow) if action is not None: data['action'] = action = unsqueeze0(action) camera_model = data["metadata"].get("conventions", {}).get("intrinsics", {}).get("model", None) 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], geometry=camera_model, ) for key in rgb.keys()} if action is not None: name = data['metadata']['name'] for key in action: if name == 'LBM': from anydata.utils.action import action_to_pose action_key = action[key]['action'] action[key] = action_to_pose(action_key) elif name == 'EgoDex': action_key = action[key]['action'] action[key] = action[key]['action'] elif name == 'AgiBotWorld': qvec = action[key]['state']['end']['orientation'][0] tvec = action[key]['state']['end']['position'][0] transf = [] for i in range(qvec.shape[0]): rot = qvec2rotmat(qvec[i]) transf.append(np.vstack([ np.hstack((rot, np.expand_dims(tvec[i], axis=1))), np.array([0.0, 0.0, 0.0, 1.0]) ]).astype(np.float32)) action[key] = torch.stack([torch.tensor(t) for t in transf], 0).unsqueeze(0) elif name == 'DROID': position = action[key]['cartesian_position'] transfs = [] for i in range(position.shape[0]): transf = np.eye(4).astype(np.float32) transf[:3, :3] = R.from_euler("xyz", position[i][3:]).as_matrix() transf[:3, 3] = position[i][:3] transfs.append(transf) action[key] = torch.stack([torch.tensor(t) for t in transf], 0).unsqueeze(0).unsqueeze(0) elif name == 'RH20T': act = action[key] if 'action' in act and 'state_tcp' in act['action']: act = act['action'] tvec = act['state_tcp'][:, :3] qvec = act['state_tcp'][:, 3:] transf = [] for i in range(qvec.shape[0]): rot = qvec2rotmat(qvec[i]) transf.append(np.vstack([ np.hstack((rot, np.expand_dims(tvec[i], axis=1))), np.array([0.0, 0.0, 0.0, 1.0]) ]).astype(np.float32)) action[key] = torch.stack([torch.tensor(t) for t in transf], 0).unsqueeze(0) else: action = None data.pop('action') break if action is not None: # NOTE(bvh): action is camera-independent (same TCP for all cameras), # so deduplicate to avoid 6x repeated points in the trajectory. sample_key = next(iter(action.keys())) if isinstance(sample_key, tuple) and len(sample_key) == 2: first_cam = min(k[1] for k in action.keys()) action_dedup = {k: v for k, v in action.items() if k[1] == first_cam} action = torch.stack( [action_dedup[k] for k in sorted(action_dedup.keys())], 1) else: action = torch.stack( [action[k] for k in sorted(action.keys())], 1) # Re-center all cameras relative to first camera base = cams[(0,0)].Twc for key in cams.keys(): cams[key].Tcw = base @ cams[key].Tcw if 'action' in data: b, l, n, x, y = action.shape action = base @ action.view(b, l * n, x, y) data['action'] = action.view(b, l, n, x, y) ego_cams = {key: val for key, val in cams.items() if key[1] == 'ego'} cams = {key: val for key, val in cams.items() if key[1] != 'ego'} cams_idx = set([key[1] for key in cams.keys()]) wh = [val.shape[-2:][::-1] for key, val in rgb.items() if key[0] == 0] wh = [int(v) for v in sorted(wh, key=lambda x: x[0])[-1]] # Reconstruct 3D points from depth points = {} if depth is not None: for key, val in cams.items(): points[key] = cams[key].reconstruct_depth_map(depth[key], euclidean=self.euclidean, to_world=True) data['normals'] = {key: calculate_normals(depth[key], cams[key], to_world=True) for key in depth.keys()} points_normals = {key: cams[key].reconstruct_depth_map(depth[key], euclidean=self.euclidean, to_world=True, world_scene_flow=data['normals'][key] / 5) for key in depth.keys()} points_normals = interleave_dict(points, points_normals) else: points_normals = None idx = [i for i in range(len(self.tasks)) if self.tasks[i] in data.keys()] keys, offsets = [self.tasks[i] for i in idx], [self.offset[i] for i in idx] return data, wh, keys, offsets, cams, ego_cams, cams_idx, points, points_normals, timestep def show_loading(self, msg='Loading...'): """Overlay a loading message on the current frame.""" import pygame wld_scr = self.draw.scr('wld') center_x = wld_scr.luwh[2] // 2 - 80 self.draw['wld'].text(msg, (center_x, 3)) pygame.display.flip() def loop(self, data=None): if data is not None: self.dataset = data data, wh, keys, offsets, cams, ego_cams, cams_idx, points, points_normals, timestep = self.process() camcv = {key: CameraCV.from_anydata(val, b=0, scale=self.scale) for key, val in cams.items()} length = sorted(set([key[0] for key in data.keys()])) zeros3 = torch.zeros((wh[1], wh[0], 3)) zeros4 = torch.zeros((wh[1] * wh[0], 4)) metadata = data['metadata'] c, t, k = 0, 0, 0 key = keys[k] change = True show_color = True show_normals = False show_all = True show_context = True show_text = True show_rays = False show_help = False save_video = False show_proj3d = False while self.draw.input(): if self.draw.SPACE: show_color = not show_color change = True if self.draw.RIGHT: change = True k = (k + 1) % len(keys) key = keys[k] if self.draw.LEFT: change = True k = (k - 1) % len(keys) key = keys[k] if self.draw.UP: change = True step = 10 if (self.draw.LSHIFT or self.draw.RSHIFT) else 1 for _ in range(step): c, t = change_key(data[key], c, t, 1) if self.draw.DOWN: change = True step = 10 if (self.draw.LSHIFT or self.draw.RSHIFT) else 1 for _ in range(step): c, t = change_key(data[key], c, t, -1) if self.draw.KEY_Q: self.scale += 0.01 camcv = {key: CameraCV.from_anydata(val, b=0, scale=self.scale) for key, val in cams.items()} ego_camcv = {key: CameraCV.from_anydata(val, b=0, scale=self.scale) for key, val in ego_cams.items()} if self.draw.KEY_W and self.scale > 0.01: self.scale -= 0.01 camcv = {key: CameraCV.from_anydata(val, b=0, scale=self.scale) for key, val in cams.items()} ego_camcv = {key: CameraCV.from_anydata(val, b=0, scale=self.scale) for key, val in ego_cams.items()} if self.draw.KEY_A and self.idx < len(self.dataset) - 1: t = 0 change = True self.idx += 1 self.show_loading(f'Loading sample {self.idx}...') data, wh, keys, offsets, cams, ego_cams, cams_idx, points, points_normals, timestep = self.process() metadata = data['metadata'] if wh != self.wh: self.create_draw(data, wh, cams, self.virtual_pose) if self.draw.KEY_Z and self.idx > 0: t = 0 change = True self.idx -= 1 self.show_loading(f'Loading sample {self.idx}...') data, wh, keys, offsets, cams, ego_cams, cams_idx, points, points_normals, timestep = self.process() metadata = data['metadata'] if wh != self.wh: self.create_draw(data, wh, cams, self.virtual_pose) if self.draw.KEY_X and points_normals is not None: show_normals = not show_normals self.draw.halt(100) if self.draw.KEY_R: show_rays = not show_rays self.draw.halt(100) if self.draw.KEY_C: show_context = not show_context self.draw.halt(100) if self.draw.KEY_P: show_proj3d = not show_proj3d self.draw.halt(100) if self.draw.KEY_F: self.filter_depth_outliers = not self.filter_depth_outliers change = True self.draw.halt(100) # TODO(bvh): upright mode (KEY_U) disabled for now, needs fixing if self.draw.KEY_T: show_text = not show_text self.draw.halt(100) if self.draw.KEY_H: show_help = not show_help self.draw.halt(100) if self.draw.KEY_V and not save_video: save_video, t, frames = True, 0, [] self.draw.halt(100) if self.draw.KEY_B: self.draw.save_frame('frame.png') self.draw.halt(100) # Bounding boxes if 'bbox2d' in data.keys(): all_bboxes2d = {} for i in cams_idx: bboxes2d = data['bbox2d'][(t,i)] if bboxes2d.shape[0] == 0: continue bboxes2d = bboxes2d[bboxes2d[:, -3] > 0] bboxes2d = [BBox2D(bboxes2d[i]) for i in range(bboxes2d.shape[0])] all_bboxes2d[i] = bboxes2d if 'bbox3d' in data.keys(): all_bboxes3d = {} proj_bboxes3d = {} for i in cams_idx: bboxes3d = data['bbox3d'][(t,i)] if bboxes3d.shape[0] == 0: continue bboxes3d = bboxes3d[bboxes3d[:, -3] > 0] h, w, l = bboxes3d[:, 0], bboxes3d[:, 1], bboxes3d[:, 2] tvec = np.array(bboxes3d[:, 3:6]) qvec = np.array(bboxes3d[:, 6:10]) rotmat = np.transpose(qvec2rotmat_batched(qvec), [2, 0, 1]) Tcw = np.repeat(np.expand_dims(np.eye(4), 0), rotmat.shape[0], axis=0) Tcw[:, :3, :3] = rotmat Tcw[:, :3, -1] = tvec 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 = cams[(t,i)].Tcw.numpy() @ Tcw bboxes3d = Tcw[:, :3, :3] @ corners + Tcw[:, :3, [-1]] if show_proj3d: corners = torch.tensor(bboxes3d).float().contiguous() n, d, c = corners.shape corners = corners.permute(0, 2, 1).reshape(n * c, d).unsqueeze(0).permute(0, 2, 1) proj = cams[(t,i)].project_points(corners, from_world=False, normalize=False).view(n, c, 2).permute(0, 2, 1) proj = proj.permute(0, 2, 1) proj = torch.stack([proj[i] for i in range(proj.shape[0]) if not (proj[i] == -2).any()], 0) proj_bboxes3d[i] = [BBox3D(proj[i].numpy()) for i in range(proj.shape[0])] bboxes3d = bboxes3d.transpose(0, 2, 1) bboxes3d = [BBox3D(bboxes3d[i]) for i in range(bboxes3d.shape[0])] all_bboxes3d[i] = bboxes3d if "map" in data.keys(): # sides: left(1), right(2), center(3) colours = {1: 'yel', 2: 'red', 3: 'cya'} widths = {1: 2, 2: 2, 3: 0.5} all_polylines = {} proj_polylines = {} for i in cams_idx: if (t,i) in data['map']: all_polylines[i] = [] proj_polylines[i] = [] Tcw: np.ndarray = cams[(t,i)].Tcw.numpy() lane_groups: List[Dict[str, np.ndarray]] = data['map'][(t,i)]["points"] for lane_group in lane_groups: for side, polyline in lane_group.items(): polyline_world = Tcw[:, :3, :3] @ polyline[:,:,None] + Tcw[:, :3, [-1]] all_polylines[i].append((colours[side], widths[side], polyline_world.squeeze(-1))) if show_proj3d: corners = torch.tensor(polyline_world).float().contiguous() n, d, c = corners.shape corners = corners.permute(0, 2, 1).reshape(n * c, d).unsqueeze(0).permute(0, 2, 1) proj = cams[(t,i)].project_points(corners, from_world=False, normalize=False, max_depth=200.0).view(n, c, 2).permute(0, 2, 1) proj = proj.squeeze(-1).numpy() proj = proj[:n // 2 * 2] invalid = (proj[:, 0] == -2) | (proj[:, 1] == -2) invalid[0::2] = invalid[0::2] | invalid[1::2] invalid[1::2] = invalid[1::2] | invalid[0::2] proj = proj[~invalid] proj_polylines[i].append((colours[side], widths[side], proj)) all_polylines[i] = [p for j, p in enumerate(all_polylines[i]) if j not in [4,7,10,13,16]] if show_proj3d: proj_polylines[i] = [p for j, p in enumerate(proj_polylines[i]) if j not in [4,7,10,13,16]] if change: change = False camcv = {key: CameraCV.from_anydata(val, b=0, scale=self.scale) for key, val in cams.items()} ego_camcv = {key: CameraCV.from_anydata(val, b=0, scale=self.scale) for key, val in ego_cams.items()} # 3D action trajectory if 'action' in data: action_xyz = data['action'][..., :3, -1] data['action3d'] = action_xyz for n in range(action_xyz.shape[2]): action_xyz_n = action_xyz[:, :, n] self.draw.addBufferf('action3d%d' % n, action_xyz_n[0]) clr = np.array(list(range(action_xyz_n.shape[1]))) self.draw.addBufferf('color3d%d' % n, jet(clr, exp=2.0)) # Per-camera image and point cloud updates data['action2d'] = dict() for i in cams_idx: # Projected 2D action trajectory if 'action' in data: action_xyz = data['action'][..., :3, -1] for n in range(action_xyz.shape[2]): action_xyz_n = action_xyz[:, :, n] data['action2d'][i] = action_uv_n = cams[(t,i)].project_points(action_xyz_n.permute(0, 2, 1), normalize=False) self.draw.addBufferf('action2d%d%d' % (i, n), action_uv_n[0]) clr = np.array(list(range(action_uv_n.shape[1]))) self.draw.addBufferf('color2d%d%d' % (i, n), jet(clr, exp=2.0)) # Update 2D image and color buffers img = data[key][(t, i)] if key == 'rgb': img = img[0] self.draw.updTexture('img%d' % i, img) self.draw.updBufferf('clr%d' % i, img) elif key == 'depth': img = viz_depth(img[0], filter_zeros=True) if self.filter_depth_outliers and 'depth' in data and (t, i) in data['depth']: d = data['depth'][(t, i)][0, 0] # (H, W) d_flat = d[d > 0] if len(d_flat) > 0: lo = np.percentile(d_flat.numpy(), self.depth_percentile_lo) hi = np.percentile(d_flat.numpy(), self.depth_percentile_hi) too_far = (d <= 0) | (d > hi) too_close = (d > 0) & (d < lo) if isinstance(img, torch.Tensor): img[too_far] = torch.tensor([0.3, 0.3, 0.3]) img[too_close] = torch.tensor([0.6, 0.6, 0.6]) else: img[too_far.numpy()] = 0.3 img[too_close.numpy()] = 0.6 self.draw.updTexture('img%d' % i, img) self.draw.updBufferf('clr%d' % i, img.reshape(-1, 3)) elif key == 'normals': img = viz_normals(img[0]) self.draw.updTexture('img%d' % i, img) self.draw.updBufferf('clr%d' % i, img.reshape(-1, 3)) elif key == 'optflow': if offsets[k] in img.keys(): img = viz_optflow(img[offsets[k]][0]) self.draw.updTexture('img%d' % i, img) self.draw.updBufferf('clr%d' % i, img.reshape(-1, 3)) else: self.draw.updTexture('img%d' % i, zeros3) self.draw.updBufferf('clr%d' % i, zeros4) elif key == 'semantic': ontology = data['metadata']['conventions']['semantic']['ontology'] img = viz_semantic(img[0], ontology) self.draw.updTexture('img%d' % i, img) self.draw.updBufferf('clr%d' % i, img.reshape(-1, 3)) # Update 3D point cloud buffers if len(points) > 0: pts = points[(t, i)][0] if self.filter_depth_outliers and 'depth' in data and (t, i) in data['depth']: d = data['depth'][(t, i)][0, 0] # (H, W) d_flat = d[d > 0] if len(d_flat) > 0: lo = np.percentile(d_flat.numpy(), self.depth_percentile_lo) hi = np.percentile(d_flat.numpy(), self.depth_percentile_hi) keep = (d > lo) & (d <= hi) pts = pts.clone() pts[:, ~keep] = 0 self.draw.updBufferf('pts%d' % i, pts) if points_normals is not None: self.draw.updBufferf('pts_nrm%d' % i, points_normals[(t, i)][0]) img_nrm = viz_normals(data['normals'][(t, i)][0]).reshape(-1, 3) img_nrm2 = np.zeros((img_nrm.shape[0] * 2, img_nrm.shape[1])) img_nrm2[::2], img_nrm2[1::2] = img_nrm, img_nrm self.draw.updBufferf('clr_nrm%d' % i, img_nrm2) rays = cams[(t,i)].get_viewdirs(normalize='unit', to_world=True)[0].permute(1, 2, 0).reshape(-1, 3) * self.scale orig = cams[(t,i)].get_origin(flatten=True).reshape(-1, 3) rays_save = np.zeros((rays.shape[0] * 2, rays.shape[1])) rays_save[::2], rays_save[1::2] = orig, orig + rays self.draw.updBufferf('rays%d' % i, rays_save) # Gather language metadata lowdim_language = data.get('language', dict()) if len(lowdim_language) > 0: lowdim_language = lowdim_language[t] # Gather metadata language (task/prompt promoted to top level by normalize_metadata) metadata_language = {} for lk in ['task', 'prompt']: if lk in metadata: metadata_language[lk] = metadata[lk] language = dict() for label in ['task','prompt','skill']: language[label] = \ lowdim_language.get(label, None) if label in lowdim_language else \ metadata_language.get(label, None) if label in metadata_language else None # ---- Rendering ---- self.draw.clear() # Draw 3D action trajectory if 'action' in data: for n in range(data['action'].shape[2]): self.draw['wld'].width(2).color('yel').strips('action3d%d' % n, 'color3d%d' % n) self.draw['wld'].size(2).color('whi').points('action3d%d' % n) actcams = [Camera(K=None, Tcw=data['action'][:, i, n], hw=(20,60)) for i in range(data['action'].shape[1])] cvactcams = [CameraCV.from_anydata(val, b=0, scale=0.005) for val in actcams] for cvactcam in cvactcams: self.draw['wld'].object(cvactcam, width=1, color='whi') pt3d = data['action3d'][0, c, n].unsqueeze(0).numpy() self.draw['wld'].size(10).color('whi').points(pt3d) # Draw cameras, point clouds, and images for i, (cam_key, cam_val) in enumerate(camcv.items()): if cam_key[0] == t or show_all: self.draw['wld'].size(2).color( self.cam_colors[cam_key[1]]).points( ('pts%d' % cam_key[1]), ('clr%d' % cam_key[1]) if show_color else None) if show_normals: self.draw['wld'].width(1).lines( 'pts_nrm%d' % cam_key[1], 'clr_nrm%d' % cam_key[1]) if show_rays: self.draw['wld'].size(3).color('yel').points( 'rays%d' % cam_key[1]) self.draw['img%d%d' % modrem(cam_key[1], 2)].image('img%d' % cam_key[1]) if 'action' in data: # Projected actions for n in range(data['action'].shape[2]): self.draw['img%d%d' % modrem(cam_key[1], 2)].width(2).strips( 'action2d%d%d' % (cam_key[1], n), 'color2d%d%d' % (cam_key[1], n)) self.draw['img%d%d' % modrem(cam_key[1], 2)].size(2).color('whi').points( 'action2d%d%d' % (cam_key[1], n)) pt2d = data['action2d'][cam_key[1]][0, c].unsqueeze(0).numpy() self.draw['img%d%d' % modrem(cam_key[1], 2)].size(5).color('whi').points(pt2d) if not show_context and cam_key[0] != t: continue clr = self.cam_colors[cam_key[1]] if cam_key[0] == t else 'gra' tex = 'img%d' % cam_key[1] if cam_key[0] == t else None self.draw['wld'].object(cam_val, color=clr, tex=tex) for i, (cam_key, cam_val) in enumerate(ego_camcv.items()): self.draw['wld'].object(cam_val, color='whi') if show_text: # Sample information (top left) text = keys[k].upper() + ('' if offsets[k] is None else '_' + offsets[k].upper()) self.draw['wld'].text(text, (0, 0)) self.draw['wld'].text(f'IDX: {str(self.idx)} / {t}', (0, 1)) # Viewer state wld_screen = self.draw.scr('wld') T = wld_screen.viewer.T cam_pos = T[:3, 3] look_at = T[:3, 3] + T[:3, 2] self.draw['wld'].text( f'Cam: ({cam_pos[0]:.2f}, {cam_pos[1]:.2f}, {cam_pos[2]:.2f})' f' -> ({look_at[0]:.2f}, {look_at[1]:.2f}, {look_at[2]:.2f})', (0, 2)) # Dataset metadata (bottom left) text = f'Dataset: {metadata["name"]}'.replace("'", "") if 'station' in metadata: text += f' ({metadata["station"]})' self.draw['wld'].text(text, (0, -9)) tar_path = metadata.get('path', '') if tar_path: parts = tar_path.split('/') short = '/'.join(parts[-4:]) if len(parts) > 4 else tar_path self.draw['wld'].text(f'Tar: {short}', (0, -10)) tags = f'Tags: {metadata["tags"]}'.replace("'", "") self.draw['wld'].text(tags, (0, -8)) labels = f'Labels: {metadata["labels"]}'.replace("'", "") self.draw['wld'].text(labels, (0, -7)) cameras = f'Cameras: {metadata["cameras"]}'.replace("'", "") self.draw['wld'].text(cameras, (0, -6)) # Language if language['task'] is not None: text = f'Task: {language["task"]}'.replace("'", "") self.draw['wld'].text(text, (0, -5)) if language['prompt'] is not None: text = f'Prompts: {language["prompt"]}'.replace("'", "") if language['skill'] is not None: text += f' ({language["skill"]})' self.draw['wld'].text(text, (0, -4)) # Timing timestep_start = float('%.3f' % data['metadata']['timestep_start'][0]) timestep_t = {key[1]: float('%.3f' % val) for key, val in timestep.items() if key[0] == t} timestep_t = reduce_dict(timestep_t) self.draw['wld'].text(f'Timestep: {timestep_t} (Start: {timestep_start})', (0, -3)) resolution = reduce_dict(metadata["resolution"]) self.draw['wld'].text(f'Resolution: {resolution}', (0, -2)) framerate = reduce_dict(metadata["framerate"]) stride = reduce_dict(metadata["stride"]) length = reduce_dict(metadata["length"]) self.draw['wld'].text(f'Framerate: {framerate} / Stride: {stride} / Length: {length}', (0, -1)) # Keyboard shortcut guide (lower right) if show_help: shortcuts = [ 'SHORTCUTS', 'SPACE Toggle color', ' Switch modality', 'v/^ Cycle timestep', 'A/Z Next/Prev sample', 'Q/W Scale cams +/-', 'Sh+v/^ Step by 10', 'X Toggle normals', 'R Toggle rays', 'C Toggle context', 'T Toggle text', 'F Filter depth outliers', # 'U Lock upright', # TODO(bvh): disabled, needs fixing '', 'LMB Pan', 'RMB Rotate', 'Wheel Zoom', 'MMB Reset view', ] else: shortcuts = [ 'Press H for Help', ] wld_scr = self.draw.scr('wld') shortcut_x = wld_scr.luwh[2] - 220 for si, s in enumerate(shortcuts): wld_scr.text(s, (shortcut_x, -(len(shortcuts) - si)), color=(0.5, 0.5, 0.5)) # Draw 2D bounding boxes if 'bbox2d' in data.keys() and not show_proj3d: for i, val in all_bboxes2d.items(): for b in val: self.draw['img%d%d' % modrem(i, 2)].object(b) # Draw 3D bounding boxes if 'bbox3d' in data.keys(): for i, val in all_bboxes3d.items(): for b in val: self.draw['wld'].object(b) if show_proj3d: for i, val in proj_bboxes3d.items(): for b in val: self.draw['img%d%d' % modrem(i, 2)].object(b, color_line='gre') # Draw map information if "map" in data.keys(): for i, val in all_polylines.items(): for j, (colour, width, polyline) in enumerate(val): self.draw['wld'].width(width).color(colour).lines(polyline) if show_proj3d: for i, val in proj_polylines.items(): for j, (colour, width, polyline) in enumerate(val): self.draw['img%d%d' % modrem(i, 2)].width(width).color(colour).lines(polyline) # Draw borders for i, cam_key in enumerate(camcv.keys()): w, h = self.wh color = self.cam_colors[cam_key[1]] tex = 'img%d%d' % modrem(cam_key[1], 2) self.draw[tex].color(color).width(10).lines(np.array([ [0,0],[0,h],[0,h],[w,h],[w,h],[w,0],[w,0],[0,0], ])) # Update drawing window self.draw.update(30) if save_video: if t < length * stride - stride: frame = self.draw.save(f'video_%03d.jpg' % t, video=True) frames.append(frame) _, t = change_key(data[key], c, t, 1) else: frame = self.draw.save(f'video_%03d.jpg' % t, video=True) frames.append(frame) save_video, t = False, 0 from anydata.utils.write import write_video frames = torch.stack([torch.tensor(np.array(f)) for f in frames], 0) write_video('video.mp4', frames[..., :3], fps=5) change = True