src/interiorfusion/utils/mesh_utils.py

"""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,
    }