From 61001a73754e21377b8c8dcedab1b0bd1d8f68af Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Tue, 28 Jul 2026 11:27:32 +0800 Subject: [PATCH 01/23] Add a new feature: gen_sim/scene_engine --- .../gen_sim/scene_engine/cli/__init__.py | 19 + .../gen_sim/scene_engine/cli/preview.py | 201 +++ embodichain/gen_sim/scene_engine/cli/start.py | 70 + .../gen_sim/scene_engine/clients/__init__.py | 19 + .../clients/geometry_generation.py | 374 ++++ .../clients/image_segmentation.py | 230 +++ .../gen_sim/scene_engine/configs/__init__.py | 19 + .../configs/scene_engine_config.json | 25 + .../gen_sim/scene_engine/core/__init__.py | 19 + .../gen_sim/scene_engine/core/asset.py | 51 + .../gen_sim/scene_engine/core/scene.py | 36 + .../gen_sim/scene_engine/core/table.py | 51 + .../gen_sim/scene_engine/llms/__init__.py | 19 + .../gen_sim/scene_engine/llms/load_config.py | 93 + .../llms/openai_compatible_client.py | 141 ++ .../gen_sim/scene_engine/pipeline/__init__.py | 19 + .../gen_sim/scene_engine/pipeline/generate.py | 118 ++ .../scene_engine/pipeline/gym_export.py | 220 +++ .../scene_engine/pipeline/scene_generation.py | 576 ++++++ .../pipeline/scene_segmentation.py | 479 +++++ .../pipeline/scene_understanding.py | 254 +++ .../scene_engine/pipeline/utils/__init__.py | 19 + .../pipeline/utils/scene_generation_utils.py | 1547 +++++++++++++++++ .../utils/scene_segmentation_utils.py | 340 ++++ .../gen_sim/scene_engine/utils/__init__.py | 19 + .../gen_sim/scene_engine/utils/logger.py | 38 + 26 files changed, 4996 insertions(+) create mode 100644 embodichain/gen_sim/scene_engine/cli/__init__.py create mode 100644 embodichain/gen_sim/scene_engine/cli/preview.py create mode 100644 embodichain/gen_sim/scene_engine/cli/start.py create mode 100644 embodichain/gen_sim/scene_engine/clients/__init__.py create mode 100644 embodichain/gen_sim/scene_engine/clients/geometry_generation.py create mode 100644 embodichain/gen_sim/scene_engine/clients/image_segmentation.py create mode 100644 embodichain/gen_sim/scene_engine/configs/__init__.py create mode 100644 embodichain/gen_sim/scene_engine/configs/scene_engine_config.json create mode 100644 embodichain/gen_sim/scene_engine/core/__init__.py create mode 100644 embodichain/gen_sim/scene_engine/core/asset.py create mode 100644 embodichain/gen_sim/scene_engine/core/scene.py create mode 100644 embodichain/gen_sim/scene_engine/core/table.py create mode 100644 embodichain/gen_sim/scene_engine/llms/__init__.py create mode 100644 embodichain/gen_sim/scene_engine/llms/load_config.py create mode 100644 embodichain/gen_sim/scene_engine/llms/openai_compatible_client.py create mode 100644 embodichain/gen_sim/scene_engine/pipeline/__init__.py create mode 100644 embodichain/gen_sim/scene_engine/pipeline/generate.py create mode 100644 embodichain/gen_sim/scene_engine/pipeline/gym_export.py create mode 100644 embodichain/gen_sim/scene_engine/pipeline/scene_generation.py create mode 100644 embodichain/gen_sim/scene_engine/pipeline/scene_segmentation.py create mode 100644 embodichain/gen_sim/scene_engine/pipeline/scene_understanding.py create mode 100644 embodichain/gen_sim/scene_engine/pipeline/utils/__init__.py create mode 100644 embodichain/gen_sim/scene_engine/pipeline/utils/scene_generation_utils.py create mode 100644 embodichain/gen_sim/scene_engine/pipeline/utils/scene_segmentation_utils.py create mode 100644 embodichain/gen_sim/scene_engine/utils/__init__.py create mode 100644 embodichain/gen_sim/scene_engine/utils/logger.py 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..e283d3831 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/cli/preview.py @@ -0,0 +1,201 @@ +# ---------------------------------------------------------------------------- +# 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 typing import Any + +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.cfg import LightCfg, MeshCfg, RigidObjectCfg + + +def preview_gym_export( + *, + output_root: str | Path, + device: str = "cpu", + headless: bool = False, +) -> None: + """Load ``gym_export/gym_config.json`` and preview its table and assets.""" + resolved_output_root = Path(output_root).expanduser().resolve() + config_path = resolved_output_root / "gym_export" / "gym_config.json" + if not config_path.is_file(): + raise FileNotFoundError(f"Gym config not found: {config_path}") + + try: + gym_config = json.loads(config_path.read_text(encoding="utf-8")) + except json.JSONDecodeError as exc: + raise ValueError(f"Gym config is not valid JSON: {config_path}") from exc + if not isinstance(gym_config, dict): + raise ValueError("Gym config must be a JSON object.") + + sim = SimulationManager( + SimulationManagerCfg( + width=1920, + height=1080, + headless=headless, + physics_dt=1.0 / 100.0, + sim_device=device, + ) + ) + try: + if sim.is_use_gpu_physics: + sim.init_gpu_physics() + _add_lights(sim) + _add_objects( + sim=sim, + entries=_config_entries(gym_config, "background"), + config_dir=config_path.parent, + label="table", + ) + _add_objects( + sim=sim, + entries=_config_entries(gym_config, "rigid_object"), + config_dir=config_path.parent, + label="asset", + ) + + if headless: + sim.update(step=1) + print(f"Loaded gym export headlessly: {config_path}") + return + + print(f"Previewing: {config_path}") + print("Close with Ctrl-C.") + sim.open_window() + while True: + time.sleep(0.1) + except KeyboardInterrupt: + print("Stopping preview.") + finally: + sim.destroy() + + +def _config_entries( + gym_config: dict[str, Any], + field_name: str, +) -> list[dict[str, Any]]: + entries = gym_config.get(field_name, []) + if not isinstance(entries, list) or not all( + isinstance(entry, dict) for entry in entries + ): + raise ValueError(f"Gym 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.""" + for entry in entries: + uid = entry.get("uid") + shape = entry.get("shape") + if not isinstance(uid, str) or not uid: + raise ValueError(f"Gym {label} has no valid uid.") + if not isinstance(shape, dict) or not isinstance(shape.get("fpath"), str): + raise ValueError(f"Gym {label} {uid!r} has no shape.fpath.") + if shape.get("shape_type") != "Mesh": + raise ValueError( + f"Gym {label} {uid!r} must use shape_type='Mesh' for preview." + ) + + mesh_path = (config_dir / shape["fpath"]).resolve() + 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, + ) + ) + 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"Gym 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"Gym config field {field_name!r} must be numeric.") from exc + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Preview a Scene Engine gym export in EmbodiChain simulation." + ) + parser.add_argument( + "output_root", + type=Path, + help="Scene Engine output root containing gym_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.", + ) + args = parser.parse_args() + preview_gym_export( + output_root=args.output_root, + device=args.device, + headless=args.headless, + ) + + +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..719f54749 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/cli/start.py @@ -0,0 +1,70 @@ +# ---------------------------------------------------------------------------- +# 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 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: + 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() -> None: + parser = argparse.ArgumentParser( + description="embodichain.gen_sim.scene_engine Scene Engine Pipeline" + ) + 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() + + 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..1181503b5 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/clients/geometry_generation.py @@ -0,0 +1,374 @@ +# ---------------------------------------------------------------------------- +# 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 +from typing import Any + +import requests + +_DEFAULT_CONFIG_PATH = ( + Path(__file__).resolve().parents[1] / "configs" / "scene_engine_config.json" +) + + +class GeometryGenerationClient: + """Manage the Geometry Generation Server connection.""" + + def __init__( + self, + *, + base_url: str, + timeout_s: int, + max_attempts: int, + health_path: str, + generate_multiple_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_multiple_objects_path = generate_multiple_objects_path + self._session = session or requests.Session() + + @classmethod + def from_config( + cls, + config_path: str | Path | None = None, + ) -> "GeometryGenerationClient": + return cls(**_load_config(config_path)) + + 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, + timeout=10, # Use a shorter timeout for avoiding long waits. + ) + response.raise_for_status() + return + except requests.RequestException 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_multiple_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 multiple objects from: + - An input image. + - A list of object masks, each with a unique object_id and a binary mask path. + """ + + # 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)) + + # Use the wrapped data structure to send the request. + response_data, response_objects = self._request_multiple_objects( + image_path=resolved_image_path, + object_masks=resolved_object_masks, + ) + + resolved_output_root = Path(output_root).expanduser().resolve() + + # 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, + ): + output_path = resolved_output_root / f"{object_id}.glb" + self._download_glb(response_object["mesh"], output_path) + return response_data, response_objects + + def _request_multiple_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: + with ExitStack() as stack: # This stack manages the context of multiple open files, ensuring they are closed after the request. + 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_multiple_objects_path), + data={"json": "1"}, + files=[ + ("image", (image_path.name, image_file)), + *[ + ("masks", (f"{object_id}.png", mask_file)) + 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_objects = ( + _parse_multiple_objects_response( # Parse the 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 _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_multiple_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 _load_config(config_path: str | Path | None) -> dict[str, Any]: + resolved_config_path = Path(config_path or _DEFAULT_CONFIG_PATH).expanduser() + resolved_config_path = resolved_config_path.resolve() + if not resolved_config_path.is_file(): + raise FileNotFoundError(f"Config not found: {resolved_config_path}") + + try: + config_data = json.loads(resolved_config_path.read_text(encoding="utf-8")) + except json.JSONDecodeError as exc: + raise ValueError(f"Config is not valid JSON: {resolved_config_path}") from exc + + config = config_data.get("geometry_generation") + if not isinstance(config, dict): + raise ValueError("Config key geometry_generation must be an object.") + + required_keys = ( + "base_url", + "timeout_s", + "max_attempts", + "health_path", + "generate_multiple_objects_path", + ) + missing = [key for key in required_keys if key not in config] + if missing: + raise ValueError(f"Missing Geometry Generation Server config keys: {missing}") + + try: + timeout_s = int(config["timeout_s"]) + except (TypeError, ValueError) as exc: + raise ValueError( + "Geometry Generation Server config timeout_s must be an integer." + ) from exc + if timeout_s < 1: + raise ValueError( + "Geometry Generation Server config timeout_s must be at least 1." + ) + + try: + max_attempts = int(config["max_attempts"]) + except (TypeError, ValueError) as exc: + raise ValueError( + "Geometry Generation Server config max_attempts must be an integer." + ) from exc + if max_attempts < 1: + raise ValueError( + "Geometry Generation Server config max_attempts must be at least 1." + ) + + string_keys = ( + "base_url", + "health_path", + "generate_multiple_objects_path", + ) + for key in string_keys: + if not isinstance(config[key], str) or not config[key].strip(): + raise ValueError( + f"Geometry Generation Server config key {key} must be a non-empty string." + ) + + return { + "base_url": config["base_url"].strip(), + "timeout_s": timeout_s, + "max_attempts": max_attempts, + "health_path": config["health_path"].strip(), + "generate_multiple_objects_path": config[ + "generate_multiple_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..083adca7d --- /dev/null +++ b/embodichain/gen_sim/scene_engine/clients/image_segmentation.py @@ -0,0 +1,230 @@ +# ---------------------------------------------------------------------------- +# 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 +from typing import Any + +import requests + +_DEFAULT_CONFIG_PATH = ( + Path(__file__).resolve().parents[1] / "configs" / "scene_engine_config.json" +) + + +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_config( + cls, + config_path: str | Path | None = None, + ) -> "ImageSegmentationClient": + config = _load_config(config_path) + return cls(**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_config(config_path: str | Path | None) -> dict[str, Any]: + resolved_config_path = Path(config_path or _DEFAULT_CONFIG_PATH).expanduser() + resolved_config_path = resolved_config_path.resolve() + if not resolved_config_path.is_file(): + raise FileNotFoundError(f"Config not found: {resolved_config_path}") + + try: + config_data = json.loads(resolved_config_path.read_text(encoding="utf-8")) + except json.JSONDecodeError as exc: + raise ValueError(f"Config is not valid JSON: {resolved_config_path}") from exc + + config = config_data.get("image_segmentation") + if not isinstance(config, dict): + raise ValueError("Config key image_segmentation must be an object.") + + required_keys = ( + "base_url", + "timeout_s", + "max_attempts", + "health_path", + "segment_single_object_path", + ) + missing = [key for key in required_keys if key not in config] + if missing: + raise ValueError(f"Missing Image Segmentation Server config keys: {missing}") + + try: + timeout_s = int(config["timeout_s"]) + except (TypeError, ValueError) as exc: + raise ValueError( + "Image Segmentation Server config timeout_s must be an integer." + ) from exc + if timeout_s < 1: + raise ValueError( + "Image Segmentation Server config timeout_s must be at least 1." + ) + + try: + max_attempts = int(config["max_attempts"]) + except (TypeError, ValueError) as exc: + raise ValueError( + "Image Segmentation Server config max_attempts must be an integer." + ) from exc + if max_attempts < 1: + raise ValueError( + "Image Segmentation Server config max_attempts must be at least 1." + ) + + string_keys = ("base_url", "health_path", "segment_single_object_path") + for key in string_keys: + if not isinstance(config[key], str) or not config[key].strip(): + raise ValueError( + f"Image Segmentation Server config key {key} must be a non-empty string." + ) + + return { + "base_url": config["base_url"].strip(), + "timeout_s": timeout_s, + "max_attempts": max_attempts, + "health_path": config["health_path"].strip(), + "segment_single_object_path": config["segment_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/scene_engine_config.json b/embodichain/gen_sim/scene_engine/configs/scene_engine_config.json new file mode 100644 index 000000000..a87c24b23 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/configs/scene_engine_config.json @@ -0,0 +1,25 @@ +{ + "llm": { + "openai_compatible": { + "api_key": "", + "model": "", + "base_url": "", + "default_query": {}, + "max_attempts": 3 + } + }, + "image_segmentation": { + "base_url": "", + "timeout_s": 30, + "max_attempts": 3, + "health_path": "/health", + "segment_single_object_path": "/predict" + }, + "geometry_generation": { + "base_url": "", + "timeout_s": 600, + "max_attempts": 3, + "health_path": "/health", + "generate_multiple_objects_path": "/generate_multiple_objects" + } +} 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/asset.py b/embodichain/gen_sim/scene_engine/core/asset.py new file mode 100644 index 000000000..81306d329 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/core/asset.py @@ -0,0 +1,51 @@ +# ---------------------------------------------------------------------------- +# 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 + + +@dataclass +class Asset: + """A scene asset identified during scene understanding.""" + + id: str + category: str + name: str + description: str + # Path to a binary mask image aligned with the input image. White pixels + # identify this asset; black pixels identify the background. + mask_path: str | None = None + # Absolute path to the canonicalized GLB used by the final simulation. + simready_glb_path: str | None = None + # Final y-up layout after scene refinement and gravity settling. + rot: list[float] | None = None + pos: list[float] | None = None + scale: list[float] | None = None + + def to_dict(self) -> dict[str, object]: + return { + "id": self.id, + "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, + } 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..ed41aa67a --- /dev/null +++ b/embodichain/gen_sim/scene_engine/core/scene.py @@ -0,0 +1,36 @@ +# ---------------------------------------------------------------------------- +# 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.asset import Asset +from embodichain.gen_sim.scene_engine.core.table import Table + + +@dataclass +class Scene: + """A scene containing a table and zero or more assets.""" + + table: Table | None = None + assets: list[Asset] = field(default_factory=list) + + def to_dict(self) -> dict[str, object]: + return { + "table": self.table.to_dict() if self.table is not None else None, + "assets": [asset.to_dict() for asset in self.assets], + } diff --git a/embodichain/gen_sim/scene_engine/core/table.py b/embodichain/gen_sim/scene_engine/core/table.py new file mode 100644 index 000000000..bab0f94fb --- /dev/null +++ b/embodichain/gen_sim/scene_engine/core/table.py @@ -0,0 +1,51 @@ +# ---------------------------------------------------------------------------- +# 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 + + +@dataclass +class Table: + """The table identified during scene understanding.""" + + id: str + category: str + name: str + description: str + # Path to a binary mask image aligned with the input image. White pixels + # identify the table; black pixels identify the background. + mask_path: str | None = None + # Absolute path to the canonicalized GLB used by the final simulation. + simready_glb_path: str | None = None + # Final y-up layout after scene refinement and gravity settling. + rot: list[float] | None = None + pos: list[float] | None = None + scale: list[float] | None = None + + def to_dict(self) -> dict[str, object]: + return { + "id": self.id, + "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, + } 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..f2a786399 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/llms/load_config.py @@ -0,0 +1,93 @@ +# ---------------------------------------------------------------------------- +# 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 +import os +from pathlib import Path +from typing import Any + +DEFAULT_LLM_CONFIG_PATH = ( + Path(__file__).resolve().parents[1] / "configs" / "scene_engine_config.json" +) + + +@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(config_path: str | Path | None = None) -> LLMConfig: + """Load LLM settings from JSON, with ``OPENAI_*`` overrides.""" + resolved_config_path = Path(config_path or DEFAULT_LLM_CONFIG_PATH).expanduser() + resolved_config_path = resolved_config_path.resolve() + if not resolved_config_path.is_file(): + raise FileNotFoundError(f"LLM config not found: {resolved_config_path}") + + try: + raw_config = json.loads(resolved_config_path.read_text(encoding="utf-8")) + except json.JSONDecodeError as exc: + raise ValueError( + f"LLM config is not valid JSON: {resolved_config_path}" + ) from exc + + llm_config = raw_config.get("llm", {}).get("openai_compatible", {}) + if not isinstance(llm_config, dict): + raise ValueError("LLM config key llm.openai_compatible must be an object.") + + api_key = os.getenv("OPENAI_API_KEY") or llm_config.get("api_key", "") + model = os.getenv("OPENAI_MODEL") or llm_config.get("model", "") + base_url = os.getenv("OPENAI_BASE_URL") or llm_config.get("base_url", "") + default_query = llm_config.get("default_query", {}) + max_attempts = os.getenv("OPENAI_MAX_ATTEMPTS") or llm_config.get("max_attempts", 3) + + if not isinstance(default_query, dict): + raise ValueError("LLM config key default_query must be an object.") + missing = [ + key + for key, value in { + "api_key": api_key, + "model": model, + "base_url": base_url, + }.items() + if not isinstance(value, str) or not value.strip() + ] + if missing: + raise ValueError(f"Missing required LLM config keys: {missing}") + + try: + parsed_max_attempts = int(max_attempts) + except (TypeError, ValueError) as exc: + raise ValueError("LLM config key max_attempts must be an integer.") from exc + if parsed_max_attempts < 1: + raise ValueError("LLM config key max_attempts must be at least 1.") + + return LLMConfig( + api_key=api_key.strip(), + model=model.strip(), + base_url=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..0b7cf3786 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/llms/openai_compatible_client.py @@ -0,0 +1,141 @@ +# ---------------------------------------------------------------------------- +# 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_config( + cls, config_path: str | Path | None = None + ) -> "OpenAICompatibleVLM": + """Create a client from the scene-engine LLM configuration.""" + return cls(load_llm_config(config_path)) + + 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..a8f6d0b70 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/generate.py @@ -0,0 +1,118 @@ +# ---------------------------------------------------------------------------- +# 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.image_segmentation import ( + ImageSegmentationClient, +) + +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.gen_sim.scene_engine.pipeline.scene_segmentation import ( + segment_scene, +) +from embodichain.gen_sim.scene_engine.utils.logger import log_stage_end, log_stage_start + +from embodichain.gen_sim.scene_engine.pipeline.scene_generation import ( + generate_scene_and_refine, +) +from embodichain.gen_sim.scene_engine.pipeline.gym_export import export_scene_to_gym + + +def generate_scene_from_image( + image_path: str | Path, + output_root: str | Path, + *, + llm_config_path: str | Path | None = None, + image_segmentation_config_path: str | Path | None = None, + geometry_generation_config_path: str | Path | None = None, +) -> 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_config(llm_config_path) + scene = Scene() + + # 1. Scene Understanding + log_stage_start("Scene Understanding") + scene = understand_scene( + scene=scene, + image_path=image_path, + output_root=resolved_output_root, + vlm_client=vlm_client, + ) + log_stage_end("Scene Understanding") + + # 2. Scene Segmentation + log_stage_start("Scene Segmentation") + # Load the config and fail if the Image Segmentation Server is unavailable. + image_segmentation_client = ImageSegmentationClient.from_config( + image_segmentation_config_path + ) + image_segmentation_client.check_health() # Error raising will happen internally. + scene = segment_scene( + image_path=image_path, + output_root=resolved_output_root, + scene=scene, + vlm_client=vlm_client, + image_segmentation_client=image_segmentation_client, + ) + image_segmentation_client.close() # Kill the session. + log_stage_end("Scene Segmentation") + + # 3. Objects + Coarse Layout Generation + log_stage_start("Objects + Coarse Layout Generation") + # Load the config and fail if the Geometry Generation Server is unavailable. + geometry_generation_client = GeometryGenerationClient.from_config( + geometry_generation_config_path + ) + geometry_generation_client.check_health() + + scene = generate_scene_and_refine( + image_path=image_path, + output_root=resolved_output_root, + scene=scene, + vlm_client=vlm_client, + geometry_generation_client=geometry_generation_client, + ) + geometry_generation_client.close() # Kill the session. + log_stage_end("Objects + Coarse Layout Generation") + + # 4. Scene Export + log_stage_start("Scene Export") + export_scene_to_gym( + scene=scene, + output_root=resolved_output_root, + table_max_convex_hull_num=16, + asset_max_convex_hull_num=16, + ) + log_stage_end("Scene Export") + + return scene diff --git a/embodichain/gen_sim/scene_engine/pipeline/gym_export.py b/embodichain/gen_sim/scene_engine/pipeline/gym_export.py new file mode 100644 index 000000000..793fe0f52 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/gym_export.py @@ -0,0 +1,220 @@ +# ---------------------------------------------------------------------------- +# 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.asset import Asset +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.table import Table + +_DEFAULT_MAX_CONVEX_HULL_NUM = 16 +_TABLE_PHYSICS_ATTRS = { + "mass": 10.0, + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.01, +} +_ASSET_PHYSICS_ATTRS = { + "mass": 0.01, + "contact_offset": 0.003, + "rest_offset": 0.001, + "restitution": 0.01, + "max_depenetration_velocity": 10.0, + "min_position_iters": 32, + "min_velocity_iters": 8, +} +_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, +) + + +def export_scene_to_gym( + *, + scene: Scene, + output_root: str | Path, + table_max_convex_hull_num: int = _DEFAULT_MAX_CONVEX_HULL_NUM, + asset_max_convex_hull_num: int = _DEFAULT_MAX_CONVEX_HULL_NUM, +) -> Path: + """Write the Gym 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. + """ + if scene.table is None: + raise ValueError("Cannot export a gym scene without a table.") + table_max_convex_hull_num = _positive_int( + table_max_convex_hull_num, + field_name="table_max_convex_hull_num", + ) + asset_max_convex_hull_num = _positive_int( + asset_max_convex_hull_num, + field_name="asset_max_convex_hull_num", + ) + + export_root = Path(output_root).expanduser().resolve() / "gym_export" + mesh_assets_root = export_root / "mesh_assets" + mesh_assets_root.mkdir(parents=True, exist_ok=True) + + scene_objects = [scene.table, *scene.assets] + object_ids = [scene_object.id for scene_object in scene_objects] + if len(set(object_ids)) != len(object_ids): + raise ValueError("Gym export requires unique table and asset ids.") + + exported_entries = { + scene_object.id: _copy_scene_object_to_gym_assets( + scene_object=scene_object, + mesh_assets_root=mesh_assets_root, + ) + for scene_object in scene_objects + } + gym_config = { + "id": f"Prompt2Scene-{int(time.time() * 1000)}-v0", + "max_episodes": 10, + "max_episode_steps": 300, + "env": {"events": {}, "observations": {}, "dataset": {}}, + "robot": {}, + "sensor": [], + "light": {}, + "background": [ + _gym_object_config( + scene_object=scene.table, + asset_relative_path=exported_entries[scene.table.id], + body_type="kinematic", + attrs=_TABLE_PHYSICS_ATTRS, + max_convex_hull_num=table_max_convex_hull_num, + ) + ], + "rigid_object": [ + _gym_object_config( + scene_object=asset, + asset_relative_path=exported_entries[asset.id], + body_type="dynamic", + attrs=_ASSET_PHYSICS_ATTRS, + max_convex_hull_num=asset_max_convex_hull_num, + ) + for asset in scene.assets + ], + } + gym_config_path = export_root / "gym_config.json" + gym_config_path.write_text( + json.dumps(gym_config, indent=2, ensure_ascii=False) + "\n", + encoding="utf-8", + ) + return gym_config_path + + +def _copy_scene_object_to_gym_assets( + *, + scene_object: Table | Asset, + 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 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( + f"SimReady GLB for scene object {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() + + +def _gym_object_config( + *, + scene_object: Table | Asset, + asset_relative_path: str, + body_type: str, + attrs: dict[str, float | int], + max_convex_hull_num: int, +) -> dict[str, object]: + """Build one z-up gym object config from a final y-up scene object.""" + pos_y_up = _scene_vector(scene_object, "pos") + rot_y_up = _scene_vector(scene_object, "rot") + scale_y_up = _scene_vector(scene_object, "scale") + + 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": attrs, + "body_type": 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": max_convex_hull_num, + } + + +def _scene_vector(scene_object: Table | Asset, 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 {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 {field_name!r}." + ) + return vector + + +def _positive_int(value: int, *, field_name: str) -> int: + result = int(value) + if result <= 0: + raise ValueError(f"{field_name} must be positive.") + return result 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..7400fcead --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/scene_generation.py @@ -0,0 +1,576 @@ +# ---------------------------------------------------------------------------- +# 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 + +from embodichain.gen_sim.scene_engine.clients.geometry_generation import ( + GeometryGenerationClient, +) +from embodichain.gen_sim.scene_engine.core.asset import Asset +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.table import Table +from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( + OpenAICompatibleVLM, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_generation_utils import ( + align_assets_group_to_table_aabb_top, + align_assets_to_table_aabb_top, # Currently be replaced by align_assets_group_to_table_aabb_top. + export_baked_layout_object_glbs, + gravity_settle_assets_on_table, + heuristic_table_largest_internal_rectangle, + heuristic_table_support_surface, + layout_object_to_transform_matrix, + make_assets_2d_aabb_inside_table_largest_rectangle, + quaternion_wxyz_to_euler_xyz_degrees, + simready_object_glb, + transform_matrix_to_layout_object, +) + +_SUPPORTED_IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png"} + + +def generate_scene_and_refine( + image_path: str | Path, + output_root: str | Path, + scene: Scene, + *, + vlm_client: OpenAICompatibleVLM, + 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. + vlm_client=vlm_client, + geometry_generation_client=geometry_generation_client, + ) + + # Geometries refinement and layout refinement. + _refine_geometries_and_layout( + image_path=resolved_image_path, + debug_output_root=debug_output_root, + coarse_geometry_output_root=coarse_geometry_output_root, + simready_geometry_output_root=simready_geometry_output_root, + scene=scene, + vlm_client=vlm_client, + ) + + # 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, + *, + vlm_client: OpenAICompatibleVLM, + 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_multiple_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 _refine_geometries_and_layout( + image_path: str | Path, + debug_output_root: str | Path, + coarse_geometry_output_root: str | Path, + simready_geometry_output_root: str | Path, + scene: Scene, + *, + vlm_client: OpenAICompatibleVLM, +) -> None: + + # Simready all the assets(includes table). + # Treat table and assets seperately. + # Notice that, currently the simready process is only + # scale + canonicalize the glb (no real-world scale, no physical attributes). + + # Load the coarse layout. + coarse_layout = _load_layout( + Path(coarse_geometry_output_root) / "coarse_layout.json" + ) + coarse_layout_by_id = { + layout_object["id"]: layout_object for layout_object in coarse_layout + } + + # Simready all the assets. + simready_assets_layout = _simready_assets( + scene=scene, + coarse_layout_by_id=coarse_layout_by_id, + coarse_geometry_output_root=coarse_geometry_output_root, + simready_geometry_output_root=simready_geometry_output_root, + ) + + # Simready the table. + simready_table_layout = _simready_table( + scene=scene, + coarse_layout_by_id=coarse_layout_by_id, + coarse_geometry_output_root=coarse_geometry_output_root, + simready_geometry_output_root=simready_geometry_output_root, + ) + # Concat then save the table info and the assets info in one JSON file. + simready_layout = [simready_table_layout, *simready_assets_layout] + (Path(simready_geometry_output_root) / "simready_layout.json").write_text( + json.dumps(simready_layout, indent=2, ensure_ascii=False) + "\n", + encoding="utf-8", + ) + # Update the scene data structure with the simready glb paths. + _update_scene_simready_glb_paths( + scene=scene, + simready_geometry_output_root=simready_geometry_output_root, + ) + + # Layout refinement will start with the table. + refined_table_layout, refined_assets_layout = _layout_refinement( + scene=scene, + 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. + vlm_client=vlm_client, # For some cases the heuristic method still faces some undeterministic issues. + ) + # 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, + ) + + # Only for debugging. + # Save the refined layout JSON. + refined_layout = [refined_table_layout, *refined_assets_layout] + (Path(debug_output_root) / "refined_layout.json").write_text( + json.dumps(refined_layout, indent=2, ensure_ascii=False) + "\n", + encoding="utf-8", + ) + # Then use export_baked_layout_object_glbs to export it for debugging. + export_baked_layout_object_glbs( + layout=refined_layout, + geometry_root=simready_geometry_output_root, + output_root=Path(debug_output_root) / "refined_baked_geometries", + ) + + return None + + +def _update_scene_simready_glb_paths( + *, + scene: Scene, + simready_geometry_output_root: str | Path, +) -> None: + """Store the canonicalized GLB path for every scene object.""" + if scene.table is None: + raise ValueError("Cannot update SimReady paths without a table.") + + geometry_root = Path(simready_geometry_output_root).expanduser().resolve() + for scene_object in [scene.table, *scene.assets]: + glb_path = geometry_root / f"{scene_object.id}.glb" + if not glb_path.is_file(): + raise FileNotFoundError(f"SimReady geometry not found: {glb_path}") + scene_object.simready_glb_path = str(glb_path) + + +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: Table | Asset, + 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, + vlm_client: OpenAICompatibleVLM, +) -> 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. + + # refined_table_layout, refined_assets_layout = align_assets_to_table_aabb_top( + # table_layout=refined_table_layout, + # assets_layout=refined_assets_layout, + # geometry_root=simready_geometry_output_root, + # ) + refined_table_layout, refined_assets_layout = align_assets_group_to_table_aabb_top( + table_layout=refined_table_layout, + assets_layout=refined_assets_layout, + geometry_root=simready_geometry_output_root, + ) + + # 4.1. Get the table's support surface info. + # Return value format: in z-up world, the 2D convex-hull boundary coordinates. + ( + table_support_surface_2d_z_up_world_boundary, + assets_aabb_2d_z_up_world_corners_by_id, + table_mesh_2d_z_up_world_projection, + ) = heuristic_table_support_surface( + table_layout=refined_table_layout, + assets_layout=refined_assets_layout, # Render each asset's 2D AABB with its own id for checking whether any asset's AABB is outside the table's support surface. + geometry_root=simready_geometry_output_root, + debug_output_root=debug_output_root, # Keep the support surface rendered image(s) for debugging. + ) + + # 4.2. Find the table's largest internal biggest rectangle. (AABB-aligned largest rectangle.) + # Notice that, this heuristic method assumes that the table does not have some big rotation angle around z-axis in z-up world. + # Render one image for debugging. + # This rectange is axis-aligned with the z-up world coordinate system. + table_largest_internal_rectangle_2d_z_up_world = heuristic_table_largest_internal_rectangle( + table_support_surface_2d_z_up_world_boundary=table_support_surface_2d_z_up_world_boundary, # For computing the largest internal rectangle + rendering. + assets_aabb_2d_z_up_world_corners_by_id=assets_aabb_2d_z_up_world_corners_by_id, # Only for rendering. + table_mesh_2d_z_up_world_projection=table_mesh_2d_z_up_world_projection, # Only for rendering. + debug_output_root=debug_output_root, + ) + + # 6. Use the table's largest internal AABB-aligned rectange as boundary to do 2D AABB optimization, + # to let all the projected 2D AABBs of the assets inside this boundary, and keep them have no overlap + # with each other. (prepare for the next step: gravity simulation.) + # The assets layout will only update their x-y pos, and keep their z pos and rot unchanged. (do not forget the + # differences between y-up and z-up!) + refined_assets_layout = make_assets_2d_aabb_inside_table_largest_rectangle( + table_id=scene.table.id, + table_support_surface_2d_z_up_world_boundary=( + table_support_surface_2d_z_up_world_boundary + ), + table_mesh_2d_z_up_world_projection=table_mesh_2d_z_up_world_projection, + table_largest_internal_rectangle_2d_z_up_world=table_largest_internal_rectangle_2d_z_up_world, + assets_aabb_2d_z_up_world_corners_by_id=assets_aabb_2d_z_up_world_corners_by_id, + debug_output_root=debug_output_root, + assets_layout=refined_assets_layout, + ) + + # 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. + refined_assets_layout = gravity_settle_assets_on_table( + table_layout=refined_table_layout, + assets_layout=refined_assets_layout, + geometry_root=simready_geometry_output_root, + ) + + return refined_table_layout, refined_assets_layout + + +def _simready_assets( + *, + scene: Scene, + coarse_layout_by_id: dict[str, dict[str, object]], + coarse_geometry_output_root: str | Path, + simready_geometry_output_root: str | Path, +) -> list[dict[str, object]]: + # Batch process all the assets in the scene. + return [ + _simready_asset( + asset_id=asset.id, + coarse_layout=coarse_layout_by_id.get(asset.id), + coarse_geometry_output_root=coarse_geometry_output_root, + simready_geometry_output_root=simready_geometry_output_root, + ) + for asset in scene.assets + ] + + +def _simready_asset( + *, + asset_id: str, + coarse_layout: dict[str, object] | None, + coarse_geometry_output_root: str | Path, + simready_geometry_output_root: str | Path, +) -> dict[str, object]: + # Hard code some asset like bottle, treat their z-axis carefully. + # For the table, treat it with the same strategy for now. + # Add asset-id-specific SimReady processing here before the generic path. + return _simready_object( + asset_id=asset_id, + coarse_layout=coarse_layout, + coarse_geometry_output_root=coarse_geometry_output_root, + simready_geometry_output_root=simready_geometry_output_root, + ) + + +def _simready_object( + *, + asset_id: str, + coarse_layout: dict[str, object] | None, + coarse_geometry_output_root: str | Path, + simready_geometry_output_root: str | Path, +) -> dict[str, object]: + if coarse_layout is None: + raise ValueError(f"Coarse layout does not contain object {asset_id!r}.") + simready_mesh, simready_transform = simready_object_glb( + Path(coarse_geometry_output_root) / f"{asset_id}.glb", + object_id=asset_id, + rot=coarse_layout.get("rot"), + pos=coarse_layout.get("pos"), + scale=coarse_layout.get("scale"), + ) + output_path = Path(simready_geometry_output_root) / f"{asset_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 geometry was not written: {output_path}" + ) + return {"id": asset_id, **simready_transform} + + +def _simready_table( + *, + scene: Scene, + coarse_layout_by_id: dict[str, dict[str, object]], + coarse_geometry_output_root: str | Path, + simready_geometry_output_root: str | Path, +) -> dict[str, object]: + # There must be a table in one scene. + if scene.table is None: + raise ValueError("Cannot SimReady a scene without a table.") + + # Using the same strategy as the normal assets first. + return _simready_object( + asset_id=scene.table.id, + coarse_layout=coarse_layout_by_id.get(scene.table.id), + coarse_geometry_output_root=coarse_geometry_output_root, + simready_geometry_output_root=simready_geometry_output_root, + ) + + +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_segmentation.py b/embodichain/gen_sim/scene_engine/pipeline/scene_segmentation.py new file mode 100644 index 000000000..1505a970f --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/scene_segmentation.py @@ -0,0 +1,479 @@ +# ---------------------------------------------------------------------------- +# 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 +from typing import Any + +from embodichain.gen_sim.scene_engine.core.asset import Asset +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.table import Table +from embodichain.gen_sim.scene_engine.clients.image_segmentation import ( + ImageSegmentationClient, +) +from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( + OpenAICompatibleVLM, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_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"} +_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 segment_scene( + image_path: str | Path, + output_root: str | Path, + scene: Scene, + *, + vlm_client: OpenAICompatibleVLM, + image_segmentation_client: ImageSegmentationClient, +) -> 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_segmentation" + if stage_output_root.exists(): + shutil.rmtree(stage_output_root) + stage_output_root.mkdir(parents=True, exist_ok=True) + debug_output_root = stage_output_root / "debug" # Keeps the mask debug images. + masks_output_root = ( + 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=resolved_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 = render_image_without_masks( + image_path=resolved_image_path, + mask_paths=asset_mask_paths, + output_path=Path(debug_output_root) / "table_validation_base.png", + ) + # Segment the table. + _segment_table( + image_path=resolved_image_path, + validation_image_path=table_validation_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, + ) + # 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 _segment_table( + image_path: str | Path, + validation_image_path: str | Path, + 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, + 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: Table, + 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 _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 table validation response has an incomplete code fence.") + return "\n".join(lines[1:-1]).strip() + + +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[Asset]] = {} + 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[Asset], + 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[Asset], + 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 + + +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 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..2990c70e5 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/scene_understanding.py @@ -0,0 +1,254 @@ +# ---------------------------------------------------------------------------- +# 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 embodichain.gen_sim.scene_engine.core.asset import Asset +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.table import Table +from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( + OpenAICompatibleVLM, +) + +_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." + + +def understand_scene( + scene: Scene, + image_path: str | Path, + output_root: str | Path, + *, + vlm_client: OpenAICompatibleVLM, + json_max_attempts: int = 3, +) -> Scene: + + if json_max_attempts < 1: + raise ValueError("json_max_attempts must be at least 1.") + + 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) + + last_validation_error: ValueError | None = None + for attempt in range(1, json_max_attempts + 1): + response_text = vlm_client.complete( + image_path=resolved_image_path, + system_prompt=_SYSTEM_PROMPT, + user_prompt=_USER_PROMPT, + ) + try: + understood_scene = validate_scene_understanding_json(response_text) + scene.table = understood_scene.table + scene.assets = understood_scene.assets + validate_scene_understanding(scene) + except ValueError as exc: + last_validation_error = exc + continue + + (stage_output_root / "scene.json").write_text( + json.dumps(scene.to_dict(), indent=2, ensure_ascii=False) + "\n", + encoding="utf-8", + ) + return scene + + assert last_validation_error is not None + raise ValueError( + "VLM returned invalid scene-understanding JSON after " + f"{json_max_attempts} attempts: {last_validation_error}" + ) from last_validation_error + + +def validate_scene_understanding_json(response_text: str) -> Scene: + """Parse a VLM response and create a core ``Scene`` with generated IDs.""" + 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 = Table( + # id=_next_id(table_fields["category"], id_counters) + # Use a fixed ID for the table. + id="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[Asset] = [] + for index, asset in enumerate(assets_value): + fields = _parse_scene_object_fields(asset, field_name=f"assets[{index}]") + assets.append( + Asset( + id=_next_id(fields["category"], id_counters), + **fields, + ) + ) + + return Scene(table=table, assets=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 true. 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}" 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/scene_generation_utils.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_generation_utils.py new file mode 100644 index 000000000..ca4034863 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_generation_utils.py @@ -0,0 +1,1547 @@ +# ---------------------------------------------------------------------------- +# 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 re +from typing import Sequence + +from embodichain.lab.sim import SimulationManager as _EmbodiSimManager +from embodichain.lab.sim import SimulationManagerCfg +from embodichain.lab.sim.cfg import RigidObjectCfg +from embodichain.lab.sim.shapes import MeshCfg +import matplotlib +import numpy as np +import open3d as o3d +from scipy.spatial import ConvexHull, QhullError +from scipy.spatial.transform import Rotation +import trimesh + +matplotlib.use("Agg") + +import matplotlib.pyplot as plt +from matplotlib.collections import PolyCollection +from matplotlib.ticker import MaxNLocator + +_UPRIGHT_CONTAINER_ID_TOKENS = frozenset({"bottle", "can", "jar", "flask", "thermos"}) + + +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_rotation_to_simulation_euler_xyz_degrees( + layout_object: dict[str, object], +) -> list[float]: + """Convert a layout's lowercase-``xyz`` Euler rotation for SimulationManager. + + Scene layouts use ``Rotation.from_euler("xyz", ...)``, whereas + ``RigidObjectCfg.init_rot`` is interpreted with uppercase ``"XYZ"``. + Convert through the rotation matrix so both represent exactly the same pose. + """ + layout_rotation = Rotation.from_euler( + "xyz", + _three_floats(layout_object.get("rot"), field_name="rot"), + degrees=True, + ) + return layout_rotation.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 align_assets_to_table_aabb_top( + *, + table_layout: dict[str, object], + assets_layout: list[dict[str, object]], + geometry_root: str | Path, + clearance: float = 0.02, # 2cm. +) -> tuple[dict[str, object], list[dict[str, object]]]: + """Place assets above a table using temporary z-up AABB height calculations. + + Input and output layouts use y-up, matching the GLBs on disk. The geometry + and layouts are converted to z-up only while measuring and changing height. + + Notice: + - The refinement pipeline currently uses the group version so it preserves + the assets' relative vertical arrangement before gravity simulation. + """ + if clearance < 0: + raise ValueError("Table clearance must be non-negative.") + + # Prepare y-up and z-up conversion matrices. + 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 = _convert_layout_coordinate_system( + table_layout, + source_to_target_matrix=y_up_to_z_up_matrix, + ) + z_up_assets_layout = [ + _convert_layout_coordinate_system( + asset_layout, + source_to_target_matrix=y_up_to_z_up_matrix, + ) + for asset_layout in assets_layout + ] + + # Get the table's top z position in z-up coordinates, and add the clearance to it. + resolved_geometry_root = Path(geometry_root).expanduser().resolve() + table_mesh = load_glb_mesh( + resolved_geometry_root / f"{z_up_table_layout['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_asset_bottom_z = table_mesh.bounds[1, 2] + clearance + + # Iterate through each asset and adjust its z position to sit above the table. + for asset_layout in z_up_assets_layout: + asset_mesh = load_glb_mesh(resolved_geometry_root / f"{asset_layout['id']}.glb") + asset_mesh.apply_transform(y_up_to_z_up_matrix) + asset_mesh.apply_transform(layout_object_to_transform_matrix(asset_layout)) + asset_bottom_z = asset_mesh.bounds[0, 2] + asset_layout["pos"][2] += target_asset_bottom_z - asset_bottom_z + + return ( + _convert_layout_coordinate_system( + z_up_table_layout, + source_to_target_matrix=z_up_to_y_up_matrix, + ), + [ + _convert_layout_coordinate_system( + asset_layout, + source_to_target_matrix=z_up_to_y_up_matrix, + ) + for asset_layout in z_up_assets_layout + ], + ) + + +def align_assets_group_to_table_aabb_top( + *, + table_layout: dict[str, object], + assets_layout: list[dict[str, object]], + geometry_root: str | Path, + clearance: float = 0.02, # 2cm. +) -> tuple[dict[str, object], list[dict[str, object]]]: + """Place all assets as one rigid vertical 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; + """ + if clearance < 0: + raise ValueError("Table clearance must be non-negative.") + if not assets_layout: + return table_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 = _convert_layout_coordinate_system( + table_layout, + source_to_target_matrix=y_up_to_z_up_matrix, + ) + z_up_assets_layout = [ + _convert_layout_coordinate_system( + asset_layout, + source_to_target_matrix=y_up_to_z_up_matrix, + ) + for asset_layout in assets_layout + ] + + resolved_geometry_root = Path(geometry_root).expanduser().resolve() + table_mesh = load_glb_mesh( + resolved_geometry_root / f"{z_up_table_layout['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] + clearance + + group_bottom_z = np.inf + for asset_layout in z_up_assets_layout: + asset_mesh = load_glb_mesh(resolved_geometry_root / f"{asset_layout['id']}.glb") + asset_mesh.apply_transform(y_up_to_z_up_matrix) + asset_mesh.apply_transform(layout_object_to_transform_matrix(asset_layout)) + group_bottom_z = min( + group_bottom_z, float(asset_mesh.bounds[0, 2]) + ) # Find the lowest z among all the assets. + + 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 + + return ( + _convert_layout_coordinate_system( + z_up_table_layout, + source_to_target_matrix=z_up_to_y_up_matrix, + ), + [ + _convert_layout_coordinate_system( + asset_layout, + source_to_target_matrix=z_up_to_y_up_matrix, + ) + for asset_layout in z_up_assets_layout + ], + ) + + +def _prepare_gravity_sim_body( + *, + layout_object: dict[str, object], + geometry_root: Path, + y_up_to_z_up_matrix: np.ndarray, +) -> tuple[ + Path, + trimesh.Trimesh, + dict[str, object], + list[float], + list[float], +]: + """Load one y-up GLB and derive its z-up rigid pose for gravity simulation.""" + object_id = str(layout_object["id"]) + source_mesh_path = geometry_root / f"{object_id}.glb" + source_mesh = load_glb_mesh(source_mesh_path) + z_up_layout = _convert_layout_coordinate_system( + layout_object, + source_to_target_matrix=y_up_to_z_up_matrix, + ) + y_up_scale = _three_floats(layout_object.get("scale"), field_name="scale") + z_up_scale = _three_floats(z_up_layout.get("scale"), field_name="scale") + z_up_rigid_layout = { + "id": object_id, + "rot": _three_floats(z_up_layout.get("rot"), field_name="rot"), + "pos": _three_floats(z_up_layout.get("pos"), field_name="pos"), + "scale": [1.0, 1.0, 1.0], + } + return ( + source_mesh_path, + source_mesh, + z_up_rigid_layout, + y_up_scale, + z_up_scale, + ) + + +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.""" + y_up_mesh.apply_transform(y_up_to_z_up_matrix) + scale_matrix = np.eye(4) + scale_matrix[:3, :3] = np.diag(z_up_scale) + y_up_mesh.apply_transform(scale_matrix) + y_up_mesh.apply_transform(layout_object_to_transform_matrix(z_up_rigid_layout)) + return y_up_mesh + + +def gravity_settle_assets_on_table( + *, + table_layout: dict[str, object], + assets_layout: list[dict[str, object]], + geometry_root: str | Path, + clearance: float = 0.02, + settle_steps: int = 300, + physics_dt: float = 1.0 / 100.0, + sim_device: str = "cpu", + max_convex_hull_num: int = 32, +) -> list[dict[str, object]]: + """Settle all assets together on a static table with z-up gravity. + + Layouts and source GLBs are y-up. The simulator automatically converts its + y-up GLB inputs to z-up, while its gravity poses are expressed in z-up. + This function therefore keeps the source meshes y-up and converts only the + layout poses for measurement and simulation. Before all dynamic assets are + added to one simulation, each asset's own lowest AABB z is placed + ``clearance`` above the table AABB top. The final rigid-body poses are + converted back to y-up layouts, with their original scales preserved. + """ + + # Check. + if clearance < 0.0: + raise ValueError("Gravity-settle clearance must be non-negative.") + if settle_steps <= 0: + raise ValueError("Gravity-settle steps must be positive.") + if physics_dt <= 0.0: + raise ValueError("Gravity-settle physics_dt must be positive.") + if max_convex_hull_num <= 0: + raise ValueError("Gravity-settle max_convex_hull_num must be positive.") + if not assets_layout: + return [] + + 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.") + asset_ids: set[str] = set() + 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_ids: + raise ValueError(f"Asset layouts contain duplicate id {asset_id!r}.") + asset_ids.add(asset_id) + + # The source GLBs/layouts are y-up, while the gravity service uses z-up. + 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() + + ( + table_mesh_path, + table_mesh, + table_rigid_layout, + table_y_up_scale, + table_z_up_scale, + ) = _prepare_gravity_sim_body( + layout_object=table_layout, + geometry_root=resolved_geometry_root, + y_up_to_z_up_matrix=y_up_to_z_up_matrix, + ) + # Match the simulator's automatic y-up-GLB conversion while measuring the + # physical z-up table top. + table_world_mesh = _mesh_to_z_up_world_for_aabb( + y_up_mesh=table_mesh, + z_up_rigid_layout=table_rigid_layout, + z_up_scale=table_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 assets_layout: + asset_id = str(asset_layout["id"]) + ( + asset_mesh_path, + asset_mesh, + asset_rigid_layout, + asset_y_up_scale, + asset_z_up_scale, + ) = _prepare_gravity_sim_body( + layout_object=asset_layout, + geometry_root=resolved_geometry_root, + y_up_to_z_up_matrix=y_up_to_z_up_matrix, + ) + asset_world_mesh = _mesh_to_z_up_world_for_aabb( + y_up_mesh=asset_mesh, + z_up_rigid_layout=asset_rigid_layout, + z_up_scale=asset_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_rigid_layout["pos"][2] += table_top_z + clearance - asset_bottom_z + prepared_assets[asset_id] = { + "mesh_path": asset_mesh_path, + "rigid_layout": asset_rigid_layout, + "y_up_scale": asset_y_up_scale, + "z_up_scale": asset_z_up_scale, + } + + sim = _EmbodiSimManager( + SimulationManagerCfg( + headless=True, + physics_dt=physics_dt, + sim_device=sim_device, + ) + ) + try: + sim.add_rigid_object( + RigidObjectCfg( + uid=table_id, + shape=MeshCfg(fpath=str(table_mesh_path)), + init_pos=tuple(table_rigid_layout["pos"]), + init_rot=tuple( + _layout_rotation_to_simulation_euler_xyz_degrees(table_rigid_layout) + ), + body_scale=tuple(table_y_up_scale), + body_type="static", + max_convex_hull_num=max_convex_hull_num, + ) + ) + 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( + _layout_rotation_to_simulation_euler_xyz_degrees(rigid_layout) + ), + body_scale=tuple(asset_info["y_up_scale"]), + body_type="dynamic", + max_convex_hull_num=max_convex_hull_num, + ) + ) + + # All assets share this one simulation, so they can collide with the + # table and with one another while settling. + sim.update(step=settle_steps) + + 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: + sim._deferred_destroy() + + settled_assets_layout = [ + settled_layout_by_id[str(asset_layout["id"])] for asset_layout in assets_layout + ] + return settled_assets_layout + + +def heuristic_table_support_surface( + *, + table_layout: dict[str, object], + assets_layout: list[dict[str, object]], + geometry_root: str | Path, + debug_output_root: str | Path, +) -> tuple[ + list[list[float]], + dict[str, list[list[float]]], + dict[str, list[list[float]] | list[list[int]]], +]: + """Return the table support boundary, asset AABBs, and table 2D mesh. + + The input table layout and its GLB use y-up. This function will convert + both to temporary z-up coordinates before extracting the support surface. + The returned convex-hull boundary is ordered counter-clockwise in the z-up + world x-y plane. Each projected rectangle is keyed by asset id and contains + four counter-clockwise x-y corners. The projected table mesh contains 2D + vertices and triangle faces, so later stages do not need to recompute it. + """ + 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.") + + resolved_geometry_root = Path(geometry_root).expanduser().resolve() + table_glb_path = resolved_geometry_root / f"{table_id}.glb" + if not table_glb_path.is_file(): + raise FileNotFoundError(f"Table geometry not found: {table_glb_path}") + + resolved_debug_output_root = Path(debug_output_root).expanduser().resolve() + resolved_debug_output_root.mkdir(parents=True, exist_ok=True) + + # 1. Load the y-up table GLB, convert its vertices and layout to z-up, then + # apply the z-up world transform to obtain the table world geometry. + 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_table_layout = _convert_layout_coordinate_system( + table_layout, + source_to_target_matrix=y_up_to_z_up_matrix, + ) + table_world_mesh = load_glb_mesh(table_glb_path) + table_world_mesh.apply_transform(y_up_to_z_up_matrix) + table_world_mesh.apply_transform( + layout_object_to_transform_matrix(z_up_table_layout) + ) + + # Prepare every asset's z-up world x-y AABB for the debug rendering. + # To check if any asset's AABB is outside the table's support surface. + assets_2d_aabbs: list[tuple[str, np.ndarray]] = ( + [] + ) # id + 2D AABB infos in z-up world x-y plane. + projected_rectangles_by_id: dict[str, list[list[float]]] = {} + 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.") + asset_glb_path = resolved_geometry_root / f"{asset_id}.glb" + if not asset_glb_path.is_file(): + raise FileNotFoundError(f"Asset geometry not found: {asset_glb_path}") + + z_up_asset_layout = _convert_layout_coordinate_system( + asset_layout, + source_to_target_matrix=y_up_to_z_up_matrix, + ) + asset_world_mesh = load_glb_mesh(asset_glb_path) + asset_world_mesh.apply_transform(y_up_to_z_up_matrix) + asset_world_mesh.apply_transform( + layout_object_to_transform_matrix(z_up_asset_layout) + ) + asset_bounds_xy = asset_world_mesh.bounds[:, :2] + asset_2d_aabb = 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]], + ] + ) + assets_2d_aabbs.append((asset_id, asset_2d_aabb)) + projected_rectangles_by_id[asset_id] = asset_2d_aabb.tolist() + + # 2. Project every table triangle into the z-up world's x-y plane. + if len(table_world_mesh.vertices) < 3 or len(table_world_mesh.faces) == 0: + raise ValueError("Table geometry must contain at least one triangle.") + projected_vertices = table_world_mesh.vertices[ + :, :2 + ] # Ignore z, for we wanna get the x-y plane projection. + try: + projected_hull = ConvexHull( + projected_vertices + ) # Compute the convex hull for the 2D projection. + # Notice that: for the L-shape table, this will return a bad result. + except QhullError as exc: + raise ValueError("Table's x-y projection is degenerate.") from exc + support_region_boundary = projected_vertices[projected_hull.vertices] + + projected_triangles = projected_vertices[table_world_mesh.faces] + # 3. Render the full projected mesh and its outer boundary for debugging. + _render_table_xy_projection( + projected_triangles=projected_triangles, # All the projection triangles, draw with blue color. + support_region_boundary=support_region_boundary, # The convex hull boundary, draw with red line. + assets_2d_aabbs=assets_2d_aabbs, # Render together for debugging. + table_id=table_id, + output_path=resolved_debug_output_root / "table_xy_projection.png", + ) + + # 4. Return the convex-hull boundary, each asset's AABB, and the table 2D mesh. + table_projected_mesh_2d: dict[str, list[list[float]] | list[list[int]]] = { + "vertices": projected_vertices.tolist(), + "faces": table_world_mesh.faces.tolist(), + } + return ( + support_region_boundary.tolist(), + projected_rectangles_by_id, + table_projected_mesh_2d, + ) + + +def _render_table_xy_projection( + *, + projected_triangles: np.ndarray, + support_region_boundary: np.ndarray, + assets_2d_aabbs: list[tuple[str, np.ndarray]], + largest_internal_rectangle: np.ndarray | None = None, + table_id: str, + output_path: str | Path, +) -> Path: + """Render a table's z-up world x-y projection with axes and tick marks.""" + resolved_output_path = Path(output_path).expanduser().resolve() + resolved_output_path.parent.mkdir(parents=True, exist_ok=True) + + figure, axes = plt.subplots(figsize=(8, 8), dpi=160) + axes.add_collection( + PolyCollection( + projected_triangles, + facecolor="steelblue", + alpha=0.08, + edgecolor="none", + ) + ) + closed_boundary = np.vstack( + [support_region_boundary, support_region_boundary[0]] + ) # Close the convex hull boundary by adding the first point to the end of the array. + axes.plot( + closed_boundary[:, 0], + closed_boundary[:, 1], + color="crimson", + linewidth=2.0, + label="2D convex-hull boundary", + ) + if largest_internal_rectangle is not None: + closed_largest_internal_rectangle = np.vstack( + [largest_internal_rectangle, largest_internal_rectangle[0]] + ) + axes.fill( + closed_largest_internal_rectangle[:, 0], + closed_largest_internal_rectangle[:, 1], + color="seagreen", + alpha=0.25, + label="largest internal x-y AABB", + ) + axes.plot( + closed_largest_internal_rectangle[:, 0], + closed_largest_internal_rectangle[:, 1], + color="seagreen", + linewidth=2.0, + ) + # Render each asset's 2D AABB with its own id for debugging. + for index, (asset_id, asset_aabb) in enumerate(assets_2d_aabbs): + closed_asset_aabb = np.vstack([asset_aabb, asset_aabb[0]]) + axes.fill( + closed_asset_aabb[:, 0], + closed_asset_aabb[:, 1], + color="darkorange", + alpha=0.16, + label="asset 2D AABB" if index == 0 else None, + ) + axes.plot( + closed_asset_aabb[:, 0], + closed_asset_aabb[:, 1], + color="darkorange", + linewidth=1.5, + ) + asset_aabb_center = asset_aabb.mean(axis=0) + axes.text( + asset_aabb_center[0], + asset_aabb_center[1], + asset_id, + color="black", + fontsize=8, + ha="center", + va="center", + bbox={"facecolor": "white", "alpha": 0.8, "edgecolor": "none"}, + ) + axes.scatter( + 0.0, + 0.0, + color="black", + marker="+", + s=100, + label="world origin", + ) + axes.update_datalim(np.array([[0.0, 0.0]])) + axes.autoscale_view() + axes.axhline(0.0, color="black", linewidth=0.8, alpha=0.55) + axes.axvline(0.0, color="black", linewidth=0.8, alpha=0.55) + + x_min, x_max = axes.get_xlim() + y_min, y_max = axes.get_ylim() + axes.annotate( + "+x", + xy=(x_max, 0.0), + xytext=(x_max - (x_max - x_min) * 0.12, (y_max - y_min) * 0.03), + arrowprops={"arrowstyle": "->", "color": "black"}, + ha="right", + va="bottom", + ) + axes.annotate( + "+y", + xy=(0.0, y_max), + xytext=((x_max - x_min) * 0.03, y_max - (y_max - y_min) * 0.12), + arrowprops={"arrowstyle": "->", "color": "black"}, + ha="left", + va="top", + ) + axes.set_aspect("equal", adjustable="box") + axes.set_xlabel("x (z-up world)") + axes.set_ylabel("y (z-up world)") + axes.set_title(f"Table 2D Projection: {table_id}") + axes.xaxis.set_major_locator(MaxNLocator(nbins=8)) + axes.yaxis.set_major_locator(MaxNLocator(nbins=8)) + axes.tick_params(axis="both", which="major", labelsize=9) + axes.legend(loc="best") + axes.grid(True, alpha=0.25) + figure.savefig(resolved_output_path, bbox_inches="tight") + plt.close(figure) + return resolved_output_path + + +def heuristic_table_largest_internal_rectangle( + *, + table_support_surface_2d_z_up_world_boundary: Sequence[Sequence[float]], + assets_aabb_2d_z_up_world_corners_by_id: dict[str, list[list[float]]], + table_mesh_2d_z_up_world_projection: dict[str, list[list[float]] | list[list[int]]], + debug_output_root: str | Path, +) -> list[list[float]]: + """Return the largest centered, x/y-aligned AABB with the table AABB aspect ratio. + + The table boundary is used to binary-search a safe uniform scale. Asset + AABBs and the table mesh projection are only reused for debug rendering. + """ + # The boundary is already in the z-up world x-y plane. + boundary = np.asarray(table_support_surface_2d_z_up_world_boundary, dtype=float) + if boundary.ndim != 2 or boundary.shape[1] != 2 or len(boundary) < 3: + raise ValueError( + "Table support-region boundary must contain at least three 2D points." + ) + if not np.all(np.isfinite(boundary)): + raise ValueError( + "Table support-region boundary must contain only finite values." + ) + if np.allclose(boundary[0], boundary[-1]): + boundary = boundary[:-1] + + # The support-surface stage has already returned this as a counter-clockwise + # convex-hull boundary, so do not compute another convex hull here. + convex_boundary = boundary + + boundary_min = convex_boundary.min(axis=0) + boundary_max = convex_boundary.max(axis=0) + # Build the smallest origin-centered 2D AABB that contains the red boundary. + boundary_half_extents = np.maximum( + np.abs(boundary_min), + np.abs(boundary_max), + ) + boundary_size = boundary_half_extents * 2.0 + if np.any(boundary_size <= 0): + raise ValueError( + "Table support-region boundary must have non-zero width and height." + ) + + # Keep the internal rectangle centered at the table/world origin. + # rectangle_center = convex_boundary.mean(axis=0) # The mean is not always 0,0. + rectangle_center = np.array([0.0, 0.0]) + coordinate_scale = max(float(boundary_size.max()), 1.0) + containment_tolerance = coordinate_scale * 1e-8 + edge_starts = convex_boundary + edge_vectors = np.roll(convex_boundary, -1, axis=0) - edge_starts + + def _rectangle_at_scale(scale: float) -> np.ndarray: + half_extents = boundary_size * scale / 2.0 + return np.array( + [ + rectangle_center - half_extents, + rectangle_center + [half_extents[0], -half_extents[1]], + rectangle_center + half_extents, + rectangle_center + [-half_extents[0], half_extents[1]], + ] + ) + + def _is_inside_boundary(rectangle: np.ndarray) -> bool: + corner_offsets = rectangle[None, :, :] - edge_starts[:, None, :] + cross_products = ( + edge_vectors[:, 0, None] * corner_offsets[:, :, 1] + - edge_vectors[:, 1, None] * corner_offsets[:, :, 0] + ) + return bool(np.all(cross_products >= -containment_tolerance)) + + # Binary-search the largest safe uniform scale in [0, 1]. + largest_safe_scale = 0.0 + smallest_unsafe_scale = 1.0 + for _ in range(32): + candidate_scale = (largest_safe_scale + smallest_unsafe_scale) / 2.0 + if _is_inside_boundary(_rectangle_at_scale(candidate_scale)): + largest_safe_scale = candidate_scale + else: + smallest_unsafe_scale = candidate_scale + if largest_safe_scale <= 1e-8: + raise ValueError("Table support-region boundary has no usable interior area.") + largest_internal_rectangle = _rectangle_at_scale(largest_safe_scale) + + # These values were created by heuristic_table_support_surface in this + # pipeline, so convert them for rendering without validating them again. + projected_vertices = np.asarray( + table_mesh_2d_z_up_world_projection["vertices"], dtype=float + ) + projected_faces = np.asarray( + table_mesh_2d_z_up_world_projection["faces"], dtype=int + ) + assets_2d_aabbs = [ + (asset_id, np.asarray(asset_aabb, dtype=float)) + for asset_id, asset_aabb in assets_aabb_2d_z_up_world_corners_by_id.items() + ] + _render_table_xy_projection( + projected_triangles=projected_vertices[projected_faces], + support_region_boundary=convex_boundary, + assets_2d_aabbs=assets_2d_aabbs, + largest_internal_rectangle=largest_internal_rectangle, + table_id="table", + output_path=( + Path(debug_output_root).expanduser().resolve() + / "table_largest_internal_rectangle.png" + ), + ) + return largest_internal_rectangle.tolist() + + +def make_assets_2d_aabb_inside_table_largest_rectangle( + *, + table_id: str, + table_support_surface_2d_z_up_world_boundary: Sequence[Sequence[float]], + table_mesh_2d_z_up_world_projection: dict[str, list[list[float]] | list[list[int]]], + table_largest_internal_rectangle_2d_z_up_world: Sequence[Sequence[float]], + assets_aabb_2d_z_up_world_corners_by_id: dict[str, list[list[float]]], + debug_output_root: str | Path, + assets_layout: list[dict[str, object]], + boundary_margin: float = 1e-6, + aabb_clearance: float = 1e-6, +) -> list[dict[str, object]]: + """Center the asset AABB union, then pack the AABBs inside the table. + + All AABB inputs are in the z-up world's x-y plane. Layouts remain y-up, so + a z-up planar offset ``(dx, dy)`` is written back as ``pos.x += dx`` and + ``pos.z -= dy``. ``boundary_margin`` and ``aabb_clearance`` are deliberately + near zero by default, but remain explicit so callers can request a gap. + The table projection inputs are used only to render the final debug image. + """ + if not assets_layout: + return [] + + # Get the table's largest internal rectangle's min and max corners in the z-up world x-y plane. + rectangle_min, rectangle_max = _aabb_2d_bounds_from_corners( + table_largest_internal_rectangle_2d_z_up_world, + name="Table largest internal rectangle", + require_nonzero_extent=True, + ) + + # Prepare asset layouts by id for validation and later lookup. + layout_by_id: dict[str, dict[str, object]] = {} + 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 layout_by_id: + raise ValueError(f"Asset layouts contain duplicate id {asset_id!r}.") + layout_by_id[asset_id] = asset_layout + + aabb_ids = set(assets_aabb_2d_z_up_world_corners_by_id) + layout_ids = set(layout_by_id) + if aabb_ids != layout_ids: + missing_aabbs = sorted(layout_ids - aabb_ids) + missing_layouts = sorted(aabb_ids - layout_ids) + raise ValueError( + "Asset layouts and 2D AABBs must have the same ids: " + f"missing AABBs={missing_aabbs}, missing layouts={missing_layouts}." + ) + + aabb_corners_by_id: dict[str, np.ndarray] = {} + aabb_bounds_by_id: dict[str, tuple[np.ndarray, np.ndarray]] = {} + for asset_id, corners in assets_aabb_2d_z_up_world_corners_by_id.items(): + corner_array = np.asarray(corners, dtype=float) + asset_min, asset_max = _aabb_2d_bounds_from_corners( + corner_array, + name=f"Asset {asset_id!r} 2D AABB", + require_nonzero_extent=False, + ) + aabb_corners_by_id[asset_id] = corner_array + aabb_bounds_by_id[asset_id] = (asset_min, asset_max) + + # Union all the assets' AABBs to find the center of the group, then offset all AABBs to be centered at the origin. + # A heuristic implementation. + union_min = np.min( + np.stack([bounds[0] for bounds in aabb_bounds_by_id.values()]), axis=0 + ) + union_max = np.max( + np.stack([bounds[1] for bounds in aabb_bounds_by_id.values()]), axis=0 + ) + union_center = (union_min + union_max) / 2.0 + union_to_origin_offset = -union_center + # Center all the AABBs by subtracting the union center from each corner. + centered_aabb_corners_by_id = { + asset_id: corners + union_to_origin_offset + for asset_id, corners in aabb_corners_by_id.items() + } + # Optimize all the asset AABBs: + # 1. Do not collide with each other. + # 2. Inside the table's region. + optimizer_offsets_by_id = _optimize_assets_2d_aabbs_in_rectangle( + rectangle_min=rectangle_min, + rectangle_max=rectangle_max, + aabb_corners_by_id=centered_aabb_corners_by_id, + boundary_margin=boundary_margin, + aabb_clearance=aabb_clearance, + ) + + # Render the final packed AABBs using the original table support-surface + # projection rather than approximating the table with its internal rectangle. + projected_vertices = np.asarray( + table_mesh_2d_z_up_world_projection["vertices"], dtype=float + ) + projected_faces = np.asarray( + table_mesh_2d_z_up_world_projection["faces"], dtype=int + ) + final_assets_2d_aabbs = [ + ( + asset_id, + centered_aabb_corners_by_id[asset_id] + optimizer_offsets_by_id[asset_id], + ) + for asset_id in sorted(centered_aabb_corners_by_id) + ] + _render_table_xy_projection( + projected_triangles=projected_vertices[projected_faces], + support_region_boundary=np.asarray( + table_support_surface_2d_z_up_world_boundary, + dtype=float, + ), + assets_2d_aabbs=final_assets_2d_aabbs, + largest_internal_rectangle=np.asarray( + table_largest_internal_rectangle_2d_z_up_world, + dtype=float, + ), + table_id=table_id, + output_path=( + Path(debug_output_root).expanduser().resolve() + / "assets_2d_aabb_optimization.png" + ), + ) + + # Update each asset layout's planar position only: z-up (x, y) maps to + # y-up (x, -z), so update layout pos.x and pos.z while preserving pos.y, + # rotation, and scale. + refined_assets_layout: list[dict[str, object]] = [] + for asset_layout in assets_layout: + asset_id = str(asset_layout["id"]) + final_z_up_xy_offset = ( + union_to_origin_offset + optimizer_offsets_by_id[asset_id] + ) + refined_layout = dict(asset_layout) + refined_pos = _three_floats(asset_layout.get("pos"), field_name="pos") + refined_pos[0] += float(final_z_up_xy_offset[0]) + refined_pos[2] -= float(final_z_up_xy_offset[1]) + refined_layout["pos"] = refined_pos + refined_assets_layout.append(refined_layout) + + return refined_assets_layout + + +def _aabb_2d_bounds_from_corners( + corners: Sequence[Sequence[float]] | np.ndarray, + *, + name: str, + require_nonzero_extent: bool, +) -> tuple[np.ndarray, np.ndarray]: + """Validate 2D AABB corners and return their minimum and maximum corners.""" + corner_array = np.asarray(corners, dtype=float) + if corner_array.shape != (4, 2) or not np.all(np.isfinite(corner_array)): + raise ValueError(f"{name} must be four finite [x, y] corners.") + minimum = corner_array.min(axis=0) + maximum = corner_array.max(axis=0) + if require_nonzero_extent and np.any(maximum <= minimum): + raise ValueError(f"{name} must have non-zero width and height.") + return minimum, maximum + + +def _aabb_pair_overlap_depths( + *, + current_mins: np.ndarray, + current_maxs: np.ndarray, + first_index: int, + second_index: int, + aabb_clearance: float, + tolerance: float, +) -> tuple[float, float] | None: + """Return x/y overlap depths, or ``None`` when two AABBs do not overlap.""" + overlap_x = ( + min(current_maxs[first_index, 0], current_maxs[second_index, 0]) + - max(current_mins[first_index, 0], current_mins[second_index, 0]) + + aabb_clearance + ) + overlap_y = ( + min(current_maxs[first_index, 1], current_maxs[second_index, 1]) + - max(current_mins[first_index, 1], current_mins[second_index, 1]) + + aabb_clearance + ) + if overlap_x <= tolerance or overlap_y <= tolerance: + return None + return overlap_x, overlap_y + + +def _find_overlapping_2d_aabb_pairs( + *, + current_mins: np.ndarray, + current_maxs: np.ndarray, + aabb_clearance: float, + tolerance: float, +) -> list[tuple[float, int, int]]: + """Return overlapping pairs, most constrained pair first.""" + overlaps: list[tuple[float, int, int]] = [] + for first_index in range(len(current_mins)): + for second_index in range(first_index + 1, len(current_mins)): + overlap_depths = _aabb_pair_overlap_depths( + current_mins=current_mins, + current_maxs=current_maxs, + first_index=first_index, + second_index=second_index, + aabb_clearance=aabb_clearance, + tolerance=tolerance, + ) + if overlap_depths is not None: + overlaps.append((min(overlap_depths), first_index, second_index)) + return sorted(overlaps, reverse=True) + + +def _aabb_pair_push_candidates( + *, + current_mins: np.ndarray, + current_maxs: np.ndarray, + first_index: int, + second_index: int, + allowed_min: np.ndarray, + allowed_max: np.ndarray, + aabb_clearance: float, + tolerance: float, +) -> list[tuple[float, int, float, float, float]] | None: + """Return feasible opposite-direction pushes, or ``None`` if already separate.""" + if ( + _aabb_pair_overlap_depths( + current_mins=current_mins, + current_maxs=current_maxs, + first_index=first_index, + second_index=second_index, + aabb_clearance=aabb_clearance, + tolerance=tolerance, + ) + is None + ): + return None + + candidates: list[tuple[float, int, float, float, float]] = [] + for axis in (0, 1): + for first_direction in (-1.0, 1.0): + second_direction = -first_direction + if first_direction < 0.0: + required_distance = ( + current_maxs[first_index, axis] + + aabb_clearance + - current_mins[second_index, axis] + ) + first_capacity = max( + 0.0, + current_mins[first_index, axis] - allowed_min[axis], + ) + second_capacity = max( + 0.0, + allowed_max[axis] - current_maxs[second_index, axis], + ) + else: + required_distance = ( + current_maxs[second_index, axis] + + aabb_clearance + - current_mins[first_index, axis] + ) + first_capacity = max( + 0.0, + allowed_max[axis] - current_maxs[first_index, axis], + ) + second_capacity = max( + 0.0, + current_mins[second_index, axis] - allowed_min[axis], + ) + if first_capacity + second_capacity < required_distance - tolerance: + continue + + # Split the required movement as evenly as possible, constrained by + # each AABB's remaining distance to the table boundary. + first_move = float( + np.clip( + required_distance / 2.0, + max(0.0, required_distance - second_capacity), + min(required_distance, first_capacity), + ) + ) + second_move = required_distance - first_move + candidates.append( + ( + first_move**2 + second_move**2, + axis, + first_direction, + first_move, + second_move, + ) + ) + return candidates + + +def _optimize_assets_2d_aabbs_in_rectangle( + *, + rectangle_min: np.ndarray, + rectangle_max: np.ndarray, + aabb_corners_by_id: dict[str, np.ndarray], + boundary_margin: float, + aabb_clearance: float, + max_rounds: int = 8, +) -> dict[str, np.ndarray]: + """Greedily pack 2D AABBs with minimum local squared displacement.""" + + # Check the inputs for validity. + if not np.isfinite(boundary_margin) or boundary_margin < 0.0: + raise ValueError("boundary_margin must be a finite non-negative number.") + if not np.isfinite(aabb_clearance) or aabb_clearance < 0.0: + raise ValueError("aabb_clearance must be a finite non-negative number.") + if max_rounds <= 0: + raise ValueError("max_rounds must be positive.") + + asset_ids = sorted(aabb_corners_by_id) + if not asset_ids: + return {} + + asset_mins: list[np.ndarray] = [] + asset_maxs: list[np.ndarray] = [] + for asset_id in asset_ids: + corners = aabb_corners_by_id[asset_id] + # Get all the asset's AABB min and max corners in the z-up world x-y plane. + asset_min, asset_max = _aabb_2d_bounds_from_corners( + corners, + name=f"Asset {asset_id!r} centered 2D AABB", + require_nonzero_extent=False, + ) + asset_mins.append(asset_min) + asset_maxs.append(asset_max) + + base_mins = np.stack(asset_mins) + base_maxs = np.stack(asset_maxs) + # Get table support surface's largest internal rectangle's min and max corners in the z-up world x-y plane. + allowed_min = rectangle_min + boundary_margin + allowed_max = rectangle_max - boundary_margin + # Compute the least and greatest offsets for each asset's AABB to stay inside the table's largest internal rectangle. + lower_offset_bounds = allowed_min - base_mins + upper_offset_bounds = allowed_max - base_maxs + + # Check if any asset's AABB is larger than the table's largest internal rectangle after applying the boundary margin. If so, raise an error. + if np.any(lower_offset_bounds > upper_offset_bounds + 1e-9): + too_large_index = int( + np.argwhere(lower_offset_bounds > upper_offset_bounds)[0, 0] + ) + asset_id = asset_ids[too_large_index] + raise ValueError( + f"Asset {asset_id!r} is larger than the table packing rectangle " + "after applying boundary_margin." + ) + + # The zero vector keeps the centered initial layout. Clamp it only when an + # AABB starts outside the table; this is the smallest boundary-only move. + offsets = np.clip( + np.zeros_like(base_mins), + lower_offset_bounds, + upper_offset_bounds, + ) + tolerance = 1e-9 + + for _ in range(max_rounds): + current_mins = base_mins + offsets + current_maxs = base_maxs + offsets + overlaps = _find_overlapping_2d_aabb_pairs( + current_mins=current_mins, + current_maxs=current_maxs, + aabb_clearance=aabb_clearance, + tolerance=tolerance, + ) + if not overlaps: + return { + asset_id: offsets[index].copy() + for index, asset_id in enumerate(asset_ids) + } + + # Process every pair found at the start of this round. A preceding pair + # move may already resolve a later pair, so recheck it before moving. + for _, first_index, second_index in overlaps: + current_mins = base_mins + offsets + current_maxs = base_maxs + offsets + candidates = _aabb_pair_push_candidates( + current_mins=current_mins, + current_maxs=current_maxs, + first_index=first_index, + second_index=second_index, + allowed_min=allowed_min, + allowed_max=allowed_max, + aabb_clearance=aabb_clearance, + tolerance=tolerance, + ) + if candidates is None: + continue + if not candidates: + raise RuntimeError( + "Cannot resolve overlapping 2D AABBs inside the table rectangle: " + f"{asset_ids[first_index]!r} and {asset_ids[second_index]!r}." + ) + + _, axis, first_direction, first_move, second_move = min(candidates) + offsets[first_index, axis] += first_direction * first_move + offsets[second_index, axis] -= first_direction * second_move + offsets = np.clip(offsets, lower_offset_bounds, upper_offset_bounds) + + unresolved = _find_overlapping_2d_aabb_pairs( + current_mins=base_mins + offsets, + current_maxs=base_maxs + offsets, + aabb_clearance=aabb_clearance, + tolerance=tolerance, + ) + if unresolved: + _, first_index, second_index = unresolved[0] + raise RuntimeError( + "2D AABB packing did not converge after " + f"{max_rounds} rounds; first remaining overlap is " + f"{asset_ids[first_index]!r} and {asset_ids[second_index]!r}." + ) + return {asset_id: offsets[index].copy() for index, asset_id in enumerate(asset_ids)} + + +def _convert_layout_coordinate_system( + layout_object: dict[str, object], + *, + source_to_target_matrix: np.ndarray, +) -> dict[str, object]: + """A helper to convert a layout object between coordinate systems using a 4x4 transform.""" + target_to_source_matrix = np.linalg.inv(source_to_target_matrix) + return transform_matrix_to_layout_object( + str(layout_object["id"]), + source_to_target_matrix + @ layout_object_to_transform_matrix(layout_object) + @ target_to_source_matrix, + ) + + +def export_baked_layout_object_glbs( + layout: list[dict[str, object]], + geometry_root: str | Path, + output_root: str | Path, +) -> list[Path]: + """Bake a layout into each object GLB and export them separately.""" + if not layout: + raise ValueError("Cannot export objects without layout objects.") + + resolved_geometry_root = Path(geometry_root).expanduser().resolve() + resolved_output_root = Path(output_root).expanduser().resolve() + resolved_output_root.mkdir(parents=True, exist_ok=True) + output_paths: list[Path] = [] + for layout_object in layout: + object_id = layout_object.get("id") + if not isinstance(object_id, str) or not object_id: + raise ValueError("Layout object id must be a non-empty string.") + mesh_path = resolved_geometry_root / f"{object_id}.glb" + if not mesh_path.is_file(): + raise FileNotFoundError(f"Geometry not found: {mesh_path}") + + loaded_mesh = trimesh.load(mesh_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 geometry is not a mesh: {mesh_path}") + + mesh.apply_transform(layout_object_to_transform_matrix(layout_object)) + output_path = resolved_output_root / f"{object_id}.glb" + mesh.export(output_path, file_type="glb") + if not output_path.is_file(): + raise FileNotFoundError( + f"Baked coarse object was not written: {output_path}" + ) + output_paths.append(output_path) + return output_paths + + +def export_baked_coarse_object_glbs( + coarse_layout: list[dict[str, object]], + coarse_geometry_root: str | Path, + output_root: str | Path, +) -> list[Path]: + """Bake the coarse layout into each object GLB and export them separately.""" + return export_baked_layout_object_glbs( + layout=coarse_layout, + geometry_root=coarse_geometry_root, + output_root=output_root, + ) + + +def simready_object_glb( + coarse_glb_path: str | Path, + *, + object_id: str, + rot: object, + pos: object, + scale: object, +) -> tuple[trimesh.Trimesh, dict[str, list[float]]]: + """Bake an object's coarse scale (from the coarse layout currently) + and canonicalize its AABB bottom center to the world's x-y plane (0, 0). + + 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 = _three_floats(rot, field_name="rot") + coarse_pos = np.asarray(_three_floats(pos, field_name="pos"), dtype=float) + coarse_scale = np.asarray(_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 _is_upright_container_id(object_id): + bottle_alignment_matrix = _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(object_id: str) -> bool: + """Return True if the object id contains tokens that indicate it is a bottle-like upright container.""" + # Example: soda_can_0 + # tokens: {"soda", "can", "0"} + # _UPRIGHT_CONTAINER_ID_TOKENS: {"bottle", "can", "jar"} + # So this would return True because "can" is in the set of upright container tokens. + tokens = set(re.findall(r"[a-z0-9]+", object_id.lower())) + return bool(tokens & _UPRIGHT_CONTAINER_ID_TOKENS) + + +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 = _convex_hull_volume(upper_points) + lower_volume = _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 + + +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 + + +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/scene_segmentation_utils.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_segmentation_utils.py new file mode 100644 index 000000000..f49befe8f --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_segmentation_utils.py @@ -0,0 +1,340 @@ +# ---------------------------------------------------------------------------- +# 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 + + +@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), +) -> Path: + """Replace all the other masks with gray color.""" + 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 + + +def render_numbered_mask_candidates( + *, + image_path: str | Path, + candidates: list[MaskCandidate], + output_path: str | Path, + mask_style: str = "fill", +) -> 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. + """ + 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 = ( # If weuse outline, then need to do some another processings. + 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) + 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." + ) + _draw_number_label( + draw=draw, + label=str(candidate.index), + center=((bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2), + font=font, + ) + + 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 _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.""" + outline_width = max(1, round(min(image_size) / 400)) + 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_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 + draw.rectangle( + ( + x - padding, + y - padding, + x + label_width + padding, + y + label_height + padding, + ), + fill=(220, 0, 0, 255), + outline=(255, 255, 255, 255), + width=max(1, padding // 3), + ) + draw.text((x, y), label, fill=(255, 255, 255, 255), font=font) diff --git a/embodichain/gen_sim/scene_engine/utils/__init__.py b/embodichain/gen_sim/scene_engine/utils/__init__.py new file mode 100644 index 000000000..015c41510 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/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/utils/logger.py b/embodichain/gen_sim/scene_engine/utils/logger.py new file mode 100644 index 000000000..a61d5aca0 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/utils/logger.py @@ -0,0 +1,38 @@ +# ---------------------------------------------------------------------------- +# 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 logging + +_LOGGER = logging.getLogger("embodichain.scene_engine") +if not _LOGGER.handlers: + handler = logging.StreamHandler() + handler.setFormatter( + logging.Formatter("%(asctime)s [EmbodiChain Scene Engine] %(message)s") + ) + _LOGGER.addHandler(handler) + _LOGGER.propagate = False +_LOGGER.setLevel(logging.INFO) + + +def log_stage_start(stage_name: str) -> None: + _LOGGER.info("Starting %s", stage_name) + + +def log_stage_end(stage_name: str) -> None: + _LOGGER.info("Completed %s", stage_name) From ad708e8d3e20c27e66c518d0bd833357d17c45ac Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Tue, 28 Jul 2026 17:59:59 +0800 Subject: [PATCH 02/23] Using vhacd by default in the simulation environment --- embodichain/gen_sim/scene_engine/cli/preview.py | 1 + .../scene_engine/pipeline/utils/scene_generation_utils.py | 2 ++ 2 files changed, 3 insertions(+) diff --git a/embodichain/gen_sim/scene_engine/cli/preview.py b/embodichain/gen_sim/scene_engine/cli/preview.py index e283d3831..49fa6006a 100644 --- a/embodichain/gen_sim/scene_engine/cli/preview.py +++ b/embodichain/gen_sim/scene_engine/cli/preview.py @@ -156,6 +156,7 @@ def _add_objects( 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}") 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 index ca4034863..9b3b20f21 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_generation_utils.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_generation_utils.py @@ -465,6 +465,7 @@ def gravity_settle_assets_on_table( body_scale=tuple(table_y_up_scale), body_type="static", max_convex_hull_num=max_convex_hull_num, + acd_method="vhacd", # Use vhacd by default. ) ) simulated_assets: dict[str, object] = {} @@ -481,6 +482,7 @@ def gravity_settle_assets_on_table( body_scale=tuple(asset_info["y_up_scale"]), body_type="dynamic", max_convex_hull_num=max_convex_hull_num, + acd_method="vhacd", # Use vhacd by default. ) ) From b664562a84a6c82fca465890f3a2a63a0399c9f9 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Wed, 29 Jul 2026 11:29:12 +0800 Subject: [PATCH 03/23] RAN black --- embodichain/gen_sim/scene_engine/cli/preview.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/embodichain/gen_sim/scene_engine/cli/preview.py b/embodichain/gen_sim/scene_engine/cli/preview.py index 49fa6006a..d81ff6111 100644 --- a/embodichain/gen_sim/scene_engine/cli/preview.py +++ b/embodichain/gen_sim/scene_engine/cli/preview.py @@ -156,7 +156,7 @@ def _add_objects( 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. + acd_method="vhacd", # Use vhacd by default. ) ) print(f"[{label}] {uid}: pos={init_pos} rot={init_rot} scale={body_scale}") From fe6c1af950b8df3e11f8b3ff4481a76216f1ee37 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Wed, 29 Jul 2026 14:19:36 +0800 Subject: [PATCH 04/23] style(geometry-generation): fix CI formatting --- .../gen_sim/scene_engine/clients/geometry_generation.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/embodichain/gen_sim/scene_engine/clients/geometry_generation.py b/embodichain/gen_sim/scene_engine/clients/geometry_generation.py index 1181503b5..ca4d5e241 100644 --- a/embodichain/gen_sim/scene_engine/clients/geometry_generation.py +++ b/embodichain/gen_sim/scene_engine/clients/geometry_generation.py @@ -142,7 +142,8 @@ def _request_multiple_objects( last_error: Exception | None = None for _ in range(self._max_attempts): try: - with ExitStack() as stack: # This stack manages the context of multiple open files, ensuring they are closed after the request. + # 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")) From 4c0cfc73aae10b1854fdb385eb399b9c93b0df62 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Thu, 30 Jul 2026 11:57:30 +0800 Subject: [PATCH 05/23] Updated geometry generation client --- .../clients/geometry_generation.py | 137 ++++++++++++++---- .../configs/scene_engine_config.json | 10 +- .../scene_engine/pipeline/scene_generation.py | 10 +- 3 files changed, 119 insertions(+), 38 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/clients/geometry_generation.py b/embodichain/gen_sim/scene_engine/clients/geometry_generation.py index ca4d5e241..6044d37e8 100644 --- a/embodichain/gen_sim/scene_engine/clients/geometry_generation.py +++ b/embodichain/gen_sim/scene_engine/clients/geometry_generation.py @@ -20,6 +20,7 @@ from contextlib import ExitStack import json from pathlib import Path +import time from typing import Any import requests @@ -39,14 +40,14 @@ def __init__( timeout_s: int, max_attempts: int, health_path: str, - generate_multiple_objects_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_multiple_objects_path = generate_multiple_objects_path + self._generate_objects_path = generate_objects_path self._session = session or requests.Session() @classmethod @@ -57,17 +58,24 @@ def from_config( return cls(**_load_config(config_path)) def check_health(self) -> None: - last_error: requests.RequestException | None = None + last_error: Exception | None = None for _ in range(self._max_attempts): try: response = self._session.get( self._url(self._health_path), - # timeout=self._timeout_s, 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 as exc: + except (requests.RequestException, ValueError, RuntimeError) as exc: last_error = exc assert last_error is not None @@ -79,16 +87,19 @@ def check_health(self) -> None: def close(self) -> None: self._session.close() - def generate_multiple_objects( + 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 multiple objects from: - - An input image. - - A list of object masks, each with a unique object_id and a binary mask path. + """Generate objects through the geometry server's mask-list endpoint. + + The SAM3D 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 two client + paths from drifting apart. """ # Check, validate then wrap each content of the request. @@ -112,13 +123,14 @@ def generate_multiple_objects( ) resolved_object_masks.append((object_id, resolved_mask_path)) - # Use the wrapped data structure to send the request. - response_data, response_objects = self._request_multiple_objects( + # Send one multipart image + masks request, matching test_sam3d_client.py. + 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. @@ -133,7 +145,7 @@ def generate_multiple_objects( self._download_glb(response_object["mesh"], output_path) return response_data, response_objects - def _request_multiple_objects( + def _request_objects( self, *, image_path: Path, @@ -150,12 +162,21 @@ def _request_multiple_objects( for _, mask_path in object_masks ] response = self._session.post( - self._url(self._generate_multiple_objects_path), - data={"json": "1"}, + self._url(self._generate_objects_path), files=[ - ("image", (image_path.name, image_file)), + ( + "image", + ( + image_path.name, + image_file, + _image_content_type(image_path), + ), + ), *[ - ("masks", (f"{object_id}.png", mask_file)) + ( + "masks", + (f"{object_id}.png", mask_file, "image/png"), + ) for (object_id, _), mask_file in zip( object_masks, mask_files, @@ -171,11 +192,10 @@ def _request_multiple_objects( raise RuntimeError( "Geometry Generation Server response is not valid JSON." ) from exc - response_objects = ( - _parse_multiple_objects_response( # Parse the response. - response_data, - object_ids=[object_id for object_id, _ in object_masks], - ) + 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: @@ -187,6 +207,65 @@ def _request_multiple_objects( f"{self._max_attempts} attempts." ) from last_error + def _wait_for_task_if_needed(self, response_data: object) -> dict[str, Any]: + """Poll a queued SAM3D job until it returns its final 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): @@ -221,7 +300,7 @@ def _url(self, path: str) -> str: return f"{self._base_url}/{path.lstrip('/')}" -def _parse_multiple_objects_response( +def _parse_objects_response( response_data: object, *, object_ids: list[str], @@ -305,6 +384,12 @@ def _parse_numeric_list( ) 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_config(config_path: str | Path | None) -> dict[str, Any]: resolved_config_path = Path(config_path or _DEFAULT_CONFIG_PATH).expanduser() resolved_config_path = resolved_config_path.resolve() @@ -325,7 +410,7 @@ def _load_config(config_path: str | Path | None) -> dict[str, Any]: "timeout_s", "max_attempts", "health_path", - "generate_multiple_objects_path", + "generate_objects_path", ) missing = [key for key in required_keys if key not in config] if missing: @@ -356,7 +441,7 @@ def _load_config(config_path: str | Path | None) -> dict[str, Any]: string_keys = ( "base_url", "health_path", - "generate_multiple_objects_path", + "generate_objects_path", ) for key in string_keys: if not isinstance(config[key], str) or not config[key].strip(): @@ -369,7 +454,5 @@ def _load_config(config_path: str | Path | None) -> dict[str, Any]: "timeout_s": timeout_s, "max_attempts": max_attempts, "health_path": config["health_path"].strip(), - "generate_multiple_objects_path": config[ - "generate_multiple_objects_path" - ].strip(), + "generate_objects_path": config["generate_objects_path"].strip(), } diff --git a/embodichain/gen_sim/scene_engine/configs/scene_engine_config.json b/embodichain/gen_sim/scene_engine/configs/scene_engine_config.json index a87c24b23..642901ab3 100644 --- a/embodichain/gen_sim/scene_engine/configs/scene_engine_config.json +++ b/embodichain/gen_sim/scene_engine/configs/scene_engine_config.json @@ -16,10 +16,10 @@ "segment_single_object_path": "/predict" }, "geometry_generation": { - "base_url": "", - "timeout_s": 600, - "max_attempts": 3, - "health_path": "/health", - "generate_multiple_objects_path": "/generate_multiple_objects" + "base_url": "", + "timeout_s": 600, + "max_attempts": 3, + "health_path": "/health", + "generate_objects_path": "/generate_multiple_objects" } } diff --git a/embodichain/gen_sim/scene_engine/pipeline/scene_generation.py b/embodichain/gen_sim/scene_engine/pipeline/scene_generation.py index 7400fcead..6db70010f 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/scene_generation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/scene_generation.py @@ -140,12 +140,10 @@ def _generate_coarse_results_from_masks( ) # 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_multiple_objects( - image_path=image_path, - object_masks=object_masks, - output_root=coarse_geometry_output_root, # Keep the coarse geometries - ) + 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. From a67ac7754abeb68ef6c217bd37a20a4d5f609e9f Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Thu, 30 Jul 2026 15:13:14 +0800 Subject: [PATCH 06/23] Modified the gym export to scene export --- .../gen_sim/scene_engine/cli/preview.py | 49 ++++++++++--------- .../gen_sim/scene_engine/pipeline/generate.py | 4 +- .../{gym_export.py => scene_export.py} | 48 +++++++++--------- 3 files changed, 53 insertions(+), 48 deletions(-) rename embodichain/gen_sim/scene_engine/pipeline/{gym_export.py => scene_export.py} (83%) diff --git a/embodichain/gen_sim/scene_engine/cli/preview.py b/embodichain/gen_sim/scene_engine/cli/preview.py index d81ff6111..88bc74e05 100644 --- a/embodichain/gen_sim/scene_engine/cli/preview.py +++ b/embodichain/gen_sim/scene_engine/cli/preview.py @@ -28,24 +28,29 @@ from embodichain.lab.sim.cfg import LightCfg, MeshCfg, RigidObjectCfg -def preview_gym_export( +def preview_scene_export( *, output_root: str | Path, device: str = "cpu", headless: bool = False, ) -> None: - """Load ``gym_export/gym_config.json`` and preview its table and assets.""" + """Load ``scene_export/scene_config.json`` and preview its table and assets.""" resolved_output_root = Path(output_root).expanduser().resolve() - config_path = resolved_output_root / "gym_export" / "gym_config.json" + config_path = resolved_output_root / "scene_export" / "scene_config.json" if not config_path.is_file(): - raise FileNotFoundError(f"Gym config not found: {config_path}") + raise FileNotFoundError(f"Scene config not found: {config_path}") try: - gym_config = json.loads(config_path.read_text(encoding="utf-8")) + scene_config = json.loads(config_path.read_text(encoding="utf-8")) except json.JSONDecodeError as exc: - raise ValueError(f"Gym config is not valid JSON: {config_path}") from exc - if not isinstance(gym_config, dict): - raise ValueError("Gym config must be a JSON object.") + 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( @@ -62,20 +67,20 @@ def preview_gym_export( _add_lights(sim) _add_objects( sim=sim, - entries=_config_entries(gym_config, "background"), + entries=_config_entries(scene_config, "background"), config_dir=config_path.parent, label="table", ) _add_objects( sim=sim, - entries=_config_entries(gym_config, "rigid_object"), + entries=_config_entries(scene_config, "rigid_object"), config_dir=config_path.parent, label="asset", ) if headless: sim.update(step=1) - print(f"Loaded gym export headlessly: {config_path}") + print(f"Loaded scene export headlessly: {config_path}") return print(f"Previewing: {config_path}") @@ -90,14 +95,14 @@ def preview_gym_export( def _config_entries( - gym_config: dict[str, Any], + scene_config: dict[str, Any], field_name: str, ) -> list[dict[str, Any]]: - entries = gym_config.get(field_name, []) + entries = scene_config.get(field_name, []) if not isinstance(entries, list) or not all( isinstance(entry, dict) for entry in entries ): - raise ValueError(f"Gym config field {field_name!r} must be a list of objects.") + raise ValueError(f"Scene config field {field_name!r} must be a list of objects.") return entries @@ -126,12 +131,12 @@ def _add_objects( uid = entry.get("uid") shape = entry.get("shape") if not isinstance(uid, str) or not uid: - raise ValueError(f"Gym {label} has no valid 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"Gym {label} {uid!r} has no shape.fpath.") + raise ValueError(f"Scene {label} {uid!r} has no shape.fpath.") if shape.get("shape_type") != "Mesh": raise ValueError( - f"Gym {label} {uid!r} must use shape_type='Mesh' for preview." + f"Scene {label} {uid!r} must use shape_type='Mesh' for preview." ) mesh_path = (config_dir / shape["fpath"]).resolve() @@ -164,21 +169,21 @@ def _add_objects( def _vector3(value: object, *, field_name: str) -> list[float]: if not isinstance(value, list) or len(value) != 3: - raise ValueError(f"Gym config field {field_name!r} must be a length-3 list.") + 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"Gym config field {field_name!r} must be numeric.") from exc + raise ValueError(f"Scene config field {field_name!r} must be numeric.") from exc def main() -> None: parser = argparse.ArgumentParser( - description="Preview a Scene Engine gym export in EmbodiChain simulation." + description="Preview a Scene Engine scene-only export in EmbodiChain simulation." ) parser.add_argument( "output_root", type=Path, - help="Scene Engine output root containing gym_export/.", + help="Scene Engine output root containing scene_export/.", ) parser.add_argument( "--device", @@ -191,7 +196,7 @@ def main() -> None: help="Load and validate the exported scene without opening a window.", ) args = parser.parse_args() - preview_gym_export( + preview_scene_export( output_root=args.output_root, device=args.device, headless=args.headless, diff --git a/embodichain/gen_sim/scene_engine/pipeline/generate.py b/embodichain/gen_sim/scene_engine/pipeline/generate.py index a8f6d0b70..f51981e3d 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generate.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generate.py @@ -41,7 +41,7 @@ from embodichain.gen_sim.scene_engine.pipeline.scene_generation import ( generate_scene_and_refine, ) -from embodichain.gen_sim.scene_engine.pipeline.gym_export import export_scene_to_gym +from embodichain.gen_sim.scene_engine.pipeline.scene_export import export_scene def generate_scene_from_image( @@ -107,7 +107,7 @@ def generate_scene_from_image( # 4. Scene Export log_stage_start("Scene Export") - export_scene_to_gym( + export_scene( scene=scene, output_root=resolved_output_root, table_max_convex_hull_num=16, diff --git a/embodichain/gen_sim/scene_engine/pipeline/gym_export.py b/embodichain/gen_sim/scene_engine/pipeline/scene_export.py similarity index 83% rename from embodichain/gen_sim/scene_engine/pipeline/gym_export.py rename to embodichain/gen_sim/scene_engine/pipeline/scene_export.py index 793fe0f52..7593a3c95 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/gym_export.py +++ b/embodichain/gen_sim/scene_engine/pipeline/scene_export.py @@ -55,22 +55,24 @@ ) -def export_scene_to_gym( +def export_scene( *, scene: Scene, output_root: str | Path, table_max_convex_hull_num: int = _DEFAULT_MAX_CONVEX_HULL_NUM, asset_max_convex_hull_num: int = _DEFAULT_MAX_CONVEX_HULL_NUM, ) -> Path: - """Write the Gym config and copy SimReady GLBs into ``mesh_assets``. + """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. + 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 its control setup. """ if scene.table is None: - raise ValueError("Cannot export a gym scene without a table.") + raise ValueError("Cannot export a scene without a table.") table_max_convex_hull_num = _positive_int( table_max_convex_hull_num, field_name="table_max_convex_hull_num", @@ -80,32 +82,30 @@ def export_scene_to_gym( field_name="asset_max_convex_hull_num", ) - export_root = Path(output_root).expanduser().resolve() / "gym_export" + export_root = Path(output_root).expanduser().resolve() / "scene_export" mesh_assets_root = export_root / "mesh_assets" mesh_assets_root.mkdir(parents=True, exist_ok=True) scene_objects = [scene.table, *scene.assets] object_ids = [scene_object.id for scene_object in scene_objects] if len(set(object_ids)) != len(object_ids): - raise ValueError("Gym export requires unique table and asset ids.") + raise ValueError("Scene export requires unique table and asset ids.") exported_entries = { - scene_object.id: _copy_scene_object_to_gym_assets( + scene_object.id: _copy_scene_object_to_assets( scene_object=scene_object, mesh_assets_root=mesh_assets_root, ) for scene_object in scene_objects } - gym_config = { - "id": f"Prompt2Scene-{int(time.time() * 1000)}-v0", - "max_episodes": 10, - "max_episode_steps": 300, - "env": {"events": {}, "observations": {}, "dataset": {}}, - "robot": {}, - "sensor": [], - "light": {}, + 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": [ - _gym_object_config( + _scene_object_config( scene_object=scene.table, asset_relative_path=exported_entries[scene.table.id], body_type="kinematic", @@ -114,7 +114,7 @@ def export_scene_to_gym( ) ], "rigid_object": [ - _gym_object_config( + _scene_object_config( scene_object=asset, asset_relative_path=exported_entries[asset.id], body_type="dynamic", @@ -124,15 +124,15 @@ def export_scene_to_gym( for asset in scene.assets ], } - gym_config_path = export_root / "gym_config.json" - gym_config_path.write_text( - json.dumps(gym_config, indent=2, ensure_ascii=False) + "\n", + scene_config_path = export_root / "scene_config.json" + scene_config_path.write_text( + json.dumps(scene_config, indent=2, ensure_ascii=False) + "\n", encoding="utf-8", ) - return gym_config_path + return scene_config_path -def _copy_scene_object_to_gym_assets( +def _copy_scene_object_to_assets( *, scene_object: Table | Asset, mesh_assets_root: Path, @@ -157,7 +157,7 @@ def _copy_scene_object_to_gym_assets( return destination_glb_path.relative_to(mesh_assets_root.parent).as_posix() -def _gym_object_config( +def _scene_object_config( *, scene_object: Table | Asset, asset_relative_path: str, @@ -165,7 +165,7 @@ def _gym_object_config( attrs: dict[str, float | int], max_convex_hull_num: int, ) -> dict[str, object]: - """Build one z-up gym object config from a final y-up scene object.""" + """Build one z-up scene-only object config from a final y-up scene object.""" pos_y_up = _scene_vector(scene_object, "pos") rot_y_up = _scene_vector(scene_object, "rot") scale_y_up = _scene_vector(scene_object, "scale") From c630f1676f2b8e4484050993ab33b7a2248a4b07 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Thu, 30 Jul 2026 16:22:15 +0800 Subject: [PATCH 07/23] Make the 2D AABB optimization more robust --- .../pipeline/utils/scene_generation_utils.py | 29 +++++++------------ 1 file changed, 11 insertions(+), 18 deletions(-) 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 index 9b3b20f21..f07dc9093 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_generation_utils.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_generation_utils.py @@ -1163,7 +1163,7 @@ def _optimize_assets_2d_aabbs_in_rectangle( aabb_corners_by_id: dict[str, np.ndarray], boundary_margin: float, aabb_clearance: float, - max_rounds: int = 8, + max_rounds: int = 64, ) -> dict[str, np.ndarray]: """Greedily pack 2D AABBs with minimum local squared displacement.""" @@ -1254,29 +1254,22 @@ def _optimize_assets_2d_aabbs_in_rectangle( if candidates is None: continue if not candidates: - raise RuntimeError( - "Cannot resolve overlapping 2D AABBs inside the table rectangle: " - f"{asset_ids[first_index]!r} and {asset_ids[second_index]!r}." - ) + # Both AABBs are already blocked by the table boundary on every + # separating axis. Keep the current boundary-safe layout and + # let the later gravity simulation handle this residual overlap. + return { + asset_id: offsets[index].copy() + for index, asset_id in enumerate(asset_ids) + } _, axis, first_direction, first_move, second_move = min(candidates) offsets[first_index, axis] += first_direction * first_move offsets[second_index, axis] -= first_direction * second_move offsets = np.clip(offsets, lower_offset_bounds, upper_offset_bounds) - unresolved = _find_overlapping_2d_aabb_pairs( - current_mins=base_mins + offsets, - current_maxs=base_maxs + offsets, - aabb_clearance=aabb_clearance, - tolerance=tolerance, - ) - if unresolved: - _, first_index, second_index = unresolved[0] - raise RuntimeError( - "2D AABB packing did not converge after " - f"{max_rounds} rounds; first remaining overlap is " - f"{asset_ids[first_index]!r} and {asset_ids[second_index]!r}." - ) + # The bounded greedy search may leave overlaps in densely packed scenes. + # Return its best boundary-safe result instead of aborting scene generation; + # the following gravity simulation can resolve remaining physical contacts. return {asset_id: offsets[index].copy() for index, asset_id in enumerate(asset_ids)} From 8ae888d130c9d1ca003ee659a29339d24e270464 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Thu, 30 Jul 2026 16:26:59 +0800 Subject: [PATCH 08/23] Add optional --config in cli --- .../gen_sim/scene_engine/cli/preview.py | 4 ++- embodichain/gen_sim/scene_engine/cli/start.py | 25 +++++++++++++++++-- 2 files changed, 26 insertions(+), 3 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/cli/preview.py b/embodichain/gen_sim/scene_engine/cli/preview.py index 88bc74e05..783dedfac 100644 --- a/embodichain/gen_sim/scene_engine/cli/preview.py +++ b/embodichain/gen_sim/scene_engine/cli/preview.py @@ -102,7 +102,9 @@ def _config_entries( 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.") + raise ValueError( + f"Scene config field {field_name!r} must be a list of objects." + ) return entries diff --git a/embodichain/gen_sim/scene_engine/cli/start.py b/embodichain/gen_sim/scene_engine/cli/start.py index 719f54749..cb05a66a9 100644 --- a/embodichain/gen_sim/scene_engine/cli/start.py +++ b/embodichain/gen_sim/scene_engine/cli/start.py @@ -24,7 +24,13 @@ _SUPPORTED_IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png"} -def cli_scene_engine(image: str | Path, output_root: str | Path) -> None: +def cli_scene_engine( + image: str | Path, + output_root: str | Path, + *, + config_path: str | Path | None = None, +) -> None: + """Generate one scene using an optional user-owned service configuration.""" resolved_image_path = Path(image).expanduser().resolve() if not resolved_image_path.exists(): raise FileNotFoundError(f"Image input not found: {resolved_image_path}") @@ -41,6 +47,12 @@ def cli_scene_engine(image: str | Path, output_root: str | Path) -> None: generate_scene_from_image( image_path=resolved_image_path, output_root=resolved_output_root, + # One Scene Engine config contains the LLM, segmentation, and geometry + # sections. Passing it through lets callers use their own service URLs + # instead of editing the package-installed default JSON. + llm_config_path=config_path, + image_segmentation_config_path=config_path, + geometry_generation_config_path=config_path, ) print("Successfully completed!") @@ -61,9 +73,18 @@ def main() -> None: required=True, help="Path to the output directory", ) + parser.add_argument( + "--config", + type=Path, + default=None, + help=( + "Optional Scene Engine JSON config containing the llm, " + "image_segmentation, and geometry_generation service settings." + ), + ) args = parser.parse_args() - cli_scene_engine(args.image, args.output_root) + cli_scene_engine(args.image, args.output_root, config_path=args.config) if __name__ == "__main__": From c38c0fc18cf64a73465685de6d5f351b06860f3c Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Thu, 30 Jul 2026 16:36:04 +0800 Subject: [PATCH 09/23] Register scene-engine and preview-scene in embodichain.__main__.COMMANDS --- embodichain/__main__.py | 10 ++++++++++ embodichain/gen_sim/scene_engine/cli/preview.py | 8 +++++--- embodichain/gen_sim/scene_engine/cli/start.py | 6 ++++-- 3 files changed, 19 insertions(+), 5 deletions(-) diff --git a/embodichain/__main__.py b/embodichain/__main__.py index fd4859d3c..e0f371a59 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.", + ), + 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/cli/preview.py b/embodichain/gen_sim/scene_engine/cli/preview.py index 783dedfac..2bc834244 100644 --- a/embodichain/gen_sim/scene_engine/cli/preview.py +++ b/embodichain/gen_sim/scene_engine/cli/preview.py @@ -22,6 +22,7 @@ import math from pathlib import Path import time +from collections.abc import Sequence from typing import Any from embodichain.lab.sim import SimulationManager, SimulationManagerCfg @@ -178,9 +179,10 @@ def _vector3(value: object, *, field_name: str) -> list[float]: raise ValueError(f"Scene config field {field_name!r} must be numeric.") from exc -def main() -> None: +def main(argv: Sequence[str] | None = None) -> None: parser = argparse.ArgumentParser( - description="Preview a Scene Engine scene-only export in EmbodiChain simulation." + prog="embodichain preview-scene", + description="Preview a Scene Engine scene export in EmbodiChain simulation.", ) parser.add_argument( "output_root", @@ -197,7 +199,7 @@ def main() -> None: action="store_true", help="Load and validate the exported scene without opening a window.", ) - args = parser.parse_args() + args = parser.parse_args(argv) preview_scene_export( output_root=args.output_root, device=args.device, diff --git a/embodichain/gen_sim/scene_engine/cli/start.py b/embodichain/gen_sim/scene_engine/cli/start.py index cb05a66a9..cc0e29586 100644 --- a/embodichain/gen_sim/scene_engine/cli/start.py +++ b/embodichain/gen_sim/scene_engine/cli/start.py @@ -17,6 +17,7 @@ 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 @@ -57,8 +58,9 @@ def cli_scene_engine( print("Successfully completed!") -def main() -> None: +def main(argv: Sequence[str] | None = None) -> None: parser = argparse.ArgumentParser( + prog="embodichain scene-engine", description="embodichain.gen_sim.scene_engine Scene Engine Pipeline" ) parser.add_argument( @@ -82,7 +84,7 @@ def main() -> None: "image_segmentation, and geometry_generation service settings." ), ) - args = parser.parse_args() + args = parser.parse_args(argv) cli_scene_engine(args.image, args.output_root, config_path=args.config) From b86c3b5e0f3b8d00f03bd1cbba10264999676701 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Thu, 30 Jul 2026 18:00:26 +0800 Subject: [PATCH 10/23] Fixed with the suggestion from Copilot --- .../gen_sim/scene_engine/cli/preview.py | 3 +- .../gen_sim/scene_engine/pipeline/generate.py | 41 ++++++++++--------- .../pipeline/scene_understanding.py | 4 +- .../pipeline/utils/scene_generation_utils.py | 3 +- .../utils/scene_segmentation_utils.py | 2 +- 5 files changed, 28 insertions(+), 25 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/cli/preview.py b/embodichain/gen_sim/scene_engine/cli/preview.py index 2bc834244..e4aefb5a2 100644 --- a/embodichain/gen_sim/scene_engine/cli/preview.py +++ b/embodichain/gen_sim/scene_engine/cli/preview.py @@ -92,7 +92,8 @@ def preview_scene_export( except KeyboardInterrupt: print("Stopping preview.") finally: - sim.destroy() + sim.destroy(exit_process=False) + _EmbodiSimManager.flush_cleanup_queue() def _config_entries( diff --git a/embodichain/gen_sim/scene_engine/pipeline/generate.py b/embodichain/gen_sim/scene_engine/pipeline/generate.py index f51981e3d..658b4000a 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generate.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generate.py @@ -76,15 +76,17 @@ def generate_scene_from_image( image_segmentation_client = ImageSegmentationClient.from_config( image_segmentation_config_path ) - image_segmentation_client.check_health() # Error raising will happen internally. - scene = segment_scene( - image_path=image_path, - output_root=resolved_output_root, - scene=scene, - vlm_client=vlm_client, - image_segmentation_client=image_segmentation_client, - ) - image_segmentation_client.close() # Kill the session. + try: + image_segmentation_client.check_health() # Error raising will happen internally. + scene = segment_scene( + image_path=image_path, + output_root=resolved_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. log_stage_end("Scene Segmentation") # 3. Objects + Coarse Layout Generation @@ -93,16 +95,17 @@ def generate_scene_from_image( geometry_generation_client = GeometryGenerationClient.from_config( geometry_generation_config_path ) - geometry_generation_client.check_health() - - scene = generate_scene_and_refine( - image_path=image_path, - output_root=resolved_output_root, - scene=scene, - vlm_client=vlm_client, - geometry_generation_client=geometry_generation_client, - ) - geometry_generation_client.close() # Kill the session. + 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, + vlm_client=vlm_client, + geometry_generation_client=geometry_generation_client, + ) + finally: + geometry_generation_client.close() # Kill the session to avoid resource leaks. log_stage_end("Objects + Coarse Layout Generation") # 4. Scene Export diff --git a/embodichain/gen_sim/scene_engine/pipeline/scene_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/scene_understanding.py index 2990c70e5..98244055a 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/scene_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/scene_understanding.py @@ -174,9 +174,7 @@ 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 true. For we hardcode the table id to "table". + if scene.table.id != "table": # Currently it will always return true. 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] 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 index f07dc9093..a1671bfee 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_generation_utils.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_generation_utils.py @@ -507,7 +507,8 @@ def gravity_settle_assets_on_table( z_up_to_y_up_matrix @ final_z_up_layout_matrix @ y_up_to_z_up_matrix, ) finally: - sim._deferred_destroy() + sim.destroy(exit_process=False) + _EmbodiSimManager.flush_cleanup_queue() settled_assets_layout = [ settled_layout_by_id[str(asset_layout["id"])] for asset_layout in assets_layout diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_segmentation_utils.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_segmentation_utils.py index f49befe8f..3685ec8ae 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_segmentation_utils.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_segmentation_utils.py @@ -212,7 +212,7 @@ def render_numbered_mask_candidates( "RGBA", image.size, colors[(candidate.index - 1) % len(colors)] ) transparent_layer = Image.new("RGBA", image.size, (0, 0, 0, 0)) - rendered_mask = ( # If weuse outline, then need to do some another processings. + rendered_mask = ( # If we use outline, then need to do some another processings. mask if mask_style == "fill" else _mask_outer_outline(mask, image.size) ) overlay.alpha_composite( From c99182caba91f583d77ea4a28e726d7f8c0d5cb2 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Thu, 30 Jul 2026 19:23:38 +0800 Subject: [PATCH 11/23] Added __init__.py and ran black. --- embodichain/gen_sim/scene_engine/__init__.py | 19 +++++++++++++++++++ embodichain/gen_sim/scene_engine/cli/start.py | 2 +- .../gen_sim/scene_engine/pipeline/generate.py | 8 ++++---- .../pipeline/scene_understanding.py | 4 +++- 4 files changed, 27 insertions(+), 6 deletions(-) create mode 100644 embodichain/gen_sim/scene_engine/__init__.py 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/start.py b/embodichain/gen_sim/scene_engine/cli/start.py index cc0e29586..427e1a2f8 100644 --- a/embodichain/gen_sim/scene_engine/cli/start.py +++ b/embodichain/gen_sim/scene_engine/cli/start.py @@ -61,7 +61,7 @@ def cli_scene_engine( def main(argv: Sequence[str] | None = None) -> None: parser = argparse.ArgumentParser( prog="embodichain scene-engine", - description="embodichain.gen_sim.scene_engine Scene Engine Pipeline" + description="embodichain.gen_sim.scene_engine Scene Engine Pipeline", ) parser.add_argument( "--image", diff --git a/embodichain/gen_sim/scene_engine/pipeline/generate.py b/embodichain/gen_sim/scene_engine/pipeline/generate.py index 658b4000a..5819ce68e 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generate.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generate.py @@ -77,7 +77,7 @@ def generate_scene_from_image( image_segmentation_config_path ) try: - image_segmentation_client.check_health() # Error raising will happen internally. + image_segmentation_client.check_health() # Error raising will happen internally. scene = segment_scene( image_path=image_path, output_root=resolved_output_root, @@ -86,7 +86,7 @@ def generate_scene_from_image( image_segmentation_client=image_segmentation_client, ) finally: - image_segmentation_client.close() # Kill the session to avoid resource leaks. + image_segmentation_client.close() # Kill the session to avoid resource leaks. log_stage_end("Scene Segmentation") # 3. Objects + Coarse Layout Generation @@ -96,7 +96,7 @@ def generate_scene_from_image( geometry_generation_config_path ) try: - geometry_generation_client.check_health() # Error raising will happen internally. + geometry_generation_client.check_health() # Error raising will happen internally. scene = generate_scene_and_refine( image_path=image_path, output_root=resolved_output_root, @@ -105,7 +105,7 @@ def generate_scene_from_image( geometry_generation_client=geometry_generation_client, ) finally: - geometry_generation_client.close() # Kill the session to avoid resource leaks. + geometry_generation_client.close() # Kill the session to avoid resource leaks. log_stage_end("Objects + Coarse Layout Generation") # 4. Scene Export diff --git a/embodichain/gen_sim/scene_engine/pipeline/scene_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/scene_understanding.py index 98244055a..2990c70e5 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/scene_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/scene_understanding.py @@ -174,7 +174,9 @@ 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 true. For we hardcode the table id to "table". + if ( + scene.table.id != "table" + ): # Currently it will always return true. 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] From 79442f50f2a3808e57ea5746700245fda215bfcf Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Fri, 31 Jul 2026 10:01:01 +0800 Subject: [PATCH 12/23] Fix import bug --- embodichain/gen_sim/scene_engine/cli/preview.py | 2 +- .../scene_engine/pipeline/utils/scene_generation_utils.py | 7 +++---- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/cli/preview.py b/embodichain/gen_sim/scene_engine/cli/preview.py index e4aefb5a2..51441cfbc 100644 --- a/embodichain/gen_sim/scene_engine/cli/preview.py +++ b/embodichain/gen_sim/scene_engine/cli/preview.py @@ -93,7 +93,7 @@ def preview_scene_export( print("Stopping preview.") finally: sim.destroy(exit_process=False) - _EmbodiSimManager.flush_cleanup_queue() + SimulationManager.flush_cleanup_queue() def _config_entries( 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 index a1671bfee..42237711a 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_generation_utils.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_generation_utils.py @@ -21,8 +21,7 @@ import re from typing import Sequence -from embodichain.lab.sim import SimulationManager as _EmbodiSimManager -from embodichain.lab.sim import SimulationManagerCfg +from embodichain.lab.sim import SimulationManagerCfg, SimulationManager from embodichain.lab.sim.cfg import RigidObjectCfg from embodichain.lab.sim.shapes import MeshCfg import matplotlib @@ -446,7 +445,7 @@ def gravity_settle_assets_on_table( "z_up_scale": asset_z_up_scale, } - sim = _EmbodiSimManager( + sim = SimulationManager( SimulationManagerCfg( headless=True, physics_dt=physics_dt, @@ -508,7 +507,7 @@ def gravity_settle_assets_on_table( ) finally: sim.destroy(exit_process=False) - _EmbodiSimManager.flush_cleanup_queue() + SimulationManager.flush_cleanup_queue() settled_assets_layout = [ settled_layout_by_id[str(asset_layout["id"])] for asset_layout in assets_layout From a0f6797905af705127dd240a3d4b9f847079268f Mon Sep 17 00:00:00 2001 From: PengXuanchao Date: Fri, 31 Jul 2026 10:17:18 +0800 Subject: [PATCH 13/23] Add Viser support --- .../gen_sim/scene_engine/cli/preview.py | 31 ++++++++++++++++--- 1 file changed, 27 insertions(+), 4 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/cli/preview.py b/embodichain/gen_sim/scene_engine/cli/preview.py index e4aefb5a2..3ae6330ca 100644 --- a/embodichain/gen_sim/scene_engine/cli/preview.py +++ b/embodichain/gen_sim/scene_engine/cli/preview.py @@ -27,6 +27,11 @@ 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( @@ -34,8 +39,16 @@ 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.""" + """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(): @@ -60,6 +73,9 @@ def preview_scene_export( headless=headless, physics_dt=1.0 / 100.0, sim_device=device, + visualization=( + VisualizationCfg() if visualization is None else visualization + ), ) ) try: @@ -79,14 +95,19 @@ def preview_scene_export( label="asset", ) - if headless: + 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 - print(f"Previewing: {config_path}") + 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.") - sim.open_window() while True: time.sleep(0.1) except KeyboardInterrupt: @@ -200,11 +221,13 @@ def main(argv: Sequence[str] | None = None) -> None: 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), ) From d1b363aa20ecdb843e0e6c53d31ce51f2670d2e0 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Fri, 31 Jul 2026 10:39:38 +0800 Subject: [PATCH 14/23] test(scene_engine): add unit coverage --- tests/gen_sim/scene_engine/test_cli.py | 96 +++++++++++++ tests/gen_sim/scene_engine/test_generate.py | 108 ++++++++++++++ .../scene_engine/test_geometry_generation.py | 132 ++++++++++++++++++ .../scene_engine/test_image_segmentation.py | 95 +++++++++++++ .../gen_sim/scene_engine/test_scene_export.py | 84 +++++++++++ .../test_scene_generation_utils.py | 102 ++++++++++++++ tests/test_main.py | 2 + 7 files changed, 619 insertions(+) create mode 100644 tests/gen_sim/scene_engine/test_cli.py create mode 100644 tests/gen_sim/scene_engine/test_generate.py create mode 100644 tests/gen_sim/scene_engine/test_geometry_generation.py create mode 100644 tests/gen_sim/scene_engine/test_image_segmentation.py create mode 100644 tests/gen_sim/scene_engine/test_scene_export.py create mode 100644 tests/gen_sim/scene_engine/test_scene_generation_utils.py diff --git a/tests/gen_sim/scene_engine/test_cli.py b/tests/gen_sim/scene_engine/test_cli.py new file mode 100644 index 000000000..afb620758 --- /dev/null +++ b/tests/gen_sim/scene_engine/test_cli.py @@ -0,0 +1,96 @@ +# ---------------------------------------------------------------------------- +# 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 + + +def test_cli_scene_engine_creates_output_and_forwards_config( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + image_path = tmp_path / "scene.png" + image_path.write_bytes(b"png") + config_path = tmp_path / "scene_engine_config.json" + config_path.write_text("{}", encoding="utf-8") + output_root = tmp_path / "generated" + received: dict[str, object] = {} + + def fake_generate_scene_from_image(**kwargs: object) -> None: + received.update(kwargs) + + monkeypatch.setattr( + start, "generate_scene_from_image", fake_generate_scene_from_image + ) + + start.cli_scene_engine( + image=image_path, + output_root=output_root, + config_path=config_path, + ) + + assert output_root.is_dir() + assert received["image_path"] == image_path.resolve() + assert received["output_root"] == output_root.resolve() + assert received["llm_config_path"] == config_path + assert received["image_segmentation_config_path"] == config_path + assert received["geometry_generation_config_path"] == config_path + + +def test_cli_scene_engine_rejects_non_image_input(tmp_path: Path) -> None: + text_path = tmp_path / "scene.txt" + text_path.write_text("not an image", encoding="utf-8") + + with pytest.raises(ValueError, match="extensions"): + start.cli_scene_engine(text_path, tmp_path / "output") + + +def test_main_forwards_parsed_arguments(monkeypatch: pytest.MonkeyPatch) -> None: + received: dict[str, object] = {} + + def fake_cli_scene_engine( + image: str | Path, + output_root: str | Path, + *, + config_path: str | Path | None, + ) -> None: + received["image"] = image + received["output_root"] = output_root + received["config_path"] = config_path + + monkeypatch.setattr(start, "cli_scene_engine", fake_cli_scene_engine) + + start.main( + [ + "--image", + "input.png", + "--output_root", + "output", + "--config", + "services.json", + ] + ) + + assert received == { + "image": "input.png", + "output_root": "output", + "config_path": Path("services.json"), + } diff --git a/tests/gen_sim/scene_engine/test_generate.py b/tests/gen_sim/scene_engine/test_generate.py new file mode 100644 index 000000000..85a71d642 --- /dev/null +++ b/tests/gen_sim/scene_engine/test_generate.py @@ -0,0 +1,108 @@ +# ---------------------------------------------------------------------------- +# 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.core.scene import Scene +from embodichain.gen_sim.scene_engine.pipeline import generate + + +class _Client: + def __init__(self) -> None: + self.closed = False + + def check_health(self) -> None: + return None + + def close(self) -> None: + self.closed = True + + +def test_segmentation_client_closes_when_segmentation_raises( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + segmentation_client = _Client() + + class FakeVLM: + @classmethod + def from_config(cls, _config_path: object) -> object: + return object() + + class FakeSegmentationClient: + @classmethod + def from_config(cls, _config_path: object) -> _Client: + return segmentation_client + + monkeypatch.setattr(generate, "OpenAICompatibleVLM", FakeVLM) + monkeypatch.setattr(generate, "ImageSegmentationClient", FakeSegmentationClient) + monkeypatch.setattr(generate, "understand_scene", lambda **_: Scene()) + + def fail_segment_scene(**_: object) -> Scene: + raise RuntimeError("segmentation failed") + + monkeypatch.setattr(generate, "segment_scene", fail_segment_scene) + + with pytest.raises(RuntimeError, match="segmentation failed"): + generate.generate_scene_from_image(tmp_path / "image.png", tmp_path / "output") + + assert segmentation_client.closed is True + + +def test_geometry_client_closes_when_refinement_raises( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + segmentation_client = _Client() + geometry_client = _Client() + + class FakeVLM: + @classmethod + def from_config(cls, _config_path: object) -> object: + return object() + + class FakeSegmentationClient: + @classmethod + def from_config(cls, _config_path: object) -> _Client: + return segmentation_client + + class FakeGeometryClient: + @classmethod + def from_config(cls, _config_path: object) -> _Client: + return geometry_client + + monkeypatch.setattr(generate, "OpenAICompatibleVLM", FakeVLM) + monkeypatch.setattr(generate, "ImageSegmentationClient", FakeSegmentationClient) + monkeypatch.setattr(generate, "GeometryGenerationClient", FakeGeometryClient) + monkeypatch.setattr(generate, "understand_scene", lambda **_: Scene()) + monkeypatch.setattr(generate, "segment_scene", lambda **kwargs: kwargs["scene"]) + + def fail_generate_scene_and_refine(**_: object) -> Scene: + raise RuntimeError("refinement failed") + + monkeypatch.setattr( + generate, "generate_scene_and_refine", fail_generate_scene_and_refine + ) + + with pytest.raises(RuntimeError, match="refinement failed"): + generate.generate_scene_from_image(tmp_path / "image.png", tmp_path / "output") + + assert segmentation_client.closed is True + assert geometry_client.closed is True diff --git a/tests/gen_sim/scene_engine/test_geometry_generation.py b/tests/gen_sim/scene_engine/test_geometry_generation.py new file mode 100644 index 000000000..2ff65dd3e --- /dev/null +++ b/tests/gen_sim/scene_engine/test_geometry_generation.py @@ -0,0 +1,132 @@ +# ---------------------------------------------------------------------------- +# 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.geometry_generation import ( + GeometryGenerationClient, + _parse_objects_response, +) + +_GLB_BYTES = b"glTF\x02\x00\x00\x00" + + +class _Response: + def __init__(self, *, payload: object | None = None, 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: + def __init__(self, *, payload: dict[str, Any], downloads: dict[str, bytes]) -> None: + self._payload = payload + self._downloads = downloads + self.post_file_names: list[tuple[str, str]] = [] + self.closed = False + + def post( + self, _url: str, *, files: list[tuple[str, tuple[Any, ...]]], **_: object + ) -> _Response: + self.post_file_names = [(field, str(value[0])) for field, value in files] + return _Response(payload=self._payload) + + def get(self, url: str, **_: object) -> _Response: + return _Response(content=self._downloads[url]) + + def close(self) -> None: + self.closed = True + + +def _object_response(object_id: str, mesh_path: str) -> dict[str, object]: + return { + "name": object_id, + "mesh": mesh_path, + "rotation_quaternion_wxyz": [1.0, 0.0, 0.0, 0.0], + "translation": [0.0, 0.0, 0.0], + "scale": [1.0, 1.0, 1.0], + } + + +def test_generate_objects_preserves_requested_mask_order(tmp_path: Path) -> None: + image_path = tmp_path / "image.png" + table_mask_path = tmp_path / "table.png" + cup_mask_path = tmp_path / "cup.png" + for path in (image_path, table_mask_path, cup_mask_path): + path.write_bytes(b"image") + response_payload = { + "ok": True, + "result": { + "objects": [ + _object_response("table", "/assets/table.glb"), + _object_response("cup", "/assets/cup.glb"), + ] + }, + } + session = _Session( + payload=response_payload, + downloads={ + "http://geometry.test/assets/table.glb": _GLB_BYTES, + "http://geometry.test/assets/cup.glb": _GLB_BYTES, + }, + ) + client = GeometryGenerationClient( + base_url="http://geometry.test", + timeout_s=1, + max_attempts=1, + health_path="/health", + generate_objects_path="/generate_objects", + session=session, + ) + output_root = tmp_path / "generated" / "meshes" + + _, objects = client.generate_objects( + image_path=image_path, + object_masks=[("table", table_mask_path), ("cup", cup_mask_path)], + output_root=output_root, + ) + + assert session.post_file_names == [ + ("image", "image.png"), + ("masks", "table.png"), + ("masks", "cup.png"), + ] + assert [object_data["mesh"] for object_data in objects] == [ + "/assets/table.glb", + "/assets/cup.glb", + ] + assert (output_root / "table.glb").read_bytes() == _GLB_BYTES + assert (output_root / "cup.glb").read_bytes() == _GLB_BYTES + + +def test_parse_objects_response_rejects_mismatched_object_name() -> None: + payload = { + "ok": True, + "result": {"objects": [_object_response("wrong", "/assets/wrong.glb")]}, + } + + with pytest.raises(RuntimeError, match="does not match"): + _parse_objects_response(payload, object_ids=["table"]) diff --git a/tests/gen_sim/scene_engine/test_image_segmentation.py b/tests/gen_sim/scene_engine/test_image_segmentation.py new file mode 100644 index 000000000..1b2f87d1b --- /dev/null +++ b/tests/gen_sim/scene_engine/test_image_segmentation.py @@ -0,0 +1,95 @@ +# ---------------------------------------------------------------------------- +# 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.image_segmentation import ( + ImageSegmentationClient, + _extract_rle_masks, +) + + +class _Response: + def __init__(self, payload: object) -> None: + self._payload = payload + + def raise_for_status(self) -> None: + return None + + def json(self) -> object: + return self._payload + + +class _Session: + def __init__(self, payload: object) -> None: + self._payload = payload + self.prompt: str | None = None + + def post(self, _url: str, *, data: dict[str, str], **_: object) -> _Response: + self.prompt = data["prompt"] + return _Response(self._payload) + + def close(self) -> None: + return None + + +def test_extract_rle_masks_accepts_instances_response() -> None: + mask = {"counts": [1, 2], "size": [2, 2]} + + masks = _extract_rle_masks({"result": {"instances": [{"mask_rle": mask}]}}) + + assert masks == [mask] + + +def test_segment_single_object_strips_prompt(tmp_path: Path) -> None: + image_path = tmp_path / "scene.png" + image_path.write_bytes(b"image") + mask = {"counts": [4], "size": [2, 2]} + session = _Session({"ok": True, "result": {"masks": [mask]}}) + client = ImageSegmentationClient( + base_url="http://segmentation.test", + timeout_s=1, + max_attempts=1, + health_path="/health", + segment_single_object_path="/segment", + session=session, + ) + + masks = client.segment_single_object(image_path=image_path, prompt=" table ") + + assert session.prompt == "table" + assert masks == [mask] + + +def test_segment_single_object_rejects_empty_prompt(tmp_path: Path) -> None: + image_path = tmp_path / "scene.png" + image_path.write_bytes(b"image") + client = ImageSegmentationClient( + base_url="http://segmentation.test", + timeout_s=1, + max_attempts=1, + health_path="/health", + segment_single_object_path="/segment", + session=_Session({"ok": True, "result": {"masks": []}}), + ) + + with pytest.raises(ValueError, match="prompt"): + client.segment_single_object(image_path=image_path, prompt=" ") diff --git a/tests/gen_sim/scene_engine/test_scene_export.py b/tests/gen_sim/scene_engine/test_scene_export.py new file mode 100644 index 000000000..c78bbbfbe --- /dev/null +++ b/tests/gen_sim/scene_engine/test_scene_export.py @@ -0,0 +1,84 @@ +# ---------------------------------------------------------------------------- +# 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 +from scipy.spatial.transform import Rotation + +from embodichain.gen_sim.scene_engine.core.asset import Asset +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.table import Table +from embodichain.gen_sim.scene_engine.pipeline import scene_export + + +def test_export_scene_copies_meshes_and_converts_y_up_layout(tmp_path: Path) -> None: + table_glb = tmp_path / "source_table.glb" + asset_glb = tmp_path / "source_cup.glb" + table_glb.write_bytes(b"glTFtable") + asset_glb.write_bytes(b"glTFasset") + table = Table( + id="table", + category="table", + name="table", + description="A table.", + simready_glb_path=str(table_glb), + rot=[0.0, 0.0, 0.0], + pos=[0.0, 0.0, 0.0], + scale=[1.0, 1.0, 1.0], + ) + asset = Asset( + id="cup", + category="cup", + name="cup", + description="A cup.", + simready_glb_path=str(asset_glb), + rot=[20.0, -35.0, 40.0], + pos=[1.0, 2.0, 3.0], + scale=[1.0, 2.0, 3.0], + ) + + config_path = scene_export.export_scene( + scene=Scene(table=table, assets=[asset]), + output_root=tmp_path / "output", + ) + config = json.loads(config_path.read_text(encoding="utf-8")) + exported_asset = config["rigid_object"][0] + + assert config["format"] == "embodichain.scene-export/v1" + assert "robot" not in config + assert "env" not in config + assert exported_asset["init_pos"] == [1.0, -3.0, 2.0] + assert exported_asset["body_scale"] == [1.0, 2.0, 3.0] + assert ( + config_path.parent / "mesh_assets" / "table" / "table.glb" + ).read_bytes() == b"glTFtable" + assert ( + config_path.parent / "mesh_assets" / "cup" / "cup.glb" + ).read_bytes() == b"glTFasset" + + expected_rotation = ( + scene_export._Y_UP_TO_Z_UP_ROTATION + @ Rotation.from_euler("xyz", asset.rot, degrees=True).as_matrix() + @ scene_export._Y_UP_TO_Z_UP_ROTATION.T + ) + actual_rotation = Rotation.from_euler( + "XYZ", exported_asset["init_rot"], degrees=True + ).as_matrix() + np.testing.assert_allclose(actual_rotation, expected_rotation, atol=1e-8) diff --git a/tests/gen_sim/scene_engine/test_scene_generation_utils.py b/tests/gen_sim/scene_engine/test_scene_generation_utils.py new file mode 100644 index 000000000..af768eb22 --- /dev/null +++ b/tests/gen_sim/scene_engine/test_scene_generation_utils.py @@ -0,0 +1,102 @@ +# ---------------------------------------------------------------------------- +# 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 embodichain.gen_sim.scene_engine.pipeline.utils import scene_generation_utils + + +def _aabb_corners( + minimum: tuple[float, float], maximum: tuple[float, float] +) -> np.ndarray: + return np.asarray( + [ + [minimum[0], minimum[1]], + [maximum[0], minimum[1]], + [maximum[0], maximum[1]], + [minimum[0], maximum[1]], + ], + dtype=float, + ) + + +def test_layout_transform_round_trip_preserves_pose_and_scale() -> None: + layout = { + "id": "cup", + "rot": [20.0, -35.0, 40.0], + "pos": [1.0, 2.0, 3.0], + "scale": [1.0, 2.0, 3.0], + } + + recovered = scene_generation_utils.transform_matrix_to_layout_object( + "cup", + scene_generation_utils.layout_object_to_transform_matrix(layout), + ) + + np.testing.assert_allclose(recovered["pos"], layout["pos"], atol=1e-8) + np.testing.assert_allclose(recovered["scale"], layout["scale"], atol=1e-8) + np.testing.assert_allclose( + scene_generation_utils.layout_object_to_transform_matrix(recovered), + scene_generation_utils.layout_object_to_transform_matrix(layout), + atol=1e-8, + ) + + +def test_aabb_optimizer_resolves_overlap_inside_boundary() -> None: + corners_by_id = { + "first": _aabb_corners((-0.75, -0.5), (0.25, 0.5)), + "second": _aabb_corners((-0.25, -0.5), (0.75, 0.5)), + } + + offsets = scene_generation_utils._optimize_assets_2d_aabbs_in_rectangle( + rectangle_min=np.asarray([-1.0, -1.0]), + rectangle_max=np.asarray([1.0, 1.0]), + aabb_corners_by_id=corners_by_id, + boundary_margin=0.0, + aabb_clearance=0.0, + ) + first_min, first_max = scene_generation_utils._aabb_2d_bounds_from_corners( + corners_by_id["first"] + offsets["first"], + name="first", + require_nonzero_extent=True, + ) + second_min, second_max = scene_generation_utils._aabb_2d_bounds_from_corners( + corners_by_id["second"] + offsets["second"], + name="second", + require_nonzero_extent=True, + ) + + assert first_min[0] >= -1.0 + assert first_max[0] <= 1.0 + assert second_min[0] >= -1.0 + assert second_max[0] <= 1.0 + assert first_max[0] <= second_min[0] or second_max[0] <= first_min[0] + + +def test_aabb_optimizer_rejects_asset_larger_than_boundary() -> None: + with pytest.raises(ValueError, match="larger than the table"): + scene_generation_utils._optimize_assets_2d_aabbs_in_rectangle( + rectangle_min=np.asarray([-1.0, -1.0]), + rectangle_max=np.asarray([1.0, 1.0]), + aabb_corners_by_id={ + "oversized": _aabb_corners((-2.0, -0.5), (2.0, 0.5)), + }, + boundary_margin=0.0, + aabb_clearance=0.0, + ) diff --git a/tests/test_main.py b/tests/test_main.py index 2c9fcd515..d6bb93882 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -29,8 +29,10 @@ "decompose-urdf", "preview-asset", "run-env", + "scene-engine", "simready", "train-rl", + "preview-scene", "workspace-cache", } From 23947e581e99b79a60510ae97acbdd9950869320 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Fri, 31 Jul 2026 10:58:51 +0800 Subject: [PATCH 15/23] Add test for the newly-modified scene-preview --- .../gen_sim/scene_engine/cli/preview.py | 3 +- tests/gen_sim/scene_engine/test_cli.py | 33 ++++++++++++++++++- tests/test_main.py | 13 ++++++++ 3 files changed, 47 insertions(+), 2 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/cli/preview.py b/embodichain/gen_sim/scene_engine/cli/preview.py index e0dcf884a..3d6a2cebf 100644 --- a/embodichain/gen_sim/scene_engine/cli/preview.py +++ b/embodichain/gen_sim/scene_engine/cli/preview.py @@ -207,8 +207,9 @@ def main(argv: Sequence[str] | None = None) -> None: description="Preview a Scene Engine scene export in EmbodiChain simulation.", ) parser.add_argument( - "output_root", + "--output_root", type=Path, + required=True, help="Scene Engine output root containing scene_export/.", ) parser.add_argument( diff --git a/tests/gen_sim/scene_engine/test_cli.py b/tests/gen_sim/scene_engine/test_cli.py index afb620758..26afbdd0e 100644 --- a/tests/gen_sim/scene_engine/test_cli.py +++ b/tests/gen_sim/scene_engine/test_cli.py @@ -20,7 +20,7 @@ import pytest -from embodichain.gen_sim.scene_engine.cli import start +from embodichain.gen_sim.scene_engine.cli import preview, start def test_cli_scene_engine_creates_output_and_forwards_config( @@ -94,3 +94,34 @@ def fake_cli_scene_engine( "output_root": "output", "config_path": Path("services.json"), } + + +def test_preview_main_forwards_output_root_and_viser_options( + monkeypatch: pytest.MonkeyPatch, +) -> None: + received: dict[str, object] = {} + + def fake_preview_scene_export(**kwargs: object) -> None: + received.update(kwargs) + + monkeypatch.setattr(preview, "preview_scene_export", fake_preview_scene_export) + + preview.main( + [ + "--output_root", + "output", + "--viser", + "--viser-host", + "0.0.0.0", + "--viser-port", + "9000", + ] + ) + + visualization = received["visualization"] + assert received["output_root"] == Path("output") + assert received["device"] == "cpu" + assert received["headless"] is False + assert visualization.backend == "viser" + assert visualization.viser_server.host == "0.0.0.0" + assert visualization.viser_server.port == 9000 diff --git a/tests/test_main.py b/tests/test_main.py index d6bb93882..b40094db3 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -88,6 +88,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: From b0774a40bad830cd64b04fc87d2b7d8936142eda Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Fri, 31 Jul 2026 11:34:57 +0800 Subject: [PATCH 16/23] docs(scene_engine): add usage documentation --- docs/source/features/generative_sim/index.rst | 1 + .../features/generative_sim/scene_engine.md | 128 ++++++++++++++++++ docs/source/guides/cli.md | 55 ++++++++ 3 files changed, 184 insertions(+) create mode 100644 docs/source/features/generative_sim/scene_engine.md 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..35e93b7d0 --- /dev/null +++ b/docs/source/features/generative_sim/scene_engine.md @@ -0,0 +1,128 @@ +# Scene Engine + +The Scene Engine converts one tabletop-scene image into a scene-only export. It +identifies a table and visible assets, generates their meshes, refines their +layout, settles them under gravity, and writes an EmbodiChain scene export. + +## Quick Start + +Install EmbodiChain with the generative-simulation dependencies. See +[Installation (gensim extra)](../../quick_start/install.md#optional-generative-simulation-gensim). + +Prepare a Scene Engine JSON config, then run: + +```bash +embodichain scene-engine \ + --image /path/to/scene.png \ + --output_root /path/to/scene_output \ + --config /path/to/scene_engine_config.json +``` + +Preview the result: + +```bash +embodichain preview-scene --output_root /path/to/scene_output +``` + +Use `--viser` for a browser-based preview, or `--headless` to validate the +export without opening a window: + +```bash +embodichain preview-scene \ + --output_root /path/to/scene_output \ + --viser +``` + +The equivalent module commands are: + +```bash +python -m embodichain.gen_sim.scene_engine.cli.start --help +python -m embodichain.gen_sim.scene_engine.cli.preview --help +``` + +## Requirements and Configuration + +The input must be one `.jpg`, `.jpeg`, or `.png` image with one main table and +visible, separate tabletop assets. The pipeline requires an OpenAI-compatible +VLM, an image-segmentation service, and a geometry-generation service. + +Pass their settings through `--config`. Keep credentials outside version +control. `OPENAI_API_KEY`, `OPENAI_MODEL`, `OPENAI_BASE_URL`, and +`OPENAI_MAX_ATTEMPTS` override the corresponding LLM settings. + +```json +{ + "llm": { + "openai_compatible": { + "api_key": "", + "model": "", + "base_url": "https://example.com/v1", + "default_query": {}, + "max_attempts": 3 + } + }, + "image_segmentation": { + "base_url": "http://segmentation-host:port", + "timeout_s": 120, + "max_attempts": 3, + "health_path": "/health", + "segment_single_object_path": "/segment_single_object" + }, + "geometry_generation": { + "base_url": "http://geometry-host:port", + "timeout_s": 600, + "max_attempts": 3, + "health_path": "/health", + "generate_objects_path": "/generate_objects" + } +} +``` + +The configured endpoint paths must match the deployed services. Geometry uses +one ordered `generate_objects` request for all masks; a single-object scene +uses the same request with one mask. + +## Output + +Each run refreshes the intermediate stage directories and writes the final +portable export: + +```text +/ +|-- scene_understanding/ +|-- scene_segmentation/ +|-- scene_generation/ +`-- scene_export/ + |-- scene_config.json + `-- mesh_assets/ + |-- /.glb + `-- /.glb +``` + +`scene_export/scene_config.json` has format +`"embodichain.scene-export/v1"`. It contains the table under `background` and +the settled assets under `rigid_object`; mesh paths are relative to +`scene_export/`. + +The internal scene layout is y-up. The exporter copies GLBs unchanged and +converts final positions and rotations to the simulator's z-up convention. +This is a scene-only export, not a `run-env` configuration: it does not define +a robot or task. + +## Python API + +Use `generate_scene_from_image` to run the full pipeline: + +```python +from embodichain.gen_sim.scene_engine.pipeline.generate import ( + generate_scene_from_image, +) + +scene = generate_scene_from_image( + image_path="scene.png", + output_root="scene_output", + llm_config_path="scene_engine_config.json", + image_segmentation_config_path="scene_engine_config.json", + geometry_generation_config_path="scene_engine_config.json", +) +``` diff --git a/docs/source/guides/cli.md b/docs/source/guides/cli.md index f4cfb4ce0..45c08708c 100644 --- a/docs/source/guides/cli.md +++ b/docs/source/guides/cli.md @@ -59,6 +59,61 @@ The generated output contains the canonical source mesh under ``asset_source/``, --- +## Scene Engine + +Generate a table-top scene from one image. The command requires a Scene Engine +JSON config for the VLM, image-segmentation, and geometry-generation services. + +```bash +embodichain scene-engine \ + --image /path/to/scene.png \ + --output_root /path/to/scene_output \ + --config /path/to/scene_engine_config.json +``` + +The generated scene-only export is written to +``/scene_export/scene_config.json``. It is intended for +``preview-scene`` and downstream scene consumers; it is not a complete +``run-env`` configuration because it does not choose or configure a robot. + +Preview the gravity-settled table and assets: + +```bash +embodichain preview-scene --output_root /path/to/scene_output +``` + +Use Viser for a browser-based preview: + +```bash +embodichain preview-scene \ + --output_root /path/to/scene_output \ + --viser +``` + +### Arguments + +``scene-engine``: + +| Argument | Default | Description | +|---|---|---| +| ``--image`` | *(required)* | Input ``.jpg``, ``.jpeg``, or ``.png`` scene image | +| ``--output_root`` | *(required)* | Directory that receives intermediate artifacts and ``scene_export/`` | +| ``--config`` | packaged config | Scene Engine JSON config containing ``llm``, ``image_segmentation``, and ``geometry_generation`` settings | + +``preview-scene``: + +| Argument | Default | Description | +|---|---|---| +| ``--output_root`` | *(required)* | Scene Engine output root containing ``scene_export/`` | +| ``--device`` | ``cpu`` | Simulation device, such as ``cpu`` or ``cuda`` | +| ``--headless`` | ``False`` | Load and validate the export without a native window | +| ``--viser`` | ``False`` | Publish the scene through Viser instead of a native window | + +For configuration, output layout, remote Viser access, and Python API usage, +see [Scene Engine](../features/generative_sim/scene_engine.md). + +--- + ## Preview Asset Preview a USD or mesh asset in the simulation without writing code. From ede8172c8f5c8e1f2c24957395303d9ac0cc51ab Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Fri, 31 Jul 2026 12:45:08 +0800 Subject: [PATCH 17/23] Change logger --- .../gen_sim/scene_engine/pipeline/generate.py | 18 ++++----- .../gen_sim/scene_engine/utils/__init__.py | 19 ---------- .../gen_sim/scene_engine/utils/logger.py | 38 ------------------- 3 files changed, 9 insertions(+), 66 deletions(-) delete mode 100644 embodichain/gen_sim/scene_engine/utils/__init__.py delete mode 100644 embodichain/gen_sim/scene_engine/utils/logger.py diff --git a/embodichain/gen_sim/scene_engine/pipeline/generate.py b/embodichain/gen_sim/scene_engine/pipeline/generate.py index 5819ce68e..92313f864 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generate.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generate.py @@ -36,7 +36,7 @@ from embodichain.gen_sim.scene_engine.pipeline.scene_segmentation import ( segment_scene, ) -from embodichain.gen_sim.scene_engine.utils.logger import log_stage_end, log_stage_start +from embodichain.utils.logger import log_info from embodichain.gen_sim.scene_engine.pipeline.scene_generation import ( generate_scene_and_refine, @@ -61,17 +61,17 @@ def generate_scene_from_image( scene = Scene() # 1. Scene Understanding - log_stage_start("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_stage_end("Scene Understanding") + log_info("Completed Scene Understanding") # 2. Scene Segmentation - log_stage_start("Scene Segmentation") + log_info("Starting Scene Segmentation") # Load the config and fail if the Image Segmentation Server is unavailable. image_segmentation_client = ImageSegmentationClient.from_config( image_segmentation_config_path @@ -87,10 +87,10 @@ def generate_scene_from_image( ) finally: image_segmentation_client.close() # Kill the session to avoid resource leaks. - log_stage_end("Scene Segmentation") + log_info("Completed Scene Segmentation") # 3. Objects + Coarse Layout Generation - log_stage_start("Objects + Coarse Layout Generation") + log_info("Starting Objects + Coarse Layout Generation") # Load the config and fail if the Geometry Generation Server is unavailable. geometry_generation_client = GeometryGenerationClient.from_config( geometry_generation_config_path @@ -106,16 +106,16 @@ def generate_scene_from_image( ) finally: geometry_generation_client.close() # Kill the session to avoid resource leaks. - log_stage_end("Objects + Coarse Layout Generation") + log_info("Completed Objects + Coarse Layout Generation") # 4. Scene Export - log_stage_start("Scene Export") + log_info("Starting Scene Export") export_scene( scene=scene, output_root=resolved_output_root, table_max_convex_hull_num=16, asset_max_convex_hull_num=16, ) - log_stage_end("Scene Export") + log_info("Completed Scene Export") return scene diff --git a/embodichain/gen_sim/scene_engine/utils/__init__.py b/embodichain/gen_sim/scene_engine/utils/__init__.py deleted file mode 100644 index 015c41510..000000000 --- a/embodichain/gen_sim/scene_engine/utils/__init__.py +++ /dev/null @@ -1,19 +0,0 @@ -# ---------------------------------------------------------------------------- -# 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/utils/logger.py b/embodichain/gen_sim/scene_engine/utils/logger.py deleted file mode 100644 index a61d5aca0..000000000 --- a/embodichain/gen_sim/scene_engine/utils/logger.py +++ /dev/null @@ -1,38 +0,0 @@ -# ---------------------------------------------------------------------------- -# 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 logging - -_LOGGER = logging.getLogger("embodichain.scene_engine") -if not _LOGGER.handlers: - handler = logging.StreamHandler() - handler.setFormatter( - logging.Formatter("%(asctime)s [EmbodiChain Scene Engine] %(message)s") - ) - _LOGGER.addHandler(handler) - _LOGGER.propagate = False -_LOGGER.setLevel(logging.INFO) - - -def log_stage_start(stage_name: str) -> None: - _LOGGER.info("Starting %s", stage_name) - - -def log_stage_end(stage_name: str) -> None: - _LOGGER.info("Completed %s", stage_name) From 73e285e1e48bb6a80f8e8a36c36ea12b86225fc1 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Fri, 31 Jul 2026 14:58:22 +0800 Subject: [PATCH 18/23] Correct the table id verification comment, make it more clear --- .../gen_sim/scene_engine/pipeline/scene_understanding.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/embodichain/gen_sim/scene_engine/pipeline/scene_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/scene_understanding.py index 2990c70e5..6e31770c3 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/scene_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/scene_understanding.py @@ -176,7 +176,7 @@ def validate_scene_understanding(scene: Scene) -> None: raise ValueError("Scene understanding must identify a table.") if ( scene.table.id != "table" - ): # Currently it will always return true. For we hardcode the table id to "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] From 244ab1bf951e92c6d017025e6967a8984fa04b91 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Fri, 31 Jul 2026 16:51:20 +0800 Subject: [PATCH 19/23] Delete a bad comment line in scene_segmentation_utils.py --- .../scene_engine/pipeline/utils/scene_segmentation_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_segmentation_utils.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_segmentation_utils.py index 3685ec8ae..7c88d62a5 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_segmentation_utils.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_segmentation_utils.py @@ -212,7 +212,7 @@ def render_numbered_mask_candidates( "RGBA", image.size, colors[(candidate.index - 1) % len(colors)] ) transparent_layer = Image.new("RGBA", image.size, (0, 0, 0, 0)) - rendered_mask = ( # If we use outline, then need to do some another processings. + rendered_mask = ( mask if mask_style == "fill" else _mask_outer_outline(mask, image.size) ) overlay.alpha_composite( From 930af9bddfa34f6a0ff9be2cf4bb5aa19a4eed57 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Fri, 31 Jul 2026 16:55:04 +0800 Subject: [PATCH 20/23] Align the doc with the scene_engine_config --- docs/source/features/generative_sim/scene_engine.md | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/docs/source/features/generative_sim/scene_engine.md b/docs/source/features/generative_sim/scene_engine.md index 35e93b7d0..81260abd2 100644 --- a/docs/source/features/generative_sim/scene_engine.md +++ b/docs/source/features/generative_sim/scene_engine.md @@ -66,21 +66,22 @@ control. `OPENAI_API_KEY`, `OPENAI_MODEL`, `OPENAI_BASE_URL`, and "timeout_s": 120, "max_attempts": 3, "health_path": "/health", - "segment_single_object_path": "/segment_single_object" + "segment_single_object_path": "/predict" }, "geometry_generation": { "base_url": "http://geometry-host:port", "timeout_s": 600, "max_attempts": 3, "health_path": "/health", - "generate_objects_path": "/generate_objects" + "generate_objects_path": "/generate_multiple_objects" } } ``` -The configured endpoint paths must match the deployed services. Geometry uses -one ordered `generate_objects` request for all masks; a single-object scene -uses the same request with one mask. +The endpoint paths above match the packaged template, but remain +service-specific placeholders: change them when the deployed services expose +different routes. Geometry uses one ordered multi-object request for all masks; +a single-object scene uses the same request with one mask. ## Output From 05f8a1a38103c2e26ece71869e355fc6155fc8cf Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Fri, 31 Jul 2026 18:33:38 +0800 Subject: [PATCH 21/23] Align the scene_engine config setup with the simready_pipeline --- .../features/generative_sim/scene_engine.md | 33 ++++- docs/source/guides/cli.md | 7 +- embodichain/gen_sim/scene_engine/cli/start.py | 8 +- .../clients/geometry_generation.py | 18 +++ .../clients/image_segmentation.py | 18 +++ setup.py | 3 + tests/gen_sim/scene_engine/test_cli.py | 22 +++ tests/gen_sim/scene_engine/test_config.py | 126 ++++++++++++++++++ .../scene_engine/test_geometry_generation.py | 16 +++ .../scene_engine/test_image_segmentation.py | 16 +++ 10 files changed, 257 insertions(+), 10 deletions(-) create mode 100644 tests/gen_sim/scene_engine/test_config.py diff --git a/docs/source/features/generative_sim/scene_engine.md b/docs/source/features/generative_sim/scene_engine.md index 81260abd2..80b4d6764 100644 --- a/docs/source/features/generative_sim/scene_engine.md +++ b/docs/source/features/generative_sim/scene_engine.md @@ -46,9 +46,36 @@ The input must be one `.jpg`, `.jpeg`, or `.png` image with one main table and visible, separate tabletop assets. The pipeline requires an OpenAI-compatible VLM, an image-segmentation service, and a geometry-generation service. -Pass their settings through `--config`. Keep credentials outside version -control. `OPENAI_API_KEY`, `OPENAI_MODEL`, `OPENAI_BASE_URL`, and -`OPENAI_MAX_ATTEMPTS` override the corresponding LLM settings. +Without `--config`, Scene Engine reads the official template at +`embodichain/gen_sim/scene_engine/configs/scene_engine_config.json`. The +checked-in template intentionally has empty service URLs and credentials. +Provide a complete user-owned JSON file with `--config`, or provide the +settings through environment variables. `--config` is an optional complete +JSON override; do not add credentials to the checked-in template. + +Keep credentials outside version control. `OPENAI_API_KEY`, `OPENAI_MODEL`, +`OPENAI_BASE_URL`, and `OPENAI_MAX_ATTEMPTS` override the corresponding LLM +settings. For example: + +```bash +export OPENAI_API_KEY="" +export OPENAI_MODEL="" +export OPENAI_BASE_URL="https://example.com/v1" +export OPENAI_MAX_ATTEMPTS="3" + +export SCENE_ENGINE_IMAGE_SEGMENTATION_BASE_URL="http://segmentation-host:port" +export SCENE_ENGINE_IMAGE_SEGMENTATION_PATH="/predict" +export SCENE_ENGINE_GEOMETRY_GENERATION_BASE_URL="http://geometry-host:port" +export SCENE_ENGINE_GEOMETRY_GENERATION_PATH="/generate_multiple_objects" +``` + +`SCENE_ENGINE_IMAGE_SEGMENTATION_TIMEOUT_S`, +`SCENE_ENGINE_IMAGE_SEGMENTATION_MAX_ATTEMPTS`, +`SCENE_ENGINE_IMAGE_SEGMENTATION_HEALTH_PATH`, +`SCENE_ENGINE_GEOMETRY_GENERATION_TIMEOUT_S`, +`SCENE_ENGINE_GEOMETRY_GENERATION_MAX_ATTEMPTS`, and +`SCENE_ENGINE_GEOMETRY_GENERATION_HEALTH_PATH` override the remaining service +fields when needed. ```json { diff --git a/docs/source/guides/cli.md b/docs/source/guides/cli.md index 45c08708c..b4b5786c9 100644 --- a/docs/source/guides/cli.md +++ b/docs/source/guides/cli.md @@ -61,8 +61,9 @@ The generated output contains the canonical source mesh under ``asset_source/``, ## Scene Engine -Generate a table-top scene from one image. The command requires a Scene Engine -JSON config for the VLM, image-segmentation, and geometry-generation services. +Generate a table-top scene from one image. Configure the VLM, +image-segmentation, and geometry-generation services with either a Scene Engine +JSON config or the documented environment variables. ```bash embodichain scene-engine \ @@ -98,7 +99,7 @@ embodichain preview-scene \ |---|---|---| | ``--image`` | *(required)* | Input ``.jpg``, ``.jpeg``, or ``.png`` scene image | | ``--output_root`` | *(required)* | Directory that receives intermediate artifacts and ``scene_export/`` | -| ``--config`` | packaged config | Scene Engine JSON config containing ``llm``, ``image_segmentation``, and ``geometry_generation`` settings | +| ``--config`` | packaged template | Optional complete Scene Engine JSON override. Without it, supply the documented service environment variables; the packaged JSON is only a template. | ``preview-scene``: diff --git a/embodichain/gen_sim/scene_engine/cli/start.py b/embodichain/gen_sim/scene_engine/cli/start.py index 427e1a2f8..4052da257 100644 --- a/embodichain/gen_sim/scene_engine/cli/start.py +++ b/embodichain/gen_sim/scene_engine/cli/start.py @@ -49,8 +49,8 @@ def cli_scene_engine( image_path=resolved_image_path, output_root=resolved_output_root, # One Scene Engine config contains the LLM, segmentation, and geometry - # sections. Passing it through lets callers use their own service URLs - # instead of editing the package-installed default JSON. + # sections. When omitted, every client reads the package template and + # applies its documented environment-variable overrides. llm_config_path=config_path, image_segmentation_config_path=config_path, geometry_generation_config_path=config_path, @@ -80,8 +80,8 @@ def main(argv: Sequence[str] | None = None) -> None: type=Path, default=None, help=( - "Optional Scene Engine JSON config containing the llm, " - "image_segmentation, and geometry_generation service settings." + "Optional Scene Engine JSON override. Without it, clients read the " + "packaged template and apply service environment-variable overrides." ), ) args = parser.parse_args(argv) diff --git a/embodichain/gen_sim/scene_engine/clients/geometry_generation.py b/embodichain/gen_sim/scene_engine/clients/geometry_generation.py index 6044d37e8..f1a0f4a08 100644 --- a/embodichain/gen_sim/scene_engine/clients/geometry_generation.py +++ b/embodichain/gen_sim/scene_engine/clients/geometry_generation.py @@ -19,6 +19,7 @@ from contextlib import ExitStack import json +import os from pathlib import Path import time from typing import Any @@ -28,6 +29,13 @@ _DEFAULT_CONFIG_PATH = ( Path(__file__).resolve().parents[1] / "configs" / "scene_engine_config.json" ) +_ENVIRONMENT_OVERRIDES = { + "base_url": "SCENE_ENGINE_GEOMETRY_GENERATION_BASE_URL", + "timeout_s": "SCENE_ENGINE_GEOMETRY_GENERATION_TIMEOUT_S", + "max_attempts": "SCENE_ENGINE_GEOMETRY_GENERATION_MAX_ATTEMPTS", + "health_path": "SCENE_ENGINE_GEOMETRY_GENERATION_HEALTH_PATH", + "generate_objects_path": "SCENE_ENGINE_GEOMETRY_GENERATION_PATH", +} class GeometryGenerationClient: @@ -404,6 +412,8 @@ def _load_config(config_path: str | Path | None) -> dict[str, Any]: config = config_data.get("geometry_generation") if not isinstance(config, dict): raise ValueError("Config key geometry_generation must be an object.") + config = dict(config) + _apply_environment_overrides(config) required_keys = ( "base_url", @@ -456,3 +466,11 @@ def _load_config(config_path: str | Path | None) -> dict[str, Any]: "health_path": config["health_path"].strip(), "generate_objects_path": config["generate_objects_path"].strip(), } + + +def _apply_environment_overrides(config: dict[str, Any]) -> None: + """Apply optional deployment-specific service settings from the environment.""" + for config_key, environment_name in _ENVIRONMENT_OVERRIDES.items(): + value = os.getenv(environment_name) + if value is not None: + config[config_key] = value diff --git a/embodichain/gen_sim/scene_engine/clients/image_segmentation.py b/embodichain/gen_sim/scene_engine/clients/image_segmentation.py index 083adca7d..d3ded7218 100644 --- a/embodichain/gen_sim/scene_engine/clients/image_segmentation.py +++ b/embodichain/gen_sim/scene_engine/clients/image_segmentation.py @@ -18,6 +18,7 @@ from __future__ import annotations import json +import os from pathlib import Path from typing import Any @@ -26,6 +27,13 @@ _DEFAULT_CONFIG_PATH = ( Path(__file__).resolve().parents[1] / "configs" / "scene_engine_config.json" ) +_ENVIRONMENT_OVERRIDES = { + "base_url": "SCENE_ENGINE_IMAGE_SEGMENTATION_BASE_URL", + "timeout_s": "SCENE_ENGINE_IMAGE_SEGMENTATION_TIMEOUT_S", + "max_attempts": "SCENE_ENGINE_IMAGE_SEGMENTATION_MAX_ATTEMPTS", + "health_path": "SCENE_ENGINE_IMAGE_SEGMENTATION_HEALTH_PATH", + "segment_single_object_path": "SCENE_ENGINE_IMAGE_SEGMENTATION_PATH", +} class ImageSegmentationClient: @@ -150,6 +158,8 @@ def _load_config(config_path: str | Path | None) -> dict[str, Any]: config = config_data.get("image_segmentation") if not isinstance(config, dict): raise ValueError("Config key image_segmentation must be an object.") + config = dict(config) + _apply_environment_overrides(config) required_keys = ( "base_url", @@ -200,6 +210,14 @@ def _load_config(config_path: str | Path | None) -> dict[str, Any]: } +def _apply_environment_overrides(config: dict[str, Any]) -> None: + """Apply optional deployment-specific service settings from the environment.""" + for config_key, environment_name in _ENVIRONMENT_OVERRIDES.items(): + value = os.getenv(environment_name) + if value is not None: + config[config_key] = value + + 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") diff --git a/setup.py b/setup.py index d3bf8fb98..17f1f055c 100644 --- a/setup.py +++ b/setup.py @@ -120,6 +120,9 @@ def main(): author="EmbodiChain Developers", description="An end-to-end, GPU-accelerated, and modular platform for building generalized Embodied Intelligence.", packages=find_packages(exclude=["docs"]), + package_data={ + "embodichain.gen_sim.scene_engine.configs": ["*.json"], + }, data_files=data_files, cmdclass=cmdclass, include_package_data=True, diff --git a/tests/gen_sim/scene_engine/test_cli.py b/tests/gen_sim/scene_engine/test_cli.py index 26afbdd0e..35446f951 100644 --- a/tests/gen_sim/scene_engine/test_cli.py +++ b/tests/gen_sim/scene_engine/test_cli.py @@ -63,6 +63,28 @@ def test_cli_scene_engine_rejects_non_image_input(tmp_path: Path) -> None: start.cli_scene_engine(text_path, tmp_path / "output") +def test_cli_scene_engine_uses_package_template_without_config( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + image_path = tmp_path / "scene.png" + image_path.write_bytes(b"png") + received: dict[str, object] = {} + + def fake_generate_scene_from_image(**kwargs: object) -> None: + received.update(kwargs) + + monkeypatch.setattr( + start, "generate_scene_from_image", fake_generate_scene_from_image + ) + + start.cli_scene_engine(image_path, tmp_path / "output") + + assert received["llm_config_path"] is None + assert received["image_segmentation_config_path"] is None + assert received["geometry_generation_config_path"] is None + + def test_main_forwards_parsed_arguments(monkeypatch: pytest.MonkeyPatch) -> None: received: dict[str, object] = {} diff --git a/tests/gen_sim/scene_engine/test_config.py b/tests/gen_sim/scene_engine/test_config.py new file mode 100644 index 000000000..cef3704de --- /dev/null +++ b/tests/gen_sim/scene_engine/test_config.py @@ -0,0 +1,126 @@ +# ---------------------------------------------------------------------------- +# 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 +from typing import Any + +import pytest + +from embodichain.gen_sim.scene_engine.clients.geometry_generation import ( + GeometryGenerationClient, +) +from embodichain.gen_sim.scene_engine.clients.image_segmentation import ( + ImageSegmentationClient, +) +from embodichain.gen_sim.scene_engine.llms.load_config import load_llm_config + +REPO_ROOT = Path(__file__).resolve().parents[3] +CONFIG_PATH = ( + REPO_ROOT + / "embodichain" + / "gen_sim" + / "scene_engine" + / "configs" + / "scene_engine_config.json" +) + + +@pytest.fixture(scope="module") +def scene_engine_config() -> dict[str, Any]: + with CONFIG_PATH.open("r", encoding="utf-8") as file: + return json.load(file) + + +def test_scene_engine_config_declares_all_service_sections( + scene_engine_config: dict[str, Any], +) -> None: + assert set(scene_engine_config) == { + "llm", + "image_segmentation", + "geometry_generation", + } + assert "openai_compatible" in scene_engine_config["llm"] + + +@pytest.mark.parametrize( + ("section_name", "path_key"), + [ + ("image_segmentation", "segment_single_object_path"), + ("geometry_generation", "generate_objects_path"), + ], +) +def test_service_template_has_valid_non_secret_defaults( + scene_engine_config: dict[str, Any], + section_name: str, + path_key: str, +) -> None: + service_config = scene_engine_config[section_name] + + assert isinstance(service_config["base_url"], str) + assert isinstance(service_config["timeout_s"], int) + assert service_config["timeout_s"] > 0 + assert isinstance(service_config["max_attempts"], int) + assert service_config["max_attempts"] > 0 + assert service_config["health_path"].startswith("/") + assert service_config[path_key].startswith("/") + + +def test_llm_environment_overrides_package_template( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("OPENAI_API_KEY", "test-api-key") + monkeypatch.setenv("OPENAI_MODEL", "test-vision-model") + monkeypatch.setenv("OPENAI_BASE_URL", "http://llm.test/v1") + monkeypatch.setenv("OPENAI_MAX_ATTEMPTS", "5") + + config = load_llm_config() + + assert config.api_key == "test-api-key" + assert config.model == "test-vision-model" + assert config.base_url == "http://llm.test/v1" + assert config.max_attempts == 5 + + +def test_package_template_reports_missing_service_configuration( + monkeypatch: pytest.MonkeyPatch, +) -> None: + for environment_name in ( + "OPENAI_API_KEY", + "OPENAI_MODEL", + "OPENAI_BASE_URL", + "OPENAI_MAX_ATTEMPTS", + "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_PATH", + "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_PATH", + ): + monkeypatch.delenv(environment_name, raising=False) + + with pytest.raises(ValueError, match="Missing required LLM config keys"): + load_llm_config() + with pytest.raises(ValueError, match="base_url must be a non-empty string"): + ImageSegmentationClient.from_config() + with pytest.raises(ValueError, match="base_url must be a non-empty string"): + GeometryGenerationClient.from_config() diff --git a/tests/gen_sim/scene_engine/test_geometry_generation.py b/tests/gen_sim/scene_engine/test_geometry_generation.py index 2ff65dd3e..9abe2fc7b 100644 --- a/tests/gen_sim/scene_engine/test_geometry_generation.py +++ b/tests/gen_sim/scene_engine/test_geometry_generation.py @@ -130,3 +130,19 @@ def test_parse_objects_response_rejects_mismatched_object_name() -> None: with pytest.raises(RuntimeError, match="does not match"): _parse_objects_response(payload, object_ids=["table"]) + + +def test_geometry_generation_environment_overrides_package_template( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv( + "SCENE_ENGINE_GEOMETRY_GENERATION_BASE_URL", + "http://geometry.test", + ) + monkeypatch.setenv("SCENE_ENGINE_GEOMETRY_GENERATION_PATH", "/generate") + + client = GeometryGenerationClient.from_config() + + assert client._base_url == "http://geometry.test" + assert client._generate_objects_path == "/generate" + client.close() diff --git a/tests/gen_sim/scene_engine/test_image_segmentation.py b/tests/gen_sim/scene_engine/test_image_segmentation.py index 1b2f87d1b..b4438f1c0 100644 --- a/tests/gen_sim/scene_engine/test_image_segmentation.py +++ b/tests/gen_sim/scene_engine/test_image_segmentation.py @@ -93,3 +93,19 @@ def test_segment_single_object_rejects_empty_prompt(tmp_path: Path) -> None: with pytest.raises(ValueError, match="prompt"): client.segment_single_object(image_path=image_path, prompt=" ") + + +def test_image_segmentation_environment_overrides_package_template( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv( + "SCENE_ENGINE_IMAGE_SEGMENTATION_BASE_URL", + "http://segmentation.test", + ) + monkeypatch.setenv("SCENE_ENGINE_IMAGE_SEGMENTATION_PATH", "/segment") + + client = ImageSegmentationClient.from_config() + + assert client._base_url == "http://segmentation.test" + assert client._segment_single_object_path == "/segment" + client.close() From d0d39cac095ab0d1aef33dd10e15f80545574d09 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Fri, 31 Jul 2026 18:46:21 +0800 Subject: [PATCH 22/23] fix(scene_engine): validate geometry output object IDs --- .../clients/geometry_generation.py | 12 ++++++- .../scene_engine/test_geometry_generation.py | 35 +++++++++++++++++++ 2 files changed, 46 insertions(+), 1 deletion(-) diff --git a/embodichain/gen_sim/scene_engine/clients/geometry_generation.py b/embodichain/gen_sim/scene_engine/clients/geometry_generation.py index f1a0f4a08..d50a3a267 100644 --- a/embodichain/gen_sim/scene_engine/clients/geometry_generation.py +++ b/embodichain/gen_sim/scene_engine/clients/geometry_generation.py @@ -149,7 +149,17 @@ def generate_objects( resolved_object_masks, response_objects, ): - output_path = resolved_output_root / f"{object_id}.glb" + 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 diff --git a/tests/gen_sim/scene_engine/test_geometry_generation.py b/tests/gen_sim/scene_engine/test_geometry_generation.py index 9abe2fc7b..2b48eded4 100644 --- a/tests/gen_sim/scene_engine/test_geometry_generation.py +++ b/tests/gen_sim/scene_engine/test_geometry_generation.py @@ -132,6 +132,41 @@ def test_parse_objects_response_rejects_mismatched_object_name() -> None: _parse_objects_response(payload, object_ids=["table"]) +@pytest.mark.parametrize("object_id", ["../outside", "nested/object", r"nested\object"]) +def test_generate_objects_rejects_unsafe_output_object_id( + tmp_path: Path, + object_id: str, +) -> None: + image_path = tmp_path / "image.png" + mask_path = tmp_path / "mask.png" + image_path.write_bytes(b"image") + mask_path.write_bytes(b"mask") + session = _Session( + payload={ + "ok": True, + "result": {"objects": [_object_response(object_id, "/assets/object.glb")]}, + }, + downloads={}, + ) + client = GeometryGenerationClient( + base_url="http://geometry.test", + timeout_s=1, + max_attempts=1, + health_path="/health", + generate_objects_path="/generate_objects", + session=session, + ) + + with pytest.raises(ValueError, match="not safe for a filename"): + client.generate_objects( + image_path=image_path, + object_masks=[(object_id, mask_path)], + output_root=tmp_path / "generated", + ) + + assert not (tmp_path / "outside.glb").exists() + + def test_geometry_generation_environment_overrides_package_template( monkeypatch: pytest.MonkeyPatch, ) -> None: From 114734d90a69e778969816de748e1dbd3d1529f2 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Fri, 31 Jul 2026 18:51:53 +0800 Subject: [PATCH 23/23] fix(scene_engine): restrict preview mesh paths --- .../gen_sim/scene_engine/cli/preview.py | 13 ++- tests/gen_sim/scene_engine/test_preview.py | 82 +++++++++++++++++++ 2 files changed, 94 insertions(+), 1 deletion(-) create mode 100644 tests/gen_sim/scene_engine/test_preview.py diff --git a/embodichain/gen_sim/scene_engine/cli/preview.py b/embodichain/gen_sim/scene_engine/cli/preview.py index 3d6a2cebf..f2d19c263 100644 --- a/embodichain/gen_sim/scene_engine/cli/preview.py +++ b/embodichain/gen_sim/scene_engine/cli/preview.py @@ -152,6 +152,7 @@ def _add_objects( 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") @@ -164,7 +165,17 @@ def _add_objects( f"Scene {label} {uid!r} must use shape_type='Mesh' for preview." ) - mesh_path = (config_dir / shape["fpath"]).resolve() + 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") diff --git a/tests/gen_sim/scene_engine/test_preview.py b/tests/gen_sim/scene_engine/test_preview.py new file mode 100644 index 000000000..2501c4c23 --- /dev/null +++ b/tests/gen_sim/scene_engine/test_preview.py @@ -0,0 +1,82 @@ +# ---------------------------------------------------------------------------- +# 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 preview + + +class _PreviewSim: + def __init__(self) -> None: + self.rigid_objects: list[object] = [] + + def add_rigid_object(self, cfg: object) -> None: + self.rigid_objects.append(cfg) + + +def test_preview_add_objects_accepts_mesh_inside_scene_export(tmp_path: Path) -> None: + config_dir = tmp_path / "scene_export" + mesh_path = config_dir / "mesh_assets" / "table" / "table.glb" + mesh_path.parent.mkdir(parents=True) + mesh_path.write_bytes(b"glTF") + sim = _PreviewSim() + + preview._add_objects( + sim=sim, + entries=[ + { + "uid": "table", + "shape": { + "shape_type": "Mesh", + "fpath": "mesh_assets/table/table.glb", + }, + "init_pos": [0.0, 0.0, 0.0], + "init_rot": [0.0, 0.0, 0.0], + } + ], + config_dir=config_dir, + label="table", + ) + + assert len(sim.rigid_objects) == 1 + + +@pytest.mark.parametrize("fpath", ["../outside.glb", "/tmp/outside.glb"]) +def test_preview_add_objects_rejects_mesh_path_outside_scene_export( + tmp_path: Path, + fpath: str, +) -> None: + config_dir = tmp_path / "scene_export" + config_dir.mkdir() + + with pytest.raises(ValueError, match="must (be a relative path|stay within)"): + preview._add_objects( + sim=_PreviewSim(), + entries=[ + { + "uid": "table", + "shape": {"shape_type": "Mesh", "fpath": fpath}, + "init_pos": [0.0, 0.0, 0.0], + "init_rot": [0.0, 0.0, 0.0], + } + ], + config_dir=config_dir, + label="table", + )