import argparse import os import sys import time import cv2 import numpy as np SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) REPO_ROOT = os.path.abspath(os.path.join(SCRIPT_DIR, "..", "..")) SRC_DIR = os.path.join(REPO_ROOT, "src") if SRC_DIR not in sys.path: sys.path.insert(0, SRC_DIR) import orisys def load_exported_data(data_dir): mesh = np.load(os.path.join(data_dir, "mesh.npz")) cam_data = np.load(os.path.join(data_dir, "camera_intrinsics.npz")) xyz_map_path = os.path.join(data_dir, "pixel_to_xyz.npz") if not os.path.exists(xyz_map_path): raise FileNotFoundError(f"Missing precomputed 3D map: {xyz_map_path}") xyz_map = np.load(xyz_map_path)["xyz_map"] width = int(cam_data["width"]) height = int(cam_data["height"]) if xyz_map.shape[:2] != (height, width): raise ValueError( f"pixel_to_xyz shape {xyz_map.shape[:2]} does not match camera size {(height, width)}" ) return { "vertices": mesh["vertices"], "faces": mesh["faces"], "uv_data": mesh["uv_data"], "loop_to_vertex": mesh["loop_to_vertex"], "K": cam_data["K"], "width": width, "height": height, "world_to_cam": np.load(os.path.join(data_dir, "world_to_cam.npy")), "xyz_map": xyz_map, } def parse_source(value): value = str(value).strip() for prefix in ("camera:", "cam:", "webcam:"): if value.lower().startswith(prefix): value = value[len(prefix):].strip() break try: return int(value) except ValueError: return value def prepare_static_points(data, display_stride): xyz_map = data["xyz_map"] height, width = xyz_map.shape[:2] valid_mask = ~np.isnan(xyz_map[:, :, 0]) valid_points = xyz_map[valid_mask].reshape(-1, 3) if len(valid_points) == 0: raise ValueError("pixel_to_xyz.npz contains no valid mesh hit points.") display_indices = np.arange(len(valid_points))[::max(1, int(display_stride))] flat_indices = np.flatnonzero(valid_mask.ravel())[display_indices] display_ys = flat_indices // width display_xs = flat_indices % width return ( valid_mask, display_indices, valid_points[display_indices], display_ys, display_xs, ) def mesh_pixel_to_world_scale(xyz_map, valid_mask): valid_points = xyz_map[valid_mask] if len(valid_points) == 0: return 1e-3 extent = np.nanmax(valid_points, axis=0) - np.nanmin(valid_points, axis=0) mesh_span = float(np.max(extent)) height, width = xyz_map.shape[:2] return mesh_span / max(height, width, 1) def sample_field_at_display_points(field, display_ys, display_xs): if field.ndim == 2: return field[display_ys, display_xs] return field[display_ys, display_xs, :] def deform_point_cloud(rest_points, displacement, pixel_to_world, deform_scale): return rest_points + displacement * pixel_to_world * float(deform_scale) def lift_flow_to_3d(flow, tangent_x, tangent_y, normals): """Lift raw 2D optical flow to tangent-plane 3D displacement.""" tx_norm = np.linalg.norm(tangent_x, axis=2, keepdims=True) ty_norm = np.linalg.norm(tangent_y, axis=2, keepdims=True) t1 = np.divide(tangent_x, tx_norm, out=np.zeros_like(tangent_x), where=tx_norm > 1e-8) t2 = np.divide(tangent_y, ty_norm, out=np.zeros_like(tangent_y), where=ty_norm > 1e-8) flow_3d = flow[:, :, 0:1] * t1 + flow[:, :, 1:2] * t2 normal_component = np.sum(flow_3d * normals, axis=2, keepdims=True) return flow_3d - normal_component * normals def align_image_to_map(img_bgr, data, valid_mask, display_indices, flip_vertical=True): target_w = int(data["width"]) target_h = int(data["height"]) h, w = img_bgr.shape[:2] if (w, h) != (target_w, target_h): img_bgr = cv2.resize(img_bgr, (target_w, target_h), interpolation=cv2.INTER_LINEAR) img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB) if flip_vertical: img_rgb = np.flipud(img_rgb).copy() colors = img_rgb.astype(np.float64) / 255.0 return colors[valid_mask].reshape(-1, 3)[display_indices] def arrow_map_reference_colors(count, color=(0.04, 0.04, 0.04)): return np.tile(np.asarray(color, dtype=np.float64), (count, 1)) def magnitude_map_colors(flow, valid_mask, display_indices, max_magnitude=None): mag = np.linalg.norm(flow, axis=2) valid_mag = mag[valid_mask].reshape(-1) display_mag = valid_mag[display_indices] if max_magnitude is None or max_magnitude <= 0: max_magnitude = max(float(np.percentile(display_mag, 95)), 1e-6) value = np.clip(display_mag / max_magnitude, 0.0, 1.0) return np.column_stack( ( value, np.clip(1.0 - np.abs(value - 0.5) * 2.0, 0.0, 1.0), 1.0 - value, ) ).astype(np.float64) def align_scalar_field_to_map(field, data, flip_vertical=True): if field.ndim == 3: field = np.linalg.norm(field, axis=2) target_w = int(data["width"]) target_h = int(data["height"]) h, w = field.shape[:2] if (w, h) != (target_w, target_h): field = cv2.resize(field, (target_w, target_h), interpolation=cv2.INTER_LINEAR) else: field = field.copy() if flip_vertical: field = np.flipud(field).copy() return field def scalar_map_colors(field, valid_mask, display_indices, value_min=None, value_max=None): valid_values = field[valid_mask].reshape(-1) display_values = valid_values[display_indices] finite_values = display_values[np.isfinite(display_values)] if len(finite_values) == 0: return arrow_map_reference_colors(len(display_indices)) if value_min is None: value_min = float(np.percentile(finite_values, 5)) if value_max is None: value_max = float(np.percentile(finite_values, 95)) if value_max <= value_min: value_max = value_min + 1e-6 value = np.clip((display_values - value_min) / (value_max - value_min), 0.0, 1.0) return np.column_stack( ( value, np.clip(1.0 - np.abs(value - 0.5) * 2.0, 0.0, 1.0), 1.0 - value, ) ).astype(np.float64) def align_vector_field_to_map(field, data, flip_vertical=True, scale_as_flow=True): target_w = int(data["width"]) target_h = int(data["height"]) h, w = field.shape[:2] if (w, h) != (target_w, target_h): sx = target_w / max(w, 1) sy = target_h / max(h, 1) field = cv2.resize(field, (target_w, target_h), interpolation=cv2.INTER_LINEAR) if scale_as_flow: field[:, :, 0] *= sx field[:, :, 1] *= sy else: field = field.copy() if flip_vertical: field = np.flipud(field).copy() field[:, :, 1] *= -1.0 return field def estimate_surface_normals(xyz_map): valid = ~np.isnan(xyz_map[:, :, 0]) filled = xyz_map.copy() filled[~valid] = 0 d_y = np.gradient(filled, axis=0) d_x = np.gradient(filled, axis=1) normals = np.cross(d_x, d_y) norm = np.linalg.norm(normals, axis=2, keepdims=True) normals = np.divide(normals, norm, out=np.zeros_like(normals), where=norm > 1e-8) normals[~valid] = 0 return normals def compute_surface_frames(xyz_map): valid = ~np.isnan(xyz_map[:, :, 0]) filled = xyz_map.copy() filled[~valid] = 0 d_y = np.gradient(filled, axis=0) d_x = np.gradient(filled, axis=1) normals = np.cross(d_x, d_y) norm = np.linalg.norm(normals, axis=2, keepdims=True) normals = np.divide(normals, norm, out=np.zeros_like(normals), where=norm > 1e-8) d_x[~valid] = 0 d_y[~valid] = 0 normals[~valid] = 0 return d_x, d_y, normals, valid def lift_deformation_to_3d(vshear, vnormal, tangent_x, tangent_y, normals, normal_gain=1.0): # Unit tangents: map 2D pixel shear to 3D direction without world-scale Jacobian blow-up. tx_norm = np.linalg.norm(tangent_x, axis=2, keepdims=True) ty_norm = np.linalg.norm(tangent_y, axis=2, keepdims=True) t1 = np.divide(tangent_x, tx_norm, out=np.zeros_like(tangent_x), where=tx_norm > 1e-8) t2 = np.divide(tangent_y, ty_norm, out=np.zeros_like(tangent_y), where=ty_norm > 1e-8) shear_3d = vshear[:, :, 0:1] * t1 + vshear[:, :, 1:2] * t2 normal_mag = np.linalg.norm(vnormal, axis=2, keepdims=True) * float(normal_gain) normal_3d = normal_mag * normals normal_component = np.sum(shear_3d * normals, axis=2, keepdims=True) shear_tangent = shear_3d - normal_component * normals return shear_tangent + normal_3d def line_colors(mag, arrow_map_style): if len(mag) == 0: return np.empty((0, 3), dtype=np.float64) if arrow_map_style: return np.tile(np.array([0.0, 0.9, 1.0], dtype=np.float64), (len(mag), 1)) norm_mag = mag / max(float(np.percentile(mag, 95)), 1e-6) norm_mag = np.clip(norm_mag, 0.0, 1.0) return np.column_stack((norm_mag, 1.0 - norm_mag, np.zeros_like(norm_mag))).astype(np.float64) def lines_from_vectors(starts, vectors, mag, flow_scale, arrow_map_style, unit_direction=False): vec_norm = np.linalg.norm(vectors, axis=1) good = np.isfinite(vec_norm) & (vec_norm > 1e-8) starts, vectors, mag = starts[good], vectors[good], mag[good] if len(starts) == 0: return np.empty((0, 3)), np.empty((0, 2), dtype=np.int32), np.empty((0, 3)) if unit_direction: vec_norm = vec_norm[:, None] vectors = np.divide(vectors, vec_norm, out=np.zeros_like(vectors), where=vec_norm > 1e-8) points = np.empty((len(starts) * 2, 3), dtype=np.float64) points[0::2] = starts points[1::2] = starts + vectors * flow_scale lines = np.column_stack((np.arange(0, len(points), 2), np.arange(1, len(points), 2))).astype(np.int32) return points, lines, line_colors(mag, arrow_map_style) def build_flow_lines( xyz_map, valid_mask, flow, arrow_stride=32, flow_scale=1.0, min_magnitude=0.5, arrow_map_style=False, ): h, w = flow.shape[:2] ys, xs = np.mgrid[0:h:arrow_stride, 0:w:arrow_stride] ys = ys.reshape(-1) xs = xs.reshape(-1) sample_flow = flow[ys, xs] mag = np.linalg.norm(sample_flow, axis=1) keep = valid_mask[ys, xs] & (mag >= min_magnitude) xs, ys, sample_flow, mag = xs[keep], ys[keep], sample_flow[keep], mag[keep] if len(xs) == 0: return np.empty((0, 3)), np.empty((0, 2), dtype=np.int32), np.empty((0, 3)) tx = np.rint(xs + sample_flow[:, 0]).astype(np.int32) ty = np.rint(ys + sample_flow[:, 1]).astype(np.int32) inside = (0 <= tx) & (tx < w) & (0 <= ty) & (ty < h) xs, ys, tx, ty, mag = xs[inside], ys[inside], tx[inside], ty[inside], mag[inside] if len(xs) == 0: return np.empty((0, 3)), np.empty((0, 2), dtype=np.int32), np.empty((0, 3)) valid_target = valid_mask[ty, tx] xs, ys, tx, ty, mag = xs[valid_target], ys[valid_target], tx[valid_target], ty[valid_target], mag[valid_target] starts = xyz_map[ys, xs] vectors = xyz_map[ty, tx] - starts return lines_from_vectors(starts, vectors, mag, flow_scale, arrow_map_style) def build_resultant_force_lines( xyz_map, valid_mask, resultant_force, arrow_stride=32, flow_scale=1.0, min_magnitude=0.01, arrow_map_style=False, ): h, w = resultant_force.shape[:2] ys, xs = np.mgrid[0:h:arrow_stride, 0:w:arrow_stride] ys = ys.reshape(-1) xs = xs.reshape(-1) vectors = resultant_force[ys, xs] mag = np.linalg.norm(vectors, axis=1) keep = valid_mask[ys, xs] & (mag >= min_magnitude) xs, ys, vectors, mag = ( xs[keep], ys[keep], vectors[keep], mag[keep], ) if len(xs) == 0: return np.empty((0, 3)), np.empty((0, 2), dtype=np.int32), np.empty((0, 3)) starts = xyz_map[ys, xs] return lines_from_vectors( starts, vectors, mag, flow_scale, arrow_map_style, unit_direction=True ) def init_viewer( cloud_points, cloud_colors, flow_points, flow_lines, flow_colors, arrow_map_style=False, show_arrows=True, ): import open3d as o3d vis = o3d.visualization.Visualizer() vis.create_window(window_name="SDK 3D Flow", width=1280, height=720) pcd = o3d.geometry.PointCloud() pcd.points = o3d.utility.Vector3dVector(cloud_points.astype(np.float64)) pcd.colors = o3d.utility.Vector3dVector(cloud_colors.astype(np.float64)) vis.add_geometry(pcd) line_set = None if show_arrows: line_set = o3d.geometry.LineSet() line_set.points = o3d.utility.Vector3dVector(flow_points.astype(np.float64)) line_set.lines = o3d.utility.Vector2iVector(flow_lines) line_set.colors = o3d.utility.Vector3dVector(flow_colors.astype(np.float64)) vis.add_geometry(line_set) render_option = vis.get_render_option() render_option.point_size = 1.0 if arrow_map_style else 2.0 render_option.line_width = 2.0 render_option.background_color = np.array([0.0, 0.0, 0.0]) vis.poll_events() vis.update_renderer() return vis, pcd, line_set, o3d def run_sdk_3d_flow(args): data = load_exported_data(args.data_dir) valid_mask, display_indices, rest_cloud_points, display_ys, display_xs = prepare_static_points( data, args.display_stride ) tangent_x, tangent_y, surface_normals, valid_mask = compute_surface_frames(data["xyz_map"]) pixel_to_world = mesh_pixel_to_world_scale(data["xyz_map"], valid_mask) show_arrows = args.display_mode == "arrows" deform_cloud = args.display_mode == "deform" sensor = orisys.Sensor( parse_source(args.source), config_name=args.config, isstitch=False, cuda=args.cuda, verbose=args.verbose, mirror=not args.no_mirror, ) vis = pcd = line_set = o3d = None frame_count = 0 t0 = time.perf_counter() try: while True: if vis is not None and not vis.poll_events(): break if sensor.get_img() is None: break needs_decomposition = args.vector_mode in ("shear-normal", "depth-map") sensor.compute_deformation(decompose=needs_decomposition, check_motion=False) if args.vector_mode == "depth-map": sensor.compute_contact() flow, img_bgr, vshear, vnormal, depth_map = sensor.read_info( sensor.info.VRAW, sensor.info.IMG, sensor.info.VSHEAR, sensor.info.VNORMAL, sensor.info.DEPTH, ) flow_map = align_vector_field_to_map(flow, data, flip_vertical=not args.no_flip_vertical) if args.vector_mode == "vraw-magnitude": cloud_colors = magnitude_map_colors(flow_map, valid_mask, display_indices, args.magnitude_max) elif args.vector_mode == "depth-map": depth_map = align_scalar_field_to_map(depth_map, data, flip_vertical=not args.no_flip_vertical) cloud_colors = scalar_map_colors( depth_map, valid_mask, display_indices, value_min=args.depth_min, value_max=args.depth_max, ) elif args.arrow_map_style: cloud_colors = arrow_map_reference_colors(len(rest_cloud_points)) else: cloud_colors = align_image_to_map( img_bgr, data, valid_mask, display_indices, flip_vertical=not args.no_flip_vertical, ) flip_vertical = not args.no_flip_vertical displacement_map = None if args.vector_mode == "shear-normal": vshear_map = align_vector_field_to_map(vshear, data, flip_vertical=flip_vertical) vnormal_map = align_vector_field_to_map(vnormal, data, flip_vertical=flip_vertical) displacement_map = lift_deformation_to_3d( vshear_map, vnormal_map, tangent_x, tangent_y, surface_normals, normal_gain=args.normal_gain, ) elif args.vector_mode == "depth-map" and deform_cloud: vnormal_map = align_vector_field_to_map(vnormal, data, flip_vertical=flip_vertical) zero_shear = np.zeros_like(vnormal_map) displacement_map = lift_deformation_to_3d( zero_shear, vnormal_map, tangent_x, tangent_y, surface_normals, normal_gain=args.normal_gain, ) elif args.vector_mode in ("vraw", "vraw-magnitude"): displacement_map = lift_flow_to_3d(flow_map, tangent_x, tangent_y, surface_normals) if deform_cloud and args.deform_color == "magnitude" and displacement_map is not None: cloud_colors = magnitude_map_colors( displacement_map, valid_mask, display_indices, args.magnitude_max ) if deform_cloud: if displacement_map is None: raise ValueError( f"display-mode deform is not supported with vector-mode {args.vector_mode!r}" ) point_displacement = sample_field_at_display_points( displacement_map, display_ys, display_xs ) cloud_points = deform_point_cloud( rest_cloud_points, point_displacement, pixel_to_world, args.deform_scale, ) flow_points = np.empty((0, 3), dtype=np.float64) flow_lines = np.empty((0, 2), dtype=np.int32) flow_colors = np.empty((0, 3), dtype=np.float64) else: cloud_points = rest_cloud_points if args.vector_mode == "shear-normal": flow_points, flow_lines, flow_colors = build_resultant_force_lines( data["xyz_map"], valid_mask, displacement_map, arrow_stride=args.arrow_stride, flow_scale=args.flow_scale, min_magnitude=args.min_flow, arrow_map_style=args.arrow_map_style, ) elif args.vector_mode in ("vraw-magnitude", "depth-map"): flow_points = np.empty((0, 3), dtype=np.float64) flow_lines = np.empty((0, 2), dtype=np.int32) flow_colors = np.empty((0, 3), dtype=np.float64) else: flow_points, flow_lines, flow_colors = build_flow_lines( data["xyz_map"], valid_mask, flow_map, arrow_stride=args.arrow_stride, flow_scale=args.flow_scale, min_magnitude=args.min_flow, arrow_map_style=args.arrow_map_style, ) if vis is None: vis, pcd, line_set, o3d = init_viewer( cloud_points, cloud_colors, flow_points, flow_lines, flow_colors, arrow_map_style=args.arrow_map_style, show_arrows=show_arrows, ) mode_label = "deformed cloud" if deform_cloud else "arrows" print(f"SDK 3D flow viewer started ({mode_label}). Close the Open3D window to stop.") else: pcd.points = o3d.utility.Vector3dVector(cloud_points.astype(np.float64)) pcd.colors = o3d.utility.Vector3dVector(cloud_colors) if line_set is not None: line_set.points = o3d.utility.Vector3dVector(flow_points.astype(np.float64)) line_set.lines = o3d.utility.Vector2iVector(flow_lines) line_set.colors = o3d.utility.Vector3dVector(flow_colors.astype(np.float64)) vis.update_geometry(line_set) vis.update_geometry(pcd) vis.update_renderer() frame_count += 1 if frame_count % 30 == 0: elapsed = max(time.perf_counter() - t0, 1e-6) arrow_count = len(flow_lines) if show_arrows else 0 print(f"frames={frame_count}, viewer FPS={frame_count / elapsed:.1f}, arrows={arrow_count}") finally: sensor.disconnect() if vis is not None: vis.destroy_window() def parse_args(): parser = argparse.ArgumentParser(description="Map SDK flow fields to a 3D point cloud.") parser.add_argument("--data-dir", default=os.path.join(SCRIPT_DIR, "exported_data")) parser.add_argument( "--source", default=os.path.join(SCRIPT_DIR, "20260602-094916.mp4"), help="Video path or live camera ID, for example 0, 1, camera:0, or cam:1.", ) parser.add_argument("--config", default=os.path.join(REPO_ROOT, "config", "dd02-ov.json")) parser.add_argument("--display-stride", type=int, default=4) parser.add_argument("--arrow-stride", type=int, default=8) parser.add_argument( "--flow-scale", type=float, default=3.0, help="Arrow length in 3D viewer units (shear-normal uses unit direction; magnitude shown by color).", ) parser.add_argument( "--min-flow", type=float, default=0.5, help="Minimum 2D-equivalent resultant magnitude before drawing an arrow.", ) parser.add_argument("--normal-gain", type=float, default=1.0) parser.add_argument( "--display-mode", choices=["arrows", "deform"], default="arrows", help="arrows: line vectors; deform: displace the point cloud by the 3D deformation field.", ) parser.add_argument( "--deform-scale", type=float, default=1.0, help="Gain on cloud displacement (multiplied by mesh pixel-to-world scale).", ) parser.add_argument( "--deform-color", choices=["texture", "magnitude"], default="texture", help="Point colors in deform mode: camera texture or deformation magnitude heat-map.", ) parser.add_argument("--magnitude-max", type=float, default=None, help="Upper VRAW magnitude for heat-map color scaling; default uses frame percentile.") parser.add_argument("--depth-min", type=float, default=None, help="Minimum depth value for heat-map color scaling; default uses frame percentile.") parser.add_argument("--depth-max", type=float, default=None, help="Maximum depth value for heat-map color scaling; default uses frame percentile.") parser.add_argument( "--vector-mode", choices=["vraw", "vraw-magnitude", "shear-normal", "depth-map"], default="vraw", help="Vector display mode. shear-normal shows 3D resultant vectors (surface-tangent shear + surface-normal component).", ) parser.add_argument( "--arrow-map-style", action="store_true", help="Hide real texture and use a dark Qt-viewer-like arrow map style.", ) parser.add_argument("--no-flip-vertical", action="store_true") parser.add_argument("--no-mirror", action="store_true") parser.add_argument("--cuda", action="store_true") parser.add_argument("--verbose", action="store_true") return parser.parse_args() if __name__ == "__main__": run_sdk_3d_flow(parse_args())