update dd02 3d viewer v1
This commit is contained in:
ykcheungab
@@ -0,0 +1,424 @@
|
||||
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):
|
||||
try:
|
||||
return int(value)
|
||||
except ValueError:
|
||||
return value
|
||||
|
||||
|
||||
def prepare_static_points(data, display_stride):
|
||||
xyz_map = data["xyz_map"]
|
||||
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))]
|
||||
return valid_mask, display_indices, valid_points[display_indices]
|
||||
|
||||
|
||||
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_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 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):
|
||||
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))
|
||||
|
||||
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_shear_normal_lines(
|
||||
xyz_map,
|
||||
valid_mask,
|
||||
vshear,
|
||||
vnormal,
|
||||
normals,
|
||||
arrow_stride=32,
|
||||
flow_scale=1.0,
|
||||
min_magnitude=0.01,
|
||||
normal_gain=1.0,
|
||||
arrow_map_style=False,
|
||||
):
|
||||
h, w = vshear.shape[:2]
|
||||
ys, xs = np.mgrid[0:h:arrow_stride, 0:w:arrow_stride]
|
||||
ys = ys.reshape(-1)
|
||||
xs = xs.reshape(-1)
|
||||
|
||||
shear = vshear[ys, xs]
|
||||
normal_field = vnormal[ys, xs]
|
||||
shear_mag = np.linalg.norm(shear, axis=1)
|
||||
normal_mag = np.linalg.norm(normal_field, axis=1) * normal_gain
|
||||
combined_mag = shear_mag + normal_mag
|
||||
keep = valid_mask[ys, xs] & (combined_mag >= min_magnitude)
|
||||
xs, ys, shear, shear_mag, normal_mag, combined_mag = (
|
||||
xs[keep],
|
||||
ys[keep],
|
||||
shear[keep],
|
||||
shear_mag[keep],
|
||||
normal_mag[keep],
|
||||
combined_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 + shear[:, 0]).astype(np.int32)
|
||||
ty = np.rint(ys + shear[:, 1]).astype(np.int32)
|
||||
inside = (0 <= tx) & (tx < w) & (0 <= ty) & (ty < h)
|
||||
xs, ys, tx, ty, shear_mag, normal_mag, combined_mag = (
|
||||
xs[inside],
|
||||
ys[inside],
|
||||
tx[inside],
|
||||
ty[inside],
|
||||
shear_mag[inside],
|
||||
normal_mag[inside],
|
||||
combined_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, shear_mag, normal_mag, combined_mag = (
|
||||
xs[valid_target],
|
||||
ys[valid_target],
|
||||
tx[valid_target],
|
||||
ty[valid_target],
|
||||
shear_mag[valid_target],
|
||||
normal_mag[valid_target],
|
||||
combined_mag[valid_target],
|
||||
)
|
||||
|
||||
starts = xyz_map[ys, xs]
|
||||
tangent = xyz_map[ty, tx] - starts
|
||||
tangent_norm = np.linalg.norm(tangent, axis=1, keepdims=True)
|
||||
tangent_unit = np.divide(tangent, tangent_norm, out=np.zeros_like(tangent), where=tangent_norm > 1e-8)
|
||||
normal_unit = normals[ys, xs]
|
||||
|
||||
shear_ratio = shear_mag / np.maximum(combined_mag, 1e-8)
|
||||
normal_ratio = normal_mag / np.maximum(combined_mag, 1e-8)
|
||||
vectors = tangent_unit * shear_ratio[:, None] + normal_unit * normal_ratio[:, None]
|
||||
vector_norm = np.linalg.norm(vectors, axis=1, keepdims=True)
|
||||
vectors = np.divide(vectors, vector_norm, out=np.zeros_like(vectors), where=vector_norm > 1e-8)
|
||||
return lines_from_vectors(starts, vectors, combined_mag, flow_scale, arrow_map_style)
|
||||
|
||||
|
||||
def init_viewer(cloud_points, cloud_colors, flow_points, flow_lines, flow_colors, arrow_map_style=False):
|
||||
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 = 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, cloud_points = prepare_static_points(data, args.display_stride)
|
||||
surface_normals = estimate_surface_normals(data["xyz_map"])
|
||||
|
||||
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
|
||||
sensor.compute_deformation(decompose=(args.vector_mode == "shear-normal"), check_motion=False)
|
||||
flow, img_bgr, vshear, vnormal = sensor.read_info(
|
||||
sensor.info.VRAW,
|
||||
sensor.info.IMG,
|
||||
sensor.info.VSHEAR,
|
||||
sensor.info.VNORMAL,
|
||||
)
|
||||
|
||||
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.arrow_map_style:
|
||||
cloud_colors = arrow_map_reference_colors(len(cloud_points))
|
||||
else:
|
||||
cloud_colors = align_image_to_map(
|
||||
img_bgr,
|
||||
data,
|
||||
valid_mask,
|
||||
display_indices,
|
||||
flip_vertical=not args.no_flip_vertical,
|
||||
)
|
||||
|
||||
if args.vector_mode == "shear-normal":
|
||||
vshear_map = align_vector_field_to_map(vshear, data, flip_vertical=not args.no_flip_vertical)
|
||||
vnormal_map = align_vector_field_to_map(vnormal, data, flip_vertical=not args.no_flip_vertical)
|
||||
flow_points, flow_lines, flow_colors = build_shear_normal_lines(
|
||||
data["xyz_map"],
|
||||
valid_mask,
|
||||
vshear_map,
|
||||
vnormal_map,
|
||||
surface_normals,
|
||||
arrow_stride=args.arrow_stride,
|
||||
flow_scale=args.flow_scale,
|
||||
min_magnitude=args.min_flow,
|
||||
normal_gain=args.normal_gain,
|
||||
arrow_map_style=args.arrow_map_style,
|
||||
)
|
||||
elif args.vector_mode == "vraw-magnitude":
|
||||
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,
|
||||
)
|
||||
print("SDK 3D flow viewer started. Close the Open3D window to stop.")
|
||||
else:
|
||||
pcd.colors = o3d.utility.Vector3dVector(cloud_colors)
|
||||
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(pcd)
|
||||
vis.update_geometry(line_set)
|
||||
vis.update_renderer()
|
||||
|
||||
frame_count += 1
|
||||
if frame_count % 30 == 0:
|
||||
elapsed = max(time.perf_counter() - t0, 1e-6)
|
||||
print(f"frames={frame_count}, viewer FPS={frame_count / elapsed:.1f}, arrows={len(flow_lines)}")
|
||||
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"))
|
||||
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)
|
||||
parser.add_argument("--min-flow", type=float, default=0.5)
|
||||
parser.add_argument("--normal-gain", type=float, default=1.0)
|
||||
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("--vector-mode", choices=["vraw", "vraw-magnitude", "shear-normal"], default="vraw")
|
||||
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())
|
||||
Reference in New Issue
Block a user