Create utils.py

utils.py

@@ -0,0 +1,265 @@
+"""
+Utility functions for point cloud processing and visualization
+"""
+
+import numpy as np
+import open3d as o3d
+from pathlib import Path
+import cv2
+
+def visualize_point_cloud(ply_file):
+    """
+    Visualize a single point cloud using Open3D
+    
+    Usage:
+        visualize_point_cloud("pointcloud_000000.ply")
+    """
+    pcd = o3d.io.read_point_cloud(ply_file)
+    
+    # Statistical outlier removal
+    pcd, _ = pcd.remove_statistical_outlier(nb_neighbors=20, std_ratio=2.0)
+    
+    # Estimate normals
+    pcd.estimate_normals(
+        search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.1, max_nn=30)
+    )
+    
+    o3d.visualization.draw_geometries(
+        [pcd],
+        window_name="Point Cloud Viewer",
+        width=1024,
+        height=768,
+        point_show_normal=False
+    )
+
+def merge_point_clouds(ply_files, output_file="merged_pointcloud.ply", subsample_rate=0.3):
+    """
+    Merge multiple point clouds into one
+    
+    Args:
+        ply_files: List of PLY file paths
+        output_file: Output merged PLY file
+        subsample_rate: Keep only this fraction of points (0.3 = 30%)
+    
+    Usage:
+        merge_point_clouds(["pc1.ply", "pc2.ply"], "merged.ply")
+    """
+    merged_pcd = o3d.geometry.PointCloud()
+    
+    for ply_file in ply_files:
+        pcd = o3d.io.read_point_cloud(str(ply_file))
+        
+        # Subsample to reduce size
+        if subsample_rate < 1.0:
+            pcd = pcd.random_down_sample(subsample_rate)
+        
+        merged_pcd += pcd
+    
+    # Remove duplicates and outliers
+    merged_pcd = merged_pcd.remove_duplicated_points()
+    merged_pcd, _ = merged_pcd.remove_statistical_outlier(nb_neighbors=20, std_ratio=2.0)
+    
+    # Save
+    o3d.io.write_point_cloud(output_file, merged_pcd)
+    print(f"Merged point cloud saved: {output_file}")
+    print(f"Total points: {len(merged_pcd.points)}")
+    
+    return output_file
+
+def create_mesh_from_pointcloud(ply_file, output_file="mesh.ply", depth=9):
+    """
+    Create mesh from point cloud using Poisson surface reconstruction
+    
+    Args:
+        ply_file: Input PLY file
+        output_file: Output mesh file
+        depth: Poisson depth (higher = more detail, slower)
+    
+    Usage:
+        create_mesh_from_pointcloud("pointcloud.ply", "mesh.ply")
+    """
+    pcd = o3d.io.read_point_cloud(ply_file)
+    
+    # Estimate normals if not present
+    if not pcd.has_normals():
+        pcd.estimate_normals(
+            search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.1, max_nn=30)
+        )
+        pcd.orient_normals_consistent_tangent_plane(k=15)
+    
+    # Poisson reconstruction
+    print("Running Poisson reconstruction...")
+    mesh, densities = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(
+        pcd, depth=depth
+    )
+    
+    # Remove low density vertices
+    vertices_to_remove = densities < np.quantile(densities, 0.01)
+    mesh.remove_vertices_by_mask(vertices_to_remove)
+    
+    # Save
+    o3d.io.write_triangle_mesh(output_file, mesh)
+    print(f"Mesh saved: {output_file}")
+    print(f"Vertices: {len(mesh.vertices)}, Triangles: {len(mesh.triangles)}")
+    
+    return output_file
+
+def animate_reconstruction(ply_directory, output_video="reconstruction.mp4", fps=10):
+    """
+    Create a video animating through the point clouds
+    
+    Args:
+        ply_directory: Directory containing PLY files
+        output_video: Output video file
+        fps: Frames per second
+    
+    Usage:
+        animate_reconstruction("./point_clouds", "animation.mp4")
+    """
+    ply_files = sorted(Path(ply_directory).glob("*.ply"))
+    
+    if not ply_files:
+        print("No PLY files found!")
+        return
+    
+    # Setup visualizer
+    vis = o3d.visualization.Visualizer()
+    vis.create_window(visible=False, width=1920, height=1080)
+    
+    # Get first point cloud for camera setup
+    pcd = o3d.io.read_point_cloud(str(ply_files[0]))
+    vis.add_geometry(pcd)
+    
+    # Setup camera
+    ctr = vis.get_view_control()
+    ctr.set_zoom(0.8)
+    
+    # Render options
+    opt = vis.get_render_option()
+    opt.point_size = 2.0
+    opt.background_color = np.asarray([0, 0, 0])
+    
+    # Capture frames
+    frames = []
+    for i, ply_file in enumerate(ply_files):
+        print(f"Rendering frame {i+1}/{len(ply_files)}")
+        
+        # Update point cloud
+        pcd = o3d.io.read_point_cloud(str(ply_file))
+        vis.clear_geometries()
+        vis.add_geometry(pcd)
+        
+        # Rotate camera slightly
+        ctr.rotate(10.0, 0.0)
+        
+        # Capture
+        vis.poll_events()
+        vis.update_renderer()
+        img = np.asarray(vis.capture_screen_float_buffer(do_render=True))
+        img = (img * 255).astype(np.uint8)
+        frames.append(cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
+    
+    vis.destroy_window()
+    
+    # Write video
+    height, width = frames[0].shape[:2]
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    out = cv2.VideoWriter(output_video, fourcc, fps, (width, height))
+    
+    for frame in frames:
+        out.write(frame)
+    
+    out.release()
+    print(f"Video saved: {output_video}")
+
+def depth_map_to_pointcloud(rgb_image_path, depth_npy_path, output_ply, subsample=0.5):
+    """
+    Convert single RGB image and depth map to point cloud
+    
+    Args:
+        rgb_image_path: Path to RGB image
+        depth_npy_path: Path to depth numpy array
+        output_ply: Output PLY file
+        subsample: Subsampling rate (0.5 = keep 50% of points)
+    
+    Usage:
+        depth_map_to_pointcloud("frame.jpg", "depth.npy", "output.ply")
+    """
+    # Load data
+    rgb = cv2.imread(rgb_image_path)
+    rgb = cv2.cvtColor(rgb, cv2.COLOR_BGR2RGB)
+    depth = np.load(depth_npy_path)
+    
+    # Normalize depth to reasonable scale
+    depth = depth.astype(np.float32)
+    depth = (depth - depth.min()) / (depth.max() - depth.min())
+    depth = depth * 10.0  # Scale to 0-10 units
+    
+    h, w = depth.shape
+    
+    # Camera intrinsics (adjust for your camera)
+    fx = fy = w * 0.7
+    cx, cy = w / 2, h / 2
+    
+    # Create 3D points
+    u, v = np.meshgrid(np.arange(w), np.arange(h))
+    z = depth
+    x = (u - cx) * z / fx
+    y = (v - cy) * z / fy
+    
+    points = np.stack([x, y, z], axis=-1).reshape(-1, 3)
+    colors = rgb.reshape(-1, 3) / 255.0
+    
+    # Subsample
+    if subsample < 1.0:
+        n_points = int(len(points) * subsample)
+        indices = np.random.choice(len(points), n_points, replace=False)
+        points = points[indices]
+        colors = colors[indices]
+    
+    # Create point cloud
+    pcd = o3d.geometry.PointCloud()
+    pcd.points = o3d.utility.Vector3dVector(points)
+    pcd.colors = o3d.utility.Vector3dVector(colors)
+    
+    # Clean up
+    pcd, _ = pcd.remove_statistical_outlier(nb_neighbors=20, std_ratio=2.0)
+    
+    # Save
+    o3d.io.write_point_cloud(output_ply, pcd)
+    print(f"Point cloud saved: {output_ply} ({len(pcd.points)} points)")
+    
+    return output_ply
+
+
+if __name__ == "__main__":
+    import sys
+    
+    if len(sys.argv) > 1:
+        command = sys.argv[1]
+        
+        if command == "visualize" and len(sys.argv) > 2:
+            visualize_point_cloud(sys.argv[2])
+        
+        elif command == "merge" and len(sys.argv) > 3:
+            ply_files = sys.argv[2:-1]
+            output = sys.argv[-1]
+            merge_point_clouds(ply_files, output)
+        
+        elif command == "mesh" and len(sys.argv) > 2:
+            output = sys.argv[3] if len(sys.argv) > 3 else "mesh.ply"
+            create_mesh_from_pointcloud(sys.argv[2], output)
+        
+        elif command == "animate" and len(sys.argv) > 2:
+            output = sys.argv[3] if len(sys.argv) > 3 else "reconstruction.mp4"
+            animate_reconstruction(sys.argv[2], output)
+        
+        else:
+            print("Usage:")
+            print("  python utils.py visualize <ply_file>")
+            print("  python utils.py merge <ply1> <ply2> ... <output>")
+            print("  python utils.py mesh <input_ply> [output_mesh]")
+            print("  python utils.py animate <ply_directory> [output_video]")
+    else:
+        print("Utility functions loaded. Import in Python:")
+        print("  from utils import visualize_point_cloud, merge_point_clouds")