src/layout_generator/assembler.py

# src/layout_generator/assembler.py

"""
Шаг пайплайна для физической сборки 3D-сцены в непрерывном пространстве.
"""
import logging
import os
from typing import Any, List

import numpy as np
import scene_synthesizer as synth
import trimesh

from dsynth.assets.ss_assets import DefaultShelf
from dsynth.scene_gen.arrangements import add_objects_to_shelf_v2, set_shelf

from .base import BaseStep, LayoutContext

logger: logging.Logger = logging.getLogger(__name__)

class AssembleSceneStep(BaseStep):
    """Шаг физического рендера 3D-сцены на основе расстановки из тензорного поля."""
    
    def process(self, context: LayoutContext) -> LayoutContext:
        """Сборка 3D-моделей по заданным координатам."""
        if not context.is_gen or not hasattr(context, 'placed_fixtures'):
            return context
            
        logger.info("🛠️ Сборка 3D-сцены и физическая выкладка инвентаря...")
        scene: synth.Scene = synth.Scene()
        actual_shelf_idx: int = 0

        for fixture in context.placed_fixtures:
            actual_shelf_idx += 1
            shelf_name: str = fixture.name
            zone_id, shelf_id = shelf_name.split('.')
            shelf_cfg: Any = context.fixed_zones[zone_id][shelf_id]
            shelf_asset_name: Any = shelf_cfg.shelf_asset

            asset_obj: Any = context.product_assets_lib.get(shelf_asset_name)
            if asset_obj is None and shelf_asset_name:
                asset_obj = context.product_assets_lib.get(f"fixtures.{shelf_asset_name}")

            shelf: Any = DefaultShelf if asset_obj is None else asset_obj.ss_asset
            
            # Извлекаем идеальные координаты и угол, посчитанные тензорным полем
            safe_x: float = fixture.x
            safe_y: float = fixture.y
            rotation_deg: int = getattr(fixture, '_rotation_deg', 0)

            # Ставим оборудование на сцену
            support_data: Any = set_shelf(
                scene, shelf, safe_x, safe_y, rotation_deg,
                f'SHELF_{actual_shelf_idx}_{shelf_name}', f'support_SHELF_{actual_shelf_idx}_{shelf_name}'
            )

            # Выкладка товаров
            if shelf_name in context.product_filling:
                placement_data: List[List[str]] = context.product_filling[shelf_name]
                if placement_data:
                    try:
                        add_objects_to_shelf_v2(
                            scene, actual_shelf_idx, placement_data, context.product_assets_lib,
                            support_data, shelf_cfg.x_gap, shelf_cfg.y_gap, shelf_cfg.delta_x, shelf_cfg.delta_y,
                            shelf_cfg.start_point_x, shelf_cfg.start_point_y, shelf_cfg.filling_type,
                            seed=context.current_seed, noise_std_x=shelf_cfg.noise_std_x, noise_std_y=shelf_cfg.noise_std_y,
                            rotation_lower=shelf_cfg.rotation_lower, rotation_upper=shelf_cfg.rotation_upper
                        )
                    except Exception as e:
                        logger.error(f"Ошибка выкладки на полку {shelf_name}: {e}")

        # Динамическая генерация пола по габаритам расставленной мебели
        margin: float = 2.0
        t: float = 0.2
        
        if context.placed_fixtures:
            all_x: List[float] = []
            all_y: List[float] = []
            for f in context.placed_fixtures:
                poly, _ = f.get_polygon()
                all_x.extend(poly[:, 0])
                all_y.extend(poly[:, 1])
            min_x, max_x = min(all_x), max(all_x)
            min_y, max_y = min(all_y), max(all_y)
        else:
            min_x, max_x = 0.0, float(context.size_n)
            min_y, max_y = 0.0, float(context.size_m)

        cx: float = (max_x + min_x) / 2.0
        cy: float = (max_y + min_y) / 2.0
        w_x: float = (max_x - min_x) + margin * 2
        w_y: float = (max_y - min_y) + margin * 2

        temp_filename: str = f"temp_scene_{os.getpid()}.glb"
        scene.export(temp_filename)
        final_scene: trimesh.Scene = trimesh.load(temp_filename, force='scene')
        
        if os.path.exists(temp_filename):
            os.remove(temp_filename)

        floor: trimesh.Trimesh = trimesh.creation.box(extents=[w_x, w_y, t])
        floor.apply_translation([cx, cy, -t/2])
        floor.visual.face_colors = [180, 180, 180, 255]
        final_scene.add_geometry(floor, node_name="floor")

        R: np.ndarray = trimesh.transformations.rotation_matrix(-np.pi / 2, [1, 0, 0])
        final_scene.apply_transform(R)

        context.final_scene = final_scene
        context.actual_shelf_idx = actual_shelf_idx
        
        return context