Upload src/interiorfusion/models/material_texture.py

src/interiorfusion/models/material_texture.py

@@ -0,0 +1,274 @@
+"""Phase 5: Material & Texture Module.
+
+Generates:
+- PBR materials (albedo, metallic, roughness, normal)
+- Texture baking from multi-view images
+- Lighting estimation for relightable scenes
+"""
+
+import os
+from typing import Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from PIL import Image
+
+
+class MaterialTextureModule(nn.Module):
+    """Generate PBR materials and bake textures onto meshes."""
+    
+    def __init__(
+        self,
+        model_size: str = "L",
+        device: str = "cuda",
+        dtype: torch.dtype = torch.float16,
+        use_pbr: bool = True,
+        cache_dir: Optional[str] = None,
+    ):
+        super().__init__()
+        self.model_size = model_size
+        self.device = device
+        self.dtype = dtype
+        self.use_pbr = use_pbr
+        self.cache_dir = cache_dir
+        
+        # Material generation model (placeholder for now)
+        self._material_model = None
+        
+        # Material type priors
+        self.material_priors = {
+            "wall": {"albedo": [0.9, 0.9, 0.9], "metallic": 0.0, "roughness": 0.8},
+            "floor_wood": {"albedo": [0.6, 0.4, 0.2], "metallic": 0.0, "roughness": 0.6},
+            "floor_tile": {"albedo": [0.8, 0.8, 0.8], "metallic": 0.1, "roughness": 0.3},
+            "floor_carpet": {"albedo": [0.5, 0.3, 0.2], "metallic": 0.0, "roughness": 0.9},
+            "ceiling": {"albedo": [0.95, 0.95, 0.95], "metallic": 0.0, "roughness": 0.9},
+            "furniture_wood": {"albedo": [0.5, 0.3, 0.15], "metallic": 0.0, "roughness": 0.5},
+            "furniture_fabric": {"albedo": [0.6, 0.5, 0.4], "metallic": 0.0, "roughness": 0.8},
+            "furniture_leather": {"albedo": [0.4, 0.2, 0.1], "metallic": 0.1, "roughness": 0.4},
+            "furniture_metal": {"albedo": [0.7, 0.7, 0.7], "metallic": 0.9, "roughness": 0.2},
+            "furniture_plastic": {"albedo": [0.8, 0.8, 0.8], "metallic": 0.0, "roughness": 0.3},
+            "furniture_glass": {"albedo": [0.9, 0.9, 0.9], "metallic": 0.0, "roughness": 0.05},
+            "default": {"albedo": [0.7, 0.7, 0.7], "metallic": 0.0, "roughness": 0.5},
+        }
+    
+    def generate_room_materials(
+        self,
+        room_shell_mesh: "trimesh.Trimesh",  # type: ignore
+        image: Image.Image,
+        semantic_seg: np.ndarray,
+    ) -> "trimesh.Trimesh":  # type: ignore
+        """
+        Generate materials for room shell (walls, floor, ceiling).
+        
+        Uses semantic segmentation to determine material types
+        and input image for color extraction.
+        """
+        if not self.use_pbr:
+            return room_shell_mesh
+        
+        # Extract dominant colors from image regions
+        img_np = np.array(image)
+        
+        # Determine material types from semantic segmentation
+        floor_region = semantic_seg == 1
+        ceiling_region = semantic_seg == 2
+        wall_regions = (semantic_seg == 3) | (semantic_seg == 4)
+        
+        # Extract colors from corresponding image regions
+        floor_color = self._extract_dominant_color(img_np, floor_region)
+        ceiling_color = self._extract_dominant_color(img_np, ceiling_region)
+        wall_color = self._extract_dominant_color(img_np, wall_regions)
+        
+        # Create materials
+        floor_mat = self._create_material("floor_wood", color=floor_color)
+        ceiling_mat = self._create_material("ceiling", color=ceiling_color)
+        wall_mat = self._create_material("wall", color=wall_color)
+        
+        # Apply materials to mesh faces
+        # In practice, this would be done per-face based on which room part the face belongs to
+        # For now, store materials as mesh metadata
+        room_shell_mesh.materials = {
+            "floor": floor_mat,
+            "ceiling": ceiling_mat,
+            "walls": wall_mat,
+        }
+        
+        return room_shell_mesh
+    
+    def generate_object_materials(
+        self,
+        object_mesh: "trimesh.Trimesh",  # type: ignore
+        multiviews: List[Image.Image],
+        object_info: Dict,
+    ) -> Tuple["trimesh.Trimesh", List[Dict]]:  # type: ignore
+        """
+        Generate PBR materials for a furniture object.
+        
+        Uses multi-view images to bake texture and infer material properties.
+        """
+        if not self.use_pbr:
+            return object_mesh, []
+        
+        class_name = object_info.get("class_name", "furniture")
+        
+        # Infer material type from class and image analysis
+        material_type = self._infer_material_type(class_name, multiviews[0])
+        
+        # Extract dominant color from multi-view images
+        colors = [self._extract_dominant_color(np.array(mv), np.ones((mv.size[1], mv.size[0]), dtype=bool))
+                  for mv in multiviews]
+        avg_color = np.mean(colors, axis=0)
+        
+        # Create material
+        material = self._create_material(material_type, color=avg_color)
+        
+        # Create simple UV atlas texture
+        texture = self._bake_texture(object_mesh, multiviews, material)
+        
+        # Attach texture to mesh
+        if texture is not None:
+            object_mesh.visual = object_mesh.visual.to_texture()
+            # In production, set actual texture image
+            object_mesh.material_override = material
+        
+        materials = [material]
+        
+        return object_mesh, materials
+    
+    def estimate_lighting(
+        self,
+        image: Image.Image,
+    ) -> Dict:
+        """
+        Estimate scene lighting from input image.
+        
+        Returns:
+            {
+                "environment_map": HDR environment map (placeholder),
+                "key_light_direction": [x, y, z],
+                "key_light_intensity": float,
+                "fill_light_intensity": float,
+                "ambient_intensity": float,
+                "color_temperature": float,  # Kelvin
+            }
+        """
+        img_np = np.array(image)
+        
+        # Simple heuristic lighting estimation
+        # In production, use trained lighting estimation network
+        
+        # Estimate brightness
+        brightness = img_np.mean()
+        
+        # Estimate color temperature from average color
+        avg_color = img_np.mean(axis=(0, 1))
+        
+        # Warm = more red, Cool = more blue
+        color_temp = 6500  # Default daylight
+        if avg_color[2] > avg_color[0] * 1.2:
+            color_temp = 8000  # Cool
+        elif avg_color[0] > avg_color[2] * 1.2:
+            color_temp = 3000  # Warm
+        
+        # Estimate light direction from shadows
+        # Placeholder: assume light from top-left
+        light_dir = np.array([0.3, 0.8, 0.2])
+        light_dir = light_dir / np.linalg.norm(light_dir)
+        
+        return {
+            "environment_map": None,  # Would generate HDR probe
+            "key_light_direction": light_dir.tolist(),
+            "key_light_intensity": float(brightness / 255.0 * 2.0),
+            "fill_light_intensity": float(brightness / 255.0 * 0.5),
+            "ambient_intensity": float(brightness / 255.0 * 0.3),
+            "color_temperature": float(color_temp),
+        }
+    
+    def _extract_dominant_color(
+        self,
+        image: np.ndarray,
+        mask: np.ndarray,
+    ) -> np.ndarray:
+        """Extract dominant color from image region."""
+        if mask.sum() == 0:
+            return np.array([0.7, 0.7, 0.7])
+        
+        masked_pixels = image[mask]
+        
+        # K-means-ish: use median for robustness
+        dominant_color = np.median(masked_pixels, axis=0) / 255.0
+        
+        return dominant_color
+    
+    def _create_material(
+        self,
+        material_type: str,
+        color: Optional[np.ndarray] = None,
+    ) -> Dict:
+        """Create PBR material from type and color."""
+        prior = self.material_priors.get(material_type, self.material_priors["default"])
+        
+        if color is not None:
+            albedo = color.tolist()
+        else:
+            albedo = prior["albedo"]
+        
+        return {
+            "type": material_type,
+            "albedo": albedo,
+            "metallic": prior["metallic"],
+            "roughness": prior["roughness"],
+            "normal_scale": 1.0,
+            "ao_scale": 1.0,
+            # Texture maps (would be actual textures in production)
+            "albedo_map": None,
+            "metallic_map": None,
+            "roughness_map": None,
+            "normal_map": None,
+            "ao_map": None,
+        }
+    
+    def _infer_material_type(
+        self,
+        class_name: str,
+        image: Image.Image,
+    ) -> str:
+        """Infer material type from object class and visual appearance."""
+        class_lower = class_name.lower()
+        
+        # Map class to material type
+        material_map = {
+            "sofa": "furniture_fabric",
+            "chair": "furniture_fabric",
+            "table": "furniture_wood",
+            "coffee_table": "furniture_wood",
+            "bed": "furniture_fabric",
+            "desk": "furniture_wood",
+            "bookshelf": "furniture_wood",
+            "lamp": "furniture_metal",
+            "wardrobe": "furniture_wood",
+            "tv_stand": "furniture_wood",
+            "rug": "floor_carpet",
+        }
+        
+        return material_map.get(class_lower, "furniture_wood")
+    
+    def _bake_texture(
+        self,
+        mesh: "trimesh.Trimesh",  # type: ignore
+        multiviews: List[Image.Image],
+        material: Dict,
+    ) -> Optional[Image.Image]:
+        """
+        Bake multi-view images into a unified UV texture.
+        
+        Uses visibility-aware projection to handle occlusions.
+        """
+        # Placeholder: in production, this would be proper UV unwrapping + projection
+        # For now, return the first multi-view as the texture
+        
+        if multiviews:
+            return multiviews[0]
+        return None