Upload interiorgen3d/core/scene_graph.py

interiorgen3d/core/scene_graph.py

@@ -0,0 +1,114 @@
+"""
+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)