src/interiorfusion/pipelines.py

"""InteriorFusion main inference pipeline."""

import os
import tempfile
from dataclasses import dataclass, field
from pathlib import Path
from typing import List, Optional, Tuple, Union

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image

from .models.scene_understanding import SceneUnderstandingModule
from .models.multiview_generation import MultiViewGenerationModule
from .models.reconstruction_3d import Reconstruction3DModule
from .models.scene_assembly import SceneAssemblyModule
from .models.material_texture import MaterialTextureModule
from .utils.mesh_utils import export_mesh
from .utils.gaussian_utils import export_gaussian_splatting


@dataclass
class InteriorFusionOutput:
    """Output container for InteriorFusion pipeline."""
    
    # 3D representations
    scene_mesh: Optional["trimesh.Trimesh"] = None  # type: ignore
    room_shell_mesh: Optional["trimesh.Trimesh"] = None  # type: ignore
    object_meshes: List["trimesh.Trimesh"] = field(default_factory=list)  # type: ignore
    gaussian_cloud: Optional[torch.Tensor] = None  # Scene Gaussians
    
    # Materials
    pbr_materials: List[dict] = field(default_factory=list)
    
    # Scene graph
    scene_graph: Optional[dict] = None
    room_layout: Optional[dict] = None
    
    # Metadata
    room_type: str = "unknown"
    style: str = "modern"
    processing_time: float = 0.0
    
    # Export paths (populated after export)
    glb_path: Optional[str] = None
    fbx_path: Optional[str] = None
    obj_path: Optional[str] = None
    usdz_path: Optional[str] = None
    ply_path: Optional[str] = None  # Gaussian splatting
    
    def export_all(self, output_dir: Union[str, Path]) -> "InteriorFusionOutput":
        """Export all formats to output directory."""
        output_dir = Path(output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)
        
        if self.scene_mesh is not None:
            self.glb_path = str(output_dir / "scene.glb")
            export_mesh(self.scene_mesh, self.glb_path, format="glb")
            
            self.fbx_path = str(output_dir / "scene.fbx")
            export_mesh(self.scene_mesh, self.fbx_path, format="fbx")
            
            self.obj_path = str(output_dir / "scene.obj")
            export_mesh(self.scene_mesh, self.obj_path, format="obj")
            
            self.usdz_path = str(output_dir / "scene.usdz")
            export_mesh(self.scene_mesh, self.usdz_path, format="usdz")
        
        if self.gaussian_cloud is not None:
            self.ply_path = str(output_dir / "scene.ply")
            export_gaussian_splatting(self.gaussian_cloud, self.ply_path)
        
        return self


class InteriorFusionPipeline:
    """
    Main inference pipeline for InteriorFusion.
    
    Orchestrates 5 phases:
    1. Scene Understanding (depth, layout, segmentation)
    2. Multi-View Generation (per-object + room shell)
    3. 3D Reconstruction (room shell + per-object)
    4. Scene Assembly (layout optimization, scale normalization)
    5. Material & Texture (PBR generation, texture baking)
    """
    
    def __init__(
        self,
        model_size: str = "L",
        device: str = "cuda",
        dtype: torch.dtype = torch.float16,
        use_scene_graph: bool = True,
        use_pbr: bool = True,
        use_gaussian_splatting: bool = True,
        cache_dir: Optional[str] = None,
    ):
        self.model_size = model_size
        self.device = device
        self.dtype = dtype
        self.use_scene_graph = use_scene_graph
        self.use_pbr = use_pbr
        self.use_gaussian_splatting = use_gaussian_splatting
        self.cache_dir = cache_dir or os.path.expanduser("~/.cache/interiorfusion")
        
        os.makedirs(self.cache_dir, exist_ok=True)
        
        # Initialize sub-modules (lazy loading)
        self._scene_understanding = None
        self._multiview_gen = None
        self._reconstruction = None
        self._scene_assembly = None
        self._material_texture = None
    
    @property
    def scene_understanding(self):
        if self._scene_understanding is None:
            self._scene_understanding = SceneUnderstandingModule(
                model_size=self.model_size,
                device=self.device,
                dtype=self.dtype,
                cache_dir=self.cache_dir,
            )
        return self._scene_understanding
    
    @property
    def multiview_gen(self):
        if self._multiview_gen is None:
            self._multiview_gen = MultiViewGenerationModule(
                model_size=self.model_size,
                device=self.device,
                dtype=self.dtype,
                cache_dir=self.cache_dir,
            )
        return self._multiview_gen
    
    @property
    def reconstruction(self):
        if self._reconstruction is None:
            self._reconstruction = Reconstruction3DModule(
                model_size=self.model_size,
                device=self.device,
                dtype=self.dtype,
                cache_dir=self.cache_dir,
            )
        return self._reconstruction
    
    @property
    def scene_assembly(self):
        if self._scene_assembly is None:
            self._scene_assembly = SceneAssemblyModule(
                device=self.device,
                dtype=self.dtype,
            )
        return self._scene_assembly
    
    @property
    def material_texture(self):
        if self._material_texture is None:
            self._material_texture = MaterialTextureModule(
                model_size=self.model_size,
                device=self.device,
                dtype=self.dtype,
                use_pbr=self.use_pbr,
                cache_dir=self.cache_dir,
            )
        return self._material_texture
    
    @torch.no_grad()
    def __call__(
        self,
        image: Union[str, Path, Image.Image, np.ndarray],
        room_type_hint: Optional[str] = None,
        style_hint: Optional[str] = None,
        output_formats: Optional[List[str]] = None,
        return_intermediates: bool = False,
    ) -> InteriorFusionOutput:
        """
        Run full InteriorFusion pipeline on a single interior image.
        
        Args:
            image: Input interior photograph
            room_type_hint: Optional room type ("living_room", "bedroom", etc.)
            style_hint: Optional style ("modern", "scandinavian", etc.)
            output_formats: List of formats to export ["glb", "fbx", "obj", "usdz", "ply"]
            return_intermediates: Whether to return intermediate stage outputs
            
        Returns:
            InteriorFusionOutput with all generated 3D content
        """
        import time
        start_time = time.time()
        
        # Convert input to PIL Image
        if isinstance(image, (str, Path)):
            image = Image.open(image).convert("RGB")
        elif isinstance(image, np.ndarray):
            image = Image.fromarray(image).convert("RGB")
        
        # ============================
        # Phase 1: Scene Understanding
        # ============================
        print("[Phase 1/5] Scene Understanding...")
        scene_info = self.scene_understanding(image)
        
        depth_map = scene_info["depth"]
        room_layout = scene_info["room_layout"]
        semantic_seg = scene_info["semantic_segmentation"]
        detected_objects = scene_info["detected_objects"]
        room_type = scene_info.get("room_type", room_type_hint or "living_room")
        style = scene_info.get("style", style_hint or "modern")
        
        # ============================
        # Phase 2: Multi-View Generation
        # ============================
        print("[Phase 2/5] Multi-View Generation...")
        
        # Per-object multi-view generation
        object_multiviews = {}
        for obj_id, obj_info in detected_objects.items():
            crop = obj_info["crop"]
            mask = obj_info["mask"]
            multiviews = self.multiview_gen.generate_object_views(
                crop, mask, depth_map, num_views=6
            )
            object_multiviews[obj_id] = multiviews
        
        # Room shell multi-view
        room_shell_views = self.multiview_gen.generate_room_shell_views(
            image, depth_map, room_layout
        )
        
        # ============================
        # Phase 3: 3D Reconstruction
        # ============================
        print("[Phase 3/5] 3D Reconstruction...")
        
        # Room shell reconstruction
        room_shell_mesh = self.reconstruction.reconstruct_room_shell(
            room_shell_views, room_layout, depth_map
        )
        
        # Per-object reconstruction
        object_meshes = []
        object_gaussians = []
        for obj_id, multiviews in object_multiviews.items():
            obj_mesh, obj_gaussians = self.reconstruction.reconstruct_object(
                multiviews,
                room_layout=room_layout,
                depth_map=depth_map,
                object_info=detected_objects[obj_id],
            )
            object_meshes.append(obj_mesh)
            object_gaussians.append(obj_gaussians)
        
        # Scene Gaussian splatting
        gaussian_cloud = None
        if self.use_gaussian_splatting:
            gaussian_cloud = self.reconstruction.build_scene_gaussians(
                room_shell_mesh, object_gaussians, object_meshes
            )
        
        # ============================
        # Phase 4: Scene Assembly
        # ============================
        print("[Phase 4/5] Scene Assembly...")
        
        assembled_scene = self.scene_assembly.assemble(
            room_shell_mesh=room_shell_mesh,
            object_meshes=object_meshes,
            room_layout=room_layout,
            detected_objects=detected_objects,
            depth_map=depth_map,
        )
        
        scene_mesh = assembled_scene["scene_mesh"]
        scene_graph = assembled_scene.get("scene_graph")
        
        # ============================
        # Phase 5: Material & Texture
        # ============================
        print("[Phase 5/5] Material & Texture...")
        
        pbr_materials = []
        if self.use_pbr:
            # Room shell materials
            room_shell_mesh = self.material_texture.generate_room_materials(
                room_shell_mesh, image, semantic_seg
            )
            
            # Per-object materials
            textured_objects = []
            for i, obj_mesh in enumerate(object_meshes):
                obj_id = list(detected_objects.keys())[i]
                textured_obj, materials = self.material_texture.generate_object_materials(
                    obj_mesh,
                    object_multiviews[obj_id],
                    detected_objects[obj_id],
                )
                textured_objects.append(textured_obj)
                pbr_materials.extend(materials)
            
            # Re-assemble with textured objects
            scene_mesh = self.scene_assembly.reassemble_with_textures(
                room_shell_mesh, textured_objects, scene_graph
            )
        
        processing_time = time.time() - start_time
        
        output = InteriorFusionOutput(
            scene_mesh=scene_mesh,
            room_shell_mesh=room_shell_mesh,
            object_meshes=object_meshes if not self.use_pbr else textured_objects,
            gaussian_cloud=gaussian_cloud,
            pbr_materials=pbr_materials,
            scene_graph=scene_graph,
            room_layout=room_layout,
            room_type=room_type,
            style=style,
            processing_time=processing_time,
        )
        
        print(f"\n✅ Generation complete in {processing_time:.1f}s")
        print(f"   Room type: {room_type}")
        print(f"   Style: {style}")
        print(f"   Objects detected: {len(detected_objects)}")
        print(f"   PBR materials: {len(pbr_materials)}")
        
        return output
    
    def edit_scene(
        self,
        scene_output: InteriorFusionOutput,
        edits: List[dict],
    ) -> InteriorFusionOutput:
        """
        Apply edits to a generated scene.
        
        Edits format:
        [
            {"action": "move", "object_id": 0, "position": [x, y, z]},
            {"action": "replace", "object_id": 1, "new_image": Image},
            {"action": "remove", "object_id": 2},
            {"action": "add", "new_image": Image, "position": [x, y, z]},
        ]
        """
        print(f"Applying {len(edits)} edits...")
        
        scene_graph = scene_output.scene_graph or {}
        object_meshes = list(scene_output.object_meshes)
        
        for edit in edits:
            action = edit["action"]
            
            if action == "move":
                obj_id = edit["object_id"]
                new_pos = edit["position"]
                # Update scene graph
                if "nodes" in scene_graph and obj_id < len(scene_graph["nodes"]):
                    scene_graph["nodes"][obj_id]["position"] = new_pos
                # Update mesh transform
                if obj_id < len(object_meshes):
                    # Apply translation
                    mesh = object_meshes[obj_id]
                    mesh.vertices += np.array(new_pos)
                    
            elif action == "replace":
                obj_id = edit["object_id"]
                new_image = edit["new_image"]
                # Generate new object from image
                new_multiviews = self.multiview_gen.generate_object_views(
                    new_image, None, None, num_views=6
                )
                new_mesh, _ = self.reconstruction.reconstruct_object(
                    new_multiviews, room_layout=scene_output.room_layout
                )
                object_meshes[obj_id] = new_mesh
                
            elif action == "remove":
                obj_id = edit["object_id"]
                if obj_id < len(object_meshes):
                    object_meshes.pop(obj_id)
                
            elif action == "add":
                new_image = edit["new_image"]
                position = edit["position"]
                new_multiviews = self.multiview_gen.generate_object_views(
                    new_image, None, None, num_views=6
                )
                new_mesh, _ = self.reconstruction.reconstruct_object(
                    new_multiviews, room_layout=scene_output.room_layout
                )
                new_mesh.vertices += np.array(position)
                object_meshes.append(new_mesh)
        
        # Re-assemble
        assembled = self.scene_assembly.reassemble_with_textures(
            scene_output.room_shell_mesh,
            object_meshes,
            scene_graph,
        )
        
        return InteriorFusionOutput(
            scene_mesh=assembled,
            room_shell_mesh=scene_output.room_shell_mesh,
            object_meshes=object_meshes,
            gaussian_cloud=scene_output.gaussian_cloud,
            pbr_materials=scene_output.pbr_materials,
            scene_graph=scene_graph,
            room_layout=scene_output.room_layout,
            room_type=scene_output.room_type,
            style=scene_output.style,
        )