foodforthought-cli 0.2.7__py3-none-any.whl → 0.3.0__py3-none-any.whl

ate/urdf/lifting.py

@@ -0,0 +1,415 @@
+"""
+2D to 3D lifting utilities for point cloud generation.
+
+This module handles Phase 2c of the pipeline:
+1. Back-project masked pixels to 3D using depth and intrinsics
+2. Aggregate per-frame point clouds
+3. Save point clouds in PLY format
+
+Uses the standard pinhole camera model for projection.
+"""
+
+import logging
+from typing import Dict, List, Optional, Tuple
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+try:
+    import numpy as np
+    NUMPY_AVAILABLE = True
+except ImportError:
+    NUMPY_AVAILABLE = False
+    np = None
+
+
+class LiftingError(Exception):
+    """Error during 2D to 3D lifting."""
+    pass
+
+
+def lift_to_3d(
+    rgb: "np.ndarray",
+    depth: "np.ndarray",
+    mask: "np.ndarray",
+    intrinsics: Tuple[float, float, float, float],
+) -> Tuple["np.ndarray", "np.ndarray"]:
+    """
+    Back-project masked pixels to 3D points.
+
+    Uses the standard pinhole camera model:
+        X = (u - cx) * Z / fx
+        Y = (v - cy) * Z / fy
+        Z = depth
+
+    Args:
+        rgb: RGB image (H, W, 3)
+        depth: Depth map (H, W) in meters
+        mask: Binary mask (H, W)
+        intrinsics: (fx, fy, cx, cy) camera intrinsics
+
+    Returns:
+        Tuple of (points, colors):
+        - points: (N, 3) XYZ coordinates in meters
+        - colors: (N, 3) RGB values [0-255]
+    """
+    if not NUMPY_AVAILABLE:
+        raise LiftingError("NumPy not available")
+
+    fx, fy, cx, cy = intrinsics
+
+    # Get masked pixel coordinates
+    v, u = np.where(mask)  # v = row = y, u = col = x
+
+    if len(u) == 0:
+        return np.zeros((0, 3)), np.zeros((0, 3), dtype=np.uint8)
+
+    # Get depth values
+    z = depth[v, u]
+
+    # Filter out invalid depths
+    valid = (z > 0.01) & (z < 100.0)  # 1cm to 100m
+    u, v, z = u[valid], v[valid], z[valid]
+
+    if len(u) == 0:
+        return np.zeros((0, 3)), np.zeros((0, 3), dtype=np.uint8)
+
+    # Back-project to 3D
+    x = (u - cx) * z / fx
+    y = (v - cy) * z / fy
+
+    points = np.stack([x, y, z], axis=1)
+
+    # Get colors
+    colors = rgb[v, u]
+
+    return points.astype(np.float32), colors.astype(np.uint8)
+
+
+def save_ply(
+    path: Path,
+    points: "np.ndarray",
+    colors: Optional["np.ndarray"] = None,
+) -> None:
+    """
+    Save point cloud to PLY file.
+
+    Args:
+        path: Output file path
+        points: (N, 3) XYZ coordinates
+        colors: Optional (N, 3) RGB colors [0-255]
+    """
+    if not NUMPY_AVAILABLE:
+        raise LiftingError("NumPy not available")
+
+    n_points = len(points)
+
+    with open(path, 'w') as f:
+        # Header
+        f.write("ply\n")
+        f.write("format ascii 1.0\n")
+        f.write(f"element vertex {n_points}\n")
+        f.write("property float x\n")
+        f.write("property float y\n")
+        f.write("property float z\n")
+        if colors is not None:
+            f.write("property uchar red\n")
+            f.write("property uchar green\n")
+            f.write("property uchar blue\n")
+        f.write("end_header\n")
+
+        # Data
+        for i in range(n_points):
+            x, y, z = points[i]
+            if colors is not None:
+                r, g, b = colors[i]
+                f.write(f"{x:.6f} {y:.6f} {z:.6f} {int(r)} {int(g)} {int(b)}\n")
+            else:
+                f.write(f"{x:.6f} {y:.6f} {z:.6f}\n")
+
+    logger.debug(f"Saved {n_points} points to {path}")
+
+
+def load_ply(path: Path) -> Tuple["np.ndarray", Optional["np.ndarray"]]:
+    """
+    Load point cloud from PLY file.
+
+    Args:
+        path: Path to PLY file
+
+    Returns:
+        Tuple of (points, colors) or (points, None)
+    """
+    if not NUMPY_AVAILABLE:
+        raise LiftingError("NumPy not available")
+
+    points = []
+    colors = []
+    has_color = False
+    in_header = True
+    n_vertices = 0
+
+    with open(path, 'r') as f:
+        for line in f:
+            line = line.strip()
+
+            if in_header:
+                if line.startswith("element vertex"):
+                    n_vertices = int(line.split()[-1])
+                elif "property" in line and "red" in line:
+                    has_color = True
+                elif line == "end_header":
+                    in_header = False
+            else:
+                parts = line.split()
+                x, y, z = float(parts[0]), float(parts[1]), float(parts[2])
+                points.append([x, y, z])
+
+                if has_color and len(parts) >= 6:
+                    r, g, b = int(parts[3]), int(parts[4]), int(parts[5])
+                    colors.append([r, g, b])
+
+    points = np.array(points, dtype=np.float32)
+    colors = np.array(colors, dtype=np.uint8) if colors else None
+
+    return points, colors
+
+
+def aggregate_point_clouds(
+    clouds: List["np.ndarray"],
+    colors_list: Optional[List["np.ndarray"]] = None,
+    voxel_size: Optional[float] = None,
+) -> Tuple["np.ndarray", Optional["np.ndarray"]]:
+    """
+    Aggregate multiple point clouds into one.
+
+    Args:
+        clouds: List of (N, 3) point clouds
+        colors_list: Optional list of (N, 3) color arrays
+        voxel_size: Optional voxel size for downsampling
+
+    Returns:
+        Tuple of (aggregated_points, aggregated_colors)
+    """
+    if not NUMPY_AVAILABLE:
+        raise LiftingError("NumPy not available")
+
+    if not clouds:
+        return np.zeros((0, 3)), None
+
+    # Concatenate all clouds
+    points = np.vstack(clouds)
+
+    colors = None
+    if colors_list:
+        colors = np.vstack(colors_list)
+
+    # Optional voxel downsampling
+    if voxel_size and voxel_size > 0:
+        points, colors = voxel_downsample(points, colors, voxel_size)
+
+    return points, colors
+
+
+def voxel_downsample(
+    points: "np.ndarray",
+    colors: Optional["np.ndarray"],
+    voxel_size: float,
+) -> Tuple["np.ndarray", Optional["np.ndarray"]]:
+    """
+    Downsample point cloud using voxel grid.
+
+    Args:
+        points: (N, 3) point cloud
+        colors: Optional (N, 3) colors
+        voxel_size: Voxel edge length
+
+    Returns:
+        Downsampled (points, colors)
+    """
+    if len(points) == 0:
+        return points, colors
+
+    # Compute voxel indices
+    voxel_indices = np.floor(points / voxel_size).astype(np.int32)
+
+    # Hash voxels
+    unique_voxels, inverse_indices = np.unique(
+        voxel_indices, axis=0, return_inverse=True
+    )
+
+    # Average points in each voxel
+    n_voxels = len(unique_voxels)
+    new_points = np.zeros((n_voxels, 3))
+    new_colors = np.zeros((n_voxels, 3)) if colors is not None else None
+    counts = np.zeros(n_voxels)
+
+    for i, idx in enumerate(inverse_indices):
+        new_points[idx] += points[i]
+        if colors is not None:
+            new_colors[idx] += colors[i]
+        counts[idx] += 1
+
+    new_points /= counts[:, np.newaxis]
+    if new_colors is not None:
+        new_colors /= counts[:, np.newaxis]
+        new_colors = new_colors.astype(np.uint8)
+
+    logger.debug(f"Voxel downsampled: {len(points)} -> {len(new_points)} points")
+    return new_points.astype(np.float32), new_colors
+
+
+def generate_link_clouds(
+    session: "ScanSession",
+    masks: Dict[str, Dict[int, "np.ndarray"]],
+    depth_maps: Dict[int, "np.ndarray"],
+    intrinsics: Optional[Tuple[float, float, float, float]] = None,
+    frame_skip: int = 1,
+    voxel_size: float = 0.005,
+    progress_callback: Optional[callable] = None,
+) -> Dict[str, Path]:
+    """
+    Generate and save point clouds for each link.
+
+    Args:
+        session: ScanSession with video
+        masks: Segmentation masks
+        depth_maps: Depth maps from depth estimation
+        intrinsics: Camera intrinsics (fx, fy, cx, cy)
+        frame_skip: Use every Nth frame
+        voxel_size: Voxel size for downsampling (meters)
+        progress_callback: Optional progress callback
+
+    Returns:
+        Dict mapping link_name -> cloud_path
+    """
+    import cv2
+
+    # Get intrinsics
+    if intrinsics is None:
+        from .scale import estimate_intrinsics_from_resolution
+        w, h = session.metadata.resolution
+        intrinsics = estimate_intrinsics_from_resolution(w, h)
+        logger.info(f"Estimated intrinsics: fx={intrinsics[0]:.1f}, fy={intrinsics[1]:.1f}")
+
+    # Open video
+    cap = cv2.VideoCapture(str(session.video_path))
+    if not cap.isOpened():
+        raise LiftingError(f"Could not open video: {session.video_path}")
+
+    # Generate clouds for each link
+    link_clouds = {name: [] for name in masks.keys()}
+    link_colors = {name: [] for name in masks.keys()}
+
+    total_frames = len(depth_maps)
+    processed = 0
+
+    try:
+        for frame_idx in sorted(depth_maps.keys()):
+            if frame_idx % frame_skip != 0:
+                continue
+
+            # Read frame
+            cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
+            ret, frame = cap.read()
+            if not ret:
+                continue
+
+            rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            depth = depth_maps[frame_idx]
+
+            # Lift each link's masked region
+            for link_name, link_masks in masks.items():
+                if frame_idx not in link_masks:
+                    continue
+
+                mask = link_masks[frame_idx]
+                points, colors = lift_to_3d(rgb, depth, mask, intrinsics)
+
+                if len(points) > 0:
+                    link_clouds[link_name].append(points)
+                    link_colors[link_name].append(colors)
+
+            processed += 1
+            if progress_callback:
+                progress_callback(processed, total_frames)
+
+    finally:
+        cap.release()
+
+    # Aggregate and save each link's cloud
+    result = {}
+    clouds_dir = session.clouds_dir
+    clouds_dir.mkdir(parents=True, exist_ok=True)
+
+    for link_name in masks.keys():
+        if not link_clouds[link_name]:
+            logger.warning(f"No points generated for link '{link_name}'")
+            continue
+
+        # Aggregate
+        points, colors = aggregate_point_clouds(
+            link_clouds[link_name],
+            link_colors[link_name],
+            voxel_size=voxel_size,
+        )
+
+        # Save
+        cloud_path = clouds_dir / f"{link_name}.ply"
+        save_ply(cloud_path, points, colors)
+        result[link_name] = cloud_path
+
+        logger.info(f"Saved cloud for '{link_name}': {len(points)} points")
+
+    return result
+
+
+def run_lifting(
+    session: "ScanSession",
+    masks: Dict[str, Dict[int, "np.ndarray"]],
+    depth_maps: Dict[int, "np.ndarray"],
+    frame_skip: int = 1,
+    voxel_size: float = 0.005,
+    progress_callback: Optional[callable] = None,
+) -> Dict[str, Path]:
+    """
+    Run the lifting phase to generate point clouds.
+
+    Args:
+        session: ScanSession
+        masks: Segmentation masks
+        depth_maps: Depth maps
+        frame_skip: Process every Nth frame
+        voxel_size: Voxel grid size for downsampling
+        progress_callback: Optional progress callback
+
+    Returns:
+        Dict mapping link_name -> cloud_path
+    """
+    cloud_paths = generate_link_clouds(
+        session,
+        masks,
+        depth_maps,
+        frame_skip=frame_skip,
+        voxel_size=voxel_size,
+        progress_callback=progress_callback,
+    )
+
+    # Update session metadata
+    session.metadata.segment_complete = True
+    session.save_metadata()
+
+    logger.info(f"Lifting complete: {len(cloud_paths)} point clouds generated")
+    return cloud_paths
+
+
+__all__ = [
+    "LiftingError",
+    "lift_to_3d",
+    "save_ply",
+    "load_ply",
+    "aggregate_point_clouds",
+    "voxel_downsample",
+    "generate_link_clouds",
+    "run_lifting",
+]

ate/urdf/meshing.py

@@ -0,0 +1,300 @@
+"""
+Mesh generation from point clouds for URDF.
+
+This module handles Phase 4a of the pipeline:
+1. Aggregate per-frame point clouds into canonical link shapes
+2. Generate visual meshes via surface reconstruction
+3. Support multiple reconstruction methods (Poisson, SuGaR)
+
+Output: Visual mesh files (OBJ format) for each link.
+"""
+
+import logging
+from typing import Dict, List, Optional, Tuple
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+try:
+    import numpy as np
+    NUMPY_AVAILABLE = True
+except ImportError:
+    NUMPY_AVAILABLE = False
+    np = None
+
+try:
+    import trimesh
+    TRIMESH_AVAILABLE = True
+except ImportError:
+    TRIMESH_AVAILABLE = False
+    trimesh = None
+
+
+class MeshingError(Exception):
+    """Error during mesh generation."""
+    pass
+
+
+def poisson_reconstruct(
+    points: "np.ndarray",
+    colors: Optional["np.ndarray"] = None,
+    depth: int = 8,
+) -> "trimesh.Trimesh":
+    """
+    Reconstruct mesh using Poisson surface reconstruction.
+
+    Requires Open3D for the actual Poisson reconstruction.
+    Falls back to convex hull if Open3D is not available.
+
+    Args:
+        points: (N, 3) point cloud
+        colors: Optional (N, 3) RGB colors
+        depth: Octree depth for Poisson (higher = more detail)
+
+    Returns:
+        Trimesh mesh object
+    """
+    if not NUMPY_AVAILABLE:
+        raise MeshingError("NumPy not available")
+
+    if len(points) < 4:
+        raise MeshingError("Need at least 4 points for mesh reconstruction")
+
+    try:
+        import open3d as o3d
+
+        # Create Open3D point cloud
+        pcd = o3d.geometry.PointCloud()
+        pcd.points = o3d.utility.Vector3dVector(points)
+
+        if colors is not None:
+            pcd.colors = o3d.utility.Vector3dVector(colors / 255.0)
+
+        # Estimate normals
+        pcd.estimate_normals(
+            search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.05, max_nn=30)
+        )
+        pcd.orient_normals_consistent_tangent_plane(k=15)
+
+        # Poisson reconstruction
+        mesh, densities = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(
+            pcd, depth=depth
+        )
+
+        # Remove low-density vertices (artifacts)
+        densities = np.asarray(densities)
+        density_threshold = np.quantile(densities, 0.1)
+        vertices_to_remove = densities < density_threshold
+        mesh.remove_vertices_by_mask(vertices_to_remove)
+
+        # Convert to trimesh
+        vertices = np.asarray(mesh.vertices)
+        faces = np.asarray(mesh.triangles)
+
+        result = trimesh.Trimesh(vertices=vertices, faces=faces)
+        logger.debug(f"Poisson mesh: {len(result.vertices)} vertices, {len(result.faces)} faces")
+        return result
+
+    except ImportError:
+        logger.warning("Open3D not available, falling back to convex hull")
+        return convex_hull_mesh(points)
+
+
+def convex_hull_mesh(points: "np.ndarray") -> "trimesh.Trimesh":
+    """
+    Generate mesh using convex hull.
+
+    Simple fallback when more sophisticated methods aren't available.
+
+    Args:
+        points: (N, 3) point cloud
+
+    Returns:
+        Trimesh mesh object
+    """
+    if not TRIMESH_AVAILABLE:
+        raise MeshingError("trimesh not available. Run: pip install trimesh")
+
+    if len(points) < 4:
+        raise MeshingError("Need at least 4 points for convex hull")
+
+    from scipy.spatial import ConvexHull
+
+    hull = ConvexHull(points)
+    mesh = trimesh.Trimesh(vertices=points[hull.vertices], faces=hull.simplices)
+
+    logger.debug(f"Convex hull mesh: {len(mesh.vertices)} vertices")
+    return mesh
+
+
+def alpha_shape_mesh(
+    points: "np.ndarray",
+    alpha: float = 0.1,
+) -> "trimesh.Trimesh":
+    """
+    Generate mesh using alpha shapes.
+
+    Good for non-convex shapes when Poisson isn't available.
+
+    Args:
+        points: (N, 3) point cloud
+        alpha: Alpha value (smaller = tighter fit)
+
+    Returns:
+        Trimesh mesh object
+    """
+    if not TRIMESH_AVAILABLE:
+        raise MeshingError("trimesh not available")
+
+    # Use trimesh's built-in alpha shape
+    cloud = trimesh.PointCloud(points)
+    mesh = cloud.convex_hull  # Start with convex hull
+
+    # TODO: Implement proper alpha shape when scipy supports it
+    # For now, fall back to convex hull
+
+    return mesh
+
+
+def simplify_mesh(
+    mesh: "trimesh.Trimesh",
+    target_faces: int = 5000,
+) -> "trimesh.Trimesh":
+    """
+    Simplify mesh to reduce face count.
+
+    Args:
+        mesh: Input mesh
+        target_faces: Target number of faces
+
+    Returns:
+        Simplified mesh
+    """
+    if not TRIMESH_AVAILABLE:
+        raise MeshingError("trimesh not available")
+
+    if len(mesh.faces) <= target_faces:
+        return mesh
+
+    try:
+        # Try quadric decimation
+        simplified = mesh.simplify_quadric_decimation(target_faces)
+        logger.debug(f"Simplified mesh: {len(mesh.faces)} -> {len(simplified.faces)} faces")
+        return simplified
+    except Exception as e:
+        logger.warning(f"Mesh simplification failed: {e}")
+        return mesh
+
+
+def generate_visual_mesh(
+    cloud_path: Path,
+    output_path: Path,
+    method: str = "auto",
+    simplify_to: int = 5000,
+) -> Path:
+    """
+    Generate visual mesh from point cloud file.
+
+    Args:
+        cloud_path: Path to PLY point cloud
+        output_path: Path to output OBJ file
+        method: "poisson", "hull", or "auto"
+        simplify_to: Target face count for simplification
+
+    Returns:
+        Path to generated mesh file
+    """
+    from .lifting import load_ply
+
+    # Load point cloud
+    points, colors = load_ply(cloud_path)
+
+    if len(points) < 4:
+        raise MeshingError(f"Insufficient points in {cloud_path}: {len(points)}")
+
+    # Generate mesh
+    if method == "auto":
+        try:
+            mesh = poisson_reconstruct(points, colors)
+        except Exception as e:
+            logger.warning(f"Poisson failed, using convex hull: {e}")
+            mesh = convex_hull_mesh(points)
+    elif method == "poisson":
+        mesh = poisson_reconstruct(points, colors)
+    elif method == "hull":
+        mesh = convex_hull_mesh(points)
+    else:
+        raise MeshingError(f"Unknown meshing method: {method}")
+
+    # Simplify if needed
+    if simplify_to > 0 and len(mesh.faces) > simplify_to:
+        mesh = simplify_mesh(mesh, simplify_to)
+
+    # Ensure mesh is watertight
+    if not mesh.is_watertight:
+        logger.warning(f"Mesh is not watertight, attempting repair")
+        mesh.fill_holes()
+
+    # Export
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    mesh.export(str(output_path))
+
+    logger.info(f"Generated visual mesh: {output_path} ({len(mesh.faces)} faces)")
+    return output_path
+
+
+def generate_all_visual_meshes(
+    session: "ScanSession",
+    method: str = "auto",
+    simplify_to: int = 5000,
+    progress_callback: Optional[callable] = None,
+) -> Dict[str, Path]:
+    """
+    Generate visual meshes for all links in a session.
+
+    Args:
+        session: ScanSession with point clouds
+        method: Meshing method
+        simplify_to: Target face count
+        progress_callback: Optional progress callback
+
+    Returns:
+        Dict mapping link_name -> mesh_path
+    """
+    clouds_dir = session.clouds_dir
+    meshes_dir = session.meshes_dir
+    meshes_dir.mkdir(parents=True, exist_ok=True)
+
+    cloud_files = list(clouds_dir.glob("*.ply"))
+    if not cloud_files:
+        raise MeshingError(f"No point clouds found in {clouds_dir}")
+
+    result = {}
+    total = len(cloud_files)
+
+    for i, cloud_path in enumerate(cloud_files):
+        link_name = cloud_path.stem
+        output_path = meshes_dir / f"{link_name}_visual.obj"
+
+        try:
+            generate_visual_mesh(cloud_path, output_path, method, simplify_to)
+            result[link_name] = output_path
+        except Exception as e:
+            logger.error(f"Failed to generate mesh for {link_name}: {e}")
+
+        if progress_callback:
+            progress_callback(i + 1, total)
+
+    logger.info(f"Generated {len(result)} visual meshes")
+    return result
+
+
+__all__ = [
+    "MeshingError",
+    "poisson_reconstruct",
+    "convex_hull_mesh",
+    "alpha_shape_mesh",
+    "simplify_mesh",
+    "generate_visual_mesh",
+    "generate_all_visual_meshes",
+]