src/interiorfusion/models/material_texture.py

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