Upload src/interiorfusion/utils/mesh_utils.py

src/interiorfusion/utils/mesh_utils.py

@@ -0,0 +1,180 @@
+"""Mesh utilities for export and manipulation."""
+
+import os
+from pathlib import Path
+from typing import Dict, List, Optional, Union
+
+import numpy as np
+
+
+def export_mesh(
+    mesh,
+    output_path: Union[str, Path],
+    format: str = "glb",
+    materials: Optional[List[Dict]] = None,
+) -> str:
+    """
+    Export mesh to various 3D formats.
+    
+    Supported formats: glb, fbx, obj, usdz, ply
+    """
+    output_path = Path(output_path)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    
+    try:
+        import trimesh
+    except ImportError:
+        raise ImportError("trimesh is required for mesh export")
+    
+    format = format.lower()
+    
+    if format == "glb" or format == "gltf":
+        # Export as GLB (GL Transmission Format)
+        mesh.export(str(output_path))
+        
+    elif format == "fbx":
+        # Export as FBX
+        # trimesh can export some formats; for FBX we may need assimp
+        mesh.export(str(output_path))
+        
+    elif format == "obj":
+        # Export as OBJ with MTL
+        mesh.export(str(output_path))
+        
+    elif format == "usdz":
+        # USDZ requires USD libraries
+        # For now, export as GLB and note conversion needed
+        glb_path = output_path.with_suffix(".glb")
+        mesh.export(str(glb_path))
+        print(f"USDZ export: converted to GLB at {glb_path}. Use Apple's usdzconvert.")
+        
+    elif format == "ply":
+        # Export as PLY (Stanford Polygon Format)
+        mesh.export(str(output_path))
+        
+    else:
+        raise ValueError(f"Unsupported format: {format}")
+    
+    return str(output_path)
+
+
+def export_gaussian_splatting(
+    gaussian_cloud: np.ndarray,
+    output_path: Union[str, Path],
+) -> str:
+    """
+    Export Gaussian Splatting representation to PLY format.
+    
+    Args:
+        gaussian_cloud: [N, 14] array with columns:
+            [x, y, z, scale_x, scale_y, scale_z,
+             rot_qx, rot_qy, rot_qz, rot_qw,
+             r, g, b, opacity]
+    """
+    output_path = Path(output_path)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    
+    try:
+        from plyfile import PlyData, PlyElement
+    except ImportError:
+        # Fallback: write simple ASCII PLY
+        _write_ascii_ply(gaussian_cloud, output_path)
+        return str(output_path)
+    
+    # Write binary PLY with Gaussian attributes
+    num_points = len(gaussian_cloud)
+    
+    # Define dtype for Gaussian splatting format
+    dtype = [
+        ('x', 'f4'), ('y', 'f4'), ('z', 'f4'),
+        ('scale_0', 'f4'), ('scale_1', 'f4'), ('scale_2', 'f4'),
+        ('rot_0', 'f4'), ('rot_1', 'f4'), ('rot_2', 'f4'), ('rot_3', 'f4'),
+        ('f_dc_0', 'f4'), ('f_dc_1', 'f4'), ('f_dc_2', 'f4'),
+        ('opacity', 'f4'),
+    ]
+    
+    # Convert RGB to spherical harmonics DC term (simplified)
+    # In production, would also include SH bands 1, 2, 3
+    
+    vertices = np.zeros(num_points, dtype=dtype)
+    vertices['x'] = gaussian_cloud[:, 0]
+    vertices['y'] = gaussian_cloud[:, 1]
+    vertices['z'] = gaussian_cloud[:, 2]
+    vertices['scale_0'] = gaussian_cloud[:, 3]
+    vertices['scale_1'] = gaussian_cloud[:, 4]
+    vertices['scale_2'] = gaussian_cloud[:, 5]
+    vertices['rot_0'] = gaussian_cloud[:, 6]
+    vertices['rot_1'] = gaussian_cloud[:, 7]
+    vertices['rot_2'] = gaussian_cloud[:, 8]
+    vertices['rot_3'] = gaussian_cloud[:, 9]
+    vertices['f_dc_0'] = gaussian_cloud[:, 10]
+    vertices['f_dc_1'] = gaussian_cloud[:, 11]
+    vertices['f_dc_2'] = gaussian_cloud[:, 12]
+    vertices['opacity'] = gaussian_cloud[:, 13]
+    
+    el = PlyElement.describe(vertices, 'vertex')
+    PlyData([el], text=True).write(str(output_path))
+    
+    return str(output_path)
+
+
+def _write_ascii_ply(gaussian_cloud: np.ndarray, output_path: Path):
+    """Write simple ASCII PLY fallback."""
+    num_points = len(gaussian_cloud)
+    
+    with open(output_path, 'w') as f:
+        f.write("ply\n")
+        f.write("format ascii 1.0\n")
+        f.write(f"element vertex {num_points}\n")
+        f.write("property float x\n")
+        f.write("property float y\n")
+        f.write("property float z\n")
+        f.write("property float scale_0\n")
+        f.write("property float scale_1\n")
+        f.write("property float scale_2\n")
+        f.write("property float rot_0\n")
+        f.write("property float rot_1\n")
+        f.write("property float rot_2\n")
+        f.write("property float rot_3\n")
+        f.write("property float f_dc_0\n")
+        f.write("property float f_dc_1\n")
+        f.write("property float f_dc_2\n")
+        f.write("property float opacity\n")
+        f.write("end_header\n")
+        
+        for point in gaussian_cloud:
+            f.write(" ".join(f"{v:.6f}" for v in point) + "\n")
+
+
+def decimate_mesh(mesh, target_faces: int = 10000):
+    """Reduce mesh complexity for mobile/VR."""
+    try:
+        import trimesh
+        if hasattr(mesh, 'simplify'):
+            return mesh.simplify(target_faces)
+    except Exception:
+        pass
+    return mesh
+
+
+def compute_uv_atlas(mesh):
+    """Compute UV atlas for mesh texture mapping."""
+    try:
+        import xatlas
+        # Placeholder: would use xatlas for UV unwrapping
+        return mesh
+    except ImportError:
+        return mesh
+
+
+def merge_materials(materials: List[Dict]) -> Dict:
+    """Merge multiple PBR materials into a single material atlas."""
+    # In production: create texture atlas
+    # For now, return first material
+    if materials:
+        return materials[0]
+    return {
+        "albedo": [0.7, 0.7, 0.7],
+        "metallic": 0.0,
+        "roughness": 0.5,
+    }