ykcheungab
188 lines
6.2 KiB
Python
188 lines
6.2 KiB
Python
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]) |