interiorgen3d/core/scene_graph.py

"""
Core data structures for InteriorGen3D scene representation.
"""
from dataclasses import dataclass, field
from typing import List, Optional, Dict, Tuple
import numpy as np


@dataclass
class BoundingBox3D:
    """3D oriented bounding box for furniture objects."""
    center: np.ndarray
    dimensions: np.ndarray
    rotation: np.ndarray
    confidence: float = 1.0
    
    @property
    def corners(self) -> np.ndarray:
        w, h, d = self.dimensions / 2
        corners_local = np.array([
            [-w, -h, -d], [w, -h, -d], [w, h, -d], [-w, h, -d],
            [-w, -h, d], [w, -h, d], [w, h, d], [-w, h, d]
        ])
        return (self.rotation @ corners_local.T).T + self.center
    
    @property
    def volume(self) -> float:
        return float(np.prod(self.dimensions))


@dataclass
class WallPlane:
    normal: np.ndarray
    offset: float
    height: float
    start_point: np.ndarray
    end_point: np.ndarray
    has_window: bool = False
    has_door: bool = False
    material_id: Optional[str] = None


@dataclass
class RoomLayout:
    walls: List[WallPlane] = field(default_factory=list)
    floor_height: float = 0.0
    ceiling_height: float = 2.7
    floor_polygon: Optional[np.ndarray] = None
    doors: List[Dict] = field(default_factory=list)
    windows: List[Dict] = field(default_factory=list)


@dataclass
class SceneObject:
    object_id: str
    category: str
    bbox: BoundingBox3D
    mesh_path: Optional[str] = None
    texture_paths: Optional[Dict[str, str]] = None
    gaussian_splat_path: Optional[str] = None
    confidence: float = 1.0
    material_type: Optional[str] = None
    is_on_floor: bool = True
    parent_object_id: Optional[str] = None


@dataclass
class SpatialRelation:
    subject_id: str
    object_id: str
    relation: str
    confidence: float = 1.0


@dataclass
class SceneGraph:
    room_layout: RoomLayout
    objects: List[SceneObject] = field(default_factory=list)
    relations: List[SpatialRelation] = field(default_factory=list)
    camera_pose: Optional[np.ndarray] = None
    camera_intrinsics: Optional[np.ndarray] = None
    metric_scale: float = 1.0
    
    def get_object(self, object_id: str) -> Optional[SceneObject]:
        for obj in self.objects:
            if obj.object_id == object_id:
                return obj
        return None
    
    def get_objects_by_category(self, category: str) -> List[SceneObject]:
        return [obj for obj in self.objects if obj.category == category]


@dataclass
class PBRMaterial:
    albedo_map: Optional[np.ndarray] = None
    metallic_map: Optional[np.ndarray] = None
    roughness_map: Optional[np.ndarray] = None
    normal_map: Optional[np.ndarray] = None
    base_color: Tuple[float, float, float] = (0.8, 0.8, 0.8)
    metallic: float = 0.0
    roughness: float = 0.5


@dataclass
class GeneratedScene:
    scene_graph: SceneGraph
    room_mesh_path: str
    object_meshes: Dict[str, str]
    materials: Dict[str, PBRMaterial]
    gaussian_splat_path: Optional[str] = None
    export_paths: Dict[str, str] = field(default_factory=dict)
    generation_time_seconds: float = 0.0
    metadata: Dict = field(default_factory=dict)