File size: 9,920 Bytes
6df8db0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 | """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
|