diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index b85deaea9..3ad17a9bc 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -147,6 +147,11 @@ jobs: --extra-index-url https://download.blender.org/pypi/ pip install pytest-xdist + - name: Verify Scene Engine gensim installation + run: | + python -c "import matplotlib, numpy, open3d, requests, scipy, shapely, trimesh; from PIL import Image; import embodichain.gen_sim.scene_engine.pipeline.generate" + pytest tests/gen_sim/scene_engine -q + - name: Run default tests run: | echo "Default test suite (GPU-marked tests are skipped)" diff --git a/docs/source/features/generative_sim/index.rst b/docs/source/features/generative_sim/index.rst index 1f7c759f7..09d041571 100644 --- a/docs/source/features/generative_sim/index.rst +++ b/docs/source/features/generative_sim/index.rst @@ -7,3 +7,4 @@ Generative Simulation collects EmbodiChain features for generating simulation-re :maxdepth: 2 SimReady Asset Pipeline + Scene Engine diff --git a/docs/source/features/generative_sim/scene_engine.md b/docs/source/features/generative_sim/scene_engine.md new file mode 100644 index 000000000..03e4b3fa4 --- /dev/null +++ b/docs/source/features/generative_sim/scene_engine.md @@ -0,0 +1,83 @@ +# Scene Engine + +Scene Engine reconstructs a table-top scene from one image. It identifies the +table and visible objects, segments their masks, generates simulation-ready +meshes, refines the object layout on the table, and exports a scene that can be +loaded by EmbodiChain. + +## Quick Start + +Install EmbodiChain with the `gensim` extra first; see +[Installation](../../quick_start/install.md#optional-generative-simulation-gensim). + +Configure the required services in `embodichain/gen_sim/.env`, then generate a +scene: + +```bash +embodichain scene-engine \ + --image /path/to/scene.png \ + --output_root /path/to/scene_output +``` + +The same command is available through the package entry point: + +```bash +python -m embodichain scene-engine \ + --image /path/to/scene.png \ + --output_root /path/to/scene_output +``` + +## Configuration + +Scene Engine reads the LLM, segmentation, and geometry-generation settings +from `embodichain/gen_sim/.env`: + +```bash +OPENAI_API_KEY="your-api-key" +OPENAI_MODEL="your-model" +OPENAI_BASE_URL="https://api.openai.com/v1" +SCENE_ENGINE_OPENAI_DEFAULT_QUERY="{}" +OPENAI_MAX_ATTEMPTS=3 + +SCENE_ENGINE_IMAGE_SEGMENTATION_BASE_URL="http://host:port" +SCENE_ENGINE_IMAGE_SEGMENTATION_TIMEOUT_S=30 +SCENE_ENGINE_IMAGE_SEGMENTATION_MAX_ATTEMPTS=3 +SCENE_ENGINE_IMAGE_SEGMENTATION_HEALTH_PATH="/health" +SCENE_ENGINE_IMAGE_SEGMENTATION_SINGLE_OBJECT_PATH="/predict" + +SCENE_ENGINE_GEOMETRY_GENERATION_BASE_URL="http://host:port" +SCENE_ENGINE_GEOMETRY_GENERATION_TIMEOUT_S=600 +SCENE_ENGINE_GEOMETRY_GENERATION_MAX_ATTEMPTS=3 +SCENE_ENGINE_GEOMETRY_GENERATION_HEALTH_PATH="/health" +SCENE_ENGINE_GEOMETRY_GENERATION_OBJECTS_PATH="/generate_multiple_objects" +``` + +## Processing Flow + +- **Scene understanding**: analyzes the image and segments the table and visible objects. +- **Scene generation**: generates meshes, prepares SimReady geometry, detects the table support surface, and refines the table-top layout. +- **Scene export**: copies the final GLBs and writes a portable z-up scene export. + +## Output and Preview + +The important final outputs are: + +```text +scene_output/ +|-- scene_understanding/ # Object analysis, masks, and stage JSON +|-- scene_generation/ # Generated, SimReady, and layout-debug artifacts +`-- scene_export/ + |-- mesh_assets/ # Final GLBs + `-- scene_config.json # Exported scene description +``` + +Validate the export without opening a window: + +```bash +embodichain preview-scene \ + --output_root /path/to/scene_output \ + --headless +``` + +For an interactive preview, omit `--headless`. Add `--viser` to publish the +scene through Viser. diff --git a/embodichain/__main__.py b/embodichain/__main__.py index 0af9abae6..975e7649d 100644 --- a/embodichain/__main__.py +++ b/embodichain/__main__.py @@ -51,6 +51,16 @@ class Command: target="embodichain.gen_sim.simready_pipeline.cli.start:main", help="Convert a raw asset directory into a SimReady asset.", ), + Command( + name="scene-engine", + target="embodichain.gen_sim.scene_engine.cli.start:main", + help="Generate a scene export from an input image using gen_sim/.env.", + ), + Command( + name="preview-scene", + target="embodichain.gen_sim.scene_engine.cli.preview:main", + help="Preview a generated Scene Engine scene export.", + ), Command( name="preview-asset", target="embodichain.lab.scripts.preview_asset:cli", diff --git a/embodichain/gen_sim/scene_engine/__init__.py b/embodichain/gen_sim/scene_engine/__init__.py new file mode 100644 index 000000000..015c41510 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/__init__.py @@ -0,0 +1,19 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +__all__: list[str] = [] diff --git a/embodichain/gen_sim/scene_engine/cli/__init__.py b/embodichain/gen_sim/scene_engine/cli/__init__.py new file mode 100644 index 000000000..015c41510 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/cli/__init__.py @@ -0,0 +1,19 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +__all__: list[str] = [] diff --git a/embodichain/gen_sim/scene_engine/cli/preview.py b/embodichain/gen_sim/scene_engine/cli/preview.py new file mode 100644 index 000000000..f2d19c263 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/cli/preview.py @@ -0,0 +1,247 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path +import time +from collections.abc import Sequence +from typing import Any + +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.cfg import LightCfg, MeshCfg, RigidObjectCfg +from embodichain.lab.visualization import ( + VisualizationCfg, + add_viser_args_to_parser, + visualization_cfg_from_args, +) + + +def preview_scene_export( + *, + output_root: str | Path, + device: str = "cpu", + headless: bool = False, + visualization: VisualizationCfg | None = None, +) -> None: + """Load ``scene_export/scene_config.json`` and preview its table and assets. + + Args: + output_root: Scene Engine output root containing ``scene_export/``. + device: Simulation device, for example ``"cpu"`` or ``"cuda"``. + headless: Load and validate the scene without an interactive preview. + visualization: Optional live-visualization configuration. + """ + resolved_output_root = Path(output_root).expanduser().resolve() + config_path = resolved_output_root / "scene_export" / "scene_config.json" + if not config_path.is_file(): + raise FileNotFoundError(f"Scene config not found: {config_path}") + + try: + scene_config = json.loads(config_path.read_text(encoding="utf-8")) + except json.JSONDecodeError as exc: + raise ValueError(f"Scene config is not valid JSON: {config_path}") from exc + if not isinstance(scene_config, dict): + raise ValueError("Scene config must be a JSON object.") + if scene_config.get("format") != "embodichain.scene-export/v1": + raise ValueError( + "Expected an EmbodiChain scene export " + "(format='embodichain.scene-export/v1')." + ) + + sim = SimulationManager( + SimulationManagerCfg( + width=1920, + height=1080, + headless=headless, + physics_dt=1.0 / 100.0, + sim_device=device, + visualization=( + VisualizationCfg() if visualization is None else visualization + ), + ) + ) + try: + if sim.is_use_gpu_physics: + sim.init_gpu_physics() + _add_lights(sim) + _add_objects( + sim=sim, + entries=_config_entries(scene_config, "background"), + config_dir=config_path.parent, + label="table", + ) + _add_objects( + sim=sim, + entries=_config_entries(scene_config, "rigid_object"), + config_dir=config_path.parent, + label="asset", + ) + + is_viser = sim.sim_config.visualization.backend == "viser" + if headless and not is_viser: + sim.update(step=1) + print(f"Loaded scene export headlessly: {config_path}") + return + + if is_viser: + sim.update(step=1) + print(f"Previewing in Viser: {config_path}") + else: + print(f"Previewing: {config_path}") + sim.open_window() + print("Close with Ctrl-C.") + while True: + time.sleep(0.1) + except KeyboardInterrupt: + print("Stopping preview.") + finally: + sim.destroy(exit_process=False) + SimulationManager.flush_cleanup_queue() + + +def _config_entries( + scene_config: dict[str, Any], + field_name: str, +) -> list[dict[str, Any]]: + entries = scene_config.get(field_name, []) + if not isinstance(entries, list) or not all( + isinstance(entry, dict) for entry in entries + ): + raise ValueError( + f"Scene config field {field_name!r} must be a list of objects." + ) + return entries + + +def _add_lights(sim: SimulationManager) -> None: + for index in range(8): + angle = 2.0 * math.pi * index / 8 + sim.add_light( + LightCfg( + uid=f"light_{index + 1}", + intensity=80.0, + radius=600, + init_pos=[5.0 * math.cos(angle), 5.0 * math.sin(angle), 8.0], + ) + ) + + +def _add_objects( + *, + sim: SimulationManager, + entries: list[dict[str, Any]], + config_dir: Path, + label: str, +) -> None: + """Add exported meshes as static bodies so previewing does not re-simulate them.""" + resolved_config_dir = config_dir.resolve() + for entry in entries: + uid = entry.get("uid") + shape = entry.get("shape") + if not isinstance(uid, str) or not uid: + raise ValueError(f"Scene {label} has no valid uid.") + if not isinstance(shape, dict) or not isinstance(shape.get("fpath"), str): + raise ValueError(f"Scene {label} {uid!r} has no shape.fpath.") + if shape.get("shape_type") != "Mesh": + raise ValueError( + f"Scene {label} {uid!r} must use shape_type='Mesh' for preview." + ) + + fpath = Path(shape["fpath"]) + if fpath.is_absolute(): + raise ValueError( + f"Scene {label} {uid!r} shape.fpath must be a relative path." + ) + mesh_path = (resolved_config_dir / fpath).resolve() + if resolved_config_dir not in mesh_path.parents: + raise ValueError( + f"Scene {label} {uid!r} shape.fpath must stay within " + f"{resolved_config_dir}." + ) + if not mesh_path.is_file(): + raise FileNotFoundError(f"Gym mesh for {uid!r} not found: {mesh_path}") + init_pos = _vector3(entry.get("init_pos"), field_name=f"{uid}.init_pos") + init_rot = _vector3(entry.get("init_rot"), field_name=f"{uid}.init_rot") + body_scale = _vector3( + entry.get("body_scale", [1.0, 1.0, 1.0]), + field_name=f"{uid}.body_scale", + ) + max_convex_hull_num = max(1, int(entry.get("max_convex_hull_num", 32))) + + sim.add_rigid_object( + RigidObjectCfg( + uid=uid, + shape=MeshCfg(fpath=str(mesh_path)), + # Keep every preview body static: exported poses are already the + # final gravity-settled poses and should not be simulated again. + body_type="static", + init_pos=tuple(init_pos), + init_rot=tuple(init_rot), + body_scale=tuple(body_scale), + max_convex_hull_num=max_convex_hull_num, + acd_method="vhacd", # Use vhacd by default. + ) + ) + print(f"[{label}] {uid}: pos={init_pos} rot={init_rot} scale={body_scale}") + + +def _vector3(value: object, *, field_name: str) -> list[float]: + if not isinstance(value, list) or len(value) != 3: + raise ValueError(f"Scene config field {field_name!r} must be a length-3 list.") + try: + return [float(item) for item in value] + except (TypeError, ValueError) as exc: + raise ValueError(f"Scene config field {field_name!r} must be numeric.") from exc + + +def main(argv: Sequence[str] | None = None) -> None: + parser = argparse.ArgumentParser( + prog="embodichain preview-scene", + description="Preview a Scene Engine scene export in EmbodiChain simulation.", + ) + parser.add_argument( + "--output_root", + type=Path, + required=True, + help="Scene Engine output root containing scene_export/.", + ) + parser.add_argument( + "--device", + default="cpu", + help="Simulation device, for example cpu or cuda.", + ) + parser.add_argument( + "--headless", + action="store_true", + help="Load and validate the exported scene without opening a window.", + ) + add_viser_args_to_parser(parser) + args = parser.parse_args(argv) + preview_scene_export( + output_root=args.output_root, + device=args.device, + headless=args.headless, + visualization=visualization_cfg_from_args(args), + ) + + +if __name__ == "__main__": + main() diff --git a/embodichain/gen_sim/scene_engine/cli/start.py b/embodichain/gen_sim/scene_engine/cli/start.py new file mode 100644 index 000000000..59454b09f --- /dev/null +++ b/embodichain/gen_sim/scene_engine/cli/start.py @@ -0,0 +1,77 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import argparse +from collections.abc import Sequence +from pathlib import Path + +from embodichain.gen_sim.scene_engine.pipeline.generate import generate_scene_from_image + +_SUPPORTED_IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png"} + + +def cli_scene_engine( + image: str | Path, + output_root: str | Path, +) -> None: + """Generate one scene using the required ``gen_sim/.env`` settings.""" + resolved_image_path = Path(image).expanduser().resolve() + if not resolved_image_path.exists(): + raise FileNotFoundError(f"Image input not found: {resolved_image_path}") + if not resolved_image_path.is_file(): + raise ValueError(f"Image input is not a file: {resolved_image_path}") + if resolved_image_path.suffix.lower() not in _SUPPORTED_IMAGE_SUFFIXES: + raise ValueError( + "Image input must have one of these extensions: .jpg, .jpeg, .png" + ) + + resolved_output_root = Path(output_root).expanduser().resolve() + resolved_output_root.mkdir(parents=True, exist_ok=True) + + generate_scene_from_image( + image_path=resolved_image_path, + output_root=resolved_output_root, + ) + print("Successfully completed!") + + +def main(argv: Sequence[str] | None = None) -> None: + parser = argparse.ArgumentParser( + prog="embodichain scene-engine", + description="Generate a Scene Engine export from one input image.", + epilog="Service settings are read from embodichain/gen_sim/.env.", + ) + parser.add_argument( + "--image", + type=str, + required=True, + help="Path to the required input image file (.jpg, .jpeg, or .png)", + ) + parser.add_argument( + "--output_root", + type=str, + required=True, + help="Path to the output directory", + ) + args = parser.parse_args(argv) + + cli_scene_engine(args.image, args.output_root) + + +if __name__ == "__main__": + main() diff --git a/embodichain/gen_sim/scene_engine/clients/__init__.py b/embodichain/gen_sim/scene_engine/clients/__init__.py new file mode 100644 index 000000000..015c41510 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/clients/__init__.py @@ -0,0 +1,19 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +__all__: list[str] = [] diff --git a/embodichain/gen_sim/scene_engine/clients/geometry_generation.py b/embodichain/gen_sim/scene_engine/clients/geometry_generation.py new file mode 100644 index 000000000..c84fe4c26 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/clients/geometry_generation.py @@ -0,0 +1,448 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from contextlib import ExitStack +import json +from pathlib import Path +import time +from typing import Any + +import requests + +from embodichain.gen_sim.scene_engine.configs.environment import ( + read_scene_engine_env_values, +) + + +class GeometryGenerationClient: + """Manage the Geometry Generation Server connection.""" + + def __init__( + self, + *, + base_url: str, + timeout_s: int, + max_attempts: int, + health_path: str, + generate_objects_path: str, + session: requests.Session | None = None, + ) -> None: + self._base_url = base_url.rstrip("/") + self._timeout_s = timeout_s + self._max_attempts = max_attempts + self._health_path = health_path + self._generate_objects_path = generate_objects_path + self._session = session or requests.Session() + + @classmethod + def from_dotenv(cls) -> "GeometryGenerationClient": + """Create a client from its required ``gen_sim/.env`` settings.""" + return cls(**_load_dotenv_config()) + + def check_health(self) -> None: + last_error: Exception | None = None + for _ in range(self._max_attempts): + try: + response = self._session.get( + self._url(self._health_path), + timeout=10, # Use a shorter timeout for avoiding long waits. + ) + response.raise_for_status() + response_data = response.json() + if ( + not isinstance(response_data, dict) + or response_data.get("ok") is not True + ): + raise RuntimeError( + "Geometry Generation Server health response does not contain ok=true." + ) + return + except (requests.RequestException, ValueError, RuntimeError) as exc: + last_error = exc + + assert last_error is not None + raise RuntimeError( + "Geometry Generation Server health check failed after " + f"{self._max_attempts} attempts." + ) from last_error + + def close(self) -> None: + self._session.close() + + def generate_objects( + self, + *, + image_path: str | Path, + object_masks: list[tuple[str, Path]], + output_root: str | Path, + ) -> tuple[dict[str, Any], list[dict[str, Any]]]: + """Generate objects through the geometry server's mask-list endpoint. + + The service represents both one-object and multi-object jobs as one + image plus a multipart ``masks`` list. The number of list items is the + only difference, so keeping one implementation prevents the client + paths from drifting apart. + """ + + # Check, validate then wrap each content of the request. + resolved_image_path = Path(image_path).expanduser().resolve() + if not resolved_image_path.is_file(): + raise FileNotFoundError( + f"Geometry generation input not found: {resolved_image_path}" + ) + if not object_masks: + raise ValueError("Geometry generation object_masks must not be empty.") + object_ids = [object_id for object_id, _ in object_masks] + if len(set(object_ids)) != len(object_ids): + raise ValueError("Geometry generation object_ids must be unique.") + + resolved_object_masks: list[tuple[str, Path]] = [] + for object_id, mask_path in object_masks: + resolved_mask_path = Path(mask_path).expanduser().resolve() + if not resolved_mask_path.is_file(): + raise FileNotFoundError( + f"Geometry generation mask not found: {resolved_mask_path}" + ) + resolved_object_masks.append((object_id, resolved_mask_path)) + + # Send one multipart image + masks request. + response_data, response_objects = self._request_objects( + image_path=resolved_image_path, + object_masks=resolved_object_masks, + ) + + resolved_output_root = Path(output_root).expanduser().resolve() + resolved_output_root.mkdir(parents=True, exist_ok=True) + + # This loop will iterate min(len(resolved_object_masks), len(response_objects)) times + # , which is safe because we validated the lengths earlier. + for ( + object_id, + _, + ), response_object in zip( # Pair each object_id with its response_object for downloading the glb. + resolved_object_masks, + response_objects, + ): + safe_object_id = Path(object_id).name + if ( + safe_object_id != object_id + or "\\" in object_id + or object_id in {"", ".", ".."} + ): + raise ValueError( + "Geometry generation object_id is not safe for a filename: " + f"{object_id!r}" + ) + output_path = resolved_output_root / f"{safe_object_id}.glb" + self._download_glb(response_object["mesh"], output_path) + return response_data, response_objects + + def _request_objects( + self, + *, + image_path: Path, + object_masks: list[tuple[str, Path]], + ) -> tuple[dict[str, Any], list[dict[str, Any]]]: + last_error: Exception | None = None + for _ in range(self._max_attempts): + try: + # This stack manages the context of multiple open files, ensuring they are closed after the request. + with ExitStack() as stack: + image_file = stack.enter_context(image_path.open("rb")) + mask_files = [ + stack.enter_context(mask_path.open("rb")) + for _, mask_path in object_masks + ] + response = self._session.post( + self._url(self._generate_objects_path), + files=[ + ( + "image", + ( + image_path.name, + image_file, + _image_content_type(image_path), + ), + ), + *[ + ( + "masks", + (f"{object_id}.png", mask_file, "image/png"), + ) + for (object_id, _), mask_file in zip( + object_masks, + mask_files, + ) + ], + ], + timeout=self._timeout_s, + ) + response.raise_for_status() + try: + response_data = response.json() + except ValueError as exc: + raise RuntimeError( + "Geometry Generation Server response is not valid JSON." + ) from exc + response_data = self._wait_for_task_if_needed(response_data) + response_objects = _parse_objects_response( + response_data, + object_ids=[object_id for object_id, _ in object_masks], + ) + return response_data, response_objects + except (requests.RequestException, RuntimeError) as exc: + last_error = exc + + assert last_error is not None + raise RuntimeError( + "Geometry Generation Server request failed after " + f"{self._max_attempts} attempts." + ) from last_error + + def _wait_for_task_if_needed(self, response_data: object) -> dict[str, Any]: + """Poll a queued geometry-generation job until it returns its result.""" + if not isinstance(response_data, dict): + raise RuntimeError( + "Geometry Generation Server response must be a JSON object." + ) + + status = response_data.get("status") + if not isinstance(status, str) or "waiting" not in status: + return response_data + + request_id = response_data.get("request_id") + if not isinstance(request_id, str) or not request_id: + raise RuntimeError( + "Geometry Generation Server queued response has no request_id." + ) + + # The server test client uses one-second polling and permits ten minutes + # for a queued job. Keep the same contract here. + for _ in range(600): + try: + response = self._session.get( + self._url(f"/tasks/{request_id}"), + timeout=10, + ) + response.raise_for_status() + task_data = response.json() + except (requests.RequestException, ValueError) as exc: + raise RuntimeError( + f"Geometry Generation Server task polling failed: {request_id}." + ) from exc + + if not isinstance(task_data, dict): + raise RuntimeError( + "Geometry Generation Server task response must be a JSON object." + ) + task_status = task_data.get("status") + if task_status == "succeeded": + return task_data + if task_status in {"failed", "cancelled"}: + raise RuntimeError( + "Geometry Generation Server task " + f"{task_status}: {task_data.get('error', 'unknown error')}" + ) + if not isinstance(task_status, str) or ( + task_status != "running" and "waiting" not in task_status + ): + raise RuntimeError( + "Geometry Generation Server returned unknown task status: " + f"{task_status!r}." + ) + + time.sleep(1) + + raise RuntimeError( + "Geometry Generation Server task timed out after 600 seconds: " + f"{request_id}." + ) + + def _download_glb(self, mesh_path: str, output_path: Path) -> None: + last_error: Exception | None = None + for _ in range(self._max_attempts): + try: + response = self._session.get( + self._mesh_url(mesh_path), + timeout=self._timeout_s, + ) + response.raise_for_status() + glb_bytes = response.content + if not glb_bytes.startswith(b"glTF"): + raise RuntimeError( + "Geometry Generation Server returned invalid GLB content." + ) + output_path.write_bytes(glb_bytes) + return + except (requests.RequestException, RuntimeError) as exc: + last_error = exc + + assert last_error is not None + raise RuntimeError( + "Geometry Generation Server GLB download failed after " + f"{self._max_attempts} attempts: {mesh_path}" + ) from last_error + + def _mesh_url(self, mesh_path: str) -> str: + if mesh_path.startswith(("http://", "https://")): + return mesh_path + return self._url(mesh_path) + + def _url(self, path: str) -> str: + return f"{self._base_url}/{path.lstrip('/')}" + + +def _parse_objects_response( + response_data: object, + *, + object_ids: list[str], +) -> list[dict[str, Any]]: + if not isinstance(response_data, dict): + raise RuntimeError("Geometry Generation Server response must be a JSON object.") + if response_data.get("ok") is not True: + raise RuntimeError( + "Geometry Generation Server request failed: " + f"{response_data.get('error', 'ok is not true')}" + ) + result = response_data.get("result") + if not isinstance(result, dict): + raise RuntimeError( + "Geometry Generation Server response must contain a result object." + ) + response_objects = result.get("objects") + if not isinstance(response_objects, list) or len(response_objects) != len( + object_ids + ): + raise RuntimeError( + "Geometry Generation Server response object count does not match masks." + ) + + parsed_objects: list[dict[str, Any]] = [] + for index, (object_id, response_object) in enumerate( + zip(object_ids, response_objects) + ): + if not isinstance(response_object, dict): + raise RuntimeError( + f"Geometry Generation Server object {index} must be a JSON object." + ) + if response_object.get("name") != object_id: + raise RuntimeError( + "Geometry Generation Server object name does not match its " + f"requested id: {object_id!r}." + ) + mesh_path = response_object.get("mesh") + if not isinstance(mesh_path, str) or not mesh_path: + raise RuntimeError( + f"Geometry Generation Server object {index} has no mesh path." + ) + parsed_objects.append( + { + "mesh": mesh_path, + "rotation_quaternion_wxyz": _parse_numeric_list( + response_object.get("rotation_quaternion_wxyz"), + expected_length=4, + field_name=f"objects[{index}].rotation_quaternion_wxyz", + ), + "translation": _parse_numeric_list( + response_object.get("translation"), + expected_length=3, + field_name=f"objects[{index}].translation", + ), + "scale": _parse_numeric_list( + response_object.get("scale"), + expected_length=3, + field_name=f"objects[{index}].scale", + ), + } + ) + return parsed_objects + + +def _parse_numeric_list( + value: object, + *, + expected_length: int, + field_name: str, +) -> list[float]: + if not isinstance(value, list) or len(value) != expected_length: + raise RuntimeError( + f"Geometry Generation Server response field {field_name} is invalid." + ) + try: + return [float(item) for item in value] + except (TypeError, ValueError) as exc: + raise RuntimeError( + f"Geometry Generation Server response field {field_name} must be numeric." + ) from exc + + +def _image_content_type(image_path: Path) -> str: + if image_path.suffix.lower() in {".jpg", ".jpeg"}: + return "image/jpeg" + return "image/png" + + +def _load_dotenv_config() -> dict[str, Any]: + values = read_scene_engine_env_values( + "SCENE_ENGINE_GEOMETRY_GENERATION_BASE_URL", + "SCENE_ENGINE_GEOMETRY_GENERATION_TIMEOUT_S", + "SCENE_ENGINE_GEOMETRY_GENERATION_MAX_ATTEMPTS", + "SCENE_ENGINE_GEOMETRY_GENERATION_HEALTH_PATH", + "SCENE_ENGINE_GEOMETRY_GENERATION_OBJECTS_PATH", + ) + try: + timeout_s = int(values["SCENE_ENGINE_GEOMETRY_GENERATION_TIMEOUT_S"]) + except (TypeError, ValueError) as exc: + raise ValueError( + "SCENE_ENGINE_GEOMETRY_GENERATION_TIMEOUT_S must be an integer." + ) from exc + if timeout_s < 1: + raise ValueError( + "SCENE_ENGINE_GEOMETRY_GENERATION_TIMEOUT_S must be at least 1." + ) + + try: + max_attempts = int(values["SCENE_ENGINE_GEOMETRY_GENERATION_MAX_ATTEMPTS"]) + except (TypeError, ValueError) as exc: + raise ValueError( + "SCENE_ENGINE_GEOMETRY_GENERATION_MAX_ATTEMPTS must be an integer." + ) from exc + if max_attempts < 1: + raise ValueError( + "SCENE_ENGINE_GEOMETRY_GENERATION_MAX_ATTEMPTS must be at least 1." + ) + + string_keys = ( + "SCENE_ENGINE_GEOMETRY_GENERATION_BASE_URL", + "SCENE_ENGINE_GEOMETRY_GENERATION_HEALTH_PATH", + "SCENE_ENGINE_GEOMETRY_GENERATION_OBJECTS_PATH", + ) + for key in string_keys: + if not values[key].strip(): + raise ValueError(f"Scene Engine .env key {key} must be a non-empty string.") + + return { + "base_url": values["SCENE_ENGINE_GEOMETRY_GENERATION_BASE_URL"].strip(), + "timeout_s": timeout_s, + "max_attempts": max_attempts, + "health_path": values["SCENE_ENGINE_GEOMETRY_GENERATION_HEALTH_PATH"].strip(), + "generate_objects_path": values[ + "SCENE_ENGINE_GEOMETRY_GENERATION_OBJECTS_PATH" + ].strip(), + } diff --git a/embodichain/gen_sim/scene_engine/clients/image_segmentation.py b/embodichain/gen_sim/scene_engine/clients/image_segmentation.py new file mode 100644 index 000000000..2c56af91a --- /dev/null +++ b/embodichain/gen_sim/scene_engine/clients/image_segmentation.py @@ -0,0 +1,212 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import requests + +from embodichain.gen_sim.scene_engine.configs.environment import ( + read_scene_engine_env_values, +) + + +class ImageSegmentationClient: + + def __init__( + self, + *, + base_url: str, + timeout_s: int, + max_attempts: int, + health_path: str, + segment_single_object_path: str, + session: requests.Session | None = None, + ) -> None: + self._base_url = base_url.rstrip("/") + self._timeout_s = timeout_s + self._max_attempts = max_attempts + self._health_path = health_path + self._segment_single_object_path = segment_single_object_path + self._session = session or requests.Session() + + @classmethod + def from_dotenv(cls) -> "ImageSegmentationClient": + """Create a client from its required ``gen_sim/.env`` settings.""" + return cls(**_load_dotenv_config()) + + def check_health(self) -> None: + last_error: requests.RequestException | None = None + for _ in range(self._max_attempts): + try: + response = self._session.get( + self._url(self._health_path), + timeout=self._timeout_s, + ) + response.raise_for_status() + return + except requests.RequestException as exc: + last_error = exc + + assert last_error is not None + raise RuntimeError( + "Image Segmentation Server health check failed after " + f"{self._max_attempts} attempts." + ) from last_error + + def close(self) -> None: + self._session.close() + + def segment_single_object( + self, + *, + image_path: str | Path, + prompt: str, + ) -> list[dict[str, Any]]: + """Segment one prompted concept and return its RLE masks. + The returned list contains only RLE dictionaries, one per mask. + """ + resolved_image_path = Path(image_path).expanduser().resolve() + if not resolved_image_path.is_file(): + raise FileNotFoundError( + f"Image segmentation input not found: {resolved_image_path}" + ) + prompt = prompt.strip() + if not prompt: + raise ValueError("Image segmentation prompt must not be empty.") + + last_error: Exception | None = None + for _ in range(self._max_attempts): + try: + with resolved_image_path.open("rb") as image_file: + response = self._session.post( + self._url(self._segment_single_object_path), + data={"prompt": prompt}, + files={"image": (resolved_image_path.name, image_file)}, + timeout=self._timeout_s, + ) + response.raise_for_status() + + try: + response_data = response.json() + except ValueError as exc: + raise RuntimeError( + "Image Segmentation Server response is not valid JSON." + ) from exc + if not isinstance(response_data, dict): + raise RuntimeError( + "Image Segmentation Server response must be a JSON object." + ) + if response_data.get("ok") is False: + raise RuntimeError( + "Image Segmentation Server request failed: " + f"{response_data.get('error', 'unknown error')}" + ) + return _extract_rle_masks(response_data) + except (requests.RequestException, RuntimeError) as exc: + last_error = exc + + assert last_error is not None + raise RuntimeError( + "Image Segmentation Server request failed after " + f"{self._max_attempts} attempts." + ) from last_error + + def _url(self, path: str) -> str: + return f"{self._base_url}/{path.lstrip('/')}" + + +def _load_dotenv_config() -> dict[str, Any]: + values = read_scene_engine_env_values( + "SCENE_ENGINE_IMAGE_SEGMENTATION_BASE_URL", + "SCENE_ENGINE_IMAGE_SEGMENTATION_TIMEOUT_S", + "SCENE_ENGINE_IMAGE_SEGMENTATION_MAX_ATTEMPTS", + "SCENE_ENGINE_IMAGE_SEGMENTATION_HEALTH_PATH", + "SCENE_ENGINE_IMAGE_SEGMENTATION_SINGLE_OBJECT_PATH", + ) + try: + timeout_s = int(values["SCENE_ENGINE_IMAGE_SEGMENTATION_TIMEOUT_S"]) + except (TypeError, ValueError) as exc: + raise ValueError( + "SCENE_ENGINE_IMAGE_SEGMENTATION_TIMEOUT_S must be an integer." + ) from exc + if timeout_s < 1: + raise ValueError( + "SCENE_ENGINE_IMAGE_SEGMENTATION_TIMEOUT_S must be at least 1." + ) + + try: + max_attempts = int(values["SCENE_ENGINE_IMAGE_SEGMENTATION_MAX_ATTEMPTS"]) + except (TypeError, ValueError) as exc: + raise ValueError( + "SCENE_ENGINE_IMAGE_SEGMENTATION_MAX_ATTEMPTS must be an integer." + ) from exc + if max_attempts < 1: + raise ValueError( + "SCENE_ENGINE_IMAGE_SEGMENTATION_MAX_ATTEMPTS must be at least 1." + ) + + string_keys = ( + "SCENE_ENGINE_IMAGE_SEGMENTATION_BASE_URL", + "SCENE_ENGINE_IMAGE_SEGMENTATION_HEALTH_PATH", + "SCENE_ENGINE_IMAGE_SEGMENTATION_SINGLE_OBJECT_PATH", + ) + for key in string_keys: + if not values[key].strip(): + raise ValueError(f"Scene Engine .env key {key} must be a non-empty string.") + + return { + "base_url": values["SCENE_ENGINE_IMAGE_SEGMENTATION_BASE_URL"].strip(), + "timeout_s": timeout_s, + "max_attempts": max_attempts, + "health_path": values["SCENE_ENGINE_IMAGE_SEGMENTATION_HEALTH_PATH"].strip(), + "segment_single_object_path": values[ + "SCENE_ENGINE_IMAGE_SEGMENTATION_SINGLE_OBJECT_PATH" + ].strip(), + } + + +def _extract_rle_masks(response_data: dict[str, Any]) -> list[dict[str, Any]]: + """Extract RLE masks from accepted Image Segmentation Server layouts.""" + result_data = response_data.get("result") or response_data.get("data") + if not isinstance(result_data, dict): + result_data = response_data + + masks = result_data.get("masks") + if isinstance(masks, list): + rle_masks = [mask for mask in masks if isinstance(mask, dict)] + if rle_masks: + return rle_masks + + instances = result_data.get("instances", []) + if isinstance(instances, list): + rle_masks: list[dict[str, Any]] = [] + for instance in instances: + if not isinstance(instance, dict): + continue + mask = ( + instance.get("mask_rle") + or instance.get("mask") + or instance.get("segmentation") + ) + if isinstance(mask, dict): + rle_masks.append(mask) + return rle_masks + + return [] diff --git a/embodichain/gen_sim/scene_engine/configs/__init__.py b/embodichain/gen_sim/scene_engine/configs/__init__.py new file mode 100644 index 000000000..015c41510 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/configs/__init__.py @@ -0,0 +1,19 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +__all__: list[str] = [] diff --git a/embodichain/gen_sim/scene_engine/configs/environment.py b/embodichain/gen_sim/scene_engine/configs/environment.py new file mode 100644 index 000000000..9b686d02f --- /dev/null +++ b/embodichain/gen_sim/scene_engine/configs/environment.py @@ -0,0 +1,50 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from pathlib import Path + +_SCENE_ENGINE_ENV_PATH = Path(__file__).resolve().parents[2] / ".env" + + +def read_scene_engine_env_values(*keys: str) -> dict[str, str]: + """Read only the requested Scene Engine settings from ``gen_sim/.env``.""" + if not _SCENE_ENGINE_ENV_PATH.is_file(): + raise FileNotFoundError( + f"Scene Engine .env file not found: {_SCENE_ENGINE_ENV_PATH}" + ) + + requested_keys = set(keys) + values: dict[str, str] = {} + for raw_line in _SCENE_ENGINE_ENV_PATH.read_text(encoding="utf-8").splitlines(): + line = raw_line.strip() + if not line or line.startswith("#") or "=" not in line: + continue + key, raw_value = line.split("=", maxsplit=1) + key = key.strip() + if key not in requested_keys: + continue + value = raw_value.strip() + if len(value) >= 2 and value[0] == value[-1] and value[0] in {'"', "'"}: + value = value[1:-1] + values[key] = value + + missing_keys = [key for key in keys if key not in values] + if missing_keys: + raise ValueError(f"Missing required Scene Engine .env keys: {missing_keys}") + return values diff --git a/embodichain/gen_sim/scene_engine/core/__init__.py b/embodichain/gen_sim/scene_engine/core/__init__.py new file mode 100644 index 000000000..015c41510 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/core/__init__.py @@ -0,0 +1,19 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +__all__: list[str] = [] diff --git a/embodichain/gen_sim/scene_engine/core/scene.py b/embodichain/gen_sim/scene_engine/core/scene.py new file mode 100644 index 000000000..07d937e41 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/core/scene.py @@ -0,0 +1,53 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from dataclasses import dataclass, field + +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject + + +@dataclass +class Scene: + """A scene containing one table object and zero or more asset objects.""" + + objects: list[SceneObject] = field(default_factory=list) + + @property + def table(self) -> SceneObject | None: + """Return the sole table object, or ``None`` before understanding.""" + tables = [ + scene_object + for scene_object in self.objects + if scene_object.kind == "table" + ] + if len(tables) > 1: + raise ValueError("A scene may contain only one table object.") + return tables[0] if tables else None + + @property + def assets(self) -> list[SceneObject]: + """Return movable asset objects in their scene order.""" + return [ + scene_object + for scene_object in self.objects + if scene_object.kind == "asset" + ] + + def to_dict(self) -> dict[str, object]: + """Serialize the canonical object collection for debugging artifacts.""" + return {"objects": [scene_object.to_dict() for scene_object in self.objects]} diff --git a/embodichain/gen_sim/scene_engine/core/scene_object.py b/embodichain/gen_sim/scene_engine/core/scene_object.py new file mode 100644 index 000000000..d9a837e87 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/core/scene_object.py @@ -0,0 +1,85 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Literal + + +@dataclass +class ObjectPhysics: + """Physics and collision settings shared by settling and scene export.""" + + body_type: Literal["dynamic", "kinematic"] # Runtime behaviour in simulation. + attrs: dict[str, float | int] # Rigid-body material and contact attributes. + max_convex_hull_num: int # Collision-decomposition hull budget. + + def __post_init__(self) -> None: + """Validate physics settings before a later stage consumes them.""" + if self.body_type not in {"dynamic", "kinematic"}: + raise ValueError("body_type must be 'dynamic' or 'kinematic'.") + if self.max_convex_hull_num <= 0: + raise ValueError("max_convex_hull_num must be positive.") + if not self.attrs: + raise ValueError("attrs must contain at least one physics attribute.") + if not all( + isinstance(name, str) and isinstance(value, (float, int)) + for name, value in self.attrs.items() + ): + raise ValueError("attrs must map strings to numeric physics values.") + + def to_dict(self) -> dict[str, object]: + """Serialize the physics settings for scene debugging artifacts.""" + return { + "body_type": self.body_type, + "attrs": self.attrs, + "max_convex_hull_num": self.max_convex_hull_num, + } + + +@dataclass +class SceneObject: + """One semantic object progressing through the Scene Engine pipeline.""" + + id: str # Stable scene-unique identifier. + kind: Literal["table", "asset"] # Table support body or movable scene asset. + category: str # Semantic category identified by scene understanding. + name: str # Human-readable visual name. + description: str # Detailed semantic and spatial description. + mask_path: str | None = None # Absolute path to the validated binary image mask. + simready_glb_path: str | None = None # Absolute path to the canonical SimReady GLB. + rot: list[float] | None = None # Final y-up Euler XYZ rotation in degrees. + pos: list[float] | None = None # Final y-up world position in metres. + scale: list[float] | None = None # Final y-up object scale. + physics: ObjectPhysics | None = None # Assigned when SimReady processing succeeds. + + def to_dict(self) -> dict[str, object]: + """Serialize this object and its currently available pipeline artifacts.""" + return { + "id": self.id, + "kind": self.kind, + "category": self.category, + "name": self.name, + "description": self.description, + "mask_path": self.mask_path, + "simready_glb_path": self.simready_glb_path, + "rot": self.rot, + "pos": self.pos, + "scale": self.scale, + "physics": self.physics.to_dict() if self.physics is not None else None, + } diff --git a/embodichain/gen_sim/scene_engine/llms/__init__.py b/embodichain/gen_sim/scene_engine/llms/__init__.py new file mode 100644 index 000000000..015c41510 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/llms/__init__.py @@ -0,0 +1,19 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +__all__: list[str] = [] diff --git a/embodichain/gen_sim/scene_engine/llms/load_config.py b/embodichain/gen_sim/scene_engine/llms/load_config.py new file mode 100644 index 000000000..8a2af22d5 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/llms/load_config.py @@ -0,0 +1,83 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from dataclasses import dataclass +import json +from typing import Any + +from embodichain.gen_sim.scene_engine.configs.environment import ( + read_scene_engine_env_values, +) + + +@dataclass(frozen=True) +class LLMConfig: + """OpenAI-compatible VLM connection settings.""" + + api_key: str + model: str + base_url: str + default_query: dict[str, Any] + max_attempts: int + + +def load_llm_config() -> LLMConfig: + """Load the required OpenAI-compatible LLM settings from ``gen_sim/.env``.""" + values = read_scene_engine_env_values( + "OPENAI_API_KEY", + "OPENAI_MODEL", + "OPENAI_BASE_URL", + "SCENE_ENGINE_OPENAI_DEFAULT_QUERY", + "OPENAI_MAX_ATTEMPTS", + ) + try: + default_query = json.loads(values["SCENE_ENGINE_OPENAI_DEFAULT_QUERY"]) + except json.JSONDecodeError as exc: + raise ValueError( + "SCENE_ENGINE_OPENAI_DEFAULT_QUERY must contain a JSON object." + ) from exc + + if not isinstance(default_query, dict): + raise ValueError("SCENE_ENGINE_OPENAI_DEFAULT_QUERY must be a JSON object.") + missing = [ + key + for key, value in { + "OPENAI_API_KEY": values["OPENAI_API_KEY"], + "OPENAI_MODEL": values["OPENAI_MODEL"], + "OPENAI_BASE_URL": values["OPENAI_BASE_URL"], + }.items() + if not value.strip() + ] + if missing: + raise ValueError(f"Missing required LLM config keys: {missing}") + + try: + parsed_max_attempts = int(values["OPENAI_MAX_ATTEMPTS"]) + except (TypeError, ValueError) as exc: + raise ValueError("OPENAI_MAX_ATTEMPTS must be an integer.") from exc + if parsed_max_attempts < 1: + raise ValueError("OPENAI_MAX_ATTEMPTS must be at least 1.") + + return LLMConfig( + api_key=values["OPENAI_API_KEY"].strip(), + model=values["OPENAI_MODEL"].strip(), + base_url=values["OPENAI_BASE_URL"].rstrip("/"), + default_query=default_query, + max_attempts=parsed_max_attempts, + ) diff --git a/embodichain/gen_sim/scene_engine/llms/openai_compatible_client.py b/embodichain/gen_sim/scene_engine/llms/openai_compatible_client.py new file mode 100644 index 000000000..f83e316aa --- /dev/null +++ b/embodichain/gen_sim/scene_engine/llms/openai_compatible_client.py @@ -0,0 +1,139 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +import base64 +import json +from pathlib import Path +from typing import Any +from urllib.error import HTTPError, URLError +from urllib.request import Request, urlopen + +from embodichain.gen_sim.scene_engine.llms.load_config import LLMConfig, load_llm_config + +_SUPPORTED_IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png"} + + +class OpenAICompatibleVLM: + """Client for multimodal OpenAI-compatible chat-completions endpoints.""" + + def __init__(self, config: LLMConfig): + self._config = config + + @classmethod + def from_dotenv(cls) -> "OpenAICompatibleVLM": + """Create a client from its required ``gen_sim/.env`` settings.""" + return cls(load_llm_config()) + + def complete( + self, + *, + system_prompt: str, + user_prompt: str, + image_path: str | Path | None = None, + ) -> str: + """Send a text or text-and-image chat-completions request.""" + user_content: str | list[dict[str, object]] = user_prompt + if image_path is not None: + resolved_image_path = Path(image_path).expanduser().resolve() + if not resolved_image_path.is_file(): + raise FileNotFoundError(f"Image input not found: {resolved_image_path}") + if resolved_image_path.suffix.lower() not in _SUPPORTED_IMAGE_SUFFIXES: + raise ValueError("Image input must be a .jpg, .jpeg, or .png file.") + user_content = [ + {"type": "text", "text": user_prompt}, + { + "type": "image_url", + "image_url": {"url": _image_data_url(resolved_image_path)}, + }, + ] + + payload = dict(self._config.default_query) + payload.update( + { + "model": self._config.model, + "messages": [ + {"role": "system", "content": system_prompt}, + { + "role": "user", + "content": user_content, + }, + ], + } + ) + return self._request_chat_completion(payload) + + def _request_chat_completion(self, payload: dict[str, Any]) -> str: + """Execute a chat-completions HTTP request with transient retries.""" + endpoint = _chat_completions_endpoint(self._config.base_url) + last_error: Exception | None = None + + for attempt in range(1, self._config.max_attempts + 1): + try: + request = Request( + endpoint, + data=json.dumps(payload).encode("utf-8"), + headers={ + "Authorization": f"Bearer {self._config.api_key}", + "Content-Type": "application/json", + }, + method="POST", + ) + with urlopen(request, timeout=120) as response: + response_payload = json.loads(response.read().decode("utf-8")) + return _extract_response_text(response_payload) + except HTTPError as exc: + details = exc.read().decode("utf-8", errors="replace") + last_error = RuntimeError( + f"VLM request failed with HTTP {exc.code}: {details}" + ) + except URLError as exc: + last_error = RuntimeError(f"VLM request failed: {exc.reason}") + except (TimeoutError, OSError) as exc: + last_error = RuntimeError(f"VLM request failed: {exc}") + except (json.JSONDecodeError, ValueError): + last_error = RuntimeError("VLM API returned a malformed response.") + + assert last_error is not None + raise last_error + + +def _image_data_url(image_path: Path) -> str: + mime_type = ( + "image/jpeg" if image_path.suffix.lower() in {".jpg", ".jpeg"} else "image/png" + ) + encoded_image = base64.b64encode(image_path.read_bytes()).decode("ascii") + return f"data:{mime_type};base64,{encoded_image}" + + +def _chat_completions_endpoint(base_url: str) -> str: + if base_url.endswith("/chat/completions"): + return base_url + return f"{base_url}/chat/completions" + + +def _extract_response_text(response_payload: object) -> str: + try: + content = response_payload["choices"][0]["message"]["content"] + except (KeyError, IndexError, TypeError) as exc: + raise ValueError( + "VLM response does not contain choices[0].message.content." + ) from exc + if not isinstance(content, str): + raise ValueError("VLM response content must be a string.") + return content diff --git a/embodichain/gen_sim/scene_engine/pipeline/__init__.py b/embodichain/gen_sim/scene_engine/pipeline/__init__.py new file mode 100644 index 000000000..015c41510 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/__init__.py @@ -0,0 +1,19 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +__all__: list[str] = [] diff --git a/embodichain/gen_sim/scene_engine/pipeline/generate.py b/embodichain/gen_sim/scene_engine/pipeline/generate.py new file mode 100644 index 000000000..773a897bf --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/generate.py @@ -0,0 +1,87 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from pathlib import Path + +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( + OpenAICompatibleVLM, +) +from embodichain.gen_sim.scene_engine.clients.geometry_generation import ( + GeometryGenerationClient, +) + +from embodichain.gen_sim.scene_engine.pipeline.scene_understanding import ( + understand_scene, +) +from embodichain.utils.logger import log_info + +from embodichain.gen_sim.scene_engine.pipeline.scene_generation import ( + generate_scene_and_refine, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_exporter import SceneExporter + + +def generate_scene_from_image( + image_path: str | Path, + output_root: str | Path, +) -> Scene: + """Generate the initial core scene state from an input image.""" + resolved_output_root = Path(output_root).expanduser().resolve() + resolved_output_root.mkdir(parents=True, exist_ok=True) + + # Initialize the VLM client and the Scene data structure. + vlm_client = OpenAICompatibleVLM.from_dotenv() + scene = Scene() + + # 1. Scene Understanding + log_info("Starting Scene Understanding") + scene = understand_scene( + scene=scene, + image_path=image_path, + output_root=resolved_output_root, + vlm_client=vlm_client, + ) + log_info("Completed Scene Understanding") + + # 2. Objects + Coarse Layout Generation + log_info("Starting Objects + Coarse Layout Generation") + # Load .env settings and fail if the Geometry Generation Server is unavailable. + geometry_generation_client = GeometryGenerationClient.from_dotenv() + try: + geometry_generation_client.check_health() # Error raising will happen internally. + scene = generate_scene_and_refine( + image_path=image_path, + output_root=resolved_output_root, + scene=scene, + geometry_generation_client=geometry_generation_client, + ) + finally: + geometry_generation_client.close() # Kill the session to avoid resource leaks. + log_info("Completed Objects + Coarse Layout Generation") + + # 3. Scene Export + log_info("Starting Scene Export") + scene_exporter = SceneExporter( + scene=scene, + output_root=resolved_output_root, + ) + scene_exporter.export() + log_info("Completed Scene Export") + + return scene diff --git a/embodichain/gen_sim/scene_engine/pipeline/scene_generation.py b/embodichain/gen_sim/scene_engine/pipeline/scene_generation.py new file mode 100644 index 000000000..d14ce364f --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/scene_generation.py @@ -0,0 +1,509 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +import json +from pathlib import Path +import shutil + +import numpy as np +import trimesh + +from embodichain.gen_sim.scene_engine.clients.geometry_generation import ( + GeometryGenerationClient, +) +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.pipeline.utils.assets_group_support_clamp import ( + AssetsGroupSupportClamp, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.assets_group_table_aligner import ( + AssetsGroupTableAligner, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.assets_group_layout_optimizer import ( + AssetsSupportLayoutOptimizer, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.assets_gravity_settler import ( + AssetsGravitySettler, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_generation_utils import ( + layout_object_to_transform_matrix, + load_glb_mesh, + quaternion_wxyz_to_euler_xyz_degrees, + transform_matrix_to_layout_object, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.simready_scene_processor import ( + SimReadySceneProcessor, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.table_support_surface import ( + TableSupportSurfaceDetector, +) +from embodichain.utils.logger import log_info + +_SUPPORTED_IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png"} + + +def generate_scene_and_refine( + image_path: str | Path, + output_root: str | Path, + scene: Scene, + *, + geometry_generation_client: GeometryGenerationClient, +) -> Scene: + + resolved_image_path = _validate_image_path(image_path) + # Create stage output directory. + stage_output_root = Path(output_root).expanduser().resolve() / "scene_generation" + if stage_output_root.exists(): + shutil.rmtree(stage_output_root) + stage_output_root.mkdir(parents=True, exist_ok=True) + # Create debug folder and the sim-ready geometry folder. + debug_output_root = ( + stage_output_root / "debug" + ) # Keeps the other files for debugging. + coarse_geometry_output_root = ( + stage_output_root / "coarse_geometry" + ) # Keeps the coarse geometries. + simready_geometry_output_root = ( + stage_output_root / "simready_geometry" + ) # Keeps the final-used geometries. + debug_output_root.mkdir() + coarse_geometry_output_root.mkdir() + simready_geometry_output_root.mkdir() + + # Coarse geometry generation and coarse layout generation. + _generate_coarse_results_from_masks( + image_path=resolved_image_path, + debug_output_root=debug_output_root, + coarse_geometry_output_root=coarse_geometry_output_root, + scene=scene, # Use the masks which are kept in the scene data structure. + geometry_generation_client=geometry_generation_client, + ) + + # Simready all the assets(includes table). + # Treat table and assets seperately. + coarse_layout = _load_layout(coarse_geometry_output_root / "coarse_layout.json") + coarse_layout_by_id = { + layout_object["id"]: layout_object for layout_object in coarse_layout + } + simready_processor = SimReadySceneProcessor( + scene=scene, + coarse_layout_by_id=coarse_layout_by_id, + coarse_geometry_root=coarse_geometry_output_root, + simready_geometry_root=simready_geometry_output_root, + ) + simready_assets_layout = simready_processor.process_assets() + simready_table_layout = simready_processor.process_table() + # Concat then save the table info and the assets info in one JSON file. + simready_layout = [simready_table_layout, *simready_assets_layout] + (simready_geometry_output_root / "simready_layout.json").write_text( + json.dumps(simready_layout, indent=2, ensure_ascii=False) + "\n", + encoding="utf-8", + ) + + # Layout refinement will start with the table. + refined_table_layout, refined_assets_layout = _layout_refinement( + scene=scene, # Update this data structure internally. + simready_geometry_output_root=simready_geometry_output_root, # Contains simready assets and their current coarse layout JSON. + debug_output_root=debug_output_root, # Keep the table support surface info + optimized layout info (render with matplotlib) for debugging. + ) + + # Write the Updated scene JSON for debugging. + (stage_output_root / "scene.json").write_text( + json.dumps(scene.to_dict(), indent=2, ensure_ascii=False) + "\n", + encoding="utf-8", + ) + return scene + + +def _generate_coarse_results_from_masks( + image_path: str | Path, + debug_output_root: str | Path, + coarse_geometry_output_root: str | Path, + scene: Scene, + *, + geometry_generation_client: GeometryGenerationClient, +) -> None: + + # Parse whether the scene has each assets' binary masks. + # The original image has already been validated. + # The table must exist, for it is the base of the scene. + if scene.table is None: + raise ValueError("Scene must contain a table before geometry generation.") + + scene_objects = [scene.table, *scene.assets] + object_masks: list[tuple[str, Path]] = [] + for scene_object in scene_objects: + if scene_object.mask_path is None: + raise ValueError( + f"Scene object {scene_object.id!r} has no binary mask path." + ) + mask_path = Path(scene_object.mask_path).expanduser().resolve() + if not mask_path.is_file(): + raise FileNotFoundError( + f"Binary mask for scene object {scene_object.id!r} not found: " + f"{mask_path}" + ) + object_masks.append( + (scene_object.id, mask_path) + ) # id + mask, for avoiding the download glbs order confusion. + + # Sent the request, wait, then save the intermediate results. + response_data, response_objects = geometry_generation_client.generate_objects( + image_path=image_path, + object_masks=object_masks, + output_root=coarse_geometry_output_root, # Keep the coarse geometries + ) + # Write the response JSON which contains all the layout info the server gave us. + # Keep original response for getting the sam3d coarse layout matrix. + (Path(debug_output_root) / "geometry_generation_response.json").write_text( + json.dumps(response_data, indent=2, ensure_ascii=False) + "\n", + encoding="utf-8", + ) + # Write the coarse layout JSON as one of the results in this step. + coarse_layout = [ + { + "id": object_id, + "rot": quaternion_wxyz_to_euler_xyz_degrees( + response_object["rotation_quaternion_wxyz"] + ), + "pos": response_object["translation"], + "scale": response_object["scale"], + } + for (object_id, _), response_object in zip(object_masks, response_objects) + ] + (Path(coarse_geometry_output_root) / "coarse_layout.json").write_text( + json.dumps(coarse_layout, indent=2, ensure_ascii=False) + "\n", + encoding="utf-8", + ) + # Nothing to be returned. + return None + + +def _update_scene_final_y_up_layout( + *, + scene: Scene, + table_layout: dict[str, object], + assets_layout: list[dict[str, object]], +) -> None: + """Copy final y-up layout values into the matching table and asset objects.""" + if scene.table is None: + raise ValueError("Cannot update a final layout without a table.") + + _copy_y_up_layout_to_scene_object(scene.table, table_layout) + assets_by_id = {asset.id: asset for asset in scene.assets} + layout_ids = set() + for asset_layout in assets_layout: + asset_id = asset_layout.get("id") + if not isinstance(asset_id, str) or asset_id not in assets_by_id: + raise ValueError(f"Final layout contains unknown asset {asset_id!r}.") + if asset_id in layout_ids: + raise ValueError(f"Final layout contains duplicate asset {asset_id!r}.") + _copy_y_up_layout_to_scene_object(assets_by_id[asset_id], asset_layout) + layout_ids.add(asset_id) + + missing_assets = set(assets_by_id) - layout_ids + if missing_assets: + raise ValueError( + f"Final layout is missing scene assets: {sorted(missing_assets)}." + ) + + +def _copy_y_up_layout_to_scene_object( + scene_object: SceneObject, + layout_object: dict[str, object], +) -> None: + """Copy one y-up layout object after validating its id and numeric vectors.""" + if layout_object.get("id") != scene_object.id: + raise ValueError( + f"Layout id {layout_object.get('id')!r} does not match scene object " + f"{scene_object.id!r}." + ) + + for field_name in ("rot", "pos", "scale"): + values = layout_object.get(field_name) + if not isinstance(values, (list, tuple)) or len(values) != 3: + raise ValueError( + f"Layout object {scene_object.id!r} has invalid {field_name!r}." + ) + vector = [float(value) for value in values] + if not np.all(np.isfinite(vector)): + raise ValueError( + f"Layout object {scene_object.id!r} has non-finite {field_name!r}." + ) + setattr(scene_object, field_name, vector) + + +def _layout_refinement( + *, + scene: Scene, + simready_geometry_output_root: str | Path, + debug_output_root: str | Path, +) -> tuple[dict[str, object], list[dict[str, object]]]: + + # 1. All layouts and geometries below are SimReady outputs. Do not mix a + # coarse layout with a SimReady GLB (or vice versa), because each object's + # SimReady canonicalization may include its own local pose compensation. + simready_layout = _load_layout( + Path(simready_geometry_output_root) / "simready_layout.json" + ) + if scene.table is None: + raise ValueError("Cannot refine a layout without a table.") + table_id = scene.table.id + table_layout = next( + ( + layout_object + for layout_object in simready_layout + if layout_object["id"] == table_id + ), + None, + ) + if table_layout is None: + raise ValueError(f"SimReady layout does not contain table {table_id!r}.") + + # Keep the intermediate layout y-up; the simulator converts final GLBs to + # z-up. Left multiplication expresses every complete asset pose (position + # and rotation) in the SimReady table frame. + simready_table_to_world_matrix = layout_object_to_transform_matrix(table_layout) + world_to_simready_table_matrix = np.linalg.inv(simready_table_to_world_matrix) + + # 2. The table defines the refined world frame, so its transform is exact + # identity instead of a numerically reconstructed inverse(table) @ table. + refined_table_layout = transform_matrix_to_layout_object( + table_layout["id"], + np.eye(4), + ) + refined_assets_layout: list[dict[str, object]] = [] + for asset_layout in simready_layout: + if asset_layout["id"] == table_layout["id"]: + continue + + simready_asset_to_world_matrix = layout_object_to_transform_matrix(asset_layout) + simready_asset_to_table_matrix = ( + world_to_simready_table_matrix @ simready_asset_to_world_matrix + ) + + # Converting an asset back through the table pose must reconstruct its + # original SimReady world pose. This catches missing rotations, wrong + # matrix order, and coarse/SimReady coordinate-system mixing early. + if not np.allclose( + simready_table_to_world_matrix @ simready_asset_to_table_matrix, + simready_asset_to_world_matrix, + atol=1e-6, + ): + raise ValueError( + "SimReady table-frame conversion failed for asset " + f"{asset_layout['id']!r}." + ) + + refined_assets_layout.append( + transform_matrix_to_layout_object( + asset_layout["id"], + simready_asset_to_table_matrix, + ) + ) + + # 3. Move all assets as one rigid group so its lowest AABB point is 2cm above + # the table. This preserves the initial relative poses for the later + # gravity simulation, which can settle individual assets physically. + + group_table_aligner = AssetsGroupTableAligner( + table_layout=refined_table_layout, + assets_layout=refined_assets_layout, + geometry_root=simready_geometry_output_root, + ) + refined_table_layout, refined_assets_layout = group_table_aligner.align() + if not refined_assets_layout: + log_info("Scene has no movable assets; skipping support-region clamping.") + return refined_table_layout, [] + + # 4. Detect the actual upward support triangles instead of projecting the + # entire table mesh to one convex hull. The result retains concavities + # (for example, an L-shaped tabletop) and is the only boundary used for + # placement below. + ( + table_world_mesh_z_up, + assets_aabb_2d_z_up_world_corners_by_id, + ) = _measure_table_and_assets_in_z_up_world( + table_layout=refined_table_layout, + assets_layout=refined_assets_layout, + geometry_root=simready_geometry_output_root, + ) + support_detector = TableSupportSurfaceDetector( + table_world_mesh=table_world_mesh_z_up, + debug_output_root=debug_output_root, + ) + table_support_region = support_detector.detect() + support_detector.save_support_surface_debug_images() + + # 5. Keep the complete clutter rigid in the table plane. A successful + # result applies one shared z-up XY delta to every AABB, so it preserves + # all existing asset-to-asset relations. It is *not* an asset packing + # pass: pre-existing overlap is deliberately left to a later optimizer. + group_clamp = AssetsGroupSupportClamp( + support_region=table_support_region.support_polygon, + assets_aabb_2d_z_up_world_corners_by_id=( + assets_aabb_2d_z_up_world_corners_by_id + ), + assets_layout=refined_assets_layout, + debug_output_root=debug_output_root, + ) + refined_assets_layout = group_clamp.clamp() + group_clamp.save_group_clamp_debug_images() + + # The clamp returns y-up layouts; measure their resulting z-up AABBs again + # so the following independent optimizer consumes the same world-frame + # geometry as every other stage. + _, clamped_assets_aabb_2d_z_up_world_corners_by_id = ( + _measure_table_and_assets_in_z_up_world( + table_layout=refined_table_layout, + assets_layout=refined_assets_layout, + geometry_root=simready_geometry_output_root, + ) + ) + + # 6. Restore the previous pairwise AABB separation stage, but constrain + # every candidate with the actual support polygon rather than the legacy + # largest internal rectangle. Assets may now move independently only as + # much as needed to remove overlap; every resulting AABB remains on the + # L-shaped, circular, or otherwise non-convex support region. + overlap_optimizer = AssetsSupportLayoutOptimizer( + support_region=table_support_region.support_polygon, + assets_aabb_2d_z_up_world_corners_by_id=( + clamped_assets_aabb_2d_z_up_world_corners_by_id + ), + assets_layout=refined_assets_layout, + debug_output_root=debug_output_root, + ) + # Render this stage separately from the rigid group clamp. The latter + # intentionally preserves pre-existing overlaps, while this figure shows + # whether independent AABB separation actually resolved them. + refined_assets_layout = overlap_optimizer.optimize() + overlap_optimizer.save_overlap_optimization_debug_images() + + # 7. Gravity simulation, to let all the assets to be stable and placed well on the table's support surface. + # Notice that: we do not consider the assets like a bottle, which should be standing on the table but laid down + # after the simulation. + gravity_settler = AssetsGravitySettler( + scene=scene, + table_layout=refined_table_layout, + assets_layout=refined_assets_layout, + geometry_root=simready_geometry_output_root, + ) + refined_assets_layout = gravity_settler.settle() + + # Update the scene data structure with the final y-up layout values. + _update_scene_final_y_up_layout( + scene=scene, + table_layout=refined_table_layout, + assets_layout=refined_assets_layout, + ) + return refined_table_layout, refined_assets_layout + + +def _measure_table_and_assets_in_z_up_world( + *, + table_layout: dict[str, object], + assets_layout: list[dict[str, object]], + geometry_root: str | Path, +) -> tuple[trimesh.Trimesh, dict[str, np.ndarray]]: + """Measure a table mesh and asset AABBs in one shared z-up world frame. + + Scene layouts and SimReady GLBs are y-up. The support detector and the + group clamp both operate in z-up world XY, so this conversion is performed + once here and the exact same measured AABBs are passed to the clamp. + """ + table_id = table_layout.get("id") + if not isinstance(table_id, str) or not table_id: + raise ValueError("Table layout must contain a non-empty string id.") + + y_up_to_z_up_matrix = np.eye(4) + y_up_to_z_up_matrix[:3, :3] = np.array( + [ + [1.0, 0.0, 0.0], + [0.0, 0.0, -1.0], + [0.0, 1.0, 0.0], + ] + ) + z_up_to_y_up_matrix = np.linalg.inv(y_up_to_z_up_matrix) + resolved_geometry_root = Path(geometry_root).expanduser().resolve() + + def _mesh_in_z_up_world(layout_object: dict[str, object]) -> trimesh.Trimesh: + object_id = layout_object.get("id") + if not isinstance(object_id, str) or not object_id: + raise ValueError("Layout object must contain a non-empty string id.") + mesh = load_glb_mesh(resolved_geometry_root / f"{object_id}.glb") + z_up_layout = transform_matrix_to_layout_object( + object_id, + y_up_to_z_up_matrix + @ layout_object_to_transform_matrix(layout_object) + @ z_up_to_y_up_matrix, + ) + mesh.apply_transform(y_up_to_z_up_matrix) + mesh.apply_transform(layout_object_to_transform_matrix(z_up_layout)) + return mesh + + table_world_mesh_z_up = _mesh_in_z_up_world(table_layout) + asset_aabbs_by_id: dict[str, np.ndarray] = {} + for asset_layout in assets_layout: + asset_id = asset_layout.get("id") + if not isinstance(asset_id, str) or not asset_id: + raise ValueError("Each asset layout must contain a non-empty string id.") + if asset_id in asset_aabbs_by_id: + raise ValueError(f"Asset layouts contain duplicate id {asset_id!r}.") + asset_bounds_xy = _mesh_in_z_up_world(asset_layout).bounds[:, :2] + asset_aabbs_by_id[asset_id] = np.array( + [ + [asset_bounds_xy[0, 0], asset_bounds_xy[0, 1]], + [asset_bounds_xy[1, 0], asset_bounds_xy[0, 1]], + [asset_bounds_xy[1, 0], asset_bounds_xy[1, 1]], + [asset_bounds_xy[0, 0], asset_bounds_xy[1, 1]], + ], + dtype=float, + ) + return table_world_mesh_z_up, asset_aabbs_by_id + + +def _load_layout(layout_path: str | Path) -> list[dict[str, object]]: + # Load and check the coarse layout JSON file. + resolved_layout_path = Path(layout_path).expanduser().resolve() + if not resolved_layout_path.is_file(): + raise FileNotFoundError(f"Layout not found: {resolved_layout_path}") + try: + layout = json.loads(resolved_layout_path.read_text(encoding="utf-8")) + except json.JSONDecodeError as exc: + raise ValueError(f"Layout is not valid JSON: {resolved_layout_path}") from exc + if not isinstance(layout, list) or not all( + isinstance(item, dict) for item in layout + ): + raise ValueError("Layout must be a JSON array of objects.") + for layout_object in layout: + if not isinstance(layout_object.get("id"), str): + raise ValueError("Each layout object must have a string id.") + return layout + + +def _validate_image_path(image_path: str | Path) -> Path: + resolved_image_path = Path(image_path).expanduser().resolve() + if not resolved_image_path.is_file(): + raise FileNotFoundError(f"Image input not found: {resolved_image_path}") + if resolved_image_path.suffix.lower() not in _SUPPORTED_IMAGE_SUFFIXES: + raise ValueError( + f"Image input must be one of the supported formats: {_SUPPORTED_IMAGE_SUFFIXES}." + ) + return resolved_image_path diff --git a/embodichain/gen_sim/scene_engine/pipeline/scene_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/scene_understanding.py new file mode 100644 index 000000000..f91c665e6 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/scene_understanding.py @@ -0,0 +1,712 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +import json +from pathlib import Path +import re +import shutil +from typing import Any + +from PIL import Image + +from embodichain.gen_sim.scene_engine.clients.image_segmentation import ( + ImageSegmentationClient, +) +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( + OpenAICompatibleVLM, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.image_segmentation_utils import ( + MaskCandidate, + build_mask_candidates, + render_image_without_masks, + render_numbered_mask_candidates, + save_binary_mask, + union_overlapping_mask_candidates, +) + +_SUPPORTED_IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png"} +_CATEGORY_PATTERN = re.compile(r"^[a-z][a-z0-9_]*$") +_LOCATION_WORD_PATTERN = re.compile( + r"\b(?:left|right|front|back|center|middle|top|bottom|upper|lower|" + r"foreground|background|near|next|beside|behind|between|on|in|inside|" + r"under|above|below|against)\b", + flags=re.IGNORECASE, +) +_SYSTEM_PROMPT = """You inspect one tabletop-scene image. +Identify the main table and every visible, physically distinct object that should +be segmented and later generated as an independent 3D asset. + +Rules: +1. Ignore people, floor, carpet, walls, ceiling, doors, tiny incidental items, + and objects cut off by the image border. +2. Merge visually or functionally unified units, such as a potted plant, a vase + with flowers, or one built-in cabinet system. +3. Do not merge objects merely resting on another object. A mug on a table and + the table are separate entries. +4. List every visible physical instance separately. If two objects look alike, + keep the same category and name, but distinguish them in description using + location. Do not add location to name. +5. category is a lower-case singular snake_case class, such as mug, book, + potted_plant, or coffee_table. It must not contain color or material. +6. name contains only color, material, texture, shape, and object description. + It must not contain position or relations, such as left, right, on, in, or + near. +7. For table, description contains only its category, material, color, texture, + shape, and visible structural details. Do not mention image coverage, image + position, camera framing, or viewpoint. For example, do not write "occupying + most of the image" or "at the center of the image". +8. For assets, description may include all visible details, including location + and spatial context. + +Return JSON only: no Markdown, comments, or prose outside this exact schema: +{ + "table": { + "category": "coffee_table", + "name": "light wood coffee table", + "description": "low rectangular light wood coffee table with a smooth wood surface" + }, + "assets": [ + { + "category": "mug", + "name": "blue ceramic mug", + "description": "small blue ceramic mug on the left side of the table" + } + ] +} +For two identical blue mugs, output two asset entries with the same category and +name, and use their descriptions to state left/right or front/back. Do not +infer objects that are not visible. Use an empty assets array when no objects +are visible. Every field must be a non-empty string.""" + +_USER_PROMPT = "Analyze the provided image and return only the required JSON object." + +_TABLE_VALIDATION_SYSTEM_PROMPT = """You select the best table mask candidate. +The image contains table-mask candidates overlaid semi-transparently on the +scene. Gray regions are already-segmented non-table assets that were +intentionally removed for this validation; ignore them. Candidate numbers only +identify masks; do not treat the number or its background as scene content. + +Choose the candidate covering the main visible table. A table candidate is +acceptable when it covers the visible tabletop and/or legs, even if some edges +are incomplete, objects on the table occlude parts of it, or it slightly +overlaps those objects. Return null only when no candidate depicts the main +table. If there is one plausible candidate, select it rather than returning +null. + +Examples: +- Candidate 1 covers the tabletop and legs but misses a narrow edge: + {"selected_mask_index": 1} +- Candidate 1 is a cup and candidate 2 covers the main table: + {"selected_mask_index": 2} +- Every candidate is an object resting on the table, not the table itself: + {"selected_mask_index": null} + +Return JSON only, with exactly one key: selected_mask_index. Use a one-based +candidate index or null. Do not include Markdown or any other text.""" +_ASSET_ASSIGNMENT_SYSTEM_PROMPT = """You assign outlined mask candidates to a group of scene assets. +The image is the original scene with numbered candidate mask outlines. The +number labels identify candidates only; they are not scene content. Use the +provided category, name, and description of every asset to match each asset to +exactly one candidate. Descriptions can distinguish visually similar assets by +location. + +Extra candidate masks are normal and may be ignored. Never force a candidate +onto an asset. If any listed asset has no correct candidate, return +{"assignments": null}. + +Examples: +- Two listed paper cups match candidate 1 and candidate 3: + {"assignments": [{"asset_id": "paper_cup_001", "mask_index": 1}, {"asset_id": "paper_cup_002", "mask_index": 3}]} +- A listed asset is absent from every candidate: + {"assignments": null} + +Return JSON only, with exactly one key: assignments. It must be null or an +array of asset_id and mask_index objects. Do not include Markdown or any other +text.""" + + +def understand_scene( + scene: Scene, + image_path: str | Path, + output_root: str | Path, + *, + vlm_client: OpenAICompatibleVLM, + json_max_attempts: int = 3, +) -> Scene: + + resolved_image_path = _validate_image_path(image_path) + # The output in this stage will keep a JSON which contains + # the Scene data structure for debugging. + stage_output_root = Path(output_root).expanduser().resolve() / "scene_understanding" + if stage_output_root.exists(): + shutil.rmtree(stage_output_root) + stage_output_root.mkdir(parents=True, exist_ok=True) + + _analyze_image_objects( # Update the scene data structure internally. + scene=scene, + image_path=resolved_image_path, + vlm_client=vlm_client, + json_max_attempts=json_max_attempts, + ) + + # Load .env settings and fail if the Image Segmentation Server is unavailable. + image_segmentation_client = ImageSegmentationClient.from_dotenv() + try: + image_segmentation_client.check_health() # Error raising will happen internally. + _segment_scene( + image_path=resolved_image_path, + stage_output_root=stage_output_root, + scene=scene, + vlm_client=vlm_client, + image_segmentation_client=image_segmentation_client, + ) + finally: + image_segmentation_client.close() # Kill the session to avoid resource leaks. + + # Write the Updated scene JSON for debugging. + (stage_output_root / "scene.json").write_text( + json.dumps(scene.to_dict(), indent=2, ensure_ascii=False) + "\n", + encoding="utf-8", + ) + return scene + + +def _analyze_image_objects( + *, + scene: Scene, + image_path: str | Path, + vlm_client: OpenAICompatibleVLM, + json_max_attempts: int = 3, +) -> None: + """Analyze one image and update ``scene`` with validated semantic objects.""" + if json_max_attempts < 1: + raise ValueError("json_max_attempts must be at least 1.") + + resolved_image_path = _validate_image_path(image_path) + last_validation_error: ValueError | None = None + for _ in range(json_max_attempts): + response_text = vlm_client.complete( + image_path=resolved_image_path, + system_prompt=_SYSTEM_PROMPT, + user_prompt=_USER_PROMPT, + ) + try: + analyzed_scene = _parse_image_object_analysis_response(response_text) + validate_scene_understanding(analyzed_scene) + except ValueError as exc: + last_validation_error = exc + continue + + scene.objects = analyzed_scene.objects + return None + + assert last_validation_error is not None + raise ValueError( + "VLM returned invalid image-object analysis JSON after " + f"{json_max_attempts} attempts: {last_validation_error}" + ) from last_validation_error + + +def _parse_image_object_analysis_response(response_text: str) -> Scene: + """Parse one VLM image-object analysis response into a semantic ``Scene``.""" + json_text = _strip_json_code_fence(response_text) + try: + payload = json.loads(json_text) + except json.JSONDecodeError as exc: + raise ValueError(f"VLM response is not valid JSON: {exc.msg}") from exc + + if not isinstance(payload, dict) or set(payload) != {"table", "assets"}: + raise ValueError("VLM JSON must contain exactly the keys: table and assets.") + + id_counters: dict[str, int] = {} + table_fields = _parse_scene_object_fields(payload["table"], field_name="table") + table = SceneObject( + # id=_next_id(table_fields["category"], id_counters) + # Use a fixed ID for the table. + id="table", + kind="table", + **table_fields, + ) + assets_value = payload["assets"] + if not isinstance(assets_value, list): + raise ValueError("VLM JSON key assets must be an array.") + assets: list[SceneObject] = [] + for index, asset in enumerate(assets_value): + fields = _parse_scene_object_fields(asset, field_name=f"assets[{index}]") + assets.append( + SceneObject( + id=_next_id(fields["category"], id_counters), + kind="asset", + **fields, + ) + ) + + return Scene(objects=[table, *assets]) + + +def validate_scene_understanding(scene: Scene) -> None: + """Validate that scene understanding produced a complete semantic scene.""" + if scene.table is None: + raise ValueError("Scene understanding must identify a table.") + if ( + scene.table.id != "table" + ): # Currently it will always return false. For we hardcode the table id to "table". + raise ValueError("Scene table id must be 'table'.") + + asset_ids = [asset.id for asset in scene.assets] + if len(asset_ids) != len(set(asset_ids)): + raise ValueError("Scene asset ids must be unique.") + + for obj in [scene.table, *scene.assets]: + if not obj.category or not obj.name or not obj.description: + raise ValueError( + "Every scene object must contain category, name, and description." + ) + + +def _strip_json_code_fence(response_text: str) -> str: + stripped = response_text.strip() + if not stripped.startswith("```"): + return stripped + lines = stripped.splitlines() + if len(lines) < 3 or not lines[-1].strip().startswith("```"): + raise ValueError("VLM response contains an incomplete JSON code fence.") + return "\n".join(lines[1:-1]).strip() + + +def _validate_image_path(image_path: str | Path) -> Path: + resolved_image_path = Path(image_path).expanduser().resolve() + if not resolved_image_path.is_file(): + raise FileNotFoundError(f"Image input not found: {resolved_image_path}") + if resolved_image_path.suffix.lower() not in _SUPPORTED_IMAGE_SUFFIXES: + raise ValueError("Image input must be a .jpg, .jpeg, or .png file.") + return resolved_image_path + + +def _parse_scene_object_fields( + value: object, + *, + field_name: str, +) -> dict[str, str]: + if not isinstance(value, dict) or set(value) != { + "category", + "name", + "description", + }: + raise ValueError( + f"VLM JSON key {field_name} must contain exactly category, name, and " + "description." + ) + + fields = {} + for key in ("category", "name", "description"): + raw_value = value[key] + if not isinstance(raw_value, str) or not raw_value.strip(): + raise ValueError( + f"VLM JSON key {field_name}.{key} must be a non-empty string." + ) + fields[key] = raw_value.strip() + + if not _CATEGORY_PATTERN.fullmatch(fields["category"]): + raise ValueError( + f"VLM JSON key {field_name}.category must be a lower-case snake_case " + "class name." + ) + if _LOCATION_WORD_PATTERN.search( + fields["name"] + ): # Check whether the name contains location. + raise ValueError( + f"VLM JSON key {field_name}.name must not contain location or " + "relationship words." + ) + return fields + + +def _next_id(category: str, counters: dict[str, int]) -> str: + """Auto increment an ID for the same category, e.g. mug_001, mug_002, etc.""" + counters[category] = counters.get(category, 0) + 1 + return f"{category}_{counters[category]:03d}" + + +def _segment_scene( + *, + image_path: str | Path, + stage_output_root: str | Path, + scene: Scene, + vlm_client: OpenAICompatibleVLM, + image_segmentation_client: ImageSegmentationClient, +) -> None: + """Add validated table and asset mask paths to a semantic scene.""" + debug_output_root = ( + Path(stage_output_root) / "debug" + ) # Keeps the mask debug images. + masks_output_root = ( + Path(stage_output_root) / "masks" + ) # Keeps the validated masked images of each assets (include the table) + debug_output_root.mkdir() + masks_output_root.mkdir() + + # Segment the table and assets with VLM validation separately. + _segment_assets( + image_path=image_path, + debug_output_root=debug_output_root, + masks_output_root=masks_output_root, + scene=scene, + vlm_client=vlm_client, + image_segmentation_client=image_segmentation_client, + ) + # Prepare an image which do not contains any asset, for the VLM validation of the table + # segmentation more easily. + asset_mask_paths: list[str] = [] + for asset in scene.assets: + if asset.mask_path is None: + raise ValueError(f"Asset {asset.id!r} has no validated mask path.") + asset_mask_paths.append(asset.mask_path) + table_validation_image_path, asset_union_mask = render_image_without_masks( + image_path=image_path, + mask_paths=asset_mask_paths, + output_path=Path(debug_output_root) / "table_validation_base.png", + ) + # Segment the table. + _segment_table( + image_path=image_path, + validation_image_path=table_validation_image_path, + label_avoid_mask=asset_union_mask, + debug_output_root=debug_output_root, + masks_output_root=masks_output_root, + scene=scene, + vlm_client=vlm_client, + image_segmentation_client=image_segmentation_client, + ) + + +def _segment_table( + image_path: str | Path, + validation_image_path: str | Path, + label_avoid_mask: Image.Image, + debug_output_root: str | Path, + masks_output_root: str | Path, + scene: Scene, + *, + vlm_client: OpenAICompatibleVLM, + image_segmentation_client: ImageSegmentationClient, +) -> None: + """Segment the table. (Now it only supports segment the complete tabletop)""" + if scene.table is None: + raise ValueError("Cannot segment a scene without a table.") + + table = scene.table + # Build the segmentation prompts for table. + for prompt_label, prompt in ( + ("name", table.name), + ("description", table.description), + ("table", "table"), + ("plane", "plane"), + ): + candidates = union_overlapping_mask_candidates( + build_mask_candidates( + image_segmentation_client.segment_single_object( + image_path=image_path, + prompt=prompt, + ) + ), + min_iou=0.8, # Union masks who have iou > 0.8 + ) + # If do not have candidate, then try segment the table with description, "table", "plane"... + # Notice that, this part could be extended with other segmentation prompt like + # a board, or newly-generated prompt from another VLM-calling etc. + if not candidates: + continue + + # Maybe the mask count = 1, but not correct; + # Maybe the mask count > 1; + # Thus, we need to validate with an VLM. + candidates_image_path = render_numbered_mask_candidates( + image_path=validation_image_path, + candidates=candidates, + label_avoid_mask=label_avoid_mask, + output_path=( + Path(debug_output_root) + / f"table_candidates_{prompt_label}.png" # Render with prompt label, for easily debug. + ), + ) + selected_mask_index = _validate_table_candidates_with_vlm( + table=table, + candidates=candidates, + candidates_image_path=candidates_image_path, + vlm_client=vlm_client, + ) + if selected_mask_index is None: + continue + + # Save result. + candidate = _candidate_by_index(candidates, selected_mask_index) + table.mask_path = str( + save_binary_mask( + candidate, + image_size=_image_size(image_path), + output_path=Path(masks_output_root) / "table_mask.png", + ) + ) + return + + raise ValueError("Unable to find a VLM-validated segmentation mask for the table.") + + +def _validate_table_candidates_with_vlm( + *, + table: SceneObject, + candidates: list[MaskCandidate], + candidates_image_path: Path, + vlm_client: OpenAICompatibleVLM, + json_max_attempts: int = 3, +) -> int | None: + + if json_max_attempts < 1: + raise ValueError("json_max_attempts must be at least 1.") + + user_prompt = ( + "Table category: " + f"{table.category}\n" + f"Table name: {table.name}\n" + f"Table description: {table.description}\n" + f"Candidate indices range from 1 to {len(candidates)}." + ) + last_error: ValueError | None = None + for _ in range(json_max_attempts): + response_text = vlm_client.complete( + image_path=candidates_image_path, + system_prompt=_TABLE_VALIDATION_SYSTEM_PROMPT, + user_prompt=user_prompt, + ) + try: + return _parse_table_validation_response(response_text, candidates) + except ValueError as exc: + last_error = exc + + assert last_error is not None + raise ValueError( + "VLM returned invalid table-segmentation validation JSON after " + f"{json_max_attempts} attempts: {last_error}" + ) from last_error + + +def _parse_table_validation_response( + response_text: str, + candidates: list[MaskCandidate], +) -> int | None: + """Validate the strict VLM response schema for table candidate selection.""" + try: + payload = json.loads(_strip_json_code_fence(response_text)) + except json.JSONDecodeError as exc: + raise ValueError("VLM table validation response is not valid JSON.") from exc + if not isinstance(payload, dict) or set(payload) != {"selected_mask_index"}: + raise ValueError( + "VLM table validation JSON must contain only selected_mask_index." + ) + + selected_mask_index = payload["selected_mask_index"] + if selected_mask_index is None: + return None + if isinstance(selected_mask_index, bool) or not isinstance( + selected_mask_index, int + ): + raise ValueError("selected_mask_index must be an integer or null.") + _candidate_by_index(candidates, selected_mask_index) + return selected_mask_index + + +def _candidate_by_index( + candidates: list[MaskCandidate], + index: int, +) -> MaskCandidate: + for candidate in candidates: + if candidate.index == index: + return candidate + raise ValueError(f"VLM selected a nonexistent mask candidate: {index}.") + + +def _image_size(image_path: str | Path) -> tuple[int, int]: + from PIL import Image + + with Image.open(image_path) as image: + return image.size + + +def _segment_assets( + image_path: str | Path, + debug_output_root: str | Path, + masks_output_root: str | Path, + scene: Scene, + *, + vlm_client: OpenAICompatibleVLM, + image_segmentation_client: ImageSegmentationClient, +) -> None: + + # Group the assets by their categories. + assets_by_category: dict[str, list[SceneObject]] = {} + for asset in scene.assets: + assets_by_category.setdefault(asset.category, []).append(asset) + + image_size = _image_size(image_path) + for category, assets in assets_by_category.items(): + mask_rles: list[dict[str, Any]] = [] + # Use categories and names as segmentation prompt. + # Use category to segment first, then use each assets' name to segment. + prompts = [category, *dict.fromkeys(asset.name for asset in assets)] + for prompt in prompts: + mask_rles.extend( + image_segmentation_client.segment_single_object( + image_path=image_path, + prompt=prompt, + ) + ) + # Union duplicated mask candidates. + candidates = union_overlapping_mask_candidates( + build_mask_candidates(mask_rles), + min_iou=0.8, + ) + # If the number of candidate is less than the grouped assets, + # raise error directly. + if len(candidates) < len(assets): + raise ValueError( + f"Asset category {category!r} has {len(assets)} assets but only " + f"{len(candidates)} segmentation candidates." + ) + + candidates_image_path = render_numbered_mask_candidates( + image_path=image_path, + candidates=candidates, + output_path=Path(debug_output_root) / f"asset_candidates_{category}.png", + mask_style="outline", + ) + assignments = _validate_asset_candidates_with_vlm( + assets=assets, + candidates=candidates, + candidates_image_path=candidates_image_path, + vlm_client=vlm_client, + ) + if assignments is None: + raise ValueError( + f"VLM could not assign every {category!r} asset to a segmentation candidate." + ) + # Save results. + for asset in assets: + asset.mask_path = str( + save_binary_mask( + _candidate_by_index(candidates, assignments[asset.id]), + image_size=image_size, + output_path=Path(masks_output_root) / f"{asset.id}_mask.png", + ) + ) + + +def _validate_asset_candidates_with_vlm( + *, + assets: list[SceneObject], + candidates: list[MaskCandidate], + candidates_image_path: Path, + vlm_client: OpenAICompatibleVLM, + json_max_attempts: int = 3, +) -> dict[str, int] | None: + """Ask the VLM for a complete one-to-one asset-to-candidate assignment.""" + if json_max_attempts < 1: + raise ValueError("json_max_attempts must be at least 1.") + + assets_text = "\n".join( + "- " + f"id: {asset.id}; category: {asset.category}; name: {asset.name}; " + f"description: {asset.description}" + for asset in assets + ) + user_prompt = ( + "Asset group:\n" + f"{assets_text}\n\n" + f"Candidate indices range from 1 to {len(candidates)}." + ) + last_error: ValueError | None = None + for _ in range(json_max_attempts): + response_text = vlm_client.complete( + image_path=candidates_image_path, + system_prompt=_ASSET_ASSIGNMENT_SYSTEM_PROMPT, + user_prompt=user_prompt, + ) + try: + return _parse_asset_assignment_response(response_text, assets, candidates) + except ValueError as exc: + last_error = exc + + assert last_error is not None + raise ValueError( + "VLM returned invalid asset-segmentation assignment JSON after " + f"{json_max_attempts} attempts: {last_error}" + ) from last_error + + +def _parse_asset_assignment_response( + response_text: str, + assets: list[SceneObject], + candidates: list[MaskCandidate], +) -> dict[str, int] | None: + """Parse a strict complete assignment, or a valid missing-asset result.""" + try: + payload = json.loads(_strip_json_code_fence(response_text)) + except json.JSONDecodeError as exc: + raise ValueError("VLM asset assignment response is not valid JSON.") from exc + if not isinstance(payload, dict) or set(payload) != {"assignments"}: + raise ValueError("VLM asset assignment JSON must contain only assignments.") + + assignment_values = payload["assignments"] + if assignment_values is None: + return None + if not isinstance(assignment_values, list): + raise ValueError("assignments must be an array or null.") + + expected_asset_ids = {asset.id for asset in assets} + assignments: dict[str, int] = {} + assigned_mask_indices: set[int] = set() + for assignment in assignment_values: + if not isinstance(assignment, dict) or set(assignment) != { + "asset_id", + "mask_index", + }: + raise ValueError( + "Each assignment must contain only asset_id and mask_index." + ) + asset_id = assignment["asset_id"] + mask_index = assignment["mask_index"] + if not isinstance(asset_id, str) or not asset_id: + raise ValueError("assignment asset_id must be a non-empty string.") + if isinstance(mask_index, bool) or not isinstance(mask_index, int): + raise ValueError("assignment mask_index must be an integer.") + if asset_id in assignments: + raise ValueError(f"VLM assigned asset {asset_id!r} more than once.") + if mask_index in assigned_mask_indices: + raise ValueError( + f"VLM assigned candidate {mask_index} to more than one asset." + ) + _candidate_by_index(candidates, mask_index) + assignments[asset_id] = mask_index + assigned_mask_indices.add(mask_index) + + if set(assignments) != expected_asset_ids: + raise ValueError("VLM assignments must cover every asset in the group.") + return assignments diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/__init__.py b/embodichain/gen_sim/scene_engine/pipeline/utils/__init__.py new file mode 100644 index 000000000..015c41510 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/__init__.py @@ -0,0 +1,19 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +__all__: list[str] = [] diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/assets_gravity_settler.py b/embodichain/gen_sim/scene_engine/pipeline/utils/assets_gravity_settler.py new file mode 100644 index 000000000..31e9e8443 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/assets_gravity_settler.py @@ -0,0 +1,349 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Sequence + +import numpy as np +from scipy.spatial.transform import Rotation +import trimesh + +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_object import ( + ObjectPhysics, + SceneObject, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_generation_utils import ( + layout_object_to_transform_matrix, + load_glb_mesh, + transform_matrix_to_layout_object, +) +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg +from embodichain.lab.sim.shapes import MeshCfg +from embodichain.utils.logger import log_info + + +@dataclass(frozen=True) +class AssetsGravitySettlerConfig: + """Physics controls for table-top asset settling.""" + + clearance_m: float = 0.02 # Initial gap between each asset and the table top. + settle_steps: int = 300 # Fixed number of simulator steps to execute. + physics_dt: float = 1.0 / 100.0 # Physics timestep in seconds. + sim_device: str = "cpu" # Simulation device requested from EmbodiChain Lab. + + +class AssetsGravitySettler: + """Settle all assets together on one kinematic table in a z-up simulation.""" + + def __init__( + self, + *, + scene: Scene, + table_layout: dict[str, object], + assets_layout: list[dict[str, object]], + geometry_root: str | Path, + config: AssetsGravitySettlerConfig | None = None, + ) -> None: + self.scene = scene + self.table_layout = table_layout + self.assets_layout = assets_layout + self.geometry_root = Path(geometry_root).expanduser().resolve() + self.settled_assets_layout: list[dict[str, object]] | None = None + self.config = config if config is not None else AssetsGravitySettlerConfig() + # Check. + if self.config.clearance_m < 0.0: + raise ValueError("Gravity-settle clearance_m must be non-negative.") + if self.config.settle_steps <= 0: + raise ValueError("Gravity-settle settle_steps must be positive.") + if self.config.physics_dt <= 0.0: + raise ValueError("Gravity-settle physics_dt must be positive.") + + def settle(self) -> list[dict[str, object]]: + """Run gravity settling and return the resulting y-up asset layouts.""" + self.settled_assets_layout = None + if not self.assets_layout: + self.settled_assets_layout = [] + log_info("Scene has no movable assets; skipping gravity settling.") + return self.settled_assets_layout + + table_id = self._require_layout_id(self.table_layout, name="Table") + table_object = self._require_scene_object(table_id, kind="table") + asset_ids: set[str] = set() + asset_objects_by_id: dict[str, SceneObject] = {} + for asset_layout in self.assets_layout: + asset_id = self._require_layout_id(asset_layout, name="Asset") + if asset_id in asset_ids: + raise ValueError(f"Asset layouts contain duplicate id {asset_id!r}.") + asset_ids.add(asset_id) + asset_objects_by_id[asset_id] = self._require_scene_object( + asset_id, kind="asset" + ) + expected_asset_ids = {asset.id for asset in self.scene.assets} + if asset_ids != expected_asset_ids: + raise ValueError( + "Gravity-settle layouts must contain exactly the scene asset ids." + ) + + y_up_to_z_up_matrix = np.eye(4) + y_up_to_z_up_matrix[:3, :3] = np.array( + [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]] + ) + z_up_to_y_up_matrix = np.linalg.inv(y_up_to_z_up_matrix) + table_info = self._prepare_sim_body( + layout_object=self.table_layout, + y_up_to_z_up_matrix=y_up_to_z_up_matrix, + ) + table_world_mesh = self._mesh_to_z_up_world_for_aabb( + y_up_mesh=table_info["mesh"], + z_up_rigid_layout=table_info["rigid_layout"], + z_up_scale=table_info["z_up_scale"], + y_up_to_z_up_matrix=y_up_to_z_up_matrix, + ) + table_top_z = float(table_world_mesh.bounds[1, 2]) + + prepared_assets: dict[str, dict[str, object]] = {} + for asset_layout in self.assets_layout: + asset_id = str(asset_layout["id"]) + asset_info = self._prepare_sim_body( + layout_object=asset_layout, + y_up_to_z_up_matrix=y_up_to_z_up_matrix, + ) + asset_world_mesh = self._mesh_to_z_up_world_for_aabb( + y_up_mesh=asset_info["mesh"], + z_up_rigid_layout=asset_info["rigid_layout"], + z_up_scale=asset_info["z_up_scale"], + y_up_to_z_up_matrix=y_up_to_z_up_matrix, + ) + asset_bottom_z = float(asset_world_mesh.bounds[0, 2]) + asset_info["rigid_layout"]["pos"][2] += ( + table_top_z + self.config.clearance_m - asset_bottom_z + ) + prepared_assets[asset_id] = asset_info + + log_info( + "Gravity settling started: " + f"assets={len(prepared_assets)}, steps={self.config.settle_steps}, " + f"physics_dt={self.config.physics_dt:.4f} s." + ) + sim = SimulationManager( + SimulationManagerCfg( + headless=True, + physics_dt=self.config.physics_dt, + sim_device=self.config.sim_device, + ) + ) + try: + # Add table. + sim.add_rigid_object( + RigidObjectCfg( + uid=table_id, + shape=MeshCfg(fpath=str(table_info["mesh_path"])), + init_pos=tuple(table_info["rigid_layout"]["pos"]), + init_rot=tuple( + self._simulation_euler_xyz_degrees(table_info["rigid_layout"]) + ), + body_scale=tuple(table_info["y_up_scale"]), + attrs=self._rigid_body_attrs(table_object.physics), + body_type=table_object.physics.body_type, + max_convex_hull_num=table_object.physics.max_convex_hull_num, + acd_method="vhacd", + ) + ) + # Add assets. + simulated_assets: dict[str, object] = {} + for asset_id, asset_info in prepared_assets.items(): + rigid_layout = asset_info["rigid_layout"] + simulated_assets[asset_id] = sim.add_rigid_object( + RigidObjectCfg( + uid=asset_id, + shape=MeshCfg(fpath=str(asset_info["mesh_path"])), + init_pos=tuple(rigid_layout["pos"]), + init_rot=tuple( + self._simulation_euler_xyz_degrees(rigid_layout) + ), + body_scale=tuple(asset_info["y_up_scale"]), + attrs=self._rigid_body_attrs( + asset_objects_by_id[asset_id].physics + ), + body_type=asset_objects_by_id[asset_id].physics.body_type, + max_convex_hull_num=( + asset_objects_by_id[asset_id].physics.max_convex_hull_num + ), + acd_method="vhacd", + ) + ) + # Run simulation to settle all assets. + sim.update(step=self.config.settle_steps) + + # Update the final layouts. + settled_layout_by_id: dict[str, dict[str, object]] = {} + for asset_id, simulated_asset in simulated_assets.items(): + final_rigid_pose_z_up = np.asarray( + simulated_asset.get_local_pose(to_matrix=True)[0] + .detach() + .cpu() + .numpy(), + dtype=float, + ) + scale_matrix = np.eye(4) + scale_matrix[:3, :3] = np.diag(prepared_assets[asset_id]["z_up_scale"]) + final_z_up_layout_matrix = final_rigid_pose_z_up @ scale_matrix + settled_layout_by_id[asset_id] = transform_matrix_to_layout_object( + asset_id, + z_up_to_y_up_matrix + @ final_z_up_layout_matrix + @ y_up_to_z_up_matrix, + ) + finally: + # Release resources. + sim.destroy(exit_process=False) + SimulationManager.flush_cleanup_queue() + + self.settled_assets_layout = [ + settled_layout_by_id[str(asset_layout["id"])] + for asset_layout in self.assets_layout + ] + log_info("Gravity settling completed for all assets.") + return self.settled_assets_layout + + def _prepare_sim_body( + self, + *, + layout_object: dict[str, object], + y_up_to_z_up_matrix: np.ndarray, + ) -> dict[str, object]: + """Load one y-up GLB and prepare its z-up simulation pose.""" + object_id = self._require_layout_id(layout_object, name="Layout object") + source_mesh_path = self.geometry_root / f"{object_id}.glb" + source_mesh = load_glb_mesh(source_mesh_path) + z_up_layout = self._convert_layout_coordinate_system( + layout_object, + source_to_target_matrix=y_up_to_z_up_matrix, + ) + return { + "mesh_path": source_mesh_path, + "mesh": source_mesh, + "rigid_layout": { + "id": object_id, + "rot": self._three_floats(z_up_layout.get("rot"), field_name="rot"), + "pos": self._three_floats(z_up_layout.get("pos"), field_name="pos"), + "scale": [1.0, 1.0, 1.0], + }, + "y_up_scale": self._three_floats( + layout_object.get("scale"), field_name="scale" + ), + "z_up_scale": self._three_floats( + z_up_layout.get("scale"), field_name="scale" + ), + } + + def _require_scene_object(self, object_id: str, *, kind: str) -> SceneObject: + """Return one physics-ready scene object with the expected semantic kind.""" + matching_objects = [ + scene_object + for scene_object in self.scene.objects + if scene_object.id == object_id + ] + if len(matching_objects) != 1: + raise ValueError( + f"Gravity settling requires exactly one scene object {object_id!r}." + ) + scene_object = matching_objects[0] + if scene_object.kind != kind: + raise ValueError( + f"Scene object {object_id!r} must have kind {kind!r} before " + "gravity settling." + ) + if scene_object.physics is None: + raise ValueError( + f"Scene object {object_id!r} has no SimReady physics settings." + ) + return scene_object + + @staticmethod + def _rigid_body_attrs(physics: ObjectPhysics | None) -> RigidBodyAttributesCfg: + """Convert persisted SceneObject physics attributes into Lab config.""" + if physics is None: + raise ValueError("Gravity settling requires SimReady physics settings.") + return RigidBodyAttributesCfg(**physics.attrs) + + @staticmethod + def _mesh_to_z_up_world_for_aabb( + *, + y_up_mesh: trimesh.Trimesh, + z_up_rigid_layout: dict[str, object], + z_up_scale: Sequence[float], + y_up_to_z_up_matrix: np.ndarray, + ) -> trimesh.Trimesh: + """Transform a y-up mesh into its z-up world pose for AABB measurement.""" + mesh = y_up_mesh.copy() + mesh.apply_transform(y_up_to_z_up_matrix) + scale_matrix = np.eye(4) + scale_matrix[:3, :3] = np.diag(z_up_scale) + mesh.apply_transform(scale_matrix) + mesh.apply_transform(layout_object_to_transform_matrix(z_up_rigid_layout)) + return mesh + + @staticmethod + def _simulation_euler_xyz_degrees(layout_object: dict[str, object]) -> list[float]: + """Convert lowercase-xyz layout rotation to SimulationManager's XYZ order.""" + layout_rotation = Rotation.from_euler( + "xyz", + AssetsGravitySettler._three_floats( + layout_object.get("rot"), field_name="rot" + ), + degrees=True, + ) + return layout_rotation.as_euler("XYZ", degrees=True).tolist() + + @staticmethod + def _convert_layout_coordinate_system( + layout_object: dict[str, object], + *, + source_to_target_matrix: np.ndarray, + ) -> dict[str, object]: + """Convert one layout object between coordinate frames through its matrix.""" + return transform_matrix_to_layout_object( + str(layout_object["id"]), + source_to_target_matrix + @ layout_object_to_transform_matrix(layout_object) + @ np.linalg.inv(source_to_target_matrix), + ) + + @staticmethod + def _require_layout_id(layout_object: dict[str, object], *, name: str) -> str: + """Check id.""" + object_id = layout_object.get("id") + if not isinstance(object_id, str) or not object_id: + raise ValueError(f"{name} layout must contain a non-empty string id.") + return object_id + + @staticmethod + def _three_floats(value: object, *, field_name: str) -> list[float]: + """Check three values.""" + if not isinstance(value, list) or len(value) != 3: + raise ValueError(f"Layout field {field_name} must contain three values.") + try: + return [float(item) for item in value] + except (TypeError, ValueError) as exc: + raise ValueError( + f"Layout field {field_name} must contain numeric values." + ) from exc diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/assets_group_layout_optimizer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/assets_group_layout_optimizer.py new file mode 100644 index 000000000..d4021fede --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/assets_group_layout_optimizer.py @@ -0,0 +1,530 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") + +import matplotlib.pyplot as plt +import numpy as np +from shapely import affinity +from shapely.geometry import MultiPolygon, Polygon + +from embodichain.gen_sim.scene_engine.pipeline.utils.assets_group_support_clamp import ( + AssetsGroupSupportClamp, + SupportGeometry, +) +from embodichain.utils.logger import log_info, log_warning + + +@dataclass(frozen=True) +class AssetsSupportLayoutOptimizerConfig: + """Controls for support-constrained pairwise AABB separation.""" + + margin_m: float = 0.0 # Required clearance between each AABB and the boundary. + aabb_clearance_m: float = 1e-6 # Required clearance between AABB pairs. + max_rounds: int = 64 # Maximum greedy pair-separation passes. + split_samples: int = 9 # Candidate splits between the two overlapping AABBs. + + +class AssetsSupportLayoutOptimizer: + """Greedily separate AABBs while retaining arbitrary support containment. + + This reuses the previous packing algorithm's pairwise strategy: detect an + overlap, try the two separating directions on both XY axes, and choose the + lowest-displacement valid push. Unlike the old path, every candidate is + validated against the actual Polygon/MultiPolygon support region instead + of a largest internal rectangle. + """ + + def __init__( + self, + *, + support_region: SupportGeometry, + assets_aabb_2d_z_up_world_corners_by_id: dict[str, np.ndarray], + assets_layout: list[dict[str, object]], + debug_output_root: str | Path | None = None, + config: AssetsSupportLayoutOptimizerConfig | None = None, + ) -> None: + self.support_region = support_region + self.assets_aabb_2d_z_up_world_corners_by_id = ( + assets_aabb_2d_z_up_world_corners_by_id + ) + self.assets_layout = assets_layout + self.debug_output_root = ( + Path(debug_output_root).expanduser().resolve() + if debug_output_root is not None + else None + ) + self.refined_assets_layout: list[dict[str, object]] | None = None + self.config = ( + config if config is not None else AssetsSupportLayoutOptimizerConfig() + ) + # Check config. + if self.config.margin_m < 0.0: + raise ValueError("margin_m must be non-negative.") + if self.config.aabb_clearance_m < 0.0: + raise ValueError("aabb_clearance_m must be non-negative.") + if self.config.max_rounds <= 0 or self.config.split_samples < 2: + raise ValueError( + "max_rounds must be positive and split_samples at least two." + ) + + def optimize(self) -> list[dict[str, object]]: + """Resolve pairwise AABB overlap and return updated y-up layouts.""" + self.refined_assets_layout = None + # Check inputs just like the previous AssetsGroupSupportClamp step would have done. + aabbs_by_id = AssetsGroupSupportClamp._validate_aabbs( + self.assets_aabb_2d_z_up_world_corners_by_id + ) + raw_support = AssetsGroupSupportClamp._coerce_support_geometry( + self.support_region + ) + if raw_support is None: + log_warning("AABB overlap optimization failed: invalid support geometry.") + raise ValueError("Support region is invalid.") + safe_support = ( + raw_support + if self.config.margin_m == 0.0 + else AssetsGroupSupportClamp._polygonal_geometry( + raw_support.buffer(-self.config.margin_m) + ) + ) + if safe_support is None or safe_support.is_empty: + log_warning( + "AABB overlap optimization failed: support region is empty after " + f"applying a {self.config.margin_m:.4f} m boundary margin." + ) + raise ValueError("Support region is empty after applying layout margin.") + + asset_ids = sorted(aabbs_by_id) + base_aabbs = np.stack([aabbs_by_id[asset_id] for asset_id in asset_ids]) + offsets = np.zeros((len(asset_ids), 2), dtype=float) + if not self._all_contained(safe_support, base_aabbs, offsets): + log_warning( + "AABB overlap optimization requires all input AABBs to be inside " + "the support region." + ) + raise ValueError( + "Overlap optimization requires AABBs already inside support; " + "run AssetsGroupSupportClamp first." + ) + initial_overlaps = self._overlaps(base_aabbs, offsets) + log_info( + "Support-constrained AABB overlap optimization started: " + f"assets={len(asset_ids)}, initial_overlaps={len(initial_overlaps)}, " + f"boundary_margin={self.config.margin_m:.4f} m, " + f"aabb_clearance={self.config.aabb_clearance_m:.4f} m, " + f"max_rounds={self.config.max_rounds}." + ) + if not initial_overlaps: # Return directly if there are no overlaps to resolve. + log_info("AABB overlap optimization succeeded without movement.") + self.refined_assets_layout = self._apply_offsets_to_y_up_layouts( + asset_ids=asset_ids, + offsets=offsets, + ) + return self.refined_assets_layout + + for round_index in range(self.config.max_rounds): + # Check whether any overlaps remain. + overlaps = self._overlaps(base_aabbs, offsets) + if not overlaps: + log_info( + "AABB overlap optimization succeeded after " + f"{round_index} rounds." + ) + self.refined_assets_layout = self._apply_offsets_to_y_up_layouts( + asset_ids=asset_ids, + offsets=offsets, + ) + return self.refined_assets_layout + log_info( + "AABB overlap optimization round " + f"{round_index + 1}/{self.config.max_rounds}: " + f"remaining_overlaps={len(overlaps)}." + ) + moved = False + # Choose a pair once. + for _, first_index, second_index in overlaps: + # If the overlaps is handeled by a previous pair, skip it. + if not self._overlaps(base_aabbs, offsets, (first_index, second_index)): + continue + candidates = self._separation_candidates( + base_aabbs=base_aabbs, + offsets=offsets, + first_index=first_index, + second_index=second_index, + safe_support=safe_support, + ) + if candidates: + # A local separation must not blindly create a new + # collision with a third asset. Prefer candidates with + # no new collisions; when every placement causes one, + # retain the least-colliding state before considering + # displacement from the generated layout. + offsets = min( + candidates, + key=lambda candidate: self._candidate_score( + base_aabbs=base_aabbs, + current_offsets=offsets, + candidate_offsets=candidate, + ), + ) # Choose the best candidate with score. + moved = True + else: + log_warning( + "No support-valid axis-aligned separation candidate for " + f"overlapping AABBs {asset_ids[first_index]!r} and " + f"{asset_ids[second_index]!r}." + ) + if not moved: + break + + unresolved_pairs = [ + f"{asset_ids[first_index]}/{asset_ids[second_index]}" + for _, first_index, second_index in self._overlaps(base_aabbs, offsets) + ] + log_warning( + "Unable to resolve all asset AABB overlaps inside the detected support " + "region; unresolved pairs=" + f"{unresolved_pairs}." + ) + raise ValueError( + "Asset AABB overlap cannot be resolved while keeping all assets " + "inside the detected table support region." + ) + + def _apply_offsets_to_y_up_layouts( + self, *, asset_ids: list[str], offsets: np.ndarray + ) -> list[dict[str, object]]: + """Write independent z-up XY offsets back to the stored y-up layouts.""" + received_ids = {str(layout.get("id")) for layout in self.assets_layout} + expected_ids = set(asset_ids) + if received_ids != expected_ids: + raise ValueError( + "Asset layouts and optimized AABBs must have identical ids." + ) + updated_layouts: list[dict[str, object]] = [] + offsets_by_id = { + asset_id: offsets[index] for index, asset_id in enumerate(asset_ids) + } + for layout in self.assets_layout: + asset_id = str(layout["id"]) + position = layout.get("pos") + if not isinstance(position, list) or len(position) != 3: + raise ValueError( + "Each asset layout must contain a three-value pos list." + ) + dx, dy = offsets_by_id[asset_id] + updated_layout = dict(layout) + updated_position = [float(value) for value in position] + updated_position[0] += float(dx) + updated_position[2] -= float(dy) + updated_layout["pos"] = updated_position + updated_layouts.append(updated_layout) + return updated_layouts + + def save_overlap_optimization_debug_images(self) -> bool: + """Optionally save diagnostics for the most recent optimization.""" + if self.refined_assets_layout is None: + self.optimize() + assert self.refined_assets_layout is not None + if self.debug_output_root is None: + raise ValueError( + "A debug_output_root is required when saving overlap-optimization " + "debug images." + ) + + initial_aabbs_by_id = AssetsGroupSupportClamp._validate_aabbs( + self.assets_aabb_2d_z_up_world_corners_by_id + ) + raw_support = AssetsGroupSupportClamp._coerce_support_geometry( + self.support_region + ) + if raw_support is None: + raise ValueError("Cannot render overlap optimization for invalid support.") + original_positions_by_id = { + str(layout["id"]): layout["pos"] for layout in self.assets_layout + } + refined_positions_by_id = { + str(layout["id"]): layout["pos"] for layout in self.refined_assets_layout + } + if set(original_positions_by_id) != set(initial_aabbs_by_id) or set( + refined_positions_by_id + ) != set(initial_aabbs_by_id): + raise ValueError( + "Asset layouts and optimized AABBs must have identical ids." + ) + + translated_aabbs_by_id: dict[str, np.ndarray] = {} + moved = False + for asset_id, corners in initial_aabbs_by_id.items(): + original_position = original_positions_by_id[asset_id] + refined_position = refined_positions_by_id[asset_id] + if ( + not isinstance(original_position, list) + or not isinstance(refined_position, list) + or len(original_position) != 3 + or len(refined_position) != 3 + ): + raise ValueError( + "Each asset layout must contain a three-value pos list." + ) + delta_xy = np.array( + [ + float(refined_position[0]) - float(original_position[0]), + float(original_position[2]) - float(refined_position[2]), + ] + ) + translated_aabbs_by_id[asset_id] = corners + delta_xy + moved = moved or not np.allclose(delta_xy, 0.0) + + path = self.debug_output_root / "assets_aabb_overlap_optimization_2d.png" + path.parent.mkdir(parents=True, exist_ok=True) + figure, axes = plt.subplots( + 1, 2, figsize=(14, 7), dpi=160, constrained_layout=True + ) + self._draw_overlap_state( + axes[0], + raw_support, + initial_aabbs_by_id, + "Before AABB overlap optimization", + ) + self._draw_overlap_state( + axes[1], + raw_support, + translated_aabbs_by_id, + ( + "After AABB overlap optimization" + if moved + else "After AABB overlap optimization (already non-overlapping)" + ), + ) + figure.savefig(path, bbox_inches="tight") + plt.close(figure) + return True + + def _separation_candidates( + self, + *, + base_aabbs: np.ndarray, + offsets: np.ndarray, + first_index: int, + second_index: int, + safe_support: Polygon | MultiPolygon, + ) -> list[np.ndarray]: + # Get current aabbs. + current_aabbs = base_aabbs + offsets[:, None, :] + minimums, maximums = current_aabbs.min(axis=1), current_aabbs.max(axis=1) + candidates: list[np.ndarray] = [] + for axis in (0, 1): + directions_and_distances = ( + ( + -1.0, + maximums[first_index, axis] + + self.config.aabb_clearance_m + - minimums[second_index, axis], + ), + ( + 1.0, + maximums[second_index, axis] + + self.config.aabb_clearance_m + - minimums[first_index, axis], + ), + ) + for first_direction, required_distance in directions_and_distances: + if required_distance <= 0.0: + continue + for fraction in np.linspace(0.0, 1.0, self.config.split_samples): + candidate = offsets.copy() + first_move = required_distance * float(fraction) + candidate[first_index, axis] += first_direction * first_move + candidate[second_index, axis] -= first_direction * ( + required_distance - first_move + ) + if not self._overlaps( + base_aabbs, candidate, (first_index, second_index) + ) and self._all_contained(safe_support, base_aabbs, candidate): + candidates.append(candidate) + return candidates + + def _overlaps( + self, + base_aabbs: np.ndarray, + offsets: np.ndarray, + only_pair: tuple[int, int] | None = None, + ) -> list[tuple[float, int, int]]: + return sorted( + [ + (min(overlap_x, overlap_y), first_index, second_index) + for overlap_x, overlap_y, first_index, second_index in self._overlap_details( + base_aabbs, offsets, only_pair + ) + ], + reverse=True, + ) + + def _overlap_details( + self, + base_aabbs: np.ndarray, + offsets: np.ndarray, + only_pair: tuple[int, int] | None = None, + ) -> list[tuple[float, float, int, int]]: + """Return positive XY penetration extents, including requested clearance.""" + current_aabbs = base_aabbs + offsets[:, None, :] + minimums, maximums = current_aabbs.min(axis=1), current_aabbs.max(axis=1) + pairs = ( + [only_pair] + if only_pair is not None + else [ + (first_index, second_index) + for first_index in range(len(current_aabbs)) + for second_index in range(first_index + 1, len(current_aabbs)) + ] + ) + overlaps: list[tuple[float, float, int, int]] = [] + for first_index, second_index in pairs: + overlap_x = ( + min(maximums[first_index, 0], maximums[second_index, 0]) + - max(minimums[first_index, 0], minimums[second_index, 0]) + + self.config.aabb_clearance_m + ) + overlap_y = ( + min(maximums[first_index, 1], maximums[second_index, 1]) + - max(minimums[first_index, 1], minimums[second_index, 1]) + + self.config.aabb_clearance_m + ) + if overlap_x > 1e-9 and overlap_y > 1e-9: + overlaps.append((overlap_x, overlap_y, first_index, second_index)) + return overlaps + + def _draw_overlap_state( + self, + axis: plt.Axes, + support: Polygon | MultiPolygon, + aabbs_by_id: dict[str, np.ndarray], + title: str, + ) -> None: + """Draw a state and highlight every AABB pair that still overlaps.""" + validated_aabbs = AssetsGroupSupportClamp._validate_aabbs(aabbs_by_id) + asset_ids = sorted(validated_aabbs) + aabbs = np.stack([validated_aabbs[asset_id] for asset_id in asset_ids]) + zero_offsets = np.zeros((len(asset_ids), 2), dtype=float) + overlaps = self._overlaps(aabbs, zero_offsets) + overlapping_indices = { + index + for _, first_index, second_index in overlaps + for index in (first_index, second_index) + } + + AssetsGroupSupportClamp._draw_support(axis, support, title, "darkorange") + for index, asset_id in enumerate(asset_ids): + corners = validated_aabbs[asset_id] + polygon = AssetsGroupSupportClamp._aabb_polygon(corners) + boundary = np.asarray(polygon.exterior.coords) + is_overlapping = index in overlapping_indices + color = "firebrick" if is_overlapping else "seagreen" + axis.fill( + boundary[:, 0], + boundary[:, 1], + facecolor=color, + edgecolor=color, + linewidth=2.0 if is_overlapping else 1.0, + alpha=0.32, + ) + axis.text( + *corners.mean(axis=0), + asset_id, + ha="center", + va="center", + fontsize=8, + bbox={"facecolor": "white", "alpha": 0.75, "edgecolor": "none"}, + ) + + for _, first_index, second_index in overlaps: + first_center = aabbs[first_index].mean(axis=0) + second_center = aabbs[second_index].mean(axis=0) + axis.plot( + [first_center[0], second_center[0]], + [first_center[1], second_center[1]], + color="firebrick", + linestyle="--", + linewidth=1.3, + ) + axis.set_title(f"{title}\nremaining AABB overlaps: {len(overlaps)}") + + def _candidate_score( + self, + *, + base_aabbs: np.ndarray, + current_offsets: np.ndarray, + candidate_offsets: np.ndarray, + ) -> tuple[int, int, float, float]: + """Rank a valid pair-separation candidate by global collision impact. + + The first term is deliberately based on *new* overlap pairs: this + keeps a pairwise correction from simply transferring its collision to + a nearby third asset. If every candidate causes a new collision, the + remaining terms prefer fewer total overlaps, less total penetration, + then less layout displacement. + """ + current_pairs = { + (first_index, second_index) + for _, _, first_index, second_index in self._overlap_details( + base_aabbs, current_offsets + ) + } + candidate_details = self._overlap_details(base_aabbs, candidate_offsets) + candidate_pairs = { + (first_index, second_index) + for _, _, first_index, second_index in candidate_details + } + new_overlap_count = len(candidate_pairs - current_pairs) + total_penetration_area = sum( + overlap_x * overlap_y for overlap_x, overlap_y, _, _ in candidate_details + ) + total_squared_displacement = float( + np.einsum("ij,ij->", candidate_offsets, candidate_offsets) + ) + return ( + new_overlap_count, + len(candidate_pairs), + total_penetration_area, + total_squared_displacement, + ) + + @staticmethod + def _all_contained( + support: Polygon | MultiPolygon, + base_aabbs: np.ndarray, + offsets: np.ndarray, + ) -> bool: + return all( + support.covers( + affinity.translate( + AssetsGroupSupportClamp._aabb_polygon(corners), + xoff=float(offset[0]), + yoff=float(offset[1]), + ) + ) + for corners, offset in zip(base_aabbs, offsets) + ) diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/assets_group_support_clamp.py b/embodichain/gen_sim/scene_engine/pipeline/utils/assets_group_support_clamp.py new file mode 100644 index 000000000..927134902 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/assets_group_support_clamp.py @@ -0,0 +1,642 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import TypeAlias + +import matplotlib + +matplotlib.use("Agg") + +import matplotlib.pyplot as plt +import numpy as np +from scipy.ndimage import binary_erosion +from shapely import affinity +from shapely.geometry import GeometryCollection, MultiPolygon, Polygon + +from embodichain.utils.logger import log_info, log_warning + +SupportGeometry: TypeAlias = Polygon | MultiPolygon + + +@dataclass(frozen=True) +class AssetsGroupSupportClampConfig: + """Numerical controls for rigid group placement on a support region.""" + + margin_m: float = 0.0 # Required clearance between each AABB and the boundary. + grid_resolution_m: float = 0.005 # Raster cell size for the coarse search. + + +@dataclass(frozen=True) +class _GridTransform: + """World/grid conversion for a centre-sampled regular XY raster.""" + + x_coordinates: np.ndarray # X coordinate of each grid-column centre. + y_coordinates: np.ndarray # Y coordinate of each grid-row centre. + resolution_m: float # Uniform spacing between neighbouring cell centres. + + @property + def shape(self) -> tuple[int, int]: + return len(self.y_coordinates), len(self.x_coordinates) + + def world_to_nearest_pixel(self, point_xy: np.ndarray) -> tuple[int, int]: + column = int(np.rint((point_xy[0] - self.x_coordinates[0]) / self.resolution_m)) + row = int(np.rint((point_xy[1] - self.y_coordinates[0]) / self.resolution_m)) + return row, column + + def pixel_to_world(self, row: int, column: int) -> np.ndarray: + return np.array([self.x_coordinates[column], self.y_coordinates[row]]) + + +class AssetsGroupSupportClamp: + """Find a small shared XY shift that places all AABBs on a support region. + + Each AABB gets a feasible-centre map obtained by binary erosion of the safe + support mask. A candidate translation is valid only when it is feasible + for every asset, then it must pass exact Shapely containment. This supports + concave polygons, holes, and disconnected ``MultiPolygon`` regions. + """ + + def __init__( + self, + *, + support_region: SupportGeometry, + assets_aabb_2d_z_up_world_corners_by_id: dict[str, np.ndarray], + assets_layout: list[dict[str, object]], + debug_output_root: str | Path | None = None, + config: AssetsGroupSupportClampConfig | None = None, + ) -> None: + self.support_region = support_region + self.assets_aabb_2d_z_up_world_corners_by_id = ( + assets_aabb_2d_z_up_world_corners_by_id + ) + self.assets_layout = assets_layout + self.debug_output_root = ( + Path(debug_output_root).expanduser().resolve() + if debug_output_root is not None + else None + ) + self.refined_assets_layout: list[dict[str, object]] | None = None + self.config = config if config is not None else AssetsGroupSupportClampConfig() + # Check. + if self.config.margin_m < 0.0: + raise ValueError("margin_m must be non-negative.") + if self.config.grid_resolution_m <= 0.0: + raise ValueError("grid_resolution_m must be positive.") + + def clamp(self) -> list[dict[str, object]]: + """Return y-up layouts after one rigid, support-valid XY translation.""" + self.refined_assets_layout = None + # Validate input aabbs. + aabbs_by_id = self._validate_aabbs(self.assets_aabb_2d_z_up_world_corners_by_id) + + # Validate and coerce the support region into a usable polygonal geometry. + raw_support = self._coerce_support_geometry(self.support_region) + if raw_support is None: + log_warning("Asset-group support clamp failed: invalid support geometry.") + raise ValueError( + "Asset-group support clamp requires a valid support region." + ) + # Re coerce the support region with a margin to get the safe support region. + safe_support = ( + raw_support + if self.config.margin_m == 0.0 + else self._polygonal_geometry(raw_support.buffer(-self.config.margin_m)) + ) + if safe_support is None or safe_support.is_empty: + log_warning( + "Asset-group support clamp failed: support region is empty after " + f"applying a {self.config.margin_m:.4f} m boundary margin." + ) + raise ValueError("Asset-group support clamp has no usable support area.") + + # Get all the translated layouts, and store them for later debug rendering. + delta_xy = self._find_clamp_delta( + safe_support=safe_support, + aabbs_by_id=aabbs_by_id, + ) + if delta_xy is None: + log_warning( + "Asset-group support clamp failed: no shared translation can place " + f"all {len(aabbs_by_id)} AABBs inside the support region." + ) + raise ValueError( + "Asset clutter cannot be placed completely on the detected table " + "support region." + ) + # Translate the layouts and store them for later debug rendering. + refined_assets_layout = self._apply_delta_to_y_up_layouts( + delta_xy=delta_xy, + expected_ids=set(aabbs_by_id), + ) + self.refined_assets_layout = refined_assets_layout + if np.allclose( + delta_xy, 0.0 + ): # Judge whether the translation is zero, if so, no need to apply optimization. + log_info( + "All asset AABBs are already fully inside the detected table " + "support region; no planar group optimization was applied." + ) + else: + log_info( + "Applied rigid asset-group support optimization with " + f"delta_xy={delta_xy.tolist()} m; all AABBs passed " + "exact support containment." + ) + return refined_assets_layout + + def _find_clamp_delta( + self, + *, + safe_support: Polygon | MultiPolygon, + aabbs_by_id: dict[str, np.ndarray], + ) -> np.ndarray | None: + """Find one common z-up XY translation, or return ``None``. + + The initial position is always checked exactly first. Otherwise, grid + candidates are considered in increasing translation distance and the + first vector-valid placement is returned. + """ + + # Get all aabbs's clutter center as anchor, to keep the internal + # relative layouts unchanged. + anchor_xy = self._group_anchor(aabbs_by_id) + log_info( + "Asset-group support clamp started: " + f"assets={len(aabbs_by_id)}, support_area={safe_support.area:.4f} m^2, " + f"boundary_margin={self.config.margin_m:.4f} m, " + f"grid_resolution={self.config.grid_resolution_m:.4f} m." + ) + + # Check whether the initial position is already valid, which is a common case. + zero_translation = np.zeros(2, dtype=float) + if self._is_exactly_contained(safe_support, aabbs_by_id, zero_translation): + log_info( + "Asset-group support clamp succeeded without movement: all AABBs " + "are exactly contained by the safe support region." + ) + return zero_translation + + # Rasterize the support region and compute feasible-centre maps for each AABB.s + transform, support_mask = self._rasterize_support(safe_support) + + # Compute feasible-centre maps with cacheing to avoid repeated binary erosion + # for identical AABB half-extents. + feasible_maps_by_id, centres_by_id = self._feasible_maps( + support_mask=support_mask, + transform=transform, + aabbs_by_id=aabbs_by_id, + ) + candidate_pixels = np.argwhere(support_mask) + if len(candidate_pixels) == 0: + return None + + candidate_world = np.asarray( + [ + transform.pixel_to_world(int(row), int(column)) + for row, column in candidate_pixels + ] + ) + candidate_deltas = candidate_world - anchor_xy + candidate_order = np.argsort( + np.einsum("ij,ij->i", candidate_deltas, candidate_deltas), kind="stable" + ) + for candidate_rank, candidate_index in enumerate(candidate_order, start=1): + delta_xy = candidate_deltas[candidate_index] + if not self._grid_translation_is_feasible( + delta_xy=delta_xy, + transform=transform, + feasible_maps_by_id=feasible_maps_by_id, + centres_by_id=centres_by_id, + ): + continue + if self._is_exactly_contained(safe_support, aabbs_by_id, delta_xy): + log_info( + "Asset-group support clamp succeeded after evaluating " + f"{candidate_rank}/{len(candidate_order)} grid candidates: " + f"delta_xy=({delta_xy[0]:+.4f}, {delta_xy[1]:+.4f}) m." + ) + return delta_xy + return None + + def _apply_delta_to_y_up_layouts( + self, *, delta_xy: np.ndarray, expected_ids: set[str] + ) -> list[dict[str, object]]: + """Apply a successful common z-up XY translation to stored layouts.""" + received_ids = {str(layout.get("id")) for layout in self.assets_layout} + if received_ids != expected_ids: + raise ValueError("Asset layouts and clamped AABBs must have identical ids.") + + dx, dy = delta_xy + translated_layouts: list[dict[str, object]] = [] + for layout in self.assets_layout: + position = layout.get("pos") + if not isinstance(position, list) or len(position) != 3: + raise ValueError( + "Each asset layout must contain a three-value pos list." + ) + translated_layout = dict(layout) + translated_position = [float(value) for value in position] + translated_position[0] += float(dx) + translated_position[2] -= float(dy) + translated_layout["pos"] = translated_position + translated_layouts.append(translated_layout) + return translated_layouts + + def save_group_clamp_debug_images(self) -> bool: + """Optionally save diagnostics for the support-valid group translation.""" + if self.refined_assets_layout is None: + self.clamp() + assert self.refined_assets_layout is not None + + initial_aabbs_by_id = self._validate_aabbs( + self.assets_aabb_2d_z_up_world_corners_by_id + ) + raw_support = self._coerce_support_geometry(self.support_region) + if raw_support is None: + raise ValueError( + "Asset-group support clamp requires a valid support region." + ) + safe_support = ( + raw_support + if self.config.margin_m == 0.0 + else self._polygonal_geometry(raw_support.buffer(-self.config.margin_m)) + ) + if safe_support is None or safe_support.is_empty: + raise ValueError("Asset-group support clamp has no usable support area.") + + original_positions_by_id = { + str(layout["id"]): layout["pos"] for layout in self.assets_layout + } + refined_positions_by_id = { + str(layout["id"]): layout["pos"] for layout in self.refined_assets_layout + } + if set(original_positions_by_id) != set(initial_aabbs_by_id) or set( + refined_positions_by_id + ) != set(initial_aabbs_by_id): + raise ValueError("Asset layouts and clamped AABBs must have identical ids.") + first_asset_id = next(iter(initial_aabbs_by_id)) + original_position = original_positions_by_id[first_asset_id] + refined_position = refined_positions_by_id[first_asset_id] + if ( + not isinstance(original_position, list) + or not isinstance(refined_position, list) + or len(original_position) != 3 + or len(refined_position) != 3 + ): + raise ValueError("Each asset layout must contain a three-value pos list.") + delta_xy = np.array( + [ + float(refined_position[0]) - float(original_position[0]), + float(original_position[2]) - float(refined_position[2]), + ] + ) + translated_aabbs_by_id = { + asset_id: corners + delta_xy + for asset_id, corners in initial_aabbs_by_id.items() + } + if self.debug_output_root is None: + raise ValueError( + "A debug_output_root is required when saving group-clamp " + "debug images." + ) + self._render_support_debug( + raw_support=raw_support, + safe_support=safe_support, + output_path=self.debug_output_root / "table_support_region_safe_2d.png", + ) + self._render_clamp_debug( + raw_support=raw_support, + initial_aabbs_by_id=initial_aabbs_by_id, + translated_aabbs_by_id=translated_aabbs_by_id, + delta_xy=delta_xy, + output_path=self.debug_output_root / "assets_group_support_clamp_2d.png", + ) + return True + + def _feasible_maps( + self, + *, + support_mask: np.ndarray, + transform: _GridTransform, + aabbs_by_id: dict[str, np.ndarray], + ) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]: + """Compute each AABB's rasterized feasible-centre map and current XY centre.""" + cached_feasible_maps: dict[tuple[int, int], np.ndarray] = {} + feasible_maps_by_id: dict[str, np.ndarray] = {} + centres_by_id: dict[str, np.ndarray] = {} + for asset_id, corners in aabbs_by_id.items(): + minimum = corners.min(axis=0) + maximum = corners.max(axis=0) + half_extent = (maximum - minimum) / 2.0 + kernel_key = tuple( + np.ceil(half_extent / transform.resolution_m).astype(int) + ) + if kernel_key not in cached_feasible_maps: + cached_feasible_maps[kernel_key] = binary_erosion( + support_mask, + structure=self._footprint_kernel(kernel_key), + border_value=0, + ) + feasible_maps_by_id[asset_id] = cached_feasible_maps[kernel_key] + centres_by_id[asset_id] = (minimum + maximum) / 2.0 + return feasible_maps_by_id, centres_by_id + + @staticmethod + def _footprint_kernel(kernel_key: tuple[int, int]) -> np.ndarray: + half_width_pixels, half_height_pixels = kernel_key + return np.ones( + (2 * half_height_pixels + 1, 2 * half_width_pixels + 1), dtype=bool + ) + + @staticmethod + def _grid_translation_is_feasible( + *, + delta_xy: np.ndarray, + transform: _GridTransform, + feasible_maps_by_id: dict[str, np.ndarray], + centres_by_id: dict[str, np.ndarray], + ) -> bool: + height, width = transform.shape + # Sorting by feasible map population is a cheap early-rejection order: + # small legal regions are most likely to reject a candidate quickly. + ordered_assets = sorted( + centres_by_id.items(), + key=lambda item: int(feasible_maps_by_id[item[0]].sum()), + ) + for asset_id, centre_xy in ordered_assets: + row, column = transform.world_to_nearest_pixel(centre_xy + delta_xy) + if row < 0 or row >= height or column < 0 or column >= width: + return False + if not feasible_maps_by_id[asset_id][row, column]: + return False + return True + + def _rasterize_support( + self, support: Polygon | MultiPolygon + ) -> tuple[_GridTransform, np.ndarray]: + """Rasterize a support region into a boolean XY grid and return the transform.""" + minimum_x, minimum_y, maximum_x, maximum_y = support.bounds + resolution = self.config.grid_resolution_m + x_coordinates = np.arange( + np.floor(minimum_x / resolution) * resolution, + np.ceil(maximum_x / resolution) * resolution + resolution / 2.0, + resolution, + ) + y_coordinates = np.arange( + np.floor(minimum_y / resolution) * resolution, + np.ceil(maximum_y / resolution) * resolution + resolution / 2.0, + resolution, + ) + x_grid, y_grid = np.meshgrid(x_coordinates, y_coordinates) + points = np.column_stack((x_grid.ravel(), y_grid.ravel())) + mask = np.zeros(len(points), dtype=bool) + for polygon in self._polygon_components(support): + component_mask = self._points_in_polygon(points, polygon.exterior.coords) + for hole in polygon.interiors: + component_mask &= ~self._points_in_polygon(points, hole.coords) + mask |= component_mask + return ( + _GridTransform( + x_coordinates, y_coordinates, resolution + ), # Keeps the transform for later world/grid conversions + mask.reshape( + x_grid.shape + ), # Keeps the boolean mask of the support region in grid form + ) + + @staticmethod + def _points_in_polygon(points: np.ndarray, coordinates: object) -> np.ndarray: + from matplotlib.path import Path as MatplotlibPath + + return MatplotlibPath(np.asarray(coordinates)).contains_points( + points, radius=1e-12 + ) + + @staticmethod + def _validate_aabbs( + aabbs_by_id: dict[str, np.ndarray], + ) -> dict[str, np.ndarray]: + """Validate the input AABB(s).""" + if not aabbs_by_id: + raise ValueError("At least one asset 2D AABB is required.") + validated: dict[str, np.ndarray] = {} + for asset_id, corners in aabbs_by_id.items(): + if not isinstance(asset_id, str) or not asset_id: + raise ValueError("Each asset AABB id must be a non-empty string.") + corner_array = np.asarray(corners, dtype=float) + if corner_array.shape != (4, 2) or not np.isfinite(corner_array).all(): + raise ValueError( + f"Asset {asset_id!r} must have four finite XY corners." + ) + validated[asset_id] = corner_array + return validated + + @classmethod + def _coerce_support_geometry( + cls, support_region: SupportGeometry + ) -> Polygon | MultiPolygon | None: + # Check the input type. + if isinstance(support_region, (Polygon, MultiPolygon)): + geometry = support_region + else: + log_warning( + "Unsupported support region type: " f"{type(support_region).__name__}." + ) + return None + return cls._polygonal_geometry(geometry) + + @staticmethod + def _polygonal_geometry(geometry: object) -> Polygon | MultiPolygon | None: + if not isinstance(geometry, (Polygon, MultiPolygon)) or geometry.is_empty: + return None + repaired = geometry if geometry.is_valid else geometry.buffer(0) + if not repaired.is_valid: + log_warning("Support polygon repair did not produce a valid geometry.") + return None + if isinstance(repaired, (Polygon, MultiPolygon)): + return repaired + if isinstance(repaired, GeometryCollection): + polygons = [item for item in repaired.geoms if isinstance(item, Polygon)] + return MultiPolygon(polygons) if polygons else None + return None + + @staticmethod + def _group_anchor(aabbs_by_id: dict[str, np.ndarray]) -> np.ndarray: + all_corners = np.concatenate(list(aabbs_by_id.values()), axis=0) + return (all_corners.min(axis=0) + all_corners.max(axis=0)) / 2.0 + + @staticmethod + def _is_exactly_contained( + support: Polygon | MultiPolygon, + aabbs_by_id: dict[str, np.ndarray], + delta_xy: np.ndarray, + ) -> bool: + return all( + support.covers( + affinity.translate( + AssetsGroupSupportClamp._aabb_polygon(corners), + xoff=float(delta_xy[0]), + yoff=float(delta_xy[1]), + ) + ) + for corners in aabbs_by_id.values() + ) + + @staticmethod + def _aabb_polygon(corners: np.ndarray) -> Polygon: + """Build a non-self-intersecting footprint regardless of corner order.""" + minimum = corners.min(axis=0) + maximum = corners.max(axis=0) + return Polygon( + [ + (minimum[0], minimum[1]), + (maximum[0], minimum[1]), + (maximum[0], maximum[1]), + (minimum[0], maximum[1]), + ] + ) + + def _render_support_debug( + self, + *, + raw_support: Polygon | MultiPolygon, + safe_support: Polygon | MultiPolygon | None, + output_path: str | Path, + ) -> Path: + path = self._resolve_png_output_path(output_path) + path.parent.mkdir(parents=True, exist_ok=True) + figure, axes = plt.subplots( + 1, 2, figsize=(14, 7), dpi=160, constrained_layout=True + ) + self._draw_support(axes[0], raw_support, "Input support region", "darkorange") + title = f"Safe support (margin={self.config.margin_m:.3f} m)" + if safe_support is None or safe_support.is_empty: + axes[1].set_title(f"{title}\n(infeasible)") + axes[1].set_aspect("equal", adjustable="box") + else: + self._draw_support(axes[1], safe_support, title, "seagreen") + figure.savefig(path, bbox_inches="tight") + plt.close(figure) + return path + + def _render_clamp_debug( + self, + *, + raw_support: Polygon | MultiPolygon, + initial_aabbs_by_id: dict[str, np.ndarray], + translated_aabbs_by_id: dict[str, np.ndarray] | None, + delta_xy: np.ndarray | None, + output_path: str | Path, + ) -> Path: + path = self._resolve_png_output_path(output_path) + path.parent.mkdir(parents=True, exist_ok=True) + figure, axes = plt.subplots( + 1, 2, figsize=(14, 7), dpi=160, constrained_layout=True + ) + self._draw_state( + axes[0], raw_support, initial_aabbs_by_id, "Before group clamp", "royalblue" + ) + if translated_aabbs_by_id is not None and delta_xy is not None: + action = "no movement" if np.allclose(delta_xy, 0.0) else "translated" + self._draw_state( + axes[1], + raw_support, + translated_aabbs_by_id, + f"After group clamp ({action})\nΔxy=({delta_xy[0]:+.3f}, {delta_xy[1]:+.3f}) m", + "seagreen", + ) + else: + self._draw_state( + axes[1], + raw_support, + initial_aabbs_by_id, + "No feasible group translation", + "firebrick", + ) + figure.savefig(path, bbox_inches="tight") + plt.close(figure) + return path + + @classmethod + def _draw_support( + cls, axis: plt.Axes, support: Polygon | MultiPolygon, title: str, color: str + ) -> None: + for polygon in cls._polygon_components(support): + exterior = np.asarray(polygon.exterior.coords) + axis.fill( + exterior[:, 0], + exterior[:, 1], + facecolor=color, + edgecolor="saddlebrown", + alpha=0.35, + ) + for hole in polygon.interiors: + hole_points = np.asarray(hole.coords) + axis.fill( + hole_points[:, 0], + hole_points[:, 1], + facecolor="white", + edgecolor="saddlebrown", + alpha=1.0, + ) + axis.set_aspect("equal", adjustable="box") + axis.set_xlabel("x (z-up world)") + axis.set_ylabel("y (z-up world)") + axis.set_title(title) + axis.autoscale_view() + + @classmethod + def _draw_state( + cls, + axis: plt.Axes, + support: Polygon | MultiPolygon, + aabbs_by_id: dict[str, np.ndarray], + title: str, + color: str, + ) -> None: + cls._draw_support(axis, support, title, "darkorange") + for asset_id, corners in sorted(aabbs_by_id.items()): + polygon = cls._aabb_polygon(corners) + boundary = np.asarray(polygon.exterior.coords) + axis.fill( + boundary[:, 0], + boundary[:, 1], + facecolor=color, + edgecolor=color, + alpha=0.35, + ) + axis.text( + *corners.mean(axis=0), + asset_id, + ha="center", + va="center", + fontsize=8, + bbox={"facecolor": "white", "alpha": 0.75, "edgecolor": "none"}, + ) + + @staticmethod + def _polygon_components(geometry: Polygon | MultiPolygon) -> list[Polygon]: + return [geometry] if isinstance(geometry, Polygon) else list(geometry.geoms) + + @staticmethod + def _resolve_png_output_path(output_path: str | Path) -> Path: + path = Path(output_path).expanduser().resolve() + return path if path.suffix.lower() == ".png" else path.with_suffix(".png") diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/assets_group_table_aligner.py b/embodichain/gen_sim/scene_engine/pipeline/utils/assets_group_table_aligner.py new file mode 100644 index 000000000..990476044 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/assets_group_table_aligner.py @@ -0,0 +1,150 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +import numpy as np + +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_generation_utils import ( + layout_object_to_transform_matrix, + load_glb_mesh, + transform_matrix_to_layout_object, +) +from embodichain.utils.logger import log_info + + +@dataclass(frozen=True) +class AssetsGroupTableAlignerConfig: + """Controls for the initial vertical gap above the table.""" + + clearance_m: float = 0.02 # Initial table-to-group gap in metres. + + +class AssetsGroupTableAligner: + """Place every asset as one rigid vertical group above a table AABB top.""" + + def __init__( + self, + *, + table_layout: dict[str, object], + assets_layout: list[dict[str, object]], + geometry_root: str | Path, + config: AssetsGroupTableAlignerConfig | None = None, + ) -> None: + self.table_layout = table_layout + self.assets_layout = assets_layout + self.geometry_root = Path(geometry_root).expanduser().resolve() + self.aligned_table_layout: dict[str, object] | None = None + self.aligned_assets_layout: list[dict[str, object]] | None = None + self.config = config if config is not None else AssetsGroupTableAlignerConfig() + # Check. + if self.config.clearance_m < 0.0: + raise ValueError("Table clearance_m must be non-negative.") + + def align(self) -> tuple[dict[str, object], list[dict[str, object]]]: + """Return y-up layouts with the complete asset group above the table. + + Input and output layouts use y-up, matching the GLBs on disk. The group + is temporarily measured in z-up coordinates and every asset receives the + same vertical translation. This preserves all asset-to-asset relative + poses. + """ + self.aligned_table_layout = None + self.aligned_assets_layout = None + if not self.assets_layout: + self.aligned_table_layout = self.table_layout + self.aligned_assets_layout = [] + log_info("Scene has no movable assets; skipping vertical group alignment.") + return self.aligned_table_layout, self.aligned_assets_layout + + y_up_to_z_up_matrix = np.eye(4) + y_up_to_z_up_matrix[:3, :3] = np.array( + [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]] + ) + z_up_to_y_up_matrix = np.linalg.inv(y_up_to_z_up_matrix) + + z_up_table_layout = self._convert_layout_coordinate_system( + self.table_layout, + source_to_target_matrix=y_up_to_z_up_matrix, + ) + z_up_assets_layout = [ + self._convert_layout_coordinate_system( + asset_layout, + source_to_target_matrix=y_up_to_z_up_matrix, + ) + for asset_layout in self.assets_layout + ] + + table_id = self._require_layout_id(z_up_table_layout, name="Table") + table_mesh = load_glb_mesh(self.geometry_root / f"{table_id}.glb") + table_mesh.apply_transform(y_up_to_z_up_matrix) + table_mesh.apply_transform(layout_object_to_transform_matrix(z_up_table_layout)) + target_group_bottom_z = table_mesh.bounds[1, 2] + self.config.clearance_m + + group_bottom_z = np.inf + for asset_layout in z_up_assets_layout: + asset_id = self._require_layout_id(asset_layout, name="Asset") + asset_mesh = load_glb_mesh(self.geometry_root / f"{asset_id}.glb") + asset_mesh.apply_transform(y_up_to_z_up_matrix) + asset_mesh.apply_transform(layout_object_to_transform_matrix(asset_layout)) + # Find the lowest z among all the assets. + group_bottom_z = min(group_bottom_z, float(asset_mesh.bounds[0, 2])) + + group_vertical_translation_z = target_group_bottom_z - group_bottom_z + for asset_layout in z_up_assets_layout: + asset_layout["pos"][2] += group_vertical_translation_z + + self.aligned_table_layout = self._convert_layout_coordinate_system( + z_up_table_layout, + source_to_target_matrix=z_up_to_y_up_matrix, + ) + self.aligned_assets_layout = [ + self._convert_layout_coordinate_system( + asset_layout, + source_to_target_matrix=z_up_to_y_up_matrix, + ) + for asset_layout in z_up_assets_layout + ] + log_info( + "Aligned the asset group above the table with " + f"delta_z={group_vertical_translation_z:.4f} m and " + f"clearance={self.config.clearance_m:.4f} m." + ) + return self.aligned_table_layout, self.aligned_assets_layout + + @staticmethod + def _convert_layout_coordinate_system( + layout_object: dict[str, object], + *, + source_to_target_matrix: np.ndarray, + ) -> dict[str, object]: + """Convert one layout object between coordinate systems through its matrix.""" + return transform_matrix_to_layout_object( + str(layout_object["id"]), + source_to_target_matrix + @ layout_object_to_transform_matrix(layout_object) + @ np.linalg.inv(source_to_target_matrix), + ) + + @staticmethod + def _require_layout_id(layout_object: dict[str, object], *, name: str) -> str: + object_id = layout_object.get("id") + if not isinstance(object_id, str) or not object_id: + raise ValueError(f"{name} layout must contain a non-empty string id.") + return object_id diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py b/embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py new file mode 100644 index 000000000..e0c3f772a --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py @@ -0,0 +1,454 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from PIL import Image, ImageChops, ImageDraw, ImageFilter, ImageFont + +from embodichain.utils.logger import log_warning + + +@dataclass(frozen=True) +class MaskCandidate: + """One numbered mask candidate returned by the Image Segmentation Server.""" + + index: int + mask_rle: dict[str, Any] + + +def build_mask_candidates(mask_rles: list[dict[str, Any]]) -> list[MaskCandidate]: + return [ + MaskCandidate(index=index, mask_rle=mask_rle) + for index, mask_rle in enumerate(mask_rles, start=1) + ] + + +def decode_rle_mask(mask_rle: dict[str, Any]) -> Image.Image: + """Decode an uncompressed RLE mask into a binary image.""" + + # Check the return value's format. + size = mask_rle.get("size") + counts = mask_rle.get("counts") + if ( + not isinstance(size, list) + or len(size) != 2 + or not all(isinstance(value, int) and value > 0 for value in size) + ): + raise ValueError("Image Segmentation Server RLE needs size=[height, width].") + if not isinstance(counts, list): + raise ValueError("Image Segmentation Server RLE counts must be a list.") + + height, width = size + pixel_count = height * width + starts_with = mask_rle.get("starts_with", 0) + if starts_with not in (0, 1, False, True): + raise ValueError("Image Segmentation Server RLE starts_with must be 0 or 1.") + + pixels = bytearray(pixel_count) + is_foreground = bool( + starts_with + ) # True for white foreground, False for black background. + offset = 0 # How many pixels have been filled so far. + for raw_count in counts: + if isinstance(raw_count, bool): + raise ValueError( + "Image Segmentation Server RLE counts must contain integers." + ) + try: + count = int(raw_count) + except (TypeError, ValueError) as exc: + raise ValueError( + "Image Segmentation Server RLE counts must contain integers." + ) from exc + if count < 0 or offset + count > pixel_count: + raise ValueError( + "Image Segmentation Server RLE counts do not match its declared size." + ) + if is_foreground: + pixels[offset : offset + count] = ( + b"\xff" * count + ) # Write white pixels for the foreground. + offset += count + is_foreground = not is_foreground + + if offset != pixel_count: + raise ValueError("Image Segmentation Server RLE does not cover the image.") + return Image.frombytes("L", (width, height), bytes(pixels)) + + +def union_overlapping_mask_candidates( + candidates: list[MaskCandidate], + *, + min_iou: float = 0.8, +) -> list[MaskCandidate]: + """Union candidate masks with IOU >= min_iou into one mask candidate.""" + if not 0 < min_iou <= 1: + raise ValueError("min_iou must be greater than 0 and at most 1.") + if not candidates: + return [] + + masks = [decode_rle_mask(candidate.mask_rle) for candidate in candidates] + image_size = masks[0].size + for mask in masks: + _require_image_size(mask, image_size) + + parents = list( + range(len(candidates)) + ) # Initialize the Union-Find data structure for candidates. + for first_index, first_mask in enumerate(masks): + for second_index in range(first_index + 1, len(masks)): + if _mask_iou(first_mask, masks[second_index]) >= min_iou: + _union_parent( + parents, first_index, second_index + ) # Union the two candidates into one. + + grouped_indices: dict[int, list[int]] = {} + for index in range(len(candidates)): + # Put all the index of the same parent into one group. + grouped_indices.setdefault(_find_parent(parents, index), []).append(index) + + merged_candidates: list[MaskCandidate] = [] + for merged_index, member_indices in enumerate(grouped_indices.values(), start=1): + merged_mask = masks[member_indices[0]] + for member_index in member_indices[1:]: + # Union the masks of the same group into one mask (lighter = union). + merged_mask = ImageChops.lighter(merged_mask, masks[member_index]) + merged_candidates.append( + MaskCandidate( + index=merged_index, + mask_rle=_encode_binary_mask_rle(merged_mask), + ) + ) + return merged_candidates + + +def save_binary_mask( + candidate: MaskCandidate, + *, + image_size: tuple[int, int], + output_path: str | Path, +) -> Path: + """Save one candidate as a white-foreground, black-background PNG mask.""" + mask = decode_rle_mask(candidate.mask_rle) + _require_image_size(mask, image_size) # Check whether the image size == mask size. + + resolved_output_path = Path(output_path).expanduser().resolve() + resolved_output_path.parent.mkdir(parents=True, exist_ok=True) + mask.save(resolved_output_path) + return resolved_output_path + + +def render_image_without_masks( + *, + image_path: str | Path, + mask_paths: list[str | Path], + output_path: str | Path, + removed_color: tuple[int, int, int] = (128, 128, 128), +) -> tuple[Path, Image.Image]: + """Gray masked regions and return the image path with their combined mask.""" + image = Image.open(image_path).convert("RGB") + ignored_mask = Image.new("L", image.size, 0) + for mask_path in mask_paths: + mask = Image.open(mask_path).convert("L") + _require_image_size(mask, image.size) + ignored_mask = ImageChops.lighter(ignored_mask, mask) + + removed_layer = Image.new("RGB", image.size, removed_color) + result = Image.composite(removed_layer, image, ignored_mask) + resolved_output_path = Path(output_path).expanduser().resolve() + resolved_output_path.parent.mkdir(parents=True, exist_ok=True) + result.save(resolved_output_path) + return resolved_output_path, ignored_mask + + +def render_numbered_mask_candidates( + *, + image_path: str | Path, + candidates: list[MaskCandidate], + output_path: str | Path, + mask_style: str = "fill", + label_avoid_mask: Image.Image | None = None, +) -> Path: + """Overlay numbered mask candidates on their source image. + Notice that: + - mask_style can be either "fill" or "outline". + - The label font and its background scale with the source image resolution. + - label_avoid_mask keeps labels outside known occluding regions. + """ + if mask_style not in {"fill", "outline"}: + raise ValueError("mask_style must be 'fill' or 'outline'.") + + image = Image.open(image_path).convert("RGBA") + overlay = Image.new("RGBA", image.size, (0, 0, 0, 0)) + colors = ( + (239, 83, 80, 160), + (66, 165, 245, 160), + (102, 187, 106, 160), + (255, 202, 40, 160), + (171, 71, 188, 160), + (38, 198, 218, 160), + ) + + decoded_masks: list[tuple[MaskCandidate, Image.Image]] = [] + for candidate in candidates: + mask = decode_rle_mask(candidate.mask_rle) + _require_image_size(mask, image.size) + decoded_masks.append((candidate, mask)) + color_layer = Image.new( + "RGBA", image.size, colors[(candidate.index - 1) % len(colors)] + ) + transparent_layer = Image.new("RGBA", image.size, (0, 0, 0, 0)) + rendered_mask = ( + mask if mask_style == "fill" else _mask_outer_outline(mask, image.size) + ) + overlay.alpha_composite( + Image.composite(color_layer, transparent_layer, rendered_mask) + ) + + draw = ImageDraw.Draw(overlay) # Initialize a draw object. + font = _load_label_font(image.size) + if label_avoid_mask is not None: + label_blocked_mask = label_avoid_mask.convert("L") + _require_image_size(label_blocked_mask, image.size) + else: + label_blocked_mask = None + label_occupied_mask = Image.new("L", image.size, 0) + for candidate, mask in decoded_masks: + bbox = mask.getbbox() + if bbox is None: + raise ValueError( + f"Image Segmentation Server candidate {candidate.index} has an empty mask." + ) + center = ( + ((bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2) + if label_blocked_mask is None + else _find_label_center_inside_bbox( + candidate_mask=mask, + candidate_bbox=bbox, + blocked_mask=ImageChops.lighter( + label_blocked_mask, label_occupied_mask + ), + label=str(candidate.index), + font=font, + ) + ) + if center is None: + # Keep every candidate selectable even when its AABB is entirely + # occluded. This is preferable to silently omitting its number. + log_warning( + "Could not place table-candidate label %s outside masked objects " + "while keeping it inside the candidate AABB; using the AABB center.", + candidate.index, + ) + center = ((bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2) + label_bounds = _number_label_bounds( + draw=draw, label=str(candidate.index), center=center, font=font + ) + _draw_number_label( + draw=draw, + label=str(candidate.index), + center=center, + font=font, + ) + if label_blocked_mask is not None: + ImageDraw.Draw(label_occupied_mask).rectangle(label_bounds, fill=255) + + resolved_output_path = Path(output_path).expanduser().resolve() + resolved_output_path.parent.mkdir(parents=True, exist_ok=True) + Image.alpha_composite(image, overlay).convert("RGB").save(resolved_output_path) + return resolved_output_path + + +def _require_image_size(mask: Image.Image, image_size: tuple[int, int]) -> None: + if mask.size != image_size: + raise ValueError( + "Image Segmentation Server mask size does not match the input image: " + f"{mask.size} != {image_size}." + ) + + +def _find_label_center_inside_bbox( + *, + candidate_mask: Image.Image, + candidate_bbox: tuple[int, int, int, int], + blocked_mask: Image.Image, + label: str, + font: ImageFont.ImageFont, +) -> tuple[float, float] | None: + """Find the nearest label centre that stays in a candidate AABB and avoids masks.""" + probe_draw = ImageDraw.Draw(Image.new("RGBA", candidate_mask.size)) + bounds_at_origin = _number_label_bounds( + draw=probe_draw, label=label, center=(0.0, 0.0), font=font + ) + minimum_x = candidate_bbox[0] - bounds_at_origin[0] + maximum_x = candidate_bbox[2] - 1 - bounds_at_origin[2] + minimum_y = candidate_bbox[1] - bounds_at_origin[1] + maximum_y = candidate_bbox[3] - 1 - bounds_at_origin[3] + if minimum_x > maximum_x or minimum_y > maximum_y: + return None + + # Expand each blocked region by the label's largest half-extent, so testing + # one candidate centre guarantees the complete label rectangle stays clear. + label_radius = max( + -bounds_at_origin[0], + bounds_at_origin[2], + -bounds_at_origin[1], + bounds_at_origin[3], + ) + blocked_centres = blocked_mask.filter(ImageFilter.MaxFilter(2 * label_radius + 1)) + bbox_center = ( + (candidate_bbox[0] + candidate_bbox[2]) / 2, + (candidate_bbox[1] + candidate_bbox[3]) / 2, + ) + for require_candidate_mask in (True, False): + for step in (max(1, round(min(candidate_mask.size) / 512)), 1): + best_center: tuple[float, float] | None = None + best_distance_squared = float("inf") + for y_coordinate in range(minimum_y, maximum_y + 1, step): + for x_coordinate in range(minimum_x, maximum_x + 1, step): + if blocked_centres.getpixel((x_coordinate, y_coordinate)): + continue + if require_candidate_mask and not candidate_mask.getpixel( + (x_coordinate, y_coordinate) + ): + continue + distance_squared = (x_coordinate - bbox_center[0]) ** 2 + ( + y_coordinate - bbox_center[1] + ) ** 2 + if distance_squared < best_distance_squared: + best_center = (float(x_coordinate), float(y_coordinate)) + best_distance_squared = distance_squared + if best_center is not None: + return best_center + return None + + +def _mask_outer_outline(mask: Image.Image, image_size: tuple[int, int]) -> Image.Image: + """Use dilation and subtraction to get the outer outline of a binary mask.""" + # Asset-candidate images use outlines only; keep them visible at common + # image resolutions without obscuring the original object appearance. + outline_width = max(2, round(min(image_size) / 200)) + dilated_mask = mask.filter(ImageFilter.MaxFilter(outline_width * 2 + 1)) + return ImageChops.subtract(dilated_mask, mask) + + +def _mask_iou(first_mask: Image.Image, second_mask: Image.Image) -> float: + """Compute the Intersection over Union (IoU) of two binary masks.""" + _require_image_size(second_mask, first_mask.size) + intersection = ImageChops.multiply(first_mask, second_mask) + union = ImageChops.lighter(first_mask, second_mask) + union_pixels = union.histogram()[255] + if union_pixels == 0: + return 0.0 + return intersection.histogram()[255] / union_pixels + + +def _encode_binary_mask_rle(mask: Image.Image) -> dict[str, Any]: + binary_mask = mask.convert("L").point( + lambda value: 255 if value else 0 + ) # Force translate an image into a binary mask. + width, height = binary_mask.size + counts: list[int] = [] + current_value = 0 + run_length = 0 + for value in binary_mask.tobytes(): + value = 255 if value else 0 + if value == current_value: + run_length += 1 + continue + counts.append(run_length) + current_value = value + run_length = 1 + counts.append(run_length) + return { + "size": [height, width], + "counts": counts, + "starts_with": 0, + } + + +def _find_parent(parents: list[int], index: int) -> int: + while parents[index] != index: + parents[index] = parents[parents[index]] + index = parents[index] + return index + + +def _union_parent(parents: list[int], first_index: int, second_index: int) -> None: + first_root = _find_parent(parents, first_index) + second_root = _find_parent(parents, second_index) + if first_root != second_root: + parents[second_root] = first_root + + +def _load_label_font(image_size: tuple[int, int]) -> ImageFont.ImageFont: + font_size = max(16, round(min(image_size) / 32)) + try: + return ImageFont.truetype("DejaVuSans-Bold.ttf", font_size) + except OSError: + return ImageFont.load_default() + + +def _draw_number_label( + *, + draw: ImageDraw.ImageDraw, + label: str, + center: tuple[float, float], + font: ImageFont.ImageFont, +) -> None: + """Draw a numbered label with red background and white text at the given center position.""" + label_bounds = _number_label_bounds( + draw=draw, label=label, center=center, font=font + ) + label_box = draw.textbbox((0, 0), label, font=font) + label_width = label_box[2] - label_box[0] + label_height = label_box[3] - label_box[1] + x = center[0] - label_width / 2 + y = center[1] - label_height / 2 + draw.rectangle( + label_bounds, + fill=(220, 0, 0, 255), + outline=(255, 255, 255, 255), + width=max(1, round(max(label_width, label_height) / 12)), + ) + draw.text((x, y), label, fill=(255, 255, 255, 255), font=font) + + +def _number_label_bounds( + *, + draw: ImageDraw.ImageDraw, + label: str, + center: tuple[float, float], + font: ImageFont.ImageFont, +) -> tuple[int, int, int, int]: + """Return the red label rectangle bounds for a label centre.""" + label_box = draw.textbbox((0, 0), label, font=font) + label_width = label_box[2] - label_box[0] + label_height = label_box[3] - label_box[1] + padding = max(4, round(max(label_width, label_height) / 4)) + x = center[0] - label_width / 2 + y = center[1] - label_height / 2 + return ( + round(x - padding), + round(y - padding), + round(x + label_width + padding), + round(y + label_height + padding), + ) diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py new file mode 100644 index 000000000..fb66c30c4 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py @@ -0,0 +1,199 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +import json +from pathlib import Path +import shutil +import time + +import numpy as np +from scipy.spatial.transform import Rotation + +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.utils.logger import log_info + +_Y_UP_TO_Z_UP_ROTATION = np.array( + [ + [1.0, 0.0, 0.0], + [0.0, 0.0, -1.0], + [0.0, 1.0, 0.0], + ], + dtype=float, +) + + +class SceneExporter: + """Write one generated scene and its SimReady meshes as a scene export.""" + + def __init__( + self, + *, + scene: Scene, + output_root: str | Path, + ) -> None: + self.scene = scene + self.output_root = Path(output_root).expanduser().resolve() + self.export_root = self.output_root / "scene_export" + self.scene_config_path: Path | None = None + + def export(self) -> Path: + """Write a scene-only config and copy SimReady GLBs into ``mesh_assets``. + + Scene layouts are y-up. The simulator automatically converts each y-up + GLB to z-up, so this exporter copies each GLB unchanged and converts + only its world position and rotation for ``init_pos`` and ``init_rot``. + ``body_scale`` remains the original y-up scale associated with the GLB. + This is not a complete ``EmbodiedEnv``/``run-env`` configuration because + a generated scene does not determine a robot, its placement, or control. + """ + if self.scene.table is None: + raise ValueError("Cannot export a scene without a table.") + + mesh_assets_root = self.export_root / "mesh_assets" + mesh_assets_root.mkdir(parents=True, exist_ok=True) + scene_objects = self.scene.objects + object_ids = [scene_object.id for scene_object in scene_objects] + if len(set(object_ids)) != len(object_ids): + raise ValueError("Scene export requires unique table and asset ids.") + + exported_entries = { + scene_object.id: self._copy_scene_object_to_assets( + scene_object=scene_object, + mesh_assets_root=mesh_assets_root, + ) + for scene_object in scene_objects + } + scene_config = { + "format": "embodichain.scene-export/v1", + # This identifies the exported scene data only. It is deliberately not + # a Gymnasium environment ID because scene exports do not register or + # instantiate an EmbodiedEnv. + "scene_id": f"scene-engine-{int(time.time() * 1000)}", + "background": [ + self._scene_object_config( + scene_object=self.scene.table, + asset_relative_path=exported_entries[self.scene.table.id], + ) + ], + "rigid_object": [ + self._scene_object_config( + scene_object=asset, + asset_relative_path=exported_entries[asset.id], + ) + for asset in self.scene.assets + ], + } + self.scene_config_path = self.export_root / "scene_config.json" + self.scene_config_path.write_text( + json.dumps(scene_config, indent=2, ensure_ascii=False) + "\n", + encoding="utf-8", + ) + log_info(f"Exported scene config: {self.scene_config_path}") + return self.scene_config_path + + @staticmethod + def _copy_scene_object_to_assets( + *, + scene_object: SceneObject, + mesh_assets_root: Path, + ) -> str: + """Copy one referenced SimReady GLB and return its config-relative path.""" + object_id = scene_object.id + if ( + Path(object_id).name != object_id + or "\\" in object_id + or object_id in {"", ".", ".."} + ): + raise ValueError( + f"Scene object id is not safe for a GLB filename: {object_id!r}" + ) + if scene_object.simready_glb_path is None: + raise ValueError(f"Scene object {object_id!r} has no SimReady GLB path.") + + source_glb_path = Path(scene_object.simready_glb_path).expanduser().resolve() + if not source_glb_path.is_file(): + raise FileNotFoundError( + "SimReady GLB for scene object " + f"{object_id!r} not found: {source_glb_path}" + ) + destination_glb_path = mesh_assets_root / object_id / f"{object_id}.glb" + destination_glb_path.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(source_glb_path, destination_glb_path) + return destination_glb_path.relative_to(mesh_assets_root.parent).as_posix() + + @staticmethod + def _scene_object_config( + *, + scene_object: SceneObject, + asset_relative_path: str, + ) -> dict[str, object]: + """Build one z-up scene-only object config from a final y-up object.""" + pos_y_up = SceneExporter._scene_vector(scene_object, "pos") + rot_y_up = SceneExporter._scene_vector(scene_object, "rot") + scale_y_up = SceneExporter._scene_vector(scene_object, "scale") + if scene_object.physics is None: + raise ValueError( + f"Scene object {scene_object.id!r} has no SimReady physics settings." + ) + + pos_z_up = _Y_UP_TO_Z_UP_ROTATION @ np.asarray(pos_y_up, dtype=float) + rotation_y_up = Rotation.from_euler("xyz", rot_y_up, degrees=True).as_matrix() + rotation_z_up = ( + _Y_UP_TO_Z_UP_ROTATION @ rotation_y_up @ _Y_UP_TO_Z_UP_ROTATION.T + ) + rot_z_up = Rotation.from_matrix(rotation_z_up).as_euler( + # RigidObjectCfg.init_rot is interpreted with uppercase XYZ. + "XYZ", + degrees=True, + ) + + return { + "uid": scene_object.id, + "description": scene_object.description, + "shape": { + "shape_type": "Mesh", + "fpath": asset_relative_path, + "compute_uv": False, + }, + "attrs": scene_object.physics.attrs, + "body_type": scene_object.physics.body_type, + "init_pos": pos_z_up.tolist(), + "init_rot": rot_z_up.tolist(), + # Do not permute this scale: it belongs to the original y-up GLB, + # which SimulationManager itself converts to z-up. + "body_scale": scale_y_up, + "max_convex_hull_num": scene_object.physics.max_convex_hull_num, + } + + @staticmethod + def _scene_vector(scene_object: SceneObject, field_name: str) -> list[float]: + """Read one finite final y-up layout vector from a scene object.""" + values = getattr(scene_object, field_name) + if not isinstance(values, list) or len(values) != 3: + raise ValueError( + f"Scene object {scene_object.id!r} has no final " + f"{field_name!r} vector." + ) + vector = [float(value) for value in values] + if not np.all(np.isfinite(vector)): + raise ValueError( + f"Scene object {scene_object.id!r} has non-finite " f"{field_name!r}." + ) + return vector diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_generation_utils.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_generation_utils.py new file mode 100644 index 000000000..2f648530a --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_generation_utils.py @@ -0,0 +1,113 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from pathlib import Path +from typing import Sequence + +import numpy as np +from scipy.spatial.transform import Rotation +import trimesh + + +def quaternion_wxyz_to_euler_xyz_degrees( + quaternion_wxyz: Sequence[float], +) -> list[float]: + """Convert a ``[w, x, y, z]`` quaternion to [roll_x, pitch_y, yaw_z] degrees.""" + if len(quaternion_wxyz) != 4: + raise ValueError("Rotation quaternion must contain exactly four values.") + + w, x, y, z = quaternion_wxyz + return Rotation.from_quat([x, y, z, w]).as_euler("xyz", degrees=True).tolist() + + +def layout_object_to_transform_matrix( + layout_object: dict[str, object], +) -> np.ndarray: + """Return the matrix that maps an object's local coordinates to world coordinates.""" + transform_matrix = np.eye(4) + transform_matrix[:3, :3] = Rotation.from_euler( + "xyz", + _three_floats(layout_object.get("rot"), field_name="rot"), + degrees=True, + ).as_matrix() @ np.diag( + _three_floats(layout_object.get("scale"), field_name="scale") + ) + transform_matrix[:3, 3] = _three_floats(layout_object.get("pos"), field_name="pos") + return transform_matrix + + +def transform_matrix_to_layout_object( + object_id: str, + transform_matrix: np.ndarray, +) -> dict[str, object]: + """Convert a non-sheared 4x4 transform matrix into one layout object.""" + if not isinstance(object_id, str) or not object_id: + raise ValueError("Layout object id must be a non-empty string.") + matrix = np.asarray(transform_matrix, dtype=float) + if matrix.shape != (4, 4) or not np.all(np.isfinite(matrix)): + raise ValueError("Transform matrix must be a finite 4x4 matrix.") + if not np.allclose(matrix[3], [0.0, 0.0, 0.0, 1.0]): + raise ValueError("Transform matrix must be affine.") + + linear_matrix = matrix[:3, :3] + scale = np.linalg.norm(linear_matrix, axis=0) + if np.any(scale <= 1e-8): + raise ValueError("Transform matrix has a zero scale axis.") + rotation_matrix = linear_matrix / scale + if not np.allclose(rotation_matrix.T @ rotation_matrix, np.eye(3), atol=1e-6): + raise ValueError("Transform matrix contains shear and cannot be decomposed.") + if np.linalg.det(rotation_matrix) <= 0: + raise ValueError( + "Transform matrix contains a reflection and cannot be decomposed." + ) + + return { + "id": object_id, + "rot": Rotation.from_matrix(rotation_matrix) + .as_euler("xyz", degrees=True) + .tolist(), + "pos": matrix[:3, 3].tolist(), + "scale": scale.tolist(), + } + + +def load_glb_mesh(glb_path: str | Path) -> trimesh.Trimesh: + """Load one GLB as a single trimesh mesh.""" + resolved_glb_path = Path(glb_path).expanduser().resolve() + if not resolved_glb_path.is_file(): + raise FileNotFoundError(f"GLB geometry not found: {resolved_glb_path}") + loaded_mesh = trimesh.load(resolved_glb_path, process=False) + if isinstance(loaded_mesh, trimesh.Scene): + return loaded_mesh.dump(concatenate=True) + if isinstance(loaded_mesh, trimesh.Trimesh): + return loaded_mesh + raise ValueError(f"GLB geometry is not a mesh: {resolved_glb_path}") + + +def _three_floats(value: object, *, field_name: str) -> list[float]: + + # Validate whether the value is a list of three numeric values. + if not isinstance(value, list) or len(value) != 3: + raise ValueError(f"Coarse layout field {field_name} must contain three values.") + try: + return [float(item) for item in value] + except (TypeError, ValueError) as exc: + raise ValueError( + f"Coarse layout field {field_name} must contain numeric values." + ) from exc diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_scene_processor.py b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_scene_processor.py new file mode 100644 index 000000000..40b4a6f94 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_scene_processor.py @@ -0,0 +1,357 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +import re + +import numpy as np +import open3d as o3d +from scipy.spatial import ConvexHull, QhullError +from scipy.spatial.transform import Rotation +import trimesh + +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_object import ( + ObjectPhysics, + SceneObject, +) +from embodichain.utils.logger import log_info + +_TABLE_PHYSICS_ATTRS = { + "mass": 10.0, # Keep the table heavy if a simulator treats it as movable. + "static_friction": 0.95, # Resist lateral sliding at table contacts. + "dynamic_friction": 0.9, # Maintain high friction during sliding contacts. + "restitution": 0.01, # Prevent a table contact from producing visible bounce. +} +_ASSET_PHYSICS_ATTRS = { + "mass": 0.01, # Use a lightweight default for unconstrained generated assets. + "contact_offset": 0.003, # Start contact detection slightly before mesh contact. + "rest_offset": 0.001, # Keep a small stable separation after contact resolution. + "restitution": 0.01, # Prevent generated assets from bouncing on the table. + "max_depenetration_velocity": 10.0, # Cap corrective separation speed. + "min_position_iters": 32, # Use extra position iterations for stable contacts. + "min_velocity_iters": 8, # Use extra velocity iterations for stable contacts. +} +_FIXED_MAX_CONVEX_HULL_NUM = 16 # Shared VHACD hull budget for settling and export. + + +@dataclass(frozen=True) +class SimReadySceneProcessorConfig: + """Object-category policy for SimReady mesh canonicalization.""" + + upright_container_id_tokens: frozenset[str] = frozenset( + {"bottle", "can", "jar", "flask", "thermos"} + ) # Object-id tokens that enable upright-container standardization. + + +class SimReadySceneProcessor: + """Create SimReady GLBs and layouts for one table and its scene assets.""" + + def __init__( + self, + *, + scene: Scene, + coarse_layout_by_id: dict[str, dict[str, object]], + coarse_geometry_root: str | Path, + simready_geometry_root: str | Path, + config: SimReadySceneProcessorConfig | None = None, + ) -> None: + self.scene = scene + self.coarse_layout_by_id = coarse_layout_by_id + self.coarse_geometry_root = Path(coarse_geometry_root).expanduser().resolve() + self.simready_geometry_root = ( + Path(simready_geometry_root).expanduser().resolve() + ) + self.simready_table_layout: dict[str, object] | None = None + self.simready_assets_layout: list[dict[str, object]] | None = None + self.config = config if config is not None else SimReadySceneProcessorConfig() + if not self.config.upright_container_id_tokens: + raise ValueError("upright_container_id_tokens must not be empty.") + + def process_table(self) -> dict[str, object]: + """Process the required scene table and return its SimReady layout.""" + if self.scene.table is None: + raise ValueError("Cannot SimReady a scene without a table.") + self.simready_table_layout = self._process_object(self.scene.table) + return self.simready_table_layout + + def process_assets(self) -> list[dict[str, object]]: + """Process every scene asset and return SimReady layouts in scene order.""" + asset_ids: set[str] = set() + processed_assets: list[dict[str, object]] = [] + for asset in self.scene.assets: + if asset.id in asset_ids: + raise ValueError(f"Scene assets contain duplicate id {asset.id!r}.") + asset_ids.add(asset.id) + processed_assets.append(self._process_object(asset)) + self.simready_assets_layout = processed_assets + return self.simready_assets_layout + + def _process_object(self, scene_object: SceneObject) -> dict[str, object]: + """Canonicalize one coarse object and write its SimReady GLB.""" + object_id = scene_object.id + object_role = scene_object.kind + if object_role not in {"table", "asset"}: + raise ValueError(f"Unsupported SimReady object role {object_role!r}.") + coarse_layout = self.coarse_layout_by_id.get(object_id) + if coarse_layout is None: + raise ValueError(f"Coarse layout does not contain object {object_id!r}.") + simready_mesh, simready_transform = self._canonicalize_object_mesh( + coarse_glb_path=self.coarse_geometry_root / f"{object_id}.glb", + object_id=object_id, + rot=coarse_layout.get("rot"), + pos=coarse_layout.get("pos"), + scale=coarse_layout.get("scale"), + ) + output_path = self.simready_geometry_root / f"{object_id}.glb" + output_path.parent.mkdir(parents=True, exist_ok=True) + simready_mesh.export(output_path, file_type="glb") + if not output_path.is_file(): + raise FileNotFoundError( + f"SimReady {object_role} geometry was not written: {output_path}" + ) + scene_object.simready_glb_path = str(output_path) + scene_object.physics = self._fixed_physics_for_kind(object_role) + log_info(f"Created SimReady {object_role}: {object_id!r}.") + return {"id": object_id, **simready_transform} + + @staticmethod + def _fixed_physics_for_kind(kind: str) -> ObjectPhysics: + """Create the fixed initial physics profile for one SimReady object.""" + if kind == "table": + return ObjectPhysics( + body_type="kinematic", + attrs=dict(_TABLE_PHYSICS_ATTRS), + max_convex_hull_num=_FIXED_MAX_CONVEX_HULL_NUM, + ) + if kind == "asset": + return ObjectPhysics( + body_type="dynamic", + attrs=dict(_ASSET_PHYSICS_ATTRS), + max_convex_hull_num=_FIXED_MAX_CONVEX_HULL_NUM, + ) + raise ValueError(f"Unsupported SceneObject kind {kind!r} for physics.") + + def _canonicalize_object_mesh( + self, + *, + coarse_glb_path: str | Path, + object_id: str, + rot: object, + pos: object, + scale: object, + ) -> tuple[trimesh.Trimesh, dict[str, list[float]]]: + """Bake coarse scale and canonicalize one mesh's AABB bottom centre. + + Return the processed mesh and its updated layout transform without writing + a GLB file. The caller owns the output path and export. + """ + resolved_coarse_glb_path = Path(coarse_glb_path).expanduser().resolve() + if not resolved_coarse_glb_path.is_file(): + raise FileNotFoundError( + f"Coarse object geometry not found: {resolved_coarse_glb_path}" + ) + loaded_mesh = trimesh.load(resolved_coarse_glb_path, process=False) + if isinstance(loaded_mesh, trimesh.Scene): + mesh = loaded_mesh.dump(concatenate=True) + elif isinstance(loaded_mesh, trimesh.Trimesh): + mesh = loaded_mesh + else: + raise ValueError( + f"Coarse object geometry is not a mesh: {resolved_coarse_glb_path}" + ) + + coarse_rot = self._three_floats(rot, field_name="rot") + coarse_pos = np.asarray(self._three_floats(pos, field_name="pos"), dtype=float) + coarse_scale = np.asarray( + self._three_floats(scale, field_name="scale"), dtype=float + ) + if np.any(coarse_scale <= 0): + raise ValueError("Coarse object scale values must be positive.") + # We need the object id to determine whether it is a bottle-like object. + # If it does, then we will do a special standardization. (Hard code) + if not isinstance(object_id, str) or not object_id: + raise ValueError("Scene object id must be a non-empty string.") + + # GLB uses y-up. Convert its vertices to z-up while processing the geometry. + y_up_to_z_up_rotation = Rotation.from_euler("x", 90.0, degrees=True) + y_up_to_z_up_matrix = y_up_to_z_up_rotation.as_matrix() + y_up_to_z_up_transform = np.eye(4) + y_up_to_z_up_transform[:3, :3] = y_up_to_z_up_matrix + mesh.apply_transform(y_up_to_z_up_transform) + + # Standardize upright containers in temporary z-up coordinates before the + # shared center, scale, and bottom-center preprocessing. + # This is to ensure the action agent can pick up the bottle or can-like objects. + bottle_alignment_matrix = np.eye(3) + if self._is_upright_container_id(object_id): + bottle_alignment_matrix = self._standardize_bottle_z_up(mesh) + bottle_alignment_transform = np.eye(4) + bottle_alignment_transform[:3, :3] = bottle_alignment_matrix + mesh.apply_transform(bottle_alignment_transform) + + # First make the object's AABB center at the origin. + original_aabb_center = mesh.bounds.mean(axis=0) + mesh.apply_translation(-original_aabb_center) + + # Scale the object with the value in the coarse layout. + scale_transform = np.eye(4) + scale_transform[:3, :3] = ( + # Actually there's no need to do so, for the scale factor is all equal + # in x, y, z axes. + bottle_alignment_matrix + @ y_up_to_z_up_matrix + @ np.diag(coarse_scale) + @ y_up_to_z_up_matrix.T + @ bottle_alignment_matrix.T + ) + mesh.apply_transform(scale_transform) + + # Move the scaled object's AABB bottom center to the world's x-y plane (z=0). + scaled_bounds = mesh.bounds + scaled_aabb_bottom_center = np.array( + [ + (scaled_bounds[0, 0] + scaled_bounds[1, 0]) / 2, + (scaled_bounds[0, 1] + scaled_bounds[1, 1]) / 2, + scaled_bounds[0, 2], + ] + ) + mesh.apply_translation(-scaled_aabb_bottom_center) + + # Convert the processed GLB back to its standard y-up coordinate system. + z_up_to_y_up_transform = np.eye(4) + z_up_to_y_up_transform[:3, :3] = y_up_to_z_up_matrix.T + mesh.apply_transform(z_up_to_y_up_transform) + + # Compensate the bottle's local rotation so that its coarse world pose does + # not change. + local_bottle_rotation = Rotation.from_matrix( + y_up_to_z_up_matrix.T @ bottle_alignment_matrix @ y_up_to_z_up_matrix + ) + coarse_rotation_matrix = Rotation.from_euler( + "xyz", coarse_rot, degrees=True + ).as_matrix() + rotation = Rotation.from_matrix( + coarse_rotation_matrix @ local_bottle_rotation.inv().as_matrix() + ) + # Update the pos. + position_offset = y_up_to_z_up_matrix.T @ ( + scale_transform[:3, :3] @ original_aabb_center + scaled_aabb_bottom_center + ) + return mesh, { + "rot": rotation.as_euler("xyz", degrees=True).tolist(), + "pos": (coarse_pos + rotation.apply(position_offset)).tolist(), + "scale": [1.0, 1.0, 1.0], + } + + def _is_upright_container_id(self, object_id: str) -> bool: + """Return whether object-id tokens indicate a bottle-like container.""" + # Example: soda_can_0 + # tokens: {"soda", "can", "0"} + # upright_container_id_tokens: {"bottle", "can", "jar"} + # So this returns True because "can" is in the configured token set. + tokens = set(re.findall(r"[a-z0-9]+", object_id.lower())) + return bool(tokens & self.config.upright_container_id_tokens) + + @staticmethod + def _standardize_bottle_z_up(mesh: trimesh.Trimesh) -> np.ndarray: + """Return a proper rotation that maps a bottle-like mesh's long axis to z-up. + + Thanks to chenjian for this idea! + """ + if len(mesh.vertices) < 4 or len(mesh.faces) < 4: + raise ValueError( + "Bottle standardization requires a non-degenerate triangle mesh." + ) + open3d_mesh = o3d.geometry.TriangleMesh( + vertices=o3d.utility.Vector3dVector(mesh.vertices), + triangles=o3d.utility.Vector3iVector(mesh.faces), + ) + sampled_points = np.asarray( + open3d_mesh.sample_points_uniformly(number_of_points=10_000).points + ) # (10000, 3) x (x, y, z) + + # Check the number of the points again, and check whether have some + # non-finite values. + if sampled_points.shape[0] < 4 or not np.all(np.isfinite(sampled_points)): + raise ValueError( + "Bottle standardization could not sample valid mesh points." + ) + + centered_points = sampled_points - sampled_points.mean(axis=0) + # SVD find the longest axis. + _, _, principal_axes = np.linalg.svd(centered_points, full_matrices=False) + if np.linalg.det(principal_axes) < 0: + principal_axes[2, :] *= -1 # in case the SVD returns a reflection. + + bottle_rotation = Rotation.from_euler( + "y", 90.0, degrees=True + ).as_matrix() # 3x3 matrix + # The first PCA axis is the longest axis; rotate it onto the temporary z axis. + bottle_rotation = bottle_rotation @ principal_axes + standardized_points = (bottle_rotation @ centered_points.T).T + + axis_min = standardized_points[:, 2].min() + axis_max = standardized_points[:, 2].max() + axis_range = axis_max - axis_min + upper_points = standardized_points[ + standardized_points[:, 2] > axis_min + axis_range * 0.8 + ] + lower_points = standardized_points[ + standardized_points[:, 2] < axis_min + axis_range * 0.2 + ] + upper_volume = SimReadySceneProcessor._convex_hull_volume(upper_points) + lower_volume = SimReadySceneProcessor._convex_hull_volume(lower_points) + + # Bottles usually have a smaller top (neck) than bottom; flip if necessary. + if upper_volume > lower_volume: + bottle_rotation = ( + Rotation.from_euler("x", 180.0, degrees=True).as_matrix() + @ bottle_rotation + ) + return bottle_rotation + + @staticmethod + def _convex_hull_volume(points: np.ndarray) -> float: + """Return the volume of a non-degenerate point set's convex hull.""" + if points.shape[0] < 4: + raise ValueError( + "Bottle standardization needs at least four points per end." + ) + try: + return float(ConvexHull(points).volume) + except QhullError as exc: + raise ValueError( + "Bottle standardization found a degenerate end volume." + ) from exc + + @staticmethod + def _three_floats(value: object, *, field_name: str) -> list[float]: + """Validate and convert a three-value layout field to floats.""" + if not isinstance(value, list) or len(value) != 3: + raise ValueError( + f"Coarse layout field {field_name} must contain three values." + ) + try: + return [float(item) for item in value] + except (TypeError, ValueError) as exc: + raise ValueError( + f"Coarse layout field {field_name} must contain numeric values." + ) from exc diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/table_support_surface.py b/embodichain/gen_sim/scene_engine/pipeline/utils/table_support_surface.py new file mode 100644 index 000000000..6c3ad3e94 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/table_support_surface.py @@ -0,0 +1,497 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from collections import deque +from dataclasses import dataclass +from pathlib import Path + +from embodichain.utils.logger import log_info, log_warning +import matplotlib +import numpy as np +from scipy.spatial import ConvexHull, QhullError +from shapely.geometry import Polygon +from shapely.ops import unary_union +import trimesh + +matplotlib.use("Agg") + +import matplotlib.pyplot as plt +from mpl_toolkits.mplot3d.art3d import Poly3DCollection + + +@dataclass(frozen=True) +class SupportSurfaceConfig: + """Parameters for conservative single-level table-top detection.""" + + normal_z_min: float = 0.95 # Minimum z component for an upward face normal. + min_surface_area_m2: float = 0.01 # Minimum projected area for a candidate level. + max_face_height_span_m: float = ( + 0.01 # Maximum within-face z variation for flatness. + ) + height_level_tolerance_m: float = 0.005 # Maximum z difference within one level. + + +@dataclass(frozen=True) +class TableSupportRegion: + """Detected main table support surface in z-up world coordinates.""" + + top_z: float # Highest z value among the selected support-surface triangles. + vertices: np.ndarray # Full z-up table vertex array referenced by ``faces``. + faces: np.ndarray # Indices of triangles selected as the main support surface. + support_polygon: Polygon # Largest valid outer support contour in z-up XY. + + +class TableSupportSurfaceDetector: + """Detect one main upward table surface and render auditable diagnostics. + + The input mesh must be standing on x-y plane of the z-up world coordinates. + This class deliberately returns only the largest outer 2D contour because + the current layout stage models one main tabletop; the original support + triangles remain available for 3D diagnostics. + """ + + def __init__( + self, + *, + table_world_mesh: trimesh.Trimesh, + debug_output_root: str | Path | None = None, + config: SupportSurfaceConfig | None = None, + ) -> None: + # Init. + self.table_world_mesh = table_world_mesh + self.debug_output_root = ( + Path(debug_output_root).expanduser().resolve() + if debug_output_root is not None + else None + ) + self.support_region: TableSupportRegion | None = None + self.config = config if config is not None else SupportSurfaceConfig() + # Check config values. + if not 0.0 < self.config.normal_z_min <= 1.0: + raise ValueError("normal_z_min must be in (0, 1].") + if self.config.min_surface_area_m2 <= 0.0: + raise ValueError("min_surface_area_m2 must be positive.") + if self.config.max_face_height_span_m <= 0.0: + raise ValueError("max_face_height_span_m must be positive.") + if self.config.height_level_tolerance_m <= 0.0: + raise ValueError("height_level_tolerance_m must be positive.") + + def detect(self) -> TableSupportRegion: + """Detect the main upward-facing support surface of a z-up table mesh.""" + table_world_mesh = self.table_world_mesh + + # Check. + if len(table_world_mesh.vertices) < 3 or len(table_world_mesh.faces) == 0: + raise ValueError("Table mesh must contain at least one triangle.") + + mesh = table_world_mesh.copy() + # Material or UV seams can duplicate vertices along one physical table top. + mesh.merge_vertices(digits_vertex=7) + + # Repair normals if possible. + self._repair_normals_if_possible(mesh) + + face_vertices = mesh.vertices[mesh.faces] + face_height_ranges = np.ptp(face_vertices[:, :, 2], axis=1) + # Check: 1. range of z values of each triangle; 2. upward-facing triangles. + candidate_face_indices = np.flatnonzero( + (mesh.face_normals[:, 2] >= self.config.normal_z_min) + & (face_height_ranges <= self.config.max_face_height_span_m) + ) + if len(candidate_face_indices) == 0: + raise ValueError( + "Table mesh has no near-horizontal upward-facing triangles that " + "satisfy max_face_height_span_m." + ) + + # Select the best support level among the candidate triangles. + selected_faces = self._select_main_support_level( + mesh=mesh, + candidate_face_indices=candidate_face_indices, + full_table_hull_area=self._convex_hull_area( + mesh.vertices[:, :2], name="Table" + ), + ) + selected_vertices = face_vertices[selected_faces] + vertices = mesh.vertices.copy() + faces = mesh.faces[selected_faces].copy() + self.support_region = TableSupportRegion( + top_z=float(selected_vertices[:, :, 2].max()), + vertices=vertices, + faces=faces, + support_polygon=self._extract_largest_support_polygon(vertices[faces, :2]), + ) + return self.support_region + + def save_support_surface_debug_images( + self, + *, + save_3d: bool = True, + save_2d: bool = True, + output_3d_path: str | Path | None = None, + output_2d_path: str | Path | None = None, + ) -> bool: + """Optionally save standard 3D and 2D support-detection diagnostics.""" + if not save_3d and not save_2d: + return True + if self.support_region is None: + self.detect() + assert self.support_region is not None + if save_3d: + self._save_support_surface_3d_image( + table_world_mesh=self.table_world_mesh, + support_region=self.support_region, + output_path=self._resolve_debug_output_path( + output_3d_path, "table_support_surface_3d.png" + ), + ) + if save_2d: + self._save_support_region_2d_image( + support_region=self.support_region, + output_path=self._resolve_debug_output_path( + output_2d_path, "table_support_region_2d.png" + ), + ) + return True + + def _resolve_debug_output_path( + self, output_path: str | Path | None, default_filename: str + ) -> Path: + if output_path is not None: + return Path(output_path).expanduser().resolve() + if self.debug_output_root is None: + raise ValueError( + "A debug_output_root or an explicit debug output path is required " + "when saving support-surface debug images." + ) + return self.debug_output_root / default_filename + + @staticmethod + def _repair_normals_if_possible(mesh: trimesh.Trimesh) -> None: + if not mesh.is_watertight: + log_warning( + "Table mesh is not watertight; skipped outward-normal repair and " + "will use its input face normals for support detection." + ) + return + if mesh.is_volume: + log_info("Table mesh normals already form a valid outward-facing volume.") + return + try: + trimesh.repair.fix_normals(mesh, multibody=True) + except Exception as exc: + log_warning( + "Table mesh normal repair raised " + f"{exc}; support detection will use the resulting face normals." + ) + return + if mesh.is_volume: + log_info("Table mesh normals were repaired into an outward-facing volume.") + return + log_warning( + "Table mesh is watertight but normals could not be repaired into a " + "valid outward-facing volume; support detection will use the resulting " + "face normals." + ) + return + + def _select_main_support_level( + self, + *, + mesh: trimesh.Trimesh, + candidate_face_indices: np.ndarray, + full_table_hull_area: float, + ) -> np.ndarray: + # Find adj. + adjacency = self._face_adjacency(mesh) + # Use BFS to group the connect components. + components = self._connected_components( + set(int(index) for index in candidate_face_indices), adjacency + ) + + # Sort components by their top z value, descending. + components_by_height = sorted( + ( + ( + float(mesh.vertices[mesh.faces[list(component)], 2].max()), + component, + ) + for component in components + ), + key=lambda item: item[0], + reverse=True, + ) + + # Group components into levels by their top z value, within the height_level_tolerance_m. + levels: list[tuple[float, list[set[int]]]] = [] + for component_top_z, component in components_by_height: + for level_index, (level_top_z, level_components) in enumerate(levels): + if ( + level_top_z - component_top_z + <= self.config.height_level_tolerance_m + ): + level_components.append(component) + levels[level_index] = (level_top_z, level_components) + break + else: + levels.append((component_top_z, [component])) + + best_level_faces: np.ndarray | None = None + best_hull_gap = np.inf + best_top_z = -np.inf + best_projected_support_area = 0.0 + hull_gap_tolerance = max(full_table_hull_area * 1e-6, 1e-9) + for level_top_z, level_components in levels: + level_indices = np.asarray( + sorted( + face_index + for component in level_components + for face_index in component + ), + dtype=int, + ) + level_triangles = mesh.vertices[mesh.faces[level_indices]] + level_projected_support_area = self._projected_triangle_area( + level_triangles[:, :, :2] + ) + if level_projected_support_area < self.config.min_surface_area_m2: + continue + level_hull_area = self._convex_hull_area( + level_triangles[:, :, :2].reshape(-1, 2), + name="Candidate support level", + ) + # Compute the gap between convex hull and triangle projection area, + # for selecting the best level among multiple candidates which avoids + # small area which have the largest z value. + level_hull_gap = max(0.0, full_table_hull_area - level_hull_area) + if ( + best_level_faces is None + or level_hull_gap < best_hull_gap - hull_gap_tolerance + or ( + abs(level_hull_gap - best_hull_gap) <= hull_gap_tolerance + and level_top_z > best_top_z + self.config.height_level_tolerance_m + ) + or ( + abs(level_hull_gap - best_hull_gap) <= hull_gap_tolerance + and abs(level_top_z - best_top_z) + <= self.config.height_level_tolerance_m + and level_projected_support_area > best_projected_support_area + ) + ): + best_level_faces = level_indices + best_hull_gap = level_hull_gap + best_top_z = level_top_z + best_projected_support_area = level_projected_support_area + if best_level_faces is None: + raise ValueError( + "No upward-facing support level meets min_surface_area_m2." + ) + return best_level_faces + + @staticmethod + def _convex_hull_area(points: np.ndarray, *, name: str) -> float: + """Compute the area of the convex hull of a set of XY points.""" + unique_points = np.unique(np.asarray(points, dtype=float), axis=0) + if ( + unique_points.ndim != 2 + or unique_points.shape[1] != 2 + or len(unique_points) < 3 + ): + raise ValueError(f"{name} must contain at least three unique XY points.") + try: + return float(ConvexHull(unique_points).volume) + except QhullError as exc: + raise ValueError(f"{name} XY projection is degenerate.") from exc + + @staticmethod + def _projected_triangle_area(triangles_xy: np.ndarray) -> float: + """Compute the total area of triangles projected onto the XY plane.""" + first_edges = triangles_xy[:, 1] - triangles_xy[:, 0] + second_edges = triangles_xy[:, 2] - triangles_xy[:, 0] + cross_products = ( + first_edges[:, 0] * second_edges[:, 1] + - first_edges[:, 1] * second_edges[:, 0] + ) + return float(np.abs(cross_products).sum() / 2.0) + + @classmethod + def _extract_largest_support_polygon(cls, triangles_xy: np.ndarray) -> Polygon: + projected_triangles = [ + Polygon(triangle) + for triangle in triangles_xy + if cls._projected_triangle_area(triangle[None, ...]) > 1e-12 + ] + if not projected_triangles: + raise ValueError( + "Selected support surface has no non-degenerate XY triangles." + ) + merged_region = unary_union(projected_triangles) + if merged_region.geom_type == "Polygon": + polygons = [merged_region] + else: + polygons = [ + geometry + for geometry in merged_region.geoms + if geometry.geom_type == "Polygon" + ] + if not polygons: + raise ValueError( + "Could not create a 2D support region from the selected triangles." + ) + if len(polygons) > 1: + log_warning( + "Detected multiple disconnected outer support contours; using only " + "the largest one for the single-contour support-region output." + ) + largest_polygon = max(polygons, key=lambda polygon: polygon.area) + if largest_polygon.is_empty or not largest_polygon.is_valid: + raise ValueError("The merged 2D support region is not a valid polygon.") + boundary_xy = np.asarray(largest_polygon.exterior.coords, dtype=float) + if len(boundary_xy) < 4 or not np.isfinite(boundary_xy).all(): + raise ValueError("The merged 2D support contour is degenerate.") + return Polygon(boundary_xy) + + @staticmethod + def _face_adjacency(mesh: trimesh.Trimesh) -> dict[int, set[int]]: + """Build a face adjacency dictionary for the mesh.""" + adjacency: dict[int, set[int]] = {} + for first, second in mesh.face_adjacency: + first_index = int(first) + second_index = int(second) + adjacency.setdefault(first_index, set()).add(second_index) + adjacency.setdefault(second_index, set()).add(first_index) + return adjacency + + @staticmethod + def _connected_components( + faces: set[int], adjacency: dict[int, set[int]] + ) -> list[set[int]]: + """Group upward-facing candidate triangles into edge-connected surface components with BFS.""" + unvisited = set(faces) + components: list[set[int]] = [] + while unvisited: + component: set[int] = set() + queue = deque([unvisited.pop()]) + while queue: + face_index = queue.popleft() + component.add(face_index) + for neighbor in adjacency.get(face_index, set()): + if neighbor in unvisited: + unvisited.remove(neighbor) + queue.append(neighbor) + components.append(component) + return components + + @classmethod + def _save_support_surface_3d_image( + cls, + *, + table_world_mesh: trimesh.Trimesh, + support_region: TableSupportRegion, + output_path: str | Path, + ) -> Path: + resolved_output_path = cls._resolve_png_output_path(output_path) + resolved_output_path.parent.mkdir(parents=True, exist_ok=True) + figure = plt.figure(figsize=(9, 8), dpi=160, constrained_layout=True) + axis = figure.add_subplot(projection="3d") + axis.add_collection3d( + Poly3DCollection( + table_world_mesh.vertices[table_world_mesh.faces], + facecolor="steelblue", + edgecolor="none", + alpha=0.18, + ) + ) + axis.add_collection3d( + Poly3DCollection( + support_region.vertices[support_region.faces], + facecolor="darkorange", + edgecolor="saddlebrown", + linewidth=0.25, + alpha=0.95, + ) + ) + lower = table_world_mesh.bounds[0].copy() + upper = table_world_mesh.bounds[1].copy() + extent = upper - lower + lower[extent <= 1e-9] -= 0.001 + upper[extent <= 1e-9] += 0.001 + axis.set( + xlim=(lower[0], upper[0]), + ylim=(lower[1], upper[1]), + zlim=(lower[2], upper[2]), + ) + axis.set_box_aspect(upper - lower) + axis.view_init(elev=25.0, azim=-55.0) + axis.set_xlabel("x (z-up world)") + axis.set_ylabel("y (z-up world)") + axis.set_zlabel("z (up)") + axis.set_title( + "Detected main table support surface\n" + f"top z={support_region.top_z:.4f} m, faces={len(support_region.faces)}" + ) + figure.savefig(resolved_output_path, bbox_inches="tight") + plt.close(figure) + return resolved_output_path + + @classmethod + def _save_support_region_2d_image( + cls, + *, + support_region: TableSupportRegion, + output_path: str | Path, + ) -> Path: + resolved_output_path = cls._resolve_png_output_path(output_path) + resolved_output_path.parent.mkdir(parents=True, exist_ok=True) + boundary_xy = np.asarray(support_region.support_polygon.exterior.coords) + figure, axis = plt.subplots(figsize=(8, 8), dpi=160, constrained_layout=True) + axis.fill( + boundary_xy[:, 0], + boundary_xy[:, 1], + facecolor="darkorange", + edgecolor="none", + alpha=0.82, + label="detected support region", + ) + axis.plot( + boundary_xy[:, 0], + boundary_xy[:, 1], + color="saddlebrown", + linewidth=2.0, + label="outer support contour", + ) + axis.autoscale_view() + axis.set_aspect("equal", adjustable="box") + axis.set_xlabel("x (z-up world)") + axis.set_ylabel("y (z-up world)") + axis.set_title( + "Detected 2D table support region\n" + f"z={support_region.top_z:.4f} m, contour vertices={len(boundary_xy) - 1}" + ) + axis.legend(loc="best") + figure.savefig(resolved_output_path, bbox_inches="tight") + plt.close(figure) + return resolved_output_path + + @staticmethod + def _resolve_png_output_path(output_path: str | Path) -> Path: + resolved_output_path = Path(output_path).expanduser().resolve() + if resolved_output_path.suffix.lower() != ".png": + return resolved_output_path.with_suffix(".png") + return resolved_output_path diff --git a/pyproject.toml b/pyproject.toml index d1daf53a1..db51c3041 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,7 +54,11 @@ dependencies = [ [project.optional-dependencies] gensim = [ "bpy", - "pyrender==0.1.45" + "pyrender==0.1.45", + "requests", + "Pillow", + "scipy", + "matplotlib", ] # cuRobo V2 is distributed from its source repository and provides separate # dependency sets for CUDA 12 and CUDA 13. Keep it optional so CPU-only and diff --git a/tests/gen_sim/scene_engine/test_clients.py b/tests/gen_sim/scene_engine/test_clients.py new file mode 100644 index 000000000..513f0c5c0 --- /dev/null +++ b/tests/gen_sim/scene_engine/test_clients.py @@ -0,0 +1,274 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import pytest + +from embodichain.gen_sim.scene_engine.clients import geometry_generation +from embodichain.gen_sim.scene_engine.clients import image_segmentation +from embodichain.gen_sim.scene_engine.llms import load_config + + +class _Response: + """Minimal successful HTTP response used by client unit tests.""" + + def __init__(self, payload: object, *, content: bytes = b"") -> None: + self._payload = payload + self.content = content + + def raise_for_status(self) -> None: + return None + + def json(self) -> object: + return self._payload + + +class _Session: + """Capture HTTP calls without contacting an external service.""" + + def __init__( + self, *, get_payload: object, post_payload: object | None = None + ) -> None: + self.get_payload = get_payload + self.post_payload = post_payload + self.get_calls: list[tuple[str, int]] = [] + self.post_call: dict[str, object] | None = None + + def get(self, url: str, *, timeout: int) -> _Response: + self.get_calls.append((url, timeout)) + return _Response(self.get_payload) + + def post(self, url: str, **kwargs: object) -> _Response: + self.post_call = {"url": url, **kwargs} + return _Response(self.post_payload) + + def close(self) -> None: + return None + + +def test_clients_load_their_required_dotenv_values( + monkeypatch: pytest.MonkeyPatch, +) -> None: + geometry_values = { + "SCENE_ENGINE_GEOMETRY_GENERATION_BASE_URL": "http://geometry/", + "SCENE_ENGINE_GEOMETRY_GENERATION_TIMEOUT_S": "60", + "SCENE_ENGINE_GEOMETRY_GENERATION_MAX_ATTEMPTS": "2", + "SCENE_ENGINE_GEOMETRY_GENERATION_HEALTH_PATH": "/health", + "SCENE_ENGINE_GEOMETRY_GENERATION_OBJECTS_PATH": "/objects", + } + segmentation_values = { + "SCENE_ENGINE_IMAGE_SEGMENTATION_BASE_URL": "http://segment/", + "SCENE_ENGINE_IMAGE_SEGMENTATION_TIMEOUT_S": "30", + "SCENE_ENGINE_IMAGE_SEGMENTATION_MAX_ATTEMPTS": "2", + "SCENE_ENGINE_IMAGE_SEGMENTATION_HEALTH_PATH": "/health", + "SCENE_ENGINE_IMAGE_SEGMENTATION_SINGLE_OBJECT_PATH": "/predict", + } + llm_values = { + "OPENAI_API_KEY": "test-key", + "OPENAI_MODEL": "test-model", + "OPENAI_BASE_URL": "http://llm/v1/", + "SCENE_ENGINE_OPENAI_DEFAULT_QUERY": '{"api-version": "1"}', + "OPENAI_MAX_ATTEMPTS": "2", + } + monkeypatch.setattr( + geometry_generation, "read_scene_engine_env_values", lambda *_: geometry_values + ) + monkeypatch.setattr( + image_segmentation, + "read_scene_engine_env_values", + lambda *_: segmentation_values, + ) + monkeypatch.setattr( + load_config, "read_scene_engine_env_values", lambda *_: llm_values + ) + + geometry_client = geometry_generation.GeometryGenerationClient.from_dotenv() + segmentation_client = image_segmentation.ImageSegmentationClient.from_dotenv() + llm_client_config = load_config.load_llm_config() + + assert geometry_client._base_url == "http://geometry" + assert geometry_client._generate_objects_path == "/objects" + assert segmentation_client._base_url == "http://segment" + assert segmentation_client._segment_single_object_path == "/predict" + assert llm_client_config.default_query == {"api-version": "1"} + assert llm_client_config.base_url == "http://llm/v1" + + +def test_geometry_dotenv_config_rejects_invalid_timeout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + values = { + "SCENE_ENGINE_GEOMETRY_GENERATION_BASE_URL": "http://geometry", + "SCENE_ENGINE_GEOMETRY_GENERATION_TIMEOUT_S": "0", + "SCENE_ENGINE_GEOMETRY_GENERATION_MAX_ATTEMPTS": "1", + "SCENE_ENGINE_GEOMETRY_GENERATION_HEALTH_PATH": "/health", + "SCENE_ENGINE_GEOMETRY_GENERATION_OBJECTS_PATH": "/objects", + } + monkeypatch.setattr( + geometry_generation, "read_scene_engine_env_values", lambda *_: values + ) + + with pytest.raises(ValueError, match="TIMEOUT_S must be at least 1"): + geometry_generation.GeometryGenerationClient.from_dotenv() + + +def test_service_health_checks_use_the_configured_health_path() -> None: + geometry_session = _Session(get_payload={"ok": True}) + geometry_client = geometry_generation.GeometryGenerationClient( + base_url="http://geometry", + timeout_s=60, + max_attempts=1, + health_path="/health", + generate_objects_path="/objects", + session=geometry_session, + ) + segmentation_session = _Session(get_payload={"ok": True}) + segmentation_client = image_segmentation.ImageSegmentationClient( + base_url="http://segment", + timeout_s=30, + max_attempts=1, + health_path="/health", + segment_single_object_path="/predict", + session=segmentation_session, + ) + + geometry_client.check_health() + segmentation_client.check_health() + + assert geometry_session.get_calls == [("http://geometry/health", 10)] + assert segmentation_session.get_calls == [("http://segment/health", 30)] + + +def test_segmentation_client_posts_prompt_and_returns_rle_masks(tmp_path: Path) -> None: + image_path = tmp_path / "scene.png" + image_path.write_bytes(b"png") + rle_mask = {"counts": [0, 1], "size": [1, 1]} + session = _Session( + get_payload={"ok": True}, + post_payload={"result": {"masks": [rle_mask]}}, + ) + client = image_segmentation.ImageSegmentationClient( + base_url="http://segment", + timeout_s=30, + max_attempts=1, + health_path="/health", + segment_single_object_path="/predict", + session=session, + ) + + assert client.segment_single_object(image_path=image_path, prompt="table") == [ + rle_mask + ] + assert session.post_call is not None + assert session.post_call["url"] == "http://segment/predict" + assert session.post_call["data"] == {"prompt": "table"} + + +def test_segmentation_client_accepts_instance_mask_response() -> None: + rle_mask = {"counts": [0, 1], "size": [1, 1]} + + masks = image_segmentation._extract_rle_masks( + {"data": {"instances": [{"mask_rle": rle_mask}]}} + ) + + assert masks == [rle_mask] + + +def test_geometry_response_requires_matching_ordered_objects() -> None: + response: dict[str, Any] = { + "ok": True, + "result": { + "objects": [ + { + "name": "table_001", + "mesh": "/results/table.glb", + "rotation_quaternion_wxyz": [1, 0, 0, 0], + "translation": [0, 1, 2], + "scale": [1, 1, 1], + } + ] + }, + } + + objects = geometry_generation._parse_objects_response( + response, + object_ids=["table_001"], + ) + + assert objects == [ + { + "mesh": "/results/table.glb", + "rotation_quaternion_wxyz": [1.0, 0.0, 0.0, 0.0], + "translation": [0.0, 1.0, 2.0], + "scale": [1.0, 1.0, 1.0], + } + ] + + +def test_geometry_client_posts_masks_and_downloads_glbs(tmp_path: Path) -> None: + class GeometrySession(_Session): + def post(self, url: str, **kwargs: object) -> _Response: + self.post_call = {"url": url, **kwargs} + return _Response( + { + "ok": True, + "result": { + "objects": [ + { + "name": "cup", + "mesh": "/results/cup.glb", + "rotation_quaternion_wxyz": [1, 0, 0, 0], + "translation": [0, 0, 0], + "scale": [1, 1, 1], + } + ] + }, + } + ) + + def get(self, url: str, *, timeout: int) -> _Response: + self.get_calls.append((url, timeout)) + return _Response({}, content=b"glTF-mesh") + + image_path = tmp_path / "scene.png" + mask_path = tmp_path / "cup.png" + image_path.write_bytes(b"png") + mask_path.write_bytes(b"png") + session = GeometrySession(get_payload={"ok": True}) + client = geometry_generation.GeometryGenerationClient( + base_url="http://geometry", + timeout_s=30, + max_attempts=1, + health_path="/health", + generate_objects_path="/objects", + session=session, + ) + + _, objects = client.generate_objects( + image_path=image_path, + object_masks=[("cup", mask_path)], + output_root=tmp_path / "output", + ) + + assert objects[0]["mesh"] == "/results/cup.glb" + assert session.post_call is not None + assert session.post_call["url"] == "http://geometry/objects" + assert (tmp_path / "output/cup.glb").read_bytes() == b"glTF-mesh" diff --git a/tests/gen_sim/scene_engine/test_group_table_aligner.py b/tests/gen_sim/scene_engine/test_group_table_aligner.py new file mode 100644 index 000000000..13f5a95ee --- /dev/null +++ b/tests/gen_sim/scene_engine/test_group_table_aligner.py @@ -0,0 +1,73 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from pathlib import Path + +import pytest +import trimesh + +from embodichain.gen_sim.scene_engine.pipeline.utils.assets_group_table_aligner import ( + AssetsGroupTableAligner, + AssetsGroupTableAlignerConfig, +) + + +def _layout(object_id: str, y: float) -> dict[str, object]: + return { + "id": object_id, + "pos": [0.0, y, 0.0], + "rot": [0.0, 0.0, 0.0], + "scale": [1.0, 1.0, 1.0], + } + + +def test_group_table_aligner_preserves_relative_vertical_offsets( + tmp_path: Path, +) -> None: + trimesh.creation.box(extents=(2.0, 1.0, 2.0)).export(tmp_path / "table.glb") + trimesh.creation.box(extents=(0.5, 1.0, 0.5)).export(tmp_path / "first.glb") + trimesh.creation.box(extents=(0.5, 1.0, 0.5)).export(tmp_path / "second.glb") + assets_layout = [_layout("first", 0.0), _layout("second", 0.3)] + + _, aligned_assets = AssetsGroupTableAligner( + table_layout=_layout("table", 0.0), + assets_layout=assets_layout, + geometry_root=tmp_path, + config=AssetsGroupTableAlignerConfig(clearance_m=0.1), + ).align() + + assert aligned_assets[0]["pos"][1] > assets_layout[0]["pos"][1] # type: ignore[index] + assert aligned_assets[1]["pos"][1] - aligned_assets[0]["pos"][1] == pytest.approx( # type: ignore[index] + 0.3 + ) + + +def test_group_table_aligner_returns_empty_assets_without_mesh_loading( + tmp_path: Path, +) -> None: + table_layout = _layout("table", 0.0) + + aligned_table, aligned_assets = AssetsGroupTableAligner( + table_layout=table_layout, + assets_layout=[], + geometry_root=tmp_path, + ).align() + + assert aligned_table is table_layout + assert aligned_assets == [] diff --git a/tests/gen_sim/scene_engine/test_scene_core_and_export.py b/tests/gen_sim/scene_engine/test_scene_core_and_export.py new file mode 100644 index 000000000..56c12af1c --- /dev/null +++ b/tests/gen_sim/scene_engine/test_scene_core_and_export.py @@ -0,0 +1,171 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import pytest + +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_object import ( + ObjectPhysics, + SceneObject, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_exporter import SceneExporter + + +def _scene_object( + *, + object_id: str, + kind: str, + glb_path: Path | None = None, + physics: ObjectPhysics | None = None, +) -> SceneObject: + return SceneObject( + id=object_id, + kind=kind, # type: ignore[arg-type] + category=kind, + name=object_id, + description=f"{kind} object", + simready_glb_path=str(glb_path) if glb_path is not None else None, + rot=[0.0, 0.0, 0.0], + pos=[1.0, 2.0, 3.0], + scale=[1.0, 2.0, 3.0], + physics=physics, + ) + + +def _physics(body_type: str) -> ObjectPhysics: + return ObjectPhysics( + body_type=body_type, # type: ignore[arg-type] + attrs={"mass": 1.0, "static_friction": 0.8}, + max_convex_hull_num=16, + ) + + +def test_scene_returns_one_table_and_ordered_assets() -> None: + table = _scene_object(object_id="table", kind="table") + asset = _scene_object(object_id="cup", kind="asset") + scene = Scene(objects=[table, asset]) + + assert scene.table is table + assert scene.assets == [asset] + assert scene.to_dict()["objects"][0]["id"] == "table" # type: ignore[index] + + +def test_scene_rejects_multiple_tables() -> None: + scene = Scene( + objects=[ + _scene_object(object_id="table_001", kind="table"), + _scene_object(object_id="table_002", kind="table"), + ] + ) + + with pytest.raises(ValueError, match="only one table"): + _ = scene.table + + +@pytest.mark.parametrize( + ("body_type", "attrs", "hulls"), + [ + ("static", {"mass": 1.0}, 1), + ("dynamic", {}, 1), + ("dynamic", {"mass": 1.0}, 0), + ], +) +def test_object_physics_rejects_invalid_values( + body_type: str, + attrs: dict[str, float], + hulls: int, +) -> None: + with pytest.raises(ValueError): + _physics = ObjectPhysics( # noqa: F841 + body_type=body_type, # type: ignore[arg-type] + attrs=attrs, + max_convex_hull_num=hulls, + ) + + +def test_scene_export_copies_meshes_and_converts_y_up_pose(tmp_path: Path) -> None: + table_glb = tmp_path / "table.glb" + asset_glb = tmp_path / "cup.glb" + table_glb.write_bytes(b"glTF-table") + asset_glb.write_bytes(b"glTF-cup") + table = _scene_object( + object_id="table", + kind="table", + glb_path=table_glb, + physics=_physics("kinematic"), + ) + asset = _scene_object( + object_id="cup", + kind="asset", + glb_path=asset_glb, + physics=_physics("dynamic"), + ) + + export_path = SceneExporter( + scene=Scene(objects=[table, asset]), + output_root=tmp_path / "output", + ).export() + exported = json.loads(export_path.read_text(encoding="utf-8")) + + assert ( + export_path.parent / "mesh_assets/table/table.glb" + ).read_bytes() == b"glTF-table" + assert (export_path.parent / "mesh_assets/cup/cup.glb").read_bytes() == b"glTF-cup" + entry = exported["rigid_object"][0] + assert entry["uid"] == "cup" + assert entry["body_type"] == "dynamic" + assert entry["init_pos"] == [1.0, -3.0, 2.0] + assert entry["body_scale"] == [1.0, 2.0, 3.0] + assert np.allclose(entry["init_rot"], [0.0, 0.0, 0.0]) + + +def test_scene_export_requires_final_physics(tmp_path: Path) -> None: + glb_path = tmp_path / "table.glb" + glb_path.write_bytes(b"glTF") + table = _scene_object(object_id="table", kind="table", glb_path=glb_path) + + with pytest.raises(ValueError, match="no SimReady physics"): + SceneExporter(scene=Scene(objects=[table]), output_root=tmp_path).export() + + +def test_scene_export_rejects_backslash_in_object_id(tmp_path: Path) -> None: + glb_path = tmp_path / "table.glb" + glb_path.write_bytes(b"glTF") + table = _scene_object( + object_id="table", + kind="table", + glb_path=glb_path, + physics=_physics("kinematic"), + ) + unsafe_asset = _scene_object( + object_id=r"..\evil", + kind="asset", + glb_path=glb_path, + physics=_physics("dynamic"), + ) + + with pytest.raises(ValueError, match="not safe for a GLB filename"): + SceneExporter( + scene=Scene(objects=[table, unsafe_asset]), + output_root=tmp_path / "output", + ).export() diff --git a/tests/gen_sim/scene_engine/test_scene_engine_config.py b/tests/gen_sim/scene_engine/test_scene_engine_config.py new file mode 100644 index 000000000..01914e144 --- /dev/null +++ b/tests/gen_sim/scene_engine/test_scene_engine_config.py @@ -0,0 +1,100 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from embodichain.gen_sim.scene_engine.cli import start +from embodichain.gen_sim.scene_engine.configs import environment + + +def test_read_scene_engine_env_values_reads_requested_keys( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + env_path = tmp_path / ".env" + env_path.write_text('OPENAI_MODEL="test-model"\nUNRELATED_VALUE=ignored\n') + monkeypatch.setattr(environment, "_SCENE_ENGINE_ENV_PATH", env_path) + + assert environment.read_scene_engine_env_values("OPENAI_MODEL") == { + "OPENAI_MODEL": "test-model" + } + + +def test_read_scene_engine_env_values_reports_missing_keys( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + env_path = tmp_path / ".env" + env_path.write_text("OPENAI_MODEL=test-model\n") + monkeypatch.setattr(environment, "_SCENE_ENGINE_ENV_PATH", env_path) + + with pytest.raises(ValueError, match="OPENAI_API_KEY"): + environment.read_scene_engine_env_values("OPENAI_MODEL", "OPENAI_API_KEY") + + +def test_scene_engine_help_exposes_only_runtime_arguments( + capsys: pytest.CaptureFixture[str], +) -> None: + with pytest.raises(SystemExit) as exc_info: + start.main(["--help"]) + + assert exc_info.value.code == 0 + output = capsys.readouterr().out + assert "--image" in output + assert "--output_root" in output + assert "gen_sim/.env" in output + assert "--config" not in output + + +def test_scene_engine_cli_forwards_validated_paths( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + image_path = tmp_path / "scene.png" + image_path.write_bytes(b"png") + captured: dict[str, Path] = {} + + def generate_scene(*, image_path: Path, output_root: Path) -> None: + captured["image_path"] = image_path + captured["output_root"] = output_root + + monkeypatch.setattr(start, "generate_scene_from_image", generate_scene) + output_root = tmp_path / "output" + + start.cli_scene_engine(image_path, output_root) + + assert captured == { + "image_path": image_path.resolve(), + "output_root": output_root.resolve(), + } + + +@pytest.mark.parametrize("image_name", ["missing.png", "scene.gif"]) +def test_scene_engine_cli_rejects_invalid_image_inputs( + tmp_path: Path, + image_name: str, +) -> None: + image_path = tmp_path / image_name + if image_path.suffix == ".gif": + image_path.write_bytes(b"gif") + + with pytest.raises((FileNotFoundError, ValueError)): + start.cli_scene_engine(image_path, tmp_path / "output") diff --git a/tests/gen_sim/scene_engine/test_scene_understanding.py b/tests/gen_sim/scene_engine/test_scene_understanding.py new file mode 100644 index 000000000..1fdb51f7b --- /dev/null +++ b/tests/gen_sim/scene_engine/test_scene_understanding.py @@ -0,0 +1,85 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.pipeline import scene_understanding + + +def _response(*, asset_name: str = "cup") -> str: + return json.dumps( + { + "table": { + "category": "dining_table", + "name": "wooden table", + "description": "A rectangular wooden table.", + }, + "assets": [ + { + "category": "cup", + "name": asset_name, + "description": "A small ceramic cup.", + } + ], + } + ) + + +def test_image_object_analysis_parses_code_fence_and_assigns_stable_ids() -> None: + scene = scene_understanding._parse_image_object_analysis_response( + f"```json\n{_response()}\n```" + ) + + assert scene.table is not None + assert scene.table.id == "table" + assert [asset.id for asset in scene.assets] == ["cup_001"] + + +def test_image_object_analysis_rejects_location_words_in_object_names() -> None: + with pytest.raises(ValueError, match="must not contain location"): + scene_understanding._parse_image_object_analysis_response( + _response(asset_name="left cup") + ) + + +def test_image_object_analysis_retries_then_updates_scene(tmp_path: Path) -> None: + class VLM: + def __init__(self) -> None: + self.responses = ["not-json", _response()] + + def complete(self, **_: object) -> str: + return self.responses.pop(0) + + image_path = tmp_path / "scene.png" + image_path.write_bytes(b"png") + scene = Scene() + + scene_understanding._analyze_image_objects( + scene=scene, + image_path=image_path, + vlm_client=VLM(), # type: ignore[arg-type] + json_max_attempts=2, + ) + + assert scene.table is not None + assert [asset.id for asset in scene.assets] == ["cup_001"] diff --git a/tests/gen_sim/scene_engine/test_support_and_layout.py b/tests/gen_sim/scene_engine/test_support_and_layout.py new file mode 100644 index 000000000..98ed30491 --- /dev/null +++ b/tests/gen_sim/scene_engine/test_support_and_layout.py @@ -0,0 +1,180 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +import numpy as np +import pytest +from shapely.geometry import Point, Polygon +import trimesh + +from embodichain.gen_sim.scene_engine.pipeline.utils.assets_group_layout_optimizer import ( + AssetsSupportLayoutOptimizer, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.assets_group_support_clamp import ( + AssetsGroupSupportClamp, + AssetsGroupSupportClampConfig, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.table_support_surface import ( + TableSupportSurfaceDetector, +) + + +def _aabb( + minimum_x: float, minimum_y: float, maximum_x: float, maximum_y: float +) -> np.ndarray: + return np.array( + [ + [minimum_x, minimum_y], + [maximum_x, minimum_y], + [maximum_x, maximum_y], + [minimum_x, maximum_y], + ], + dtype=float, + ) + + +def _layout(object_id: str, x: float, y: float) -> dict[str, object]: + return {"id": object_id, "pos": [x, 0.0, -y]} + + +def _top_mesh( + vertices_xy: list[tuple[float, float]], faces: list[list[int]], z: float +) -> trimesh.Trimesh: + return trimesh.Trimesh( + vertices=np.array([[x, y, z] for x, y in vertices_xy], dtype=float), + faces=np.array(faces, dtype=int), + process=False, + ) + + +def test_support_detector_preserves_an_l_shaped_support_contour() -> None: + mesh = _top_mesh( + [(0, 0), (2, 0), (2, 1), (1, 1), (1, 2), (0, 2)], + [[0, 1, 3], [1, 2, 3], [0, 3, 5], [3, 4, 5]], + z=1.0, + ) + + region = TableSupportSurfaceDetector(table_world_mesh=mesh).detect() + + assert region.top_z == pytest.approx(1.0) + assert region.support_polygon.area == pytest.approx(3.0) + assert region.support_polygon.covers(Point(0.5, 1.5)) + assert not region.support_polygon.covers(Point(1.5, 1.5)) + + +def test_support_detector_prefers_main_tabletop_over_small_higher_piece() -> None: + main = _top_mesh([(0, 0), (2, 0), (2, 2), (0, 2)], [[0, 1, 2], [0, 2, 3]], z=1.0) + decoration = _top_mesh( + [(0.25, 0.25), (0.75, 0.25), (0.75, 0.75), (0.25, 0.75)], + [[0, 1, 2], [0, 2, 3]], + z=1.2, + ) + mesh = trimesh.util.concatenate([main, decoration]) + + region = TableSupportSurfaceDetector(table_world_mesh=mesh).detect() + + assert region.top_z == pytest.approx(1.0) + assert region.support_polygon.area == pytest.approx(4.0) + + +def test_group_clamp_preserves_relative_layout_while_moving_inside_support() -> None: + support = Polygon([(0, 0), (4, 0), (4, 4), (0, 4)]) + aabbs = { + "first": _aabb(-0.5, 1.0, 0.5, 2.0), + "second": _aabb(1.0, 1.0, 2.0, 2.0), + } + layouts = [_layout("first", 0.0, 1.5), _layout("second", 1.5, 1.5)] + + refined = AssetsGroupSupportClamp( + support_region=support, + assets_aabb_2d_z_up_world_corners_by_id=aabbs, + assets_layout=layouts, + config=AssetsGroupSupportClampConfig(grid_resolution_m=0.05), + ).clamp() + + first, second = refined + assert first["pos"][0] > layouts[0]["pos"][0] # type: ignore[index] + assert second["pos"][0] - first["pos"][0] == pytest.approx(1.5) # type: ignore[index] + assert second["pos"][2] - first["pos"][2] == pytest.approx(0.0) # type: ignore[index] + + +def test_group_clamp_returns_unchanged_layout_when_already_contained() -> None: + layout = _layout("cup", 1.5, 1.5) + refined = AssetsGroupSupportClamp( + support_region=Polygon([(0, 0), (3, 0), (3, 3), (0, 3)]), + assets_aabb_2d_z_up_world_corners_by_id={"cup": _aabb(1.0, 1.0, 2.0, 2.0)}, + assets_layout=[layout], + ).clamp() + + assert refined == [layout] + + +def test_group_clamp_reports_infeasible_oversized_group() -> None: + clamp = AssetsGroupSupportClamp( + support_region=Polygon([(0, 0), (2, 0), (2, 2), (0, 2)]), + assets_aabb_2d_z_up_world_corners_by_id={"large": _aabb(0, 0, 3, 3)}, + assets_layout=[_layout("large", 1.5, 1.5)], + config=AssetsGroupSupportClampConfig(grid_resolution_m=0.1), + ) + + with pytest.raises(ValueError, match="cannot be placed"): + clamp.clamp() + + +def test_layout_optimizer_resolves_a_simple_pair_overlap() -> None: + optimizer = AssetsSupportLayoutOptimizer( + support_region=Polygon([(0, 0), (5, 0), (5, 5), (0, 5)]), + assets_aabb_2d_z_up_world_corners_by_id={ + "first": _aabb(1.0, 1.0, 2.0, 2.0), + "second": _aabb(1.5, 1.0, 2.5, 2.0), + }, + assets_layout=[_layout("first", 1.5, 1.5), _layout("second", 2.0, 1.5)], + ) + + refined = optimizer.optimize() + + refined_offsets = np.array( + [ + [ + refined[index]["pos"][0] - optimizer.assets_layout[index]["pos"][0], # type: ignore[index] + optimizer.assets_layout[index]["pos"][2] - refined[index]["pos"][2], + ] # type: ignore[index] + for index in range(2) + ] + ) + base_aabbs = np.stack( + [ + optimizer.assets_aabb_2d_z_up_world_corners_by_id["first"], + optimizer.assets_aabb_2d_z_up_world_corners_by_id["second"], + ] + ) + assert not optimizer._overlaps(base_aabbs, refined_offsets) + + +def test_layout_optimizer_rejects_unresolvable_overlap() -> None: + optimizer = AssetsSupportLayoutOptimizer( + support_region=Polygon([(0, 0), (2, 0), (2, 2), (0, 2)]), + assets_aabb_2d_z_up_world_corners_by_id={ + "first": _aabb(0, 0, 2, 2), + "second": _aabb(0, 0, 2, 2), + }, + assets_layout=[_layout("first", 1.0, 1.0), _layout("second", 1.0, 1.0)], + ) + + with pytest.raises(ValueError, match="cannot be resolved"): + optimizer.optimize() diff --git a/tests/test_main.py b/tests/test_main.py index 607ffceb3..b24e3680b 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -30,8 +30,10 @@ "decompose-urdf", "preview-asset", "run-env", + "scene-engine", "simready", "train-rl", + "preview-scene", "workspace-cache", } @@ -87,6 +89,19 @@ def test_subcommand_help_uses_complete_command_parser( assert "--category" in output +def test_preview_scene_help_includes_output_and_viser_options( + capsys: pytest.CaptureFixture[str], +) -> None: + """Preview Scene should expose its required path and optional Viser settings.""" + with pytest.raises(SystemExit) as exc_info: + cli.main(["preview-scene", "--help"]) + + assert exc_info.value.code == 0 + output = capsys.readouterr().out + assert "--output_root" in output + assert "--viser" in output + + def test_nested_benchmark_help_uses_suite_parser( capsys: pytest.CaptureFixture[str], ) -> None: