ykcheungab
411 lines
14 KiB
Python
411 lines
14 KiB
Python
import argparse
|
||
import os
|
||
import time
|
||
|
||
import cv2
|
||
import numpy as np
|
||
|
||
|
||
def load_exported_data(data_dir):
|
||
"""Load Blender-exported mesh/camera files."""
|
||
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"Real-time video viewing requires precomputed pixel_to_xyz.npz: {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 prepare_static_points(data, display_stride=4):
|
||
"""
|
||
Prepare fixed 3D points from the precomputed pixel-to-3D map.
|
||
|
||
The geometry does not change between video frames, so real-time viewing only
|
||
needs to update point colors.
|
||
"""
|
||
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_stride = max(1, int(display_stride))
|
||
display_indices = np.arange(len(valid_points))[::display_stride]
|
||
return valid_mask, display_indices, valid_points[display_indices]
|
||
|
||
|
||
def frame_to_rgb_colors(frame_bgr, data, valid_mask, display_indices, flip_vertical=True):
|
||
"""Convert a video frame to RGB point colors aligned with the exported map."""
|
||
target_w = int(data["width"])
|
||
target_h = int(data["height"])
|
||
h, w = frame_bgr.shape[:2]
|
||
if (w, h) != (target_w, target_h):
|
||
frame_bgr = cv2.resize(frame_bgr, (target_w, target_h), interpolation=cv2.INTER_LINEAR)
|
||
|
||
frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
|
||
if flip_vertical:
|
||
frame_rgb = np.flipud(frame_rgb).copy()
|
||
|
||
colors = frame_rgb.astype(np.float64) / 255.0
|
||
valid_colors = colors[valid_mask].reshape(-1, 3)
|
||
return valid_colors[display_indices]
|
||
|
||
|
||
def init_open3d_viewer(points, colors):
|
||
import open3d as o3d
|
||
|
||
vis = o3d.visualization.Visualizer()
|
||
vis.create_window(window_name="Real-time 3D Point Cloud", width=1280, height=720)
|
||
|
||
pcd = o3d.geometry.PointCloud()
|
||
pcd.points = o3d.utility.Vector3dVector(points.astype(np.float64))
|
||
pcd.colors = o3d.utility.Vector3dVector(colors.astype(np.float64))
|
||
vis.add_geometry(pcd)
|
||
|
||
render_option = vis.get_render_option()
|
||
render_option.point_size = 2.0
|
||
render_option.background_color = np.array([0.0, 0.0, 0.0])
|
||
|
||
vis.poll_events()
|
||
vis.update_renderer()
|
||
return vis, pcd, o3d
|
||
|
||
|
||
def save_colored_pointcloud(points, colors, output_path):
|
||
"""Save displayed points as a PLY file."""
|
||
valid_colors = (colors * 255).clip(0, 255).astype(np.uint8)
|
||
with open(output_path, "w", encoding="utf-8") as f:
|
||
f.write("ply\n")
|
||
f.write("format ascii 1.0\n")
|
||
f.write(f"element vertex {len(points)}\n")
|
||
f.write("property float x\n")
|
||
f.write("property float y\n")
|
||
f.write("property float z\n")
|
||
f.write("property uchar red\n")
|
||
f.write("property uchar green\n")
|
||
f.write("property uchar blue\n")
|
||
f.write("end_header\n")
|
||
for pt, col in zip(points, valid_colors):
|
||
f.write(f"{pt[0]} {pt[1]} {pt[2]} {col[0]} {col[1]} {col[2]}\n")
|
||
print(f"Saved point cloud to {output_path}")
|
||
|
||
|
||
def view_video_pointcloud(
|
||
video_path,
|
||
data,
|
||
display_stride=4,
|
||
frame_step=1,
|
||
flip_vertical=True,
|
||
loop=False,
|
||
save_snapshot_path=None,
|
||
):
|
||
"""Open a video and show it as a real-time colored 3D point cloud."""
|
||
cap = cv2.VideoCapture(video_path)
|
||
if not cap.isOpened():
|
||
raise FileNotFoundError(f"Could not open video: {video_path}")
|
||
|
||
valid_mask, display_indices, display_points = prepare_static_points(
|
||
data,
|
||
display_stride=display_stride,
|
||
)
|
||
|
||
ok, frame = cap.read()
|
||
if not ok:
|
||
cap.release()
|
||
raise RuntimeError(f"Could not read first frame from video: {video_path}")
|
||
|
||
display_colors = frame_to_rgb_colors(
|
||
frame,
|
||
data,
|
||
valid_mask,
|
||
display_indices,
|
||
flip_vertical=flip_vertical,
|
||
)
|
||
vis, pcd, o3d = init_open3d_viewer(display_points, display_colors)
|
||
|
||
print("Real-time point cloud viewer started.")
|
||
print("Close the Open3D window or press Ctrl+C in the terminal to stop.")
|
||
print(f"Displayed points: {len(display_points)} (stride={display_stride})")
|
||
|
||
frame_idx = 0
|
||
shown_frames = 0
|
||
t0 = time.perf_counter()
|
||
|
||
try:
|
||
while True:
|
||
alive = vis.poll_events()
|
||
if not alive:
|
||
break
|
||
|
||
ok, frame = cap.read()
|
||
if not ok:
|
||
if loop:
|
||
cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
|
||
continue
|
||
break
|
||
|
||
frame_idx += 1
|
||
if frame_idx % max(1, frame_step) != 0:
|
||
continue
|
||
|
||
display_colors = frame_to_rgb_colors(
|
||
frame,
|
||
data,
|
||
valid_mask,
|
||
display_indices,
|
||
flip_vertical=flip_vertical,
|
||
)
|
||
pcd.colors = o3d.utility.Vector3dVector(display_colors)
|
||
vis.update_geometry(pcd)
|
||
vis.update_renderer()
|
||
|
||
shown_frames += 1
|
||
if shown_frames % 30 == 0:
|
||
elapsed = max(time.perf_counter() - t0, 1e-6)
|
||
print(f"Displayed {shown_frames} frames, viewer FPS={shown_frames / elapsed:.1f}")
|
||
|
||
finally:
|
||
if save_snapshot_path:
|
||
save_colored_pointcloud(display_points, display_colors, save_snapshot_path)
|
||
cap.release()
|
||
vis.destroy_window()
|
||
|
||
|
||
def parse_args():
|
||
parser = argparse.ArgumentParser(description="View a video as a real-time 3D point cloud.")
|
||
parser.add_argument("--data-dir", default="./exported_data", help="Blender exported_data folder")
|
||
parser.add_argument("--video", default="20260602-094916.mp4", help="Input video path")
|
||
parser.add_argument("--display-stride", type=int, default=4,
|
||
help="Show every Nth valid point. Increase for faster viewing.")
|
||
parser.add_argument("--frame-step", type=int, default=1,
|
||
help="Process every Nth video frame.")
|
||
parser.add_argument("--no-flip-vertical", action="store_true",
|
||
help="Disable vertical flip before color mapping.")
|
||
parser.add_argument("--loop", action="store_true", help="Loop video playback.")
|
||
parser.add_argument("--save-snapshot", default="",
|
||
help="Optional PLY path for saving the last displayed frame.")
|
||
return parser.parse_args()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
args = parse_args()
|
||
exported = load_exported_data(args.data_dir)
|
||
view_video_pointcloud(
|
||
args.video,
|
||
exported,
|
||
display_stride=args.display_stride,
|
||
frame_step=args.frame_step,
|
||
flip_vertical=not args.no_flip_vertical,
|
||
loop=args.loop,
|
||
save_snapshot_path=args.save_snapshot or None,
|
||
)
|
||
import numpy as np
|
||
import cv2
|
||
import os
|
||
|
||
def load_exported_data(data_dir):
|
||
"""Load all exported files."""
|
||
# Mesh
|
||
mesh = np.load(os.path.join(data_dir, "mesh.npz"))
|
||
vertices = mesh['vertices']
|
||
faces = mesh['faces']
|
||
uv_data = mesh['uv_data']
|
||
loop_to_vertex = mesh['loop_to_vertex']
|
||
|
||
# Camera intrinsics
|
||
cam_data = np.load(os.path.join(data_dir, "camera_intrinsics.npz"))
|
||
K = cam_data['K']
|
||
width = int(cam_data['width'])
|
||
height = int(cam_data['height'])
|
||
|
||
# Camera pose (world-to-camera)
|
||
world_to_cam = np.load(os.path.join(data_dir, "world_to_cam.npy"))
|
||
|
||
# Precomputed pixel-to-3D map (optional)
|
||
xyz_map_path = os.path.join(data_dir, "pixel_to_xyz.npz")
|
||
if os.path.exists(xyz_map_path):
|
||
xyz_map = np.load(xyz_map_path)['xyz_map']
|
||
else:
|
||
xyz_map = None
|
||
|
||
return {
|
||
'vertices': vertices,
|
||
'faces': faces,
|
||
'uv_data': uv_data,
|
||
'loop_to_vertex': loop_to_vertex,
|
||
'K': K,
|
||
'width': width,
|
||
'height': height,
|
||
'world_to_cam': world_to_cam,
|
||
'xyz_map': xyz_map
|
||
}
|
||
|
||
def ray_intersect_triangle(ray_origin, ray_dir, v0, v1, v2):
|
||
"""Möller–Trumbore intersection algorithm. Returns t (distance) or None."""
|
||
epsilon = 1e-8
|
||
edge1 = v1 - v0
|
||
edge2 = v2 - v0
|
||
h = np.cross(ray_dir, edge2)
|
||
a = np.dot(edge1, h)
|
||
if abs(a) < epsilon:
|
||
return None
|
||
f = 1.0 / a
|
||
s = ray_origin - v0
|
||
u = f * np.dot(s, h)
|
||
if u < 0.0 or u > 1.0:
|
||
return None
|
||
q = np.cross(s, edge1)
|
||
v = f * np.dot(ray_dir, q)
|
||
if v < 0.0 or u + v > 1.0:
|
||
return None
|
||
t = f * np.dot(edge2, q)
|
||
if t > epsilon:
|
||
return t
|
||
return None
|
||
|
||
def pixel_to_3d_ray(u, v, K, world_to_cam):
|
||
"""
|
||
Given pixel coordinates (u,v) and camera intrinsics K,
|
||
return a ray in world coordinates: origin and direction.
|
||
"""
|
||
# Invert K to get normalized image coordinates
|
||
fx = K[0,0]; fy = K[1,1]; cx = K[0,2]; cy = K[1,2]
|
||
x_norm = (u - cx) / fx
|
||
y_norm = (v - cy) / fy
|
||
|
||
# Ray in camera space: (x_norm, y_norm, 1.0)
|
||
ray_cam = np.array([x_norm, y_norm, 1.0], dtype=np.float64)
|
||
|
||
# Transform to world
|
||
cam_to_world = np.linalg.inv(world_to_cam)
|
||
R = cam_to_world[:3, :3]
|
||
t = cam_to_world[:3, 3]
|
||
ray_world_dir = R @ ray_cam
|
||
ray_world_dir /= np.linalg.norm(ray_world_dir)
|
||
ray_world_origin = t
|
||
|
||
return ray_world_origin, ray_world_dir
|
||
|
||
def map_image_to_3d(image_path, data, use_precomputed=True, flip_vertical=True):
|
||
"""
|
||
Map each pixel of the image to a 3D world point.
|
||
If use_precomputed and xyz_map exists, return that directly.
|
||
If `flip_vertical` is True, flip the loaded image top-to-bottom so its
|
||
row order matches the exported Blender camera/pixel mapping.
|
||
Otherwise, perform raycasting on the mesh.
|
||
Returns:
|
||
points: (H, W, 3) array of 3D points (NaN for misses).
|
||
colors: (H, W, 3) array of RGB colors from image.
|
||
"""
|
||
img = cv2.imread(image_path)
|
||
if img is None:
|
||
raise FileNotFoundError(f"Could not read image: {image_path}")
|
||
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # to RGB
|
||
if flip_vertical:
|
||
img = np.flipud(img).copy()
|
||
h, w = img.shape[:2]
|
||
|
||
if use_precomputed and data['xyz_map'] is not None:
|
||
# Use precomputed map (must match resolution)
|
||
if h != data['height'] or w != data['width']:
|
||
raise ValueError("Image resolution does not match precomputed map.")
|
||
points = data['xyz_map'].copy()
|
||
colors = img.astype(np.float32) / 255.0
|
||
return points, colors
|
||
|
||
# Otherwise, perform raycasting on the mesh
|
||
vertices = data['vertices']
|
||
faces = data['faces']
|
||
K = data['K']
|
||
world_to_cam = data['world_to_cam']
|
||
|
||
points = np.full((h, w, 3), np.nan, dtype=np.float32)
|
||
colors = img.astype(np.float32) / 255.0
|
||
|
||
# For speed, we can use a bounding volume hierarchy (BVH).
|
||
# Here we implement a simple brute-force over all faces (slow for large meshes).
|
||
# In practice, precomputing the xyz_map is strongly recommended.
|
||
|
||
for y in range(h):
|
||
for x in range(w):
|
||
origin, direction = pixel_to_3d_ray(x, y, K, world_to_cam)
|
||
best_t = np.inf
|
||
best_hit = None
|
||
for face in faces:
|
||
v0 = vertices[face[0]]
|
||
v1 = vertices[face[1]]
|
||
v2 = vertices[face[2]]
|
||
t = ray_intersect_triangle(origin, direction, v0, v1, v2)
|
||
if t is not None and t < best_t:
|
||
best_t = t
|
||
best_hit = origin + direction * t
|
||
if best_hit is not None:
|
||
points[y, x] = best_hit
|
||
|
||
return points, colors
|
||
|
||
def save_colored_pointcloud(points, colors, output_path):
|
||
"""Save as PLY file (colored point cloud)."""
|
||
# Flatten valid points
|
||
mask = ~np.isnan(points[:,:,0])
|
||
valid_points = points[mask].reshape(-1, 3)
|
||
valid_colors = (colors[mask].reshape(-1, 3) * 255).astype(np.uint8)
|
||
|
||
with open(output_path, 'w') as f:
|
||
f.write("ply\n")
|
||
f.write("format ascii 1.0\n")
|
||
f.write(f"element vertex {len(valid_points)}\n")
|
||
f.write("property float x\n")
|
||
f.write("property float y\n")
|
||
f.write("property float z\n")
|
||
f.write("property uchar red\n")
|
||
f.write("property uchar green\n")
|
||
f.write("property uchar blue\n")
|
||
f.write("end_header\n")
|
||
for pt, col in zip(valid_points, valid_colors):
|
||
f.write(f"{pt[0]} {pt[1]} {pt[2]} {col[0]} {col[1]} {col[2]}\n")
|
||
print(f"Saved point cloud to {output_path}")
|
||
|
||
# if __name__ == "__main__":
|
||
# data_dir = "./exported_data" # Path to exported files
|
||
# image_path = "captured_image_1.jpg" # Input camera image
|
||
# output_ply = "mapped_points.ply"
|
||
|
||
# data = load_exported_data(data_dir)
|
||
# points, colors = map_image_to_3d(
|
||
# image_path,
|
||
# data,
|
||
# use_precomputed=True,
|
||
# flip_vertical=True,
|
||
# )
|
||
# save_colored_pointcloud(points, colors, output_ply)
|
||
|
||
# # Optional: Visualize using Open3D
|
||
# import open3d as o3d
|
||
# pcd = o3d.geometry.PointCloud()
|
||
# mask = ~np.isnan(points[:,:,0])
|
||
# pcd.points = o3d.utility.Vector3dVector(points[mask].reshape(-1,3))
|
||
# pcd.colors = o3d.utility.Vector3dVector(colors[mask].reshape(-1,3))
|
||
# o3d.visualization.draw_geometries([pcd]) |