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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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", + ) From a6a3219cb6165e9138b552b1ea2de69e366ee886 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Mon, 3 Aug 2026 19:05:35 +0800 Subject: [PATCH 24/41] Changed support region detection algo + Changed the 2d-aabbs-optimization algo + Reformatted the code --- .../scene_engine/pipeline/scene_generation.py | 170 +++- .../utils/assets_group_layout_optimizer.py | 530 +++++++++++ .../utils/assets_group_support_clamp.py | 642 +++++++++++++ .../pipeline/utils/scene_generation_utils.py | 840 ------------------ .../pipeline/utils/table_support_surface.py | 497 +++++++++++ 5 files changed, 1800 insertions(+), 879 deletions(-) create mode 100644 embodichain/gen_sim/scene_engine/pipeline/utils/assets_group_layout_optimizer.py create mode 100644 embodichain/gen_sim/scene_engine/pipeline/utils/assets_group_support_clamp.py create mode 100644 embodichain/gen_sim/scene_engine/pipeline/utils/table_support_surface.py diff --git a/embodichain/gen_sim/scene_engine/pipeline/scene_generation.py b/embodichain/gen_sim/scene_engine/pipeline/scene_generation.py index 6db70010f..048d75cb6 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/scene_generation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/scene_generation.py @@ -22,6 +22,7 @@ import shutil import numpy as np +import trimesh from embodichain.gen_sim.scene_engine.clients.geometry_generation import ( GeometryGenerationClient, @@ -32,19 +33,26 @@ from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( OpenAICompatibleVLM, ) +from embodichain.gen_sim.scene_engine.pipeline.utils.assets_group_support_clamp import ( + AssetsGroupSupportClamp, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.assets_group_layout_optimizer import ( + AssetsSupportLayoutOptimizer, +) 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, + load_glb_mesh, quaternion_wxyz_to_euler_xyz_degrees, simready_object_glb, transform_matrix_to_layout_object, ) +from embodichain.gen_sim.scene_engine.pipeline.utils.table_support_surface import ( + TableSupportSurfaceDetector, +) +from embodichain.utils.logger import log_info _SUPPORTED_IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png"} @@ -397,57 +405,78 @@ def _layout_refinement( # 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. + if not refined_assets_layout: + log_info("Scene has no movable assets; skipping support-region clamping.") + return refined_table_layout, [] + + # 4. Detect the actual upward support triangles instead of projecting the + # entire table mesh to one convex hull. The result retains concavities + # (for example, an L-shaped tabletop) and is the only boundary used for + # placement below. ( - table_support_surface_2d_z_up_world_boundary, + table_world_mesh_z_up, assets_aabb_2d_z_up_world_corners_by_id, - table_mesh_2d_z_up_world_projection, - ) = heuristic_table_support_surface( + ) = _measure_table_and_assets_in_z_up_world( 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. + assets_layout=refined_assets_layout, 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. + support_detector = TableSupportSurfaceDetector( + table_world_mesh=table_world_mesh_z_up, 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_support_region = support_detector.detect() + support_detector.save_support_surface_debug_images() + + # 5. Keep the complete clutter rigid in the table plane. A successful + # result applies one shared z-up XY delta to every AABB, so it preserves + # all existing asset-to-asset relations. It is *not* an asset packing + # pass: pre-existing overlap is deliberately left to a later optimizer. + group_clamp = AssetsGroupSupportClamp( + support_region=table_support_region.support_polygon, + assets_aabb_2d_z_up_world_corners_by_id=( + assets_aabb_2d_z_up_world_corners_by_id ), - 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, + assets_layout=refined_assets_layout, debug_output_root=debug_output_root, + ) + refined_assets_layout = group_clamp.clamp() + group_clamp.save_group_clamp_debug_images() + + # The clamp returns y-up layouts; measure their resulting z-up AABBs again + # so the following independent optimizer consumes the same world-frame + # geometry as every other stage. + _, clamped_assets_aabb_2d_z_up_world_corners_by_id = ( + _measure_table_and_assets_in_z_up_world( + table_layout=refined_table_layout, + assets_layout=refined_assets_layout, + geometry_root=simready_geometry_output_root, + ) + ) + + # 6. Restore the previous pairwise AABB separation stage, but constrain + # every candidate with the actual support polygon rather than the legacy + # largest internal rectangle. Assets may now move independently only as + # much as needed to remove overlap; every resulting AABB remains on the + # L-shaped, circular, or otherwise non-convex support region. + overlap_optimizer = AssetsSupportLayoutOptimizer( + support_region=table_support_region.support_polygon, + assets_aabb_2d_z_up_world_corners_by_id=( + clamped_assets_aabb_2d_z_up_world_corners_by_id + ), assets_layout=refined_assets_layout, + debug_output_root=debug_output_root, ) + # Render this stage separately from the rigid group clamp. The latter + # intentionally preserves pre-existing overlaps, while this figure shows + # whether independent AABB separation actually resolved them. + refined_assets_layout = overlap_optimizer.optimize() + overlap_optimizer.save_overlap_optimization_debug_images() # 7. Gravity simulation, to let all the assets to be stable and placed well on the table's support surface. # Notice that: we do not consider the assets like a bottle, which should be standing on the table but laid down @@ -461,6 +490,69 @@ def _layout_refinement( return refined_table_layout, refined_assets_layout +def _measure_table_and_assets_in_z_up_world( + *, + table_layout: dict[str, object], + assets_layout: list[dict[str, object]], + geometry_root: str | Path, +) -> tuple[trimesh.Trimesh, dict[str, np.ndarray]]: + """Measure a table mesh and asset AABBs in one shared z-up world frame. + + Scene layouts and SimReady GLBs are y-up. The support detector and the + group clamp both operate in z-up world XY, so this conversion is performed + once here and the exact same measured AABBs are passed to the clamp. + """ + table_id = table_layout.get("id") + if not isinstance(table_id, str) or not table_id: + raise ValueError("Table layout must contain a non-empty string id.") + + y_up_to_z_up_matrix = np.eye(4) + y_up_to_z_up_matrix[:3, :3] = np.array( + [ + [1.0, 0.0, 0.0], + [0.0, 0.0, -1.0], + [0.0, 1.0, 0.0], + ] + ) + z_up_to_y_up_matrix = np.linalg.inv(y_up_to_z_up_matrix) + resolved_geometry_root = Path(geometry_root).expanduser().resolve() + + def _mesh_in_z_up_world(layout_object: dict[str, object]) -> trimesh.Trimesh: + object_id = layout_object.get("id") + if not isinstance(object_id, str) or not object_id: + raise ValueError("Layout object must contain a non-empty string id.") + mesh = load_glb_mesh(resolved_geometry_root / f"{object_id}.glb") + z_up_layout = transform_matrix_to_layout_object( + object_id, + y_up_to_z_up_matrix + @ layout_object_to_transform_matrix(layout_object) + @ z_up_to_y_up_matrix, + ) + mesh.apply_transform(y_up_to_z_up_matrix) + mesh.apply_transform(layout_object_to_transform_matrix(z_up_layout)) + return mesh + + table_world_mesh_z_up = _mesh_in_z_up_world(table_layout) + asset_aabbs_by_id: dict[str, np.ndarray] = {} + for asset_layout in assets_layout: + asset_id = asset_layout.get("id") + if not isinstance(asset_id, str) or not asset_id: + raise ValueError("Each asset layout must contain a non-empty string id.") + if asset_id in asset_aabbs_by_id: + raise ValueError(f"Asset layouts contain duplicate id {asset_id!r}.") + asset_bounds_xy = _mesh_in_z_up_world(asset_layout).bounds[:, :2] + asset_aabbs_by_id[asset_id] = np.array( + [ + [asset_bounds_xy[0, 0], asset_bounds_xy[0, 1]], + [asset_bounds_xy[1, 0], asset_bounds_xy[0, 1]], + [asset_bounds_xy[1, 0], asset_bounds_xy[1, 1]], + [asset_bounds_xy[0, 0], asset_bounds_xy[1, 1]], + ], + dtype=float, + ) + return table_world_mesh_z_up, asset_aabbs_by_id + + def _simready_assets( *, scene: Scene, diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/assets_group_layout_optimizer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/assets_group_layout_optimizer.py new file mode 100644 index 000000000..d4021fede --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/assets_group_layout_optimizer.py @@ -0,0 +1,530 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") + +import matplotlib.pyplot as plt +import numpy as np +from shapely import affinity +from shapely.geometry import MultiPolygon, Polygon + +from embodichain.gen_sim.scene_engine.pipeline.utils.assets_group_support_clamp import ( + AssetsGroupSupportClamp, + SupportGeometry, +) +from embodichain.utils.logger import log_info, log_warning + + +@dataclass(frozen=True) +class AssetsSupportLayoutOptimizerConfig: + """Controls for support-constrained pairwise AABB separation.""" + + margin_m: float = 0.0 # Required clearance between each AABB and the boundary. + aabb_clearance_m: float = 1e-6 # Required clearance between AABB pairs. + max_rounds: int = 64 # Maximum greedy pair-separation passes. + split_samples: int = 9 # Candidate splits between the two overlapping AABBs. + + +class AssetsSupportLayoutOptimizer: + """Greedily separate AABBs while retaining arbitrary support containment. + + This reuses the previous packing algorithm's pairwise strategy: detect an + overlap, try the two separating directions on both XY axes, and choose the + lowest-displacement valid push. Unlike the old path, every candidate is + validated against the actual Polygon/MultiPolygon support region instead + of a largest internal rectangle. + """ + + def __init__( + self, + *, + support_region: SupportGeometry, + assets_aabb_2d_z_up_world_corners_by_id: dict[str, np.ndarray], + assets_layout: list[dict[str, object]], + debug_output_root: str | Path | None = None, + config: AssetsSupportLayoutOptimizerConfig | None = None, + ) -> None: + self.support_region = support_region + self.assets_aabb_2d_z_up_world_corners_by_id = ( + assets_aabb_2d_z_up_world_corners_by_id + ) + self.assets_layout = assets_layout + self.debug_output_root = ( + Path(debug_output_root).expanduser().resolve() + if debug_output_root is not None + else None + ) + self.refined_assets_layout: list[dict[str, object]] | None = None + self.config = ( + config if config is not None else AssetsSupportLayoutOptimizerConfig() + ) + # Check config. + if self.config.margin_m < 0.0: + raise ValueError("margin_m must be non-negative.") + if self.config.aabb_clearance_m < 0.0: + raise ValueError("aabb_clearance_m must be non-negative.") + if self.config.max_rounds <= 0 or self.config.split_samples < 2: + raise ValueError( + "max_rounds must be positive and split_samples at least two." + ) + + def optimize(self) -> list[dict[str, object]]: + """Resolve pairwise AABB overlap and return updated y-up layouts.""" + self.refined_assets_layout = None + # Check inputs just like the previous AssetsGroupSupportClamp step would have done. + aabbs_by_id = AssetsGroupSupportClamp._validate_aabbs( + self.assets_aabb_2d_z_up_world_corners_by_id + ) + raw_support = AssetsGroupSupportClamp._coerce_support_geometry( + self.support_region + ) + if raw_support is None: + log_warning("AABB overlap optimization failed: invalid support geometry.") + raise ValueError("Support region is invalid.") + safe_support = ( + raw_support + if self.config.margin_m == 0.0 + else AssetsGroupSupportClamp._polygonal_geometry( + raw_support.buffer(-self.config.margin_m) + ) + ) + if safe_support is None or safe_support.is_empty: + log_warning( + "AABB overlap optimization failed: support region is empty after " + f"applying a {self.config.margin_m:.4f} m boundary margin." + ) + raise ValueError("Support region is empty after applying layout margin.") + + asset_ids = sorted(aabbs_by_id) + base_aabbs = np.stack([aabbs_by_id[asset_id] for asset_id in asset_ids]) + offsets = np.zeros((len(asset_ids), 2), dtype=float) + if not self._all_contained(safe_support, base_aabbs, offsets): + log_warning( + "AABB overlap optimization requires all input AABBs to be inside " + "the support region." + ) + raise ValueError( + "Overlap optimization requires AABBs already inside support; " + "run AssetsGroupSupportClamp first." + ) + initial_overlaps = self._overlaps(base_aabbs, offsets) + log_info( + "Support-constrained AABB overlap optimization started: " + f"assets={len(asset_ids)}, initial_overlaps={len(initial_overlaps)}, " + f"boundary_margin={self.config.margin_m:.4f} m, " + f"aabb_clearance={self.config.aabb_clearance_m:.4f} m, " + f"max_rounds={self.config.max_rounds}." + ) + if not initial_overlaps: # Return directly if there are no overlaps to resolve. + log_info("AABB overlap optimization succeeded without movement.") + self.refined_assets_layout = self._apply_offsets_to_y_up_layouts( + asset_ids=asset_ids, + offsets=offsets, + ) + return self.refined_assets_layout + + for round_index in range(self.config.max_rounds): + # Check whether any overlaps remain. + overlaps = self._overlaps(base_aabbs, offsets) + if not overlaps: + log_info( + "AABB overlap optimization succeeded after " + f"{round_index} rounds." + ) + self.refined_assets_layout = self._apply_offsets_to_y_up_layouts( + asset_ids=asset_ids, + offsets=offsets, + ) + return self.refined_assets_layout + log_info( + "AABB overlap optimization round " + f"{round_index + 1}/{self.config.max_rounds}: " + f"remaining_overlaps={len(overlaps)}." + ) + moved = False + # Choose a pair once. + for _, first_index, second_index in overlaps: + # If the overlaps is handeled by a previous pair, skip it. + if not self._overlaps(base_aabbs, offsets, (first_index, second_index)): + continue + candidates = self._separation_candidates( + base_aabbs=base_aabbs, + offsets=offsets, + first_index=first_index, + second_index=second_index, + safe_support=safe_support, + ) + if candidates: + # A local separation must not blindly create a new + # collision with a third asset. Prefer candidates with + # no new collisions; when every placement causes one, + # retain the least-colliding state before considering + # displacement from the generated layout. + offsets = min( + candidates, + key=lambda candidate: self._candidate_score( + base_aabbs=base_aabbs, + current_offsets=offsets, + candidate_offsets=candidate, + ), + ) # Choose the best candidate with score. + moved = True + else: + log_warning( + "No support-valid axis-aligned separation candidate for " + f"overlapping AABBs {asset_ids[first_index]!r} and " + f"{asset_ids[second_index]!r}." + ) + if not moved: + break + + unresolved_pairs = [ + f"{asset_ids[first_index]}/{asset_ids[second_index]}" + for _, first_index, second_index in self._overlaps(base_aabbs, offsets) + ] + log_warning( + "Unable to resolve all asset AABB overlaps inside the detected support " + "region; unresolved pairs=" + f"{unresolved_pairs}." + ) + raise ValueError( + "Asset AABB overlap cannot be resolved while keeping all assets " + "inside the detected table support region." + ) + + def _apply_offsets_to_y_up_layouts( + self, *, asset_ids: list[str], offsets: np.ndarray + ) -> list[dict[str, object]]: + """Write independent z-up XY offsets back to the stored y-up layouts.""" + received_ids = {str(layout.get("id")) for layout in self.assets_layout} + expected_ids = set(asset_ids) + if received_ids != expected_ids: + raise ValueError( + "Asset layouts and optimized AABBs must have identical ids." + ) + updated_layouts: list[dict[str, object]] = [] + offsets_by_id = { + asset_id: offsets[index] for index, asset_id in enumerate(asset_ids) + } + for layout in self.assets_layout: + asset_id = str(layout["id"]) + position = layout.get("pos") + if not isinstance(position, list) or len(position) != 3: + raise ValueError( + "Each asset layout must contain a three-value pos list." + ) + dx, dy = offsets_by_id[asset_id] + updated_layout = dict(layout) + updated_position = [float(value) for value in position] + updated_position[0] += float(dx) + updated_position[2] -= float(dy) + updated_layout["pos"] = updated_position + updated_layouts.append(updated_layout) + return updated_layouts + + def save_overlap_optimization_debug_images(self) -> bool: + """Optionally save diagnostics for the most recent optimization.""" + if self.refined_assets_layout is None: + self.optimize() + assert self.refined_assets_layout is not None + if self.debug_output_root is None: + raise ValueError( + "A debug_output_root is required when saving overlap-optimization " + "debug images." + ) + + initial_aabbs_by_id = AssetsGroupSupportClamp._validate_aabbs( + self.assets_aabb_2d_z_up_world_corners_by_id + ) + raw_support = AssetsGroupSupportClamp._coerce_support_geometry( + self.support_region + ) + if raw_support is None: + raise ValueError("Cannot render overlap optimization for invalid support.") + original_positions_by_id = { + str(layout["id"]): layout["pos"] for layout in self.assets_layout + } + refined_positions_by_id = { + str(layout["id"]): layout["pos"] for layout in self.refined_assets_layout + } + if set(original_positions_by_id) != set(initial_aabbs_by_id) or set( + refined_positions_by_id + ) != set(initial_aabbs_by_id): + raise ValueError( + "Asset layouts and optimized AABBs must have identical ids." + ) + + translated_aabbs_by_id: dict[str, np.ndarray] = {} + moved = False + for asset_id, corners in initial_aabbs_by_id.items(): + original_position = original_positions_by_id[asset_id] + refined_position = refined_positions_by_id[asset_id] + if ( + not isinstance(original_position, list) + or not isinstance(refined_position, list) + or len(original_position) != 3 + or len(refined_position) != 3 + ): + raise ValueError( + "Each asset layout must contain a three-value pos list." + ) + delta_xy = np.array( + [ + float(refined_position[0]) - float(original_position[0]), + float(original_position[2]) - float(refined_position[2]), + ] + ) + translated_aabbs_by_id[asset_id] = corners + delta_xy + moved = moved or not np.allclose(delta_xy, 0.0) + + path = self.debug_output_root / "assets_aabb_overlap_optimization_2d.png" + path.parent.mkdir(parents=True, exist_ok=True) + figure, axes = plt.subplots( + 1, 2, figsize=(14, 7), dpi=160, constrained_layout=True + ) + self._draw_overlap_state( + axes[0], + raw_support, + initial_aabbs_by_id, + "Before AABB overlap optimization", + ) + self._draw_overlap_state( + axes[1], + raw_support, + translated_aabbs_by_id, + ( + "After AABB overlap optimization" + if moved + else "After AABB overlap optimization (already non-overlapping)" + ), + ) + figure.savefig(path, bbox_inches="tight") + plt.close(figure) + return True + + def _separation_candidates( + self, + *, + base_aabbs: np.ndarray, + offsets: np.ndarray, + first_index: int, + second_index: int, + safe_support: Polygon | MultiPolygon, + ) -> list[np.ndarray]: + # Get current aabbs. + current_aabbs = base_aabbs + offsets[:, None, :] + minimums, maximums = current_aabbs.min(axis=1), current_aabbs.max(axis=1) + candidates: list[np.ndarray] = [] + for axis in (0, 1): + directions_and_distances = ( + ( + -1.0, + maximums[first_index, axis] + + self.config.aabb_clearance_m + - minimums[second_index, axis], + ), + ( + 1.0, + maximums[second_index, axis] + + self.config.aabb_clearance_m + - minimums[first_index, axis], + ), + ) + for first_direction, required_distance in directions_and_distances: + if required_distance <= 0.0: + continue + for fraction in np.linspace(0.0, 1.0, self.config.split_samples): + candidate = offsets.copy() + first_move = required_distance * float(fraction) + candidate[first_index, axis] += first_direction * first_move + candidate[second_index, axis] -= first_direction * ( + required_distance - first_move + ) + if not self._overlaps( + base_aabbs, candidate, (first_index, second_index) + ) and self._all_contained(safe_support, base_aabbs, candidate): + candidates.append(candidate) + return candidates + + def _overlaps( + self, + base_aabbs: np.ndarray, + offsets: np.ndarray, + only_pair: tuple[int, int] | None = None, + ) -> list[tuple[float, int, int]]: + return sorted( + [ + (min(overlap_x, overlap_y), first_index, second_index) + for overlap_x, overlap_y, first_index, second_index in self._overlap_details( + base_aabbs, offsets, only_pair + ) + ], + reverse=True, + ) + + def _overlap_details( + self, + base_aabbs: np.ndarray, + offsets: np.ndarray, + only_pair: tuple[int, int] | None = None, + ) -> list[tuple[float, float, int, int]]: + """Return positive XY penetration extents, including requested clearance.""" + current_aabbs = base_aabbs + offsets[:, None, :] + minimums, maximums = current_aabbs.min(axis=1), current_aabbs.max(axis=1) + pairs = ( + [only_pair] + if only_pair is not None + else [ + (first_index, second_index) + for first_index in range(len(current_aabbs)) + for second_index in range(first_index + 1, len(current_aabbs)) + ] + ) + overlaps: list[tuple[float, float, int, int]] = [] + for first_index, second_index in pairs: + overlap_x = ( + min(maximums[first_index, 0], maximums[second_index, 0]) + - max(minimums[first_index, 0], minimums[second_index, 0]) + + self.config.aabb_clearance_m + ) + overlap_y = ( + min(maximums[first_index, 1], maximums[second_index, 1]) + - max(minimums[first_index, 1], minimums[second_index, 1]) + + self.config.aabb_clearance_m + ) + if overlap_x > 1e-9 and overlap_y > 1e-9: + overlaps.append((overlap_x, overlap_y, first_index, second_index)) + return overlaps + + def _draw_overlap_state( + self, + axis: plt.Axes, + support: Polygon | MultiPolygon, + aabbs_by_id: dict[str, np.ndarray], + title: str, + ) -> None: + """Draw a state and highlight every AABB pair that still overlaps.""" + validated_aabbs = AssetsGroupSupportClamp._validate_aabbs(aabbs_by_id) + asset_ids = sorted(validated_aabbs) + aabbs = np.stack([validated_aabbs[asset_id] for asset_id in asset_ids]) + zero_offsets = np.zeros((len(asset_ids), 2), dtype=float) + overlaps = self._overlaps(aabbs, zero_offsets) + overlapping_indices = { + index + for _, first_index, second_index in overlaps + for index in (first_index, second_index) + } + + AssetsGroupSupportClamp._draw_support(axis, support, title, "darkorange") + for index, asset_id in enumerate(asset_ids): + corners = validated_aabbs[asset_id] + polygon = AssetsGroupSupportClamp._aabb_polygon(corners) + boundary = np.asarray(polygon.exterior.coords) + is_overlapping = index in overlapping_indices + color = "firebrick" if is_overlapping else "seagreen" + axis.fill( + boundary[:, 0], + boundary[:, 1], + facecolor=color, + edgecolor=color, + linewidth=2.0 if is_overlapping else 1.0, + alpha=0.32, + ) + axis.text( + *corners.mean(axis=0), + asset_id, + ha="center", + va="center", + fontsize=8, + bbox={"facecolor": "white", "alpha": 0.75, "edgecolor": "none"}, + ) + + for _, first_index, second_index in overlaps: + first_center = aabbs[first_index].mean(axis=0) + second_center = aabbs[second_index].mean(axis=0) + axis.plot( + [first_center[0], second_center[0]], + [first_center[1], second_center[1]], + color="firebrick", + linestyle="--", + linewidth=1.3, + ) + axis.set_title(f"{title}\nremaining AABB overlaps: {len(overlaps)}") + + def _candidate_score( + self, + *, + base_aabbs: np.ndarray, + current_offsets: np.ndarray, + candidate_offsets: np.ndarray, + ) -> tuple[int, int, float, float]: + """Rank a valid pair-separation candidate by global collision impact. + + The first term is deliberately based on *new* overlap pairs: this + keeps a pairwise correction from simply transferring its collision to + a nearby third asset. If every candidate causes a new collision, the + remaining terms prefer fewer total overlaps, less total penetration, + then less layout displacement. + """ + current_pairs = { + (first_index, second_index) + for _, _, first_index, second_index in self._overlap_details( + base_aabbs, current_offsets + ) + } + candidate_details = self._overlap_details(base_aabbs, candidate_offsets) + candidate_pairs = { + (first_index, second_index) + for _, _, first_index, second_index in candidate_details + } + new_overlap_count = len(candidate_pairs - current_pairs) + total_penetration_area = sum( + overlap_x * overlap_y for overlap_x, overlap_y, _, _ in candidate_details + ) + total_squared_displacement = float( + np.einsum("ij,ij->", candidate_offsets, candidate_offsets) + ) + return ( + new_overlap_count, + len(candidate_pairs), + total_penetration_area, + total_squared_displacement, + ) + + @staticmethod + def _all_contained( + support: Polygon | MultiPolygon, + base_aabbs: np.ndarray, + offsets: np.ndarray, + ) -> bool: + return all( + support.covers( + affinity.translate( + AssetsGroupSupportClamp._aabb_polygon(corners), + xoff=float(offset[0]), + yoff=float(offset[1]), + ) + ) + for corners, offset in zip(base_aabbs, offsets) + ) diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/assets_group_support_clamp.py b/embodichain/gen_sim/scene_engine/pipeline/utils/assets_group_support_clamp.py new file mode 100644 index 000000000..927134902 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/assets_group_support_clamp.py @@ -0,0 +1,642 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import TypeAlias + +import matplotlib + +matplotlib.use("Agg") + +import matplotlib.pyplot as plt +import numpy as np +from scipy.ndimage import binary_erosion +from shapely import affinity +from shapely.geometry import GeometryCollection, MultiPolygon, Polygon + +from embodichain.utils.logger import log_info, log_warning + +SupportGeometry: TypeAlias = Polygon | MultiPolygon + + +@dataclass(frozen=True) +class AssetsGroupSupportClampConfig: + """Numerical controls for rigid group placement on a support region.""" + + margin_m: float = 0.0 # Required clearance between each AABB and the boundary. + grid_resolution_m: float = 0.005 # Raster cell size for the coarse search. + + +@dataclass(frozen=True) +class _GridTransform: + """World/grid conversion for a centre-sampled regular XY raster.""" + + x_coordinates: np.ndarray # X coordinate of each grid-column centre. + y_coordinates: np.ndarray # Y coordinate of each grid-row centre. + resolution_m: float # Uniform spacing between neighbouring cell centres. + + @property + def shape(self) -> tuple[int, int]: + return len(self.y_coordinates), len(self.x_coordinates) + + def world_to_nearest_pixel(self, point_xy: np.ndarray) -> tuple[int, int]: + column = int(np.rint((point_xy[0] - self.x_coordinates[0]) / self.resolution_m)) + row = int(np.rint((point_xy[1] - self.y_coordinates[0]) / self.resolution_m)) + return row, column + + def pixel_to_world(self, row: int, column: int) -> np.ndarray: + return np.array([self.x_coordinates[column], self.y_coordinates[row]]) + + +class AssetsGroupSupportClamp: + """Find a small shared XY shift that places all AABBs on a support region. + + Each AABB gets a feasible-centre map obtained by binary erosion of the safe + support mask. A candidate translation is valid only when it is feasible + for every asset, then it must pass exact Shapely containment. This supports + concave polygons, holes, and disconnected ``MultiPolygon`` regions. + """ + + def __init__( + self, + *, + support_region: SupportGeometry, + assets_aabb_2d_z_up_world_corners_by_id: dict[str, np.ndarray], + assets_layout: list[dict[str, object]], + debug_output_root: str | Path | None = None, + config: AssetsGroupSupportClampConfig | None = None, + ) -> None: + self.support_region = support_region + self.assets_aabb_2d_z_up_world_corners_by_id = ( + assets_aabb_2d_z_up_world_corners_by_id + ) + self.assets_layout = assets_layout + self.debug_output_root = ( + Path(debug_output_root).expanduser().resolve() + if debug_output_root is not None + else None + ) + self.refined_assets_layout: list[dict[str, object]] | None = None + self.config = config if config is not None else AssetsGroupSupportClampConfig() + # Check. + if self.config.margin_m < 0.0: + raise ValueError("margin_m must be non-negative.") + if self.config.grid_resolution_m <= 0.0: + raise ValueError("grid_resolution_m must be positive.") + + def clamp(self) -> list[dict[str, object]]: + """Return y-up layouts after one rigid, support-valid XY translation.""" + self.refined_assets_layout = None + # Validate input aabbs. + aabbs_by_id = self._validate_aabbs(self.assets_aabb_2d_z_up_world_corners_by_id) + + # Validate and coerce the support region into a usable polygonal geometry. + raw_support = self._coerce_support_geometry(self.support_region) + if raw_support is None: + log_warning("Asset-group support clamp failed: invalid support geometry.") + raise ValueError( + "Asset-group support clamp requires a valid support region." + ) + # Re coerce the support region with a margin to get the safe support region. + safe_support = ( + raw_support + if self.config.margin_m == 0.0 + else self._polygonal_geometry(raw_support.buffer(-self.config.margin_m)) + ) + if safe_support is None or safe_support.is_empty: + log_warning( + "Asset-group support clamp failed: support region is empty after " + f"applying a {self.config.margin_m:.4f} m boundary margin." + ) + raise ValueError("Asset-group support clamp has no usable support area.") + + # Get all the translated layouts, and store them for later debug rendering. + delta_xy = self._find_clamp_delta( + safe_support=safe_support, + aabbs_by_id=aabbs_by_id, + ) + if delta_xy is None: + log_warning( + "Asset-group support clamp failed: no shared translation can place " + f"all {len(aabbs_by_id)} AABBs inside the support region." + ) + raise ValueError( + "Asset clutter cannot be placed completely on the detected table " + "support region." + ) + # Translate the layouts and store them for later debug rendering. + refined_assets_layout = self._apply_delta_to_y_up_layouts( + delta_xy=delta_xy, + expected_ids=set(aabbs_by_id), + ) + self.refined_assets_layout = refined_assets_layout + if np.allclose( + delta_xy, 0.0 + ): # Judge whether the translation is zero, if so, no need to apply optimization. + log_info( + "All asset AABBs are already fully inside the detected table " + "support region; no planar group optimization was applied." + ) + else: + log_info( + "Applied rigid asset-group support optimization with " + f"delta_xy={delta_xy.tolist()} m; all AABBs passed " + "exact support containment." + ) + return refined_assets_layout + + def _find_clamp_delta( + self, + *, + safe_support: Polygon | MultiPolygon, + aabbs_by_id: dict[str, np.ndarray], + ) -> np.ndarray | None: + """Find one common z-up XY translation, or return ``None``. + + The initial position is always checked exactly first. Otherwise, grid + candidates are considered in increasing translation distance and the + first vector-valid placement is returned. + """ + + # Get all aabbs's clutter center as anchor, to keep the internal + # relative layouts unchanged. + anchor_xy = self._group_anchor(aabbs_by_id) + log_info( + "Asset-group support clamp started: " + f"assets={len(aabbs_by_id)}, support_area={safe_support.area:.4f} m^2, " + f"boundary_margin={self.config.margin_m:.4f} m, " + f"grid_resolution={self.config.grid_resolution_m:.4f} m." + ) + + # Check whether the initial position is already valid, which is a common case. + zero_translation = np.zeros(2, dtype=float) + if self._is_exactly_contained(safe_support, aabbs_by_id, zero_translation): + log_info( + "Asset-group support clamp succeeded without movement: all AABBs " + "are exactly contained by the safe support region." + ) + return zero_translation + + # Rasterize the support region and compute feasible-centre maps for each AABB.s + transform, support_mask = self._rasterize_support(safe_support) + + # Compute feasible-centre maps with cacheing to avoid repeated binary erosion + # for identical AABB half-extents. + feasible_maps_by_id, centres_by_id = self._feasible_maps( + support_mask=support_mask, + transform=transform, + aabbs_by_id=aabbs_by_id, + ) + candidate_pixels = np.argwhere(support_mask) + if len(candidate_pixels) == 0: + return None + + candidate_world = np.asarray( + [ + transform.pixel_to_world(int(row), int(column)) + for row, column in candidate_pixels + ] + ) + candidate_deltas = candidate_world - anchor_xy + candidate_order = np.argsort( + np.einsum("ij,ij->i", candidate_deltas, candidate_deltas), kind="stable" + ) + for candidate_rank, candidate_index in enumerate(candidate_order, start=1): + delta_xy = candidate_deltas[candidate_index] + if not self._grid_translation_is_feasible( + delta_xy=delta_xy, + transform=transform, + feasible_maps_by_id=feasible_maps_by_id, + centres_by_id=centres_by_id, + ): + continue + if self._is_exactly_contained(safe_support, aabbs_by_id, delta_xy): + log_info( + "Asset-group support clamp succeeded after evaluating " + f"{candidate_rank}/{len(candidate_order)} grid candidates: " + f"delta_xy=({delta_xy[0]:+.4f}, {delta_xy[1]:+.4f}) m." + ) + return delta_xy + return None + + def _apply_delta_to_y_up_layouts( + self, *, delta_xy: np.ndarray, expected_ids: set[str] + ) -> list[dict[str, object]]: + """Apply a successful common z-up XY translation to stored layouts.""" + received_ids = {str(layout.get("id")) for layout in self.assets_layout} + if received_ids != expected_ids: + raise ValueError("Asset layouts and clamped AABBs must have identical ids.") + + dx, dy = delta_xy + translated_layouts: list[dict[str, object]] = [] + for layout in self.assets_layout: + position = layout.get("pos") + if not isinstance(position, list) or len(position) != 3: + raise ValueError( + "Each asset layout must contain a three-value pos list." + ) + translated_layout = dict(layout) + translated_position = [float(value) for value in position] + translated_position[0] += float(dx) + translated_position[2] -= float(dy) + translated_layout["pos"] = translated_position + translated_layouts.append(translated_layout) + return translated_layouts + + def save_group_clamp_debug_images(self) -> bool: + """Optionally save diagnostics for the support-valid group translation.""" + if self.refined_assets_layout is None: + self.clamp() + assert self.refined_assets_layout is not None + + initial_aabbs_by_id = self._validate_aabbs( + self.assets_aabb_2d_z_up_world_corners_by_id + ) + raw_support = self._coerce_support_geometry(self.support_region) + if raw_support is None: + raise ValueError( + "Asset-group support clamp requires a valid support region." + ) + safe_support = ( + raw_support + if self.config.margin_m == 0.0 + else self._polygonal_geometry(raw_support.buffer(-self.config.margin_m)) + ) + if safe_support is None or safe_support.is_empty: + raise ValueError("Asset-group support clamp has no usable support area.") + + original_positions_by_id = { + str(layout["id"]): layout["pos"] for layout in self.assets_layout + } + refined_positions_by_id = { + str(layout["id"]): layout["pos"] for layout in self.refined_assets_layout + } + if set(original_positions_by_id) != set(initial_aabbs_by_id) or set( + refined_positions_by_id + ) != set(initial_aabbs_by_id): + raise ValueError("Asset layouts and clamped AABBs must have identical ids.") + first_asset_id = next(iter(initial_aabbs_by_id)) + original_position = original_positions_by_id[first_asset_id] + refined_position = refined_positions_by_id[first_asset_id] + if ( + not isinstance(original_position, list) + or not isinstance(refined_position, list) + or len(original_position) != 3 + or len(refined_position) != 3 + ): + raise ValueError("Each asset layout must contain a three-value pos list.") + delta_xy = np.array( + [ + float(refined_position[0]) - float(original_position[0]), + float(original_position[2]) - float(refined_position[2]), + ] + ) + translated_aabbs_by_id = { + asset_id: corners + delta_xy + for asset_id, corners in initial_aabbs_by_id.items() + } + if self.debug_output_root is None: + raise ValueError( + "A debug_output_root is required when saving group-clamp " + "debug images." + ) + self._render_support_debug( + raw_support=raw_support, + safe_support=safe_support, + output_path=self.debug_output_root / "table_support_region_safe_2d.png", + ) + self._render_clamp_debug( + raw_support=raw_support, + initial_aabbs_by_id=initial_aabbs_by_id, + translated_aabbs_by_id=translated_aabbs_by_id, + delta_xy=delta_xy, + output_path=self.debug_output_root / "assets_group_support_clamp_2d.png", + ) + return True + + def _feasible_maps( + self, + *, + support_mask: np.ndarray, + transform: _GridTransform, + aabbs_by_id: dict[str, np.ndarray], + ) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]: + """Compute each AABB's rasterized feasible-centre map and current XY centre.""" + cached_feasible_maps: dict[tuple[int, int], np.ndarray] = {} + feasible_maps_by_id: dict[str, np.ndarray] = {} + centres_by_id: dict[str, np.ndarray] = {} + for asset_id, corners in aabbs_by_id.items(): + minimum = corners.min(axis=0) + maximum = corners.max(axis=0) + half_extent = (maximum - minimum) / 2.0 + kernel_key = tuple( + np.ceil(half_extent / transform.resolution_m).astype(int) + ) + if kernel_key not in cached_feasible_maps: + cached_feasible_maps[kernel_key] = binary_erosion( + support_mask, + structure=self._footprint_kernel(kernel_key), + border_value=0, + ) + feasible_maps_by_id[asset_id] = cached_feasible_maps[kernel_key] + centres_by_id[asset_id] = (minimum + maximum) / 2.0 + return feasible_maps_by_id, centres_by_id + + @staticmethod + def _footprint_kernel(kernel_key: tuple[int, int]) -> np.ndarray: + half_width_pixels, half_height_pixels = kernel_key + return np.ones( + (2 * half_height_pixels + 1, 2 * half_width_pixels + 1), dtype=bool + ) + + @staticmethod + def _grid_translation_is_feasible( + *, + delta_xy: np.ndarray, + transform: _GridTransform, + feasible_maps_by_id: dict[str, np.ndarray], + centres_by_id: dict[str, np.ndarray], + ) -> bool: + height, width = transform.shape + # Sorting by feasible map population is a cheap early-rejection order: + # small legal regions are most likely to reject a candidate quickly. + ordered_assets = sorted( + centres_by_id.items(), + key=lambda item: int(feasible_maps_by_id[item[0]].sum()), + ) + for asset_id, centre_xy in ordered_assets: + row, column = transform.world_to_nearest_pixel(centre_xy + delta_xy) + if row < 0 or row >= height or column < 0 or column >= width: + return False + if not feasible_maps_by_id[asset_id][row, column]: + return False + return True + + def _rasterize_support( + self, support: Polygon | MultiPolygon + ) -> tuple[_GridTransform, np.ndarray]: + """Rasterize a support region into a boolean XY grid and return the transform.""" + minimum_x, minimum_y, maximum_x, maximum_y = support.bounds + resolution = self.config.grid_resolution_m + x_coordinates = np.arange( + np.floor(minimum_x / resolution) * resolution, + np.ceil(maximum_x / resolution) * resolution + resolution / 2.0, + resolution, + ) + y_coordinates = np.arange( + np.floor(minimum_y / resolution) * resolution, + np.ceil(maximum_y / resolution) * resolution + resolution / 2.0, + resolution, + ) + x_grid, y_grid = np.meshgrid(x_coordinates, y_coordinates) + points = np.column_stack((x_grid.ravel(), y_grid.ravel())) + mask = np.zeros(len(points), dtype=bool) + for polygon in self._polygon_components(support): + component_mask = self._points_in_polygon(points, polygon.exterior.coords) + for hole in polygon.interiors: + component_mask &= ~self._points_in_polygon(points, hole.coords) + mask |= component_mask + return ( + _GridTransform( + x_coordinates, y_coordinates, resolution + ), # Keeps the transform for later world/grid conversions + mask.reshape( + x_grid.shape + ), # Keeps the boolean mask of the support region in grid form + ) + + @staticmethod + def _points_in_polygon(points: np.ndarray, coordinates: object) -> np.ndarray: + from matplotlib.path import Path as MatplotlibPath + + return MatplotlibPath(np.asarray(coordinates)).contains_points( + points, radius=1e-12 + ) + + @staticmethod + def _validate_aabbs( + aabbs_by_id: dict[str, np.ndarray], + ) -> dict[str, np.ndarray]: + """Validate the input AABB(s).""" + if not aabbs_by_id: + raise ValueError("At least one asset 2D AABB is required.") + validated: dict[str, np.ndarray] = {} + for asset_id, corners in aabbs_by_id.items(): + if not isinstance(asset_id, str) or not asset_id: + raise ValueError("Each asset AABB id must be a non-empty string.") + corner_array = np.asarray(corners, dtype=float) + if corner_array.shape != (4, 2) or not np.isfinite(corner_array).all(): + raise ValueError( + f"Asset {asset_id!r} must have four finite XY corners." + ) + validated[asset_id] = corner_array + return validated + + @classmethod + def _coerce_support_geometry( + cls, support_region: SupportGeometry + ) -> Polygon | MultiPolygon | None: + # Check the input type. + if isinstance(support_region, (Polygon, MultiPolygon)): + geometry = support_region + else: + log_warning( + "Unsupported support region type: " f"{type(support_region).__name__}." + ) + return None + return cls._polygonal_geometry(geometry) + + @staticmethod + def _polygonal_geometry(geometry: object) -> Polygon | MultiPolygon | None: + if not isinstance(geometry, (Polygon, MultiPolygon)) or geometry.is_empty: + return None + repaired = geometry if geometry.is_valid else geometry.buffer(0) + if not repaired.is_valid: + log_warning("Support polygon repair did not produce a valid geometry.") + return None + if isinstance(repaired, (Polygon, MultiPolygon)): + return repaired + if isinstance(repaired, GeometryCollection): + polygons = [item for item in repaired.geoms if isinstance(item, Polygon)] + return MultiPolygon(polygons) if polygons else None + return None + + @staticmethod + def _group_anchor(aabbs_by_id: dict[str, np.ndarray]) -> np.ndarray: + all_corners = np.concatenate(list(aabbs_by_id.values()), axis=0) + return (all_corners.min(axis=0) + all_corners.max(axis=0)) / 2.0 + + @staticmethod + def _is_exactly_contained( + support: Polygon | MultiPolygon, + aabbs_by_id: dict[str, np.ndarray], + delta_xy: np.ndarray, + ) -> bool: + return all( + support.covers( + affinity.translate( + AssetsGroupSupportClamp._aabb_polygon(corners), + xoff=float(delta_xy[0]), + yoff=float(delta_xy[1]), + ) + ) + for corners in aabbs_by_id.values() + ) + + @staticmethod + def _aabb_polygon(corners: np.ndarray) -> Polygon: + """Build a non-self-intersecting footprint regardless of corner order.""" + minimum = corners.min(axis=0) + maximum = corners.max(axis=0) + return Polygon( + [ + (minimum[0], minimum[1]), + (maximum[0], minimum[1]), + (maximum[0], maximum[1]), + (minimum[0], maximum[1]), + ] + ) + + def _render_support_debug( + self, + *, + raw_support: Polygon | MultiPolygon, + safe_support: Polygon | MultiPolygon | None, + output_path: str | Path, + ) -> Path: + path = self._resolve_png_output_path(output_path) + path.parent.mkdir(parents=True, exist_ok=True) + figure, axes = plt.subplots( + 1, 2, figsize=(14, 7), dpi=160, constrained_layout=True + ) + self._draw_support(axes[0], raw_support, "Input support region", "darkorange") + title = f"Safe support (margin={self.config.margin_m:.3f} m)" + if safe_support is None or safe_support.is_empty: + axes[1].set_title(f"{title}\n(infeasible)") + axes[1].set_aspect("equal", adjustable="box") + else: + self._draw_support(axes[1], safe_support, title, "seagreen") + figure.savefig(path, bbox_inches="tight") + plt.close(figure) + return path + + def _render_clamp_debug( + self, + *, + raw_support: Polygon | MultiPolygon, + initial_aabbs_by_id: dict[str, np.ndarray], + translated_aabbs_by_id: dict[str, np.ndarray] | None, + delta_xy: np.ndarray | None, + output_path: str | Path, + ) -> Path: + path = self._resolve_png_output_path(output_path) + path.parent.mkdir(parents=True, exist_ok=True) + figure, axes = plt.subplots( + 1, 2, figsize=(14, 7), dpi=160, constrained_layout=True + ) + self._draw_state( + axes[0], raw_support, initial_aabbs_by_id, "Before group clamp", "royalblue" + ) + if translated_aabbs_by_id is not None and delta_xy is not None: + action = "no movement" if np.allclose(delta_xy, 0.0) else "translated" + self._draw_state( + axes[1], + raw_support, + translated_aabbs_by_id, + f"After group clamp ({action})\nΔxy=({delta_xy[0]:+.3f}, {delta_xy[1]:+.3f}) m", + "seagreen", + ) + else: + self._draw_state( + axes[1], + raw_support, + initial_aabbs_by_id, + "No feasible group translation", + "firebrick", + ) + figure.savefig(path, bbox_inches="tight") + plt.close(figure) + return path + + @classmethod + def _draw_support( + cls, axis: plt.Axes, support: Polygon | MultiPolygon, title: str, color: str + ) -> None: + for polygon in cls._polygon_components(support): + exterior = np.asarray(polygon.exterior.coords) + axis.fill( + exterior[:, 0], + exterior[:, 1], + facecolor=color, + edgecolor="saddlebrown", + alpha=0.35, + ) + for hole in polygon.interiors: + hole_points = np.asarray(hole.coords) + axis.fill( + hole_points[:, 0], + hole_points[:, 1], + facecolor="white", + edgecolor="saddlebrown", + alpha=1.0, + ) + axis.set_aspect("equal", adjustable="box") + axis.set_xlabel("x (z-up world)") + axis.set_ylabel("y (z-up world)") + axis.set_title(title) + axis.autoscale_view() + + @classmethod + def _draw_state( + cls, + axis: plt.Axes, + support: Polygon | MultiPolygon, + aabbs_by_id: dict[str, np.ndarray], + title: str, + color: str, + ) -> None: + cls._draw_support(axis, support, title, "darkorange") + for asset_id, corners in sorted(aabbs_by_id.items()): + polygon = cls._aabb_polygon(corners) + boundary = np.asarray(polygon.exterior.coords) + axis.fill( + boundary[:, 0], + boundary[:, 1], + facecolor=color, + edgecolor=color, + alpha=0.35, + ) + axis.text( + *corners.mean(axis=0), + asset_id, + ha="center", + va="center", + fontsize=8, + bbox={"facecolor": "white", "alpha": 0.75, "edgecolor": "none"}, + ) + + @staticmethod + def _polygon_components(geometry: Polygon | MultiPolygon) -> list[Polygon]: + return [geometry] if isinstance(geometry, Polygon) else list(geometry.geoms) + + @staticmethod + def _resolve_png_output_path(output_path: str | Path) -> Path: + path = Path(output_path).expanduser().resolve() + return path if path.suffix.lower() == ".png" else path.with_suffix(".png") diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_generation_utils.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_generation_utils.py index 42237711a..c7ddf8f5a 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 @@ -24,19 +24,12 @@ from embodichain.lab.sim import SimulationManagerCfg, SimulationManager 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"}) @@ -132,81 +125,6 @@ def load_glb_mesh(glb_path: str | Path) -> trimesh.Trimesh: 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], @@ -515,764 +433,6 @@ def gravity_settle_assets_on_table( 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 = 64, -) -> 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: - # 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) - - # 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)} - - def _convert_layout_coordinate_system( layout_object: dict[str, object], *, diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/table_support_surface.py b/embodichain/gen_sim/scene_engine/pipeline/utils/table_support_surface.py new file mode 100644 index 000000000..6c3ad3e94 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/table_support_surface.py @@ -0,0 +1,497 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from collections import deque +from dataclasses import dataclass +from pathlib import Path + +from embodichain.utils.logger import log_info, log_warning +import matplotlib +import numpy as np +from scipy.spatial import ConvexHull, QhullError +from shapely.geometry import Polygon +from shapely.ops import unary_union +import trimesh + +matplotlib.use("Agg") + +import matplotlib.pyplot as plt +from mpl_toolkits.mplot3d.art3d import Poly3DCollection + + +@dataclass(frozen=True) +class SupportSurfaceConfig: + """Parameters for conservative single-level table-top detection.""" + + normal_z_min: float = 0.95 # Minimum z component for an upward face normal. + min_surface_area_m2: float = 0.01 # Minimum projected area for a candidate level. + max_face_height_span_m: float = ( + 0.01 # Maximum within-face z variation for flatness. + ) + height_level_tolerance_m: float = 0.005 # Maximum z difference within one level. + + +@dataclass(frozen=True) +class TableSupportRegion: + """Detected main table support surface in z-up world coordinates.""" + + top_z: float # Highest z value among the selected support-surface triangles. + vertices: np.ndarray # Full z-up table vertex array referenced by ``faces``. + faces: np.ndarray # Indices of triangles selected as the main support surface. + support_polygon: Polygon # Largest valid outer support contour in z-up XY. + + +class TableSupportSurfaceDetector: + """Detect one main upward table surface and render auditable diagnostics. + + The input mesh must be standing on x-y plane of the z-up world coordinates. + This class deliberately returns only the largest outer 2D contour because + the current layout stage models one main tabletop; the original support + triangles remain available for 3D diagnostics. + """ + + def __init__( + self, + *, + table_world_mesh: trimesh.Trimesh, + debug_output_root: str | Path | None = None, + config: SupportSurfaceConfig | None = None, + ) -> None: + # Init. + self.table_world_mesh = table_world_mesh + self.debug_output_root = ( + Path(debug_output_root).expanduser().resolve() + if debug_output_root is not None + else None + ) + self.support_region: TableSupportRegion | None = None + self.config = config if config is not None else SupportSurfaceConfig() + # Check config values. + if not 0.0 < self.config.normal_z_min <= 1.0: + raise ValueError("normal_z_min must be in (0, 1].") + if self.config.min_surface_area_m2 <= 0.0: + raise ValueError("min_surface_area_m2 must be positive.") + if self.config.max_face_height_span_m <= 0.0: + raise ValueError("max_face_height_span_m must be positive.") + if self.config.height_level_tolerance_m <= 0.0: + raise ValueError("height_level_tolerance_m must be positive.") + + def detect(self) -> TableSupportRegion: + """Detect the main upward-facing support surface of a z-up table mesh.""" + table_world_mesh = self.table_world_mesh + + # Check. + if len(table_world_mesh.vertices) < 3 or len(table_world_mesh.faces) == 0: + raise ValueError("Table mesh must contain at least one triangle.") + + mesh = table_world_mesh.copy() + # Material or UV seams can duplicate vertices along one physical table top. + mesh.merge_vertices(digits_vertex=7) + + # Repair normals if possible. + self._repair_normals_if_possible(mesh) + + face_vertices = mesh.vertices[mesh.faces] + face_height_ranges = np.ptp(face_vertices[:, :, 2], axis=1) + # Check: 1. range of z values of each triangle; 2. upward-facing triangles. + candidate_face_indices = np.flatnonzero( + (mesh.face_normals[:, 2] >= self.config.normal_z_min) + & (face_height_ranges <= self.config.max_face_height_span_m) + ) + if len(candidate_face_indices) == 0: + raise ValueError( + "Table mesh has no near-horizontal upward-facing triangles that " + "satisfy max_face_height_span_m." + ) + + # Select the best support level among the candidate triangles. + selected_faces = self._select_main_support_level( + mesh=mesh, + candidate_face_indices=candidate_face_indices, + full_table_hull_area=self._convex_hull_area( + mesh.vertices[:, :2], name="Table" + ), + ) + selected_vertices = face_vertices[selected_faces] + vertices = mesh.vertices.copy() + faces = mesh.faces[selected_faces].copy() + self.support_region = TableSupportRegion( + top_z=float(selected_vertices[:, :, 2].max()), + vertices=vertices, + faces=faces, + support_polygon=self._extract_largest_support_polygon(vertices[faces, :2]), + ) + return self.support_region + + def save_support_surface_debug_images( + self, + *, + save_3d: bool = True, + save_2d: bool = True, + output_3d_path: str | Path | None = None, + output_2d_path: str | Path | None = None, + ) -> bool: + """Optionally save standard 3D and 2D support-detection diagnostics.""" + if not save_3d and not save_2d: + return True + if self.support_region is None: + self.detect() + assert self.support_region is not None + if save_3d: + self._save_support_surface_3d_image( + table_world_mesh=self.table_world_mesh, + support_region=self.support_region, + output_path=self._resolve_debug_output_path( + output_3d_path, "table_support_surface_3d.png" + ), + ) + if save_2d: + self._save_support_region_2d_image( + support_region=self.support_region, + output_path=self._resolve_debug_output_path( + output_2d_path, "table_support_region_2d.png" + ), + ) + return True + + def _resolve_debug_output_path( + self, output_path: str | Path | None, default_filename: str + ) -> Path: + if output_path is not None: + return Path(output_path).expanduser().resolve() + if self.debug_output_root is None: + raise ValueError( + "A debug_output_root or an explicit debug output path is required " + "when saving support-surface debug images." + ) + return self.debug_output_root / default_filename + + @staticmethod + def _repair_normals_if_possible(mesh: trimesh.Trimesh) -> None: + if not mesh.is_watertight: + log_warning( + "Table mesh is not watertight; skipped outward-normal repair and " + "will use its input face normals for support detection." + ) + return + if mesh.is_volume: + log_info("Table mesh normals already form a valid outward-facing volume.") + return + try: + trimesh.repair.fix_normals(mesh, multibody=True) + except Exception as exc: + log_warning( + "Table mesh normal repair raised " + f"{exc}; support detection will use the resulting face normals." + ) + return + if mesh.is_volume: + log_info("Table mesh normals were repaired into an outward-facing volume.") + return + log_warning( + "Table mesh is watertight but normals could not be repaired into a " + "valid outward-facing volume; support detection will use the resulting " + "face normals." + ) + return + + def _select_main_support_level( + self, + *, + mesh: trimesh.Trimesh, + candidate_face_indices: np.ndarray, + full_table_hull_area: float, + ) -> np.ndarray: + # Find adj. + adjacency = self._face_adjacency(mesh) + # Use BFS to group the connect components. + components = self._connected_components( + set(int(index) for index in candidate_face_indices), adjacency + ) + + # Sort components by their top z value, descending. + components_by_height = sorted( + ( + ( + float(mesh.vertices[mesh.faces[list(component)], 2].max()), + component, + ) + for component in components + ), + key=lambda item: item[0], + reverse=True, + ) + + # Group components into levels by their top z value, within the height_level_tolerance_m. + levels: list[tuple[float, list[set[int]]]] = [] + for component_top_z, component in components_by_height: + for level_index, (level_top_z, level_components) in enumerate(levels): + if ( + level_top_z - component_top_z + <= self.config.height_level_tolerance_m + ): + level_components.append(component) + levels[level_index] = (level_top_z, level_components) + break + else: + levels.append((component_top_z, [component])) + + best_level_faces: np.ndarray | None = None + best_hull_gap = np.inf + best_top_z = -np.inf + best_projected_support_area = 0.0 + hull_gap_tolerance = max(full_table_hull_area * 1e-6, 1e-9) + for level_top_z, level_components in levels: + level_indices = np.asarray( + sorted( + face_index + for component in level_components + for face_index in component + ), + dtype=int, + ) + level_triangles = mesh.vertices[mesh.faces[level_indices]] + level_projected_support_area = self._projected_triangle_area( + level_triangles[:, :, :2] + ) + if level_projected_support_area < self.config.min_surface_area_m2: + continue + level_hull_area = self._convex_hull_area( + level_triangles[:, :, :2].reshape(-1, 2), + name="Candidate support level", + ) + # Compute the gap between convex hull and triangle projection area, + # for selecting the best level among multiple candidates which avoids + # small area which have the largest z value. + level_hull_gap = max(0.0, full_table_hull_area - level_hull_area) + if ( + best_level_faces is None + or level_hull_gap < best_hull_gap - hull_gap_tolerance + or ( + abs(level_hull_gap - best_hull_gap) <= hull_gap_tolerance + and level_top_z > best_top_z + self.config.height_level_tolerance_m + ) + or ( + abs(level_hull_gap - best_hull_gap) <= hull_gap_tolerance + and abs(level_top_z - best_top_z) + <= self.config.height_level_tolerance_m + and level_projected_support_area > best_projected_support_area + ) + ): + best_level_faces = level_indices + best_hull_gap = level_hull_gap + best_top_z = level_top_z + best_projected_support_area = level_projected_support_area + if best_level_faces is None: + raise ValueError( + "No upward-facing support level meets min_surface_area_m2." + ) + return best_level_faces + + @staticmethod + def _convex_hull_area(points: np.ndarray, *, name: str) -> float: + """Compute the area of the convex hull of a set of XY points.""" + unique_points = np.unique(np.asarray(points, dtype=float), axis=0) + if ( + unique_points.ndim != 2 + or unique_points.shape[1] != 2 + or len(unique_points) < 3 + ): + raise ValueError(f"{name} must contain at least three unique XY points.") + try: + return float(ConvexHull(unique_points).volume) + except QhullError as exc: + raise ValueError(f"{name} XY projection is degenerate.") from exc + + @staticmethod + def _projected_triangle_area(triangles_xy: np.ndarray) -> float: + """Compute the total area of triangles projected onto the XY plane.""" + first_edges = triangles_xy[:, 1] - triangles_xy[:, 0] + second_edges = triangles_xy[:, 2] - triangles_xy[:, 0] + cross_products = ( + first_edges[:, 0] * second_edges[:, 1] + - first_edges[:, 1] * second_edges[:, 0] + ) + return float(np.abs(cross_products).sum() / 2.0) + + @classmethod + def _extract_largest_support_polygon(cls, triangles_xy: np.ndarray) -> Polygon: + projected_triangles = [ + Polygon(triangle) + for triangle in triangles_xy + if cls._projected_triangle_area(triangle[None, ...]) > 1e-12 + ] + if not projected_triangles: + raise ValueError( + "Selected support surface has no non-degenerate XY triangles." + ) + merged_region = unary_union(projected_triangles) + if merged_region.geom_type == "Polygon": + polygons = [merged_region] + else: + polygons = [ + geometry + for geometry in merged_region.geoms + if geometry.geom_type == "Polygon" + ] + if not polygons: + raise ValueError( + "Could not create a 2D support region from the selected triangles." + ) + if len(polygons) > 1: + log_warning( + "Detected multiple disconnected outer support contours; using only " + "the largest one for the single-contour support-region output." + ) + largest_polygon = max(polygons, key=lambda polygon: polygon.area) + if largest_polygon.is_empty or not largest_polygon.is_valid: + raise ValueError("The merged 2D support region is not a valid polygon.") + boundary_xy = np.asarray(largest_polygon.exterior.coords, dtype=float) + if len(boundary_xy) < 4 or not np.isfinite(boundary_xy).all(): + raise ValueError("The merged 2D support contour is degenerate.") + return Polygon(boundary_xy) + + @staticmethod + def _face_adjacency(mesh: trimesh.Trimesh) -> dict[int, set[int]]: + """Build a face adjacency dictionary for the mesh.""" + adjacency: dict[int, set[int]] = {} + for first, second in mesh.face_adjacency: + first_index = int(first) + second_index = int(second) + adjacency.setdefault(first_index, set()).add(second_index) + adjacency.setdefault(second_index, set()).add(first_index) + return adjacency + + @staticmethod + def _connected_components( + faces: set[int], adjacency: dict[int, set[int]] + ) -> list[set[int]]: + """Group upward-facing candidate triangles into edge-connected surface components with BFS.""" + unvisited = set(faces) + components: list[set[int]] = [] + while unvisited: + component: set[int] = set() + queue = deque([unvisited.pop()]) + while queue: + face_index = queue.popleft() + component.add(face_index) + for neighbor in adjacency.get(face_index, set()): + if neighbor in unvisited: + unvisited.remove(neighbor) + queue.append(neighbor) + components.append(component) + return components + + @classmethod + def _save_support_surface_3d_image( + cls, + *, + table_world_mesh: trimesh.Trimesh, + support_region: TableSupportRegion, + output_path: str | Path, + ) -> Path: + resolved_output_path = cls._resolve_png_output_path(output_path) + resolved_output_path.parent.mkdir(parents=True, exist_ok=True) + figure = plt.figure(figsize=(9, 8), dpi=160, constrained_layout=True) + axis = figure.add_subplot(projection="3d") + axis.add_collection3d( + Poly3DCollection( + table_world_mesh.vertices[table_world_mesh.faces], + facecolor="steelblue", + edgecolor="none", + alpha=0.18, + ) + ) + axis.add_collection3d( + Poly3DCollection( + support_region.vertices[support_region.faces], + facecolor="darkorange", + edgecolor="saddlebrown", + linewidth=0.25, + alpha=0.95, + ) + ) + lower = table_world_mesh.bounds[0].copy() + upper = table_world_mesh.bounds[1].copy() + extent = upper - lower + lower[extent <= 1e-9] -= 0.001 + upper[extent <= 1e-9] += 0.001 + axis.set( + xlim=(lower[0], upper[0]), + ylim=(lower[1], upper[1]), + zlim=(lower[2], upper[2]), + ) + axis.set_box_aspect(upper - lower) + axis.view_init(elev=25.0, azim=-55.0) + axis.set_xlabel("x (z-up world)") + axis.set_ylabel("y (z-up world)") + axis.set_zlabel("z (up)") + axis.set_title( + "Detected main table support surface\n" + f"top z={support_region.top_z:.4f} m, faces={len(support_region.faces)}" + ) + figure.savefig(resolved_output_path, bbox_inches="tight") + plt.close(figure) + return resolved_output_path + + @classmethod + def _save_support_region_2d_image( + cls, + *, + support_region: TableSupportRegion, + output_path: str | Path, + ) -> Path: + resolved_output_path = cls._resolve_png_output_path(output_path) + resolved_output_path.parent.mkdir(parents=True, exist_ok=True) + boundary_xy = np.asarray(support_region.support_polygon.exterior.coords) + figure, axis = plt.subplots(figsize=(8, 8), dpi=160, constrained_layout=True) + axis.fill( + boundary_xy[:, 0], + boundary_xy[:, 1], + facecolor="darkorange", + edgecolor="none", + alpha=0.82, + label="detected support region", + ) + axis.plot( + boundary_xy[:, 0], + boundary_xy[:, 1], + color="saddlebrown", + linewidth=2.0, + label="outer support contour", + ) + axis.autoscale_view() + axis.set_aspect("equal", adjustable="box") + axis.set_xlabel("x (z-up world)") + axis.set_ylabel("y (z-up world)") + axis.set_title( + "Detected 2D table support region\n" + f"z={support_region.top_z:.4f} m, contour vertices={len(boundary_xy) - 1}" + ) + axis.legend(loc="best") + figure.savefig(resolved_output_path, bbox_inches="tight") + plt.close(figure) + return resolved_output_path + + @staticmethod + def _resolve_png_output_path(output_path: str | Path) -> Path: + resolved_output_path = Path(output_path).expanduser().resolve() + if resolved_output_path.suffix.lower() != ".png": + return resolved_output_path.with_suffix(".png") + return resolved_output_path From 3b1abc28a2ff4af6de368609553a3f5b204b56b8 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Mon, 3 Aug 2026 19:14:01 +0800 Subject: [PATCH 25/41] Reformatted the gravity settlement --- .../scene_engine/pipeline/scene_generation.py | 7 +- .../pipeline/utils/assets_gravity_settler.py | 300 ++++++++++++++++++ .../pipeline/utils/scene_generation_utils.py | 251 --------------- 3 files changed, 305 insertions(+), 253 deletions(-) create mode 100644 embodichain/gen_sim/scene_engine/pipeline/utils/assets_gravity_settler.py diff --git a/embodichain/gen_sim/scene_engine/pipeline/scene_generation.py b/embodichain/gen_sim/scene_engine/pipeline/scene_generation.py index 048d75cb6..b18e2bd3e 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/scene_generation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/scene_generation.py @@ -39,10 +39,12 @@ from embodichain.gen_sim.scene_engine.pipeline.utils.assets_group_layout_optimizer import ( AssetsSupportLayoutOptimizer, ) +from embodichain.gen_sim.scene_engine.pipeline.utils.assets_gravity_settler import ( + AssetsGravitySettler, +) from embodichain.gen_sim.scene_engine.pipeline.utils.scene_generation_utils import ( align_assets_group_to_table_aabb_top, export_baked_layout_object_glbs, - gravity_settle_assets_on_table, layout_object_to_transform_matrix, load_glb_mesh, quaternion_wxyz_to_euler_xyz_degrees, @@ -481,11 +483,12 @@ def _layout_refinement( # 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( + gravity_settler = AssetsGravitySettler( table_layout=refined_table_layout, assets_layout=refined_assets_layout, geometry_root=simready_geometry_output_root, ) + refined_assets_layout = gravity_settler.settle() return refined_table_layout, refined_assets_layout diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/assets_gravity_settler.py b/embodichain/gen_sim/scene_engine/pipeline/utils/assets_gravity_settler.py new file mode 100644 index 000000000..3f68508cb --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/assets_gravity_settler.py @@ -0,0 +1,300 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Sequence + +import numpy as np +from scipy.spatial.transform import Rotation +import trimesh + +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_generation_utils import ( + layout_object_to_transform_matrix, + load_glb_mesh, + transform_matrix_to_layout_object, +) +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.cfg import RigidObjectCfg +from embodichain.lab.sim.shapes import MeshCfg +from embodichain.utils.logger import log_info + + +@dataclass(frozen=True) +class AssetsGravitySettlerConfig: + """Physics controls for table-top asset settling.""" + + clearance_m: float = 0.02 # Initial gap between each asset and the table top. + settle_steps: int = 300 # Fixed number of simulator steps to execute. + physics_dt: float = 1.0 / 100.0 # Physics timestep in seconds. + sim_device: str = "cpu" # Simulation device requested from EmbodiChain Lab. + max_convex_hull_num: int = 32 # VHACD hull budget for each collision mesh. + + +class AssetsGravitySettler: + """Settle all assets together on one static table in a z-up simulation.""" + + def __init__( + self, + *, + table_layout: dict[str, object], + assets_layout: list[dict[str, object]], + geometry_root: str | Path, + config: AssetsGravitySettlerConfig | None = None, + ) -> None: + self.table_layout = table_layout + self.assets_layout = assets_layout + self.geometry_root = Path(geometry_root).expanduser().resolve() + self.settled_assets_layout: list[dict[str, object]] | None = None + self.config = config if config is not None else AssetsGravitySettlerConfig() + # Check. + if self.config.clearance_m < 0.0: + raise ValueError("Gravity-settle clearance_m must be non-negative.") + if self.config.settle_steps <= 0: + raise ValueError("Gravity-settle settle_steps must be positive.") + if self.config.physics_dt <= 0.0: + raise ValueError("Gravity-settle physics_dt must be positive.") + if self.config.max_convex_hull_num <= 0: + raise ValueError("Gravity-settle max_convex_hull_num must be positive.") + + def settle(self) -> list[dict[str, object]]: + """Run gravity settling and return the resulting y-up asset layouts.""" + self.settled_assets_layout = None + if not self.assets_layout: + self.settled_assets_layout = [] + log_info("Scene has no movable assets; skipping gravity settling.") + return self.settled_assets_layout + + table_id = self._require_layout_id(self.table_layout, name="Table") + asset_ids: set[str] = set() + for asset_layout in self.assets_layout: + asset_id = self._require_layout_id(asset_layout, name="Asset") + if asset_id in asset_ids: + raise ValueError(f"Asset layouts contain duplicate id {asset_id!r}.") + asset_ids.add(asset_id) + + y_up_to_z_up_matrix = np.eye(4) + y_up_to_z_up_matrix[:3, :3] = np.array( + [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]] + ) + z_up_to_y_up_matrix = np.linalg.inv(y_up_to_z_up_matrix) + table_info = self._prepare_sim_body( + layout_object=self.table_layout, + y_up_to_z_up_matrix=y_up_to_z_up_matrix, + ) + table_world_mesh = self._mesh_to_z_up_world_for_aabb( + y_up_mesh=table_info["mesh"], + z_up_rigid_layout=table_info["rigid_layout"], + z_up_scale=table_info["z_up_scale"], + y_up_to_z_up_matrix=y_up_to_z_up_matrix, + ) + table_top_z = float(table_world_mesh.bounds[1, 2]) + + prepared_assets: dict[str, dict[str, object]] = {} + for asset_layout in self.assets_layout: + asset_id = str(asset_layout["id"]) + asset_info = self._prepare_sim_body( + layout_object=asset_layout, + y_up_to_z_up_matrix=y_up_to_z_up_matrix, + ) + asset_world_mesh = self._mesh_to_z_up_world_for_aabb( + y_up_mesh=asset_info["mesh"], + z_up_rigid_layout=asset_info["rigid_layout"], + z_up_scale=asset_info["z_up_scale"], + y_up_to_z_up_matrix=y_up_to_z_up_matrix, + ) + asset_bottom_z = float(asset_world_mesh.bounds[0, 2]) + asset_info["rigid_layout"]["pos"][2] += ( + table_top_z + self.config.clearance_m - asset_bottom_z + ) + prepared_assets[asset_id] = asset_info + + log_info( + "Gravity settling started: " + f"assets={len(prepared_assets)}, steps={self.config.settle_steps}, " + f"physics_dt={self.config.physics_dt:.4f} s, " + f"max_convex_hulls={self.config.max_convex_hull_num}." + ) + sim = SimulationManager( + SimulationManagerCfg( + headless=True, + physics_dt=self.config.physics_dt, + sim_device=self.config.sim_device, + ) + ) + try: + # Add table. + sim.add_rigid_object( + RigidObjectCfg( + uid=table_id, + shape=MeshCfg(fpath=str(table_info["mesh_path"])), + init_pos=tuple(table_info["rigid_layout"]["pos"]), + init_rot=tuple( + self._simulation_euler_xyz_degrees(table_info["rigid_layout"]) + ), + body_scale=tuple(table_info["y_up_scale"]), + body_type="static", + max_convex_hull_num=self.config.max_convex_hull_num, + acd_method="vhacd", + ) + ) + # Add assets. + simulated_assets: dict[str, object] = {} + for asset_id, asset_info in prepared_assets.items(): + rigid_layout = asset_info["rigid_layout"] + simulated_assets[asset_id] = sim.add_rigid_object( + RigidObjectCfg( + uid=asset_id, + shape=MeshCfg(fpath=str(asset_info["mesh_path"])), + init_pos=tuple(rigid_layout["pos"]), + init_rot=tuple( + self._simulation_euler_xyz_degrees(rigid_layout) + ), + body_scale=tuple(asset_info["y_up_scale"]), + body_type="dynamic", + max_convex_hull_num=self.config.max_convex_hull_num, + acd_method="vhacd", + ) + ) + # Run simulation to settle all assets. + sim.update(step=self.config.settle_steps) + + # Update the final layouts. + settled_layout_by_id: dict[str, dict[str, object]] = {} + for asset_id, simulated_asset in simulated_assets.items(): + final_rigid_pose_z_up = np.asarray( + simulated_asset.get_local_pose(to_matrix=True)[0] + .detach() + .cpu() + .numpy(), + dtype=float, + ) + scale_matrix = np.eye(4) + scale_matrix[:3, :3] = np.diag(prepared_assets[asset_id]["z_up_scale"]) + final_z_up_layout_matrix = final_rigid_pose_z_up @ scale_matrix + settled_layout_by_id[asset_id] = transform_matrix_to_layout_object( + asset_id, + z_up_to_y_up_matrix + @ final_z_up_layout_matrix + @ y_up_to_z_up_matrix, + ) + finally: + # Release resources. + sim.destroy(exit_process=False) + SimulationManager.flush_cleanup_queue() + + self.settled_assets_layout = [ + settled_layout_by_id[str(asset_layout["id"])] + for asset_layout in self.assets_layout + ] + log_info("Gravity settling completed for all assets.") + return self.settled_assets_layout + + def _prepare_sim_body( + self, + *, + layout_object: dict[str, object], + y_up_to_z_up_matrix: np.ndarray, + ) -> dict[str, object]: + """Load one y-up GLB and prepare its z-up simulation pose.""" + object_id = self._require_layout_id(layout_object, name="Layout object") + source_mesh_path = self.geometry_root / f"{object_id}.glb" + source_mesh = load_glb_mesh(source_mesh_path) + z_up_layout = self._convert_layout_coordinate_system( + layout_object, + source_to_target_matrix=y_up_to_z_up_matrix, + ) + return { + "mesh_path": source_mesh_path, + "mesh": source_mesh, + "rigid_layout": { + "id": object_id, + "rot": self._three_floats(z_up_layout.get("rot"), field_name="rot"), + "pos": self._three_floats(z_up_layout.get("pos"), field_name="pos"), + "scale": [1.0, 1.0, 1.0], + }, + "y_up_scale": self._three_floats( + layout_object.get("scale"), field_name="scale" + ), + "z_up_scale": self._three_floats( + z_up_layout.get("scale"), field_name="scale" + ), + } + + @staticmethod + def _mesh_to_z_up_world_for_aabb( + *, + y_up_mesh: trimesh.Trimesh, + z_up_rigid_layout: dict[str, object], + z_up_scale: Sequence[float], + y_up_to_z_up_matrix: np.ndarray, + ) -> trimesh.Trimesh: + """Transform a y-up mesh into its z-up world pose for AABB measurement.""" + mesh = y_up_mesh.copy() + mesh.apply_transform(y_up_to_z_up_matrix) + scale_matrix = np.eye(4) + scale_matrix[:3, :3] = np.diag(z_up_scale) + mesh.apply_transform(scale_matrix) + mesh.apply_transform(layout_object_to_transform_matrix(z_up_rigid_layout)) + return mesh + + @staticmethod + def _simulation_euler_xyz_degrees(layout_object: dict[str, object]) -> list[float]: + """Convert lowercase-xyz layout rotation to SimulationManager's XYZ order.""" + layout_rotation = Rotation.from_euler( + "xyz", + AssetsGravitySettler._three_floats( + layout_object.get("rot"), field_name="rot" + ), + degrees=True, + ) + return layout_rotation.as_euler("XYZ", degrees=True).tolist() + + @staticmethod + def _convert_layout_coordinate_system( + layout_object: dict[str, object], + *, + source_to_target_matrix: np.ndarray, + ) -> dict[str, object]: + """Convert one layout object between coordinate frames through its matrix.""" + return transform_matrix_to_layout_object( + str(layout_object["id"]), + source_to_target_matrix + @ layout_object_to_transform_matrix(layout_object) + @ np.linalg.inv(source_to_target_matrix), + ) + + @staticmethod + def _require_layout_id(layout_object: dict[str, object], *, name: str) -> str: + """Check id.""" + object_id = layout_object.get("id") + if not isinstance(object_id, str) or not object_id: + raise ValueError(f"{name} layout must contain a non-empty string id.") + return object_id + + @staticmethod + def _three_floats(value: object, *, field_name: str) -> list[float]: + """Check three values.""" + if not isinstance(value, list) or len(value) != 3: + raise ValueError(f"Layout field {field_name} must contain three values.") + try: + return [float(item) for item in value] + except (TypeError, ValueError) as exc: + raise ValueError( + f"Layout field {field_name} must contain numeric values." + ) from exc diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_generation_utils.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_generation_utils.py index c7ddf8f5a..bf432232d 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,9 +21,6 @@ import re from typing import Sequence -from embodichain.lab.sim import SimulationManagerCfg, SimulationManager -from embodichain.lab.sim.cfg import RigidObjectCfg -from embodichain.lab.sim.shapes import MeshCfg import numpy as np import open3d as o3d from scipy.spatial import ConvexHull, QhullError @@ -44,23 +41,6 @@ def quaternion_wxyz_to_euler_xyz_degrees( 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: @@ -202,237 +182,6 @@ def align_assets_group_to_table_aabb_top( ) -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 = SimulationManager( - 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, - acd_method="vhacd", # Use vhacd by default. - ) - ) - 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, - acd_method="vhacd", # Use vhacd by default. - ) - ) - - # 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.destroy(exit_process=False) - SimulationManager.flush_cleanup_queue() - - settled_assets_layout = [ - settled_layout_by_id[str(asset_layout["id"])] for asset_layout in assets_layout - ] - return settled_assets_layout - - def _convert_layout_coordinate_system( layout_object: dict[str, object], *, From dd66de8c2302bd2df45e4dfa84517d7aba115d1c Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Mon, 3 Aug 2026 21:28:04 +0800 Subject: [PATCH 26/41] Reformatted assets group table aligner + Reformatted simready tools --- .../scene_engine/pipeline/scene_generation.py | 111 +----- .../utils/assets_group_table_aligner.py | 150 ++++++++ .../pipeline/utils/scene_generation_utils.py | 285 ---------------- .../utils/simready_scene_processor.py | 320 ++++++++++++++++++ 4 files changed, 483 insertions(+), 383 deletions(-) create mode 100644 embodichain/gen_sim/scene_engine/pipeline/utils/assets_group_table_aligner.py create mode 100644 embodichain/gen_sim/scene_engine/pipeline/utils/simready_scene_processor.py diff --git a/embodichain/gen_sim/scene_engine/pipeline/scene_generation.py b/embodichain/gen_sim/scene_engine/pipeline/scene_generation.py index b18e2bd3e..b8174f7bd 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/scene_generation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/scene_generation.py @@ -36,6 +36,9 @@ from embodichain.gen_sim.scene_engine.pipeline.utils.assets_group_support_clamp import ( AssetsGroupSupportClamp, ) +from embodichain.gen_sim.scene_engine.pipeline.utils.assets_group_table_aligner import ( + AssetsGroupTableAligner, +) from embodichain.gen_sim.scene_engine.pipeline.utils.assets_group_layout_optimizer import ( AssetsSupportLayoutOptimizer, ) @@ -43,14 +46,15 @@ AssetsGravitySettler, ) from embodichain.gen_sim.scene_engine.pipeline.utils.scene_generation_utils import ( - align_assets_group_to_table_aabb_top, export_baked_layout_object_glbs, layout_object_to_transform_matrix, load_glb_mesh, quaternion_wxyz_to_euler_xyz_degrees, - simready_object_glb, transform_matrix_to_layout_object, ) +from embodichain.gen_sim.scene_engine.pipeline.utils.simready_scene_processor import ( + SimReadySceneProcessor, +) from embodichain.gen_sim.scene_engine.pipeline.utils.table_support_surface import ( TableSupportSurfaceDetector, ) @@ -204,21 +208,14 @@ def _refine_geometries_and_layout( layout_object["id"]: layout_object for layout_object in coarse_layout } - # Simready all the assets. - simready_assets_layout = _simready_assets( + simready_processor = SimReadySceneProcessor( 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, + coarse_geometry_root=coarse_geometry_output_root, + simready_geometry_root=simready_geometry_output_root, ) + simready_assets_layout = simready_processor.process_assets() + simready_table_layout = simready_processor.process_table() # Concat then save the table info and the assets info in one JSON file. simready_layout = [simready_table_layout, *simready_assets_layout] (Path(simready_geometry_output_root) / "simready_layout.json").write_text( @@ -407,11 +404,12 @@ def _layout_refinement( # 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_group_to_table_aabb_top( + group_table_aligner = AssetsGroupTableAligner( table_layout=refined_table_layout, assets_layout=refined_assets_layout, geometry_root=simready_geometry_output_root, ) + refined_table_layout, refined_assets_layout = group_table_aligner.align() if not refined_assets_layout: log_info("Scene has no movable assets; skipping support-region clamping.") return refined_table_layout, [] @@ -556,89 +554,6 @@ def _mesh_in_z_up_world(layout_object: dict[str, object]) -> trimesh.Trimesh: return table_world_mesh_z_up, asset_aabbs_by_id -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() diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/assets_group_table_aligner.py b/embodichain/gen_sim/scene_engine/pipeline/utils/assets_group_table_aligner.py new file mode 100644 index 000000000..990476044 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/assets_group_table_aligner.py @@ -0,0 +1,150 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +import numpy as np + +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_generation_utils import ( + layout_object_to_transform_matrix, + load_glb_mesh, + transform_matrix_to_layout_object, +) +from embodichain.utils.logger import log_info + + +@dataclass(frozen=True) +class AssetsGroupTableAlignerConfig: + """Controls for the initial vertical gap above the table.""" + + clearance_m: float = 0.02 # Initial table-to-group gap in metres. + + +class AssetsGroupTableAligner: + """Place every asset as one rigid vertical group above a table AABB top.""" + + def __init__( + self, + *, + table_layout: dict[str, object], + assets_layout: list[dict[str, object]], + geometry_root: str | Path, + config: AssetsGroupTableAlignerConfig | None = None, + ) -> None: + self.table_layout = table_layout + self.assets_layout = assets_layout + self.geometry_root = Path(geometry_root).expanduser().resolve() + self.aligned_table_layout: dict[str, object] | None = None + self.aligned_assets_layout: list[dict[str, object]] | None = None + self.config = config if config is not None else AssetsGroupTableAlignerConfig() + # Check. + if self.config.clearance_m < 0.0: + raise ValueError("Table clearance_m must be non-negative.") + + def align(self) -> tuple[dict[str, object], list[dict[str, object]]]: + """Return y-up layouts with the complete asset group above the table. + + Input and output layouts use y-up, matching the GLBs on disk. The group + is temporarily measured in z-up coordinates and every asset receives the + same vertical translation. This preserves all asset-to-asset relative + poses. + """ + self.aligned_table_layout = None + self.aligned_assets_layout = None + if not self.assets_layout: + self.aligned_table_layout = self.table_layout + self.aligned_assets_layout = [] + log_info("Scene has no movable assets; skipping vertical group alignment.") + return self.aligned_table_layout, self.aligned_assets_layout + + y_up_to_z_up_matrix = np.eye(4) + y_up_to_z_up_matrix[:3, :3] = np.array( + [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]] + ) + z_up_to_y_up_matrix = np.linalg.inv(y_up_to_z_up_matrix) + + z_up_table_layout = self._convert_layout_coordinate_system( + self.table_layout, + source_to_target_matrix=y_up_to_z_up_matrix, + ) + z_up_assets_layout = [ + self._convert_layout_coordinate_system( + asset_layout, + source_to_target_matrix=y_up_to_z_up_matrix, + ) + for asset_layout in self.assets_layout + ] + + table_id = self._require_layout_id(z_up_table_layout, name="Table") + table_mesh = load_glb_mesh(self.geometry_root / f"{table_id}.glb") + table_mesh.apply_transform(y_up_to_z_up_matrix) + table_mesh.apply_transform(layout_object_to_transform_matrix(z_up_table_layout)) + target_group_bottom_z = table_mesh.bounds[1, 2] + self.config.clearance_m + + group_bottom_z = np.inf + for asset_layout in z_up_assets_layout: + asset_id = self._require_layout_id(asset_layout, name="Asset") + asset_mesh = load_glb_mesh(self.geometry_root / f"{asset_id}.glb") + asset_mesh.apply_transform(y_up_to_z_up_matrix) + asset_mesh.apply_transform(layout_object_to_transform_matrix(asset_layout)) + # Find the lowest z among all the assets. + group_bottom_z = min(group_bottom_z, float(asset_mesh.bounds[0, 2])) + + group_vertical_translation_z = target_group_bottom_z - group_bottom_z + for asset_layout in z_up_assets_layout: + asset_layout["pos"][2] += group_vertical_translation_z + + self.aligned_table_layout = self._convert_layout_coordinate_system( + z_up_table_layout, + source_to_target_matrix=z_up_to_y_up_matrix, + ) + self.aligned_assets_layout = [ + self._convert_layout_coordinate_system( + asset_layout, + source_to_target_matrix=z_up_to_y_up_matrix, + ) + for asset_layout in z_up_assets_layout + ] + log_info( + "Aligned the asset group above the table with " + f"delta_z={group_vertical_translation_z:.4f} m and " + f"clearance={self.config.clearance_m:.4f} m." + ) + return self.aligned_table_layout, self.aligned_assets_layout + + @staticmethod + def _convert_layout_coordinate_system( + layout_object: dict[str, object], + *, + source_to_target_matrix: np.ndarray, + ) -> dict[str, object]: + """Convert one layout object between coordinate systems through its matrix.""" + return transform_matrix_to_layout_object( + str(layout_object["id"]), + source_to_target_matrix + @ layout_object_to_transform_matrix(layout_object) + @ np.linalg.inv(source_to_target_matrix), + ) + + @staticmethod + def _require_layout_id(layout_object: dict[str, object], *, name: str) -> str: + object_id = layout_object.get("id") + if not isinstance(object_id, str) or not object_id: + raise ValueError(f"{name} layout must contain a non-empty string id.") + return object_id diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_generation_utils.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_generation_utils.py index bf432232d..49e5da138 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 @@ -18,17 +18,12 @@ from __future__ import annotations from pathlib import Path -import re from typing import Sequence import numpy as np -import open3d as o3d -from scipy.spatial import ConvexHull, QhullError from scipy.spatial.transform import Rotation import trimesh -_UPRIGHT_CONTAINER_ID_TOKENS = frozenset({"bottle", "can", "jar", "flask", "thermos"}) - def quaternion_wxyz_to_euler_xyz_degrees( quaternion_wxyz: Sequence[float], @@ -105,98 +100,6 @@ def load_glb_mesh(glb_path: str | Path) -> trimesh.Trimesh: raise ValueError(f"GLB geometry is not a mesh: {resolved_glb_path}") -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 _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, @@ -250,194 +153,6 @@ def export_baked_coarse_object_glbs( ) -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. diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_scene_processor.py b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_scene_processor.py new file mode 100644 index 000000000..33ee26a11 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_scene_processor.py @@ -0,0 +1,320 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +import re + +import numpy as np +import open3d as o3d +from scipy.spatial import ConvexHull, QhullError +from scipy.spatial.transform import Rotation +import trimesh + +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.utils.logger import log_info + + +@dataclass(frozen=True) +class SimReadySceneProcessorConfig: + """Object-category policy for SimReady mesh canonicalization.""" + + upright_container_id_tokens: frozenset[str] = frozenset( + {"bottle", "can", "jar", "flask", "thermos"} + ) # Object-id tokens that enable upright-container standardization. + + +class SimReadySceneProcessor: + """Create SimReady GLBs and layouts for one table and its scene assets.""" + + def __init__( + self, + *, + scene: Scene, + coarse_layout_by_id: dict[str, dict[str, object]], + coarse_geometry_root: str | Path, + simready_geometry_root: str | Path, + config: SimReadySceneProcessorConfig | None = None, + ) -> None: + self.scene = scene + self.coarse_layout_by_id = coarse_layout_by_id + self.coarse_geometry_root = Path(coarse_geometry_root).expanduser().resolve() + self.simready_geometry_root = ( + Path(simready_geometry_root).expanduser().resolve() + ) + self.simready_table_layout: dict[str, object] | None = None + self.simready_assets_layout: list[dict[str, object]] | None = None + self.config = config if config is not None else SimReadySceneProcessorConfig() + if not self.config.upright_container_id_tokens: + raise ValueError("upright_container_id_tokens must not be empty.") + + def process_table(self) -> dict[str, object]: + """Process the required scene table and return its SimReady layout.""" + if self.scene.table is None: + raise ValueError("Cannot SimReady a scene without a table.") + self.simready_table_layout = self._process_object( + object_id=self.scene.table.id, + object_role="table", + ) + return self.simready_table_layout + + def process_assets(self) -> list[dict[str, object]]: + """Process every scene asset and return SimReady layouts in scene order.""" + asset_ids: set[str] = set() + processed_assets: list[dict[str, object]] = [] + for asset in self.scene.assets: + if asset.id in asset_ids: + raise ValueError(f"Scene assets contain duplicate id {asset.id!r}.") + asset_ids.add(asset.id) + processed_assets.append( + self._process_object(object_id=asset.id, object_role="asset") + ) + self.simready_assets_layout = processed_assets + return self.simready_assets_layout + + def _process_object(self, *, object_id: str, object_role: str) -> dict[str, object]: + """Canonicalize one coarse object and write its SimReady GLB.""" + if object_role not in {"table", "asset"}: + raise ValueError(f"Unsupported SimReady object role {object_role!r}.") + coarse_layout = self.coarse_layout_by_id.get(object_id) + if coarse_layout is None: + raise ValueError(f"Coarse layout does not contain object {object_id!r}.") + simready_mesh, simready_transform = self._canonicalize_object_mesh( + coarse_glb_path=self.coarse_geometry_root / f"{object_id}.glb", + object_id=object_id, + rot=coarse_layout.get("rot"), + pos=coarse_layout.get("pos"), + scale=coarse_layout.get("scale"), + ) + output_path = self.simready_geometry_root / f"{object_id}.glb" + output_path.parent.mkdir(parents=True, exist_ok=True) + simready_mesh.export(output_path, file_type="glb") + if not output_path.is_file(): + raise FileNotFoundError( + f"SimReady {object_role} geometry was not written: {output_path}" + ) + log_info(f"Created SimReady {object_role}: {object_id!r}.") + return {"id": object_id, **simready_transform} + + def _canonicalize_object_mesh( + self, + *, + coarse_glb_path: str | Path, + object_id: str, + rot: object, + pos: object, + scale: object, + ) -> tuple[trimesh.Trimesh, dict[str, list[float]]]: + """Bake coarse scale and canonicalize one mesh's AABB bottom centre. + + Return the processed mesh and its updated layout transform without writing + a GLB file. The caller owns the output path and export. + """ + resolved_coarse_glb_path = Path(coarse_glb_path).expanduser().resolve() + if not resolved_coarse_glb_path.is_file(): + raise FileNotFoundError( + f"Coarse object geometry not found: {resolved_coarse_glb_path}" + ) + loaded_mesh = trimesh.load(resolved_coarse_glb_path, process=False) + if isinstance(loaded_mesh, trimesh.Scene): + mesh = loaded_mesh.dump(concatenate=True) + elif isinstance(loaded_mesh, trimesh.Trimesh): + mesh = loaded_mesh + else: + raise ValueError( + f"Coarse object geometry is not a mesh: {resolved_coarse_glb_path}" + ) + + coarse_rot = self._three_floats(rot, field_name="rot") + coarse_pos = np.asarray(self._three_floats(pos, field_name="pos"), dtype=float) + coarse_scale = np.asarray( + self._three_floats(scale, field_name="scale"), dtype=float + ) + if np.any(coarse_scale <= 0): + raise ValueError("Coarse object scale values must be positive.") + # We need the object id to determine whether it is a bottle-like object. + # If it does, then we will do a special standardization. (Hard code) + if not isinstance(object_id, str) or not object_id: + raise ValueError("Scene object id must be a non-empty string.") + + # GLB uses y-up. Convert its vertices to z-up while processing the geometry. + y_up_to_z_up_rotation = Rotation.from_euler("x", 90.0, degrees=True) + y_up_to_z_up_matrix = y_up_to_z_up_rotation.as_matrix() + y_up_to_z_up_transform = np.eye(4) + y_up_to_z_up_transform[:3, :3] = y_up_to_z_up_matrix + mesh.apply_transform(y_up_to_z_up_transform) + + # Standardize upright containers in temporary z-up coordinates before the + # shared center, scale, and bottom-center preprocessing. + # This is to ensure the action agent can pick up the bottle or can-like objects. + bottle_alignment_matrix = np.eye(3) + if self._is_upright_container_id(object_id): + bottle_alignment_matrix = self._standardize_bottle_z_up(mesh) + bottle_alignment_transform = np.eye(4) + bottle_alignment_transform[:3, :3] = bottle_alignment_matrix + mesh.apply_transform(bottle_alignment_transform) + + # First make the object's AABB center at the origin. + original_aabb_center = mesh.bounds.mean(axis=0) + mesh.apply_translation(-original_aabb_center) + + # Scale the object with the value in the coarse layout. + scale_transform = np.eye(4) + scale_transform[:3, :3] = ( + # Actually there's no need to do so, for the scale factor is all equal + # in x, y, z axes. + bottle_alignment_matrix + @ y_up_to_z_up_matrix + @ np.diag(coarse_scale) + @ y_up_to_z_up_matrix.T + @ bottle_alignment_matrix.T + ) + mesh.apply_transform(scale_transform) + + # Move the scaled object's AABB bottom center to the world's x-y plane (z=0). + scaled_bounds = mesh.bounds + scaled_aabb_bottom_center = np.array( + [ + (scaled_bounds[0, 0] + scaled_bounds[1, 0]) / 2, + (scaled_bounds[0, 1] + scaled_bounds[1, 1]) / 2, + scaled_bounds[0, 2], + ] + ) + mesh.apply_translation(-scaled_aabb_bottom_center) + + # Convert the processed GLB back to its standard y-up coordinate system. + z_up_to_y_up_transform = np.eye(4) + z_up_to_y_up_transform[:3, :3] = y_up_to_z_up_matrix.T + mesh.apply_transform(z_up_to_y_up_transform) + + # Compensate the bottle's local rotation so that its coarse world pose does + # not change. + local_bottle_rotation = Rotation.from_matrix( + y_up_to_z_up_matrix.T @ bottle_alignment_matrix @ y_up_to_z_up_matrix + ) + coarse_rotation_matrix = Rotation.from_euler( + "xyz", coarse_rot, degrees=True + ).as_matrix() + rotation = Rotation.from_matrix( + coarse_rotation_matrix @ local_bottle_rotation.inv().as_matrix() + ) + # Update the pos. + position_offset = y_up_to_z_up_matrix.T @ ( + scale_transform[:3, :3] @ original_aabb_center + scaled_aabb_bottom_center + ) + return mesh, { + "rot": rotation.as_euler("xyz", degrees=True).tolist(), + "pos": (coarse_pos + rotation.apply(position_offset)).tolist(), + "scale": [1.0, 1.0, 1.0], + } + + def _is_upright_container_id(self, object_id: str) -> bool: + """Return whether object-id tokens indicate a bottle-like container.""" + # Example: soda_can_0 + # tokens: {"soda", "can", "0"} + # upright_container_id_tokens: {"bottle", "can", "jar"} + # So this returns True because "can" is in the configured token set. + tokens = set(re.findall(r"[a-z0-9]+", object_id.lower())) + return bool(tokens & self.config.upright_container_id_tokens) + + @staticmethod + def _standardize_bottle_z_up(mesh: trimesh.Trimesh) -> np.ndarray: + """Return a proper rotation that maps a bottle-like mesh's long axis to z-up. + + Thanks to chenjian for this idea! + """ + if len(mesh.vertices) < 4 or len(mesh.faces) < 4: + raise ValueError( + "Bottle standardization requires a non-degenerate triangle mesh." + ) + open3d_mesh = o3d.geometry.TriangleMesh( + vertices=o3d.utility.Vector3dVector(mesh.vertices), + triangles=o3d.utility.Vector3iVector(mesh.faces), + ) + sampled_points = np.asarray( + open3d_mesh.sample_points_uniformly(number_of_points=10_000).points + ) # (10000, 3) x (x, y, z) + + # Check the number of the points again, and check whether have some + # non-finite values. + if sampled_points.shape[0] < 4 or not np.all(np.isfinite(sampled_points)): + raise ValueError( + "Bottle standardization could not sample valid mesh points." + ) + + centered_points = sampled_points - sampled_points.mean(axis=0) + # SVD find the longest axis. + _, _, principal_axes = np.linalg.svd(centered_points, full_matrices=False) + if np.linalg.det(principal_axes) < 0: + principal_axes[2, :] *= -1 # in case the SVD returns a reflection. + + bottle_rotation = Rotation.from_euler( + "y", 90.0, degrees=True + ).as_matrix() # 3x3 matrix + # The first PCA axis is the longest axis; rotate it onto the temporary z axis. + bottle_rotation = bottle_rotation @ principal_axes + standardized_points = (bottle_rotation @ centered_points.T).T + + axis_min = standardized_points[:, 2].min() + axis_max = standardized_points[:, 2].max() + axis_range = axis_max - axis_min + upper_points = standardized_points[ + standardized_points[:, 2] > axis_min + axis_range * 0.8 + ] + lower_points = standardized_points[ + standardized_points[:, 2] < axis_min + axis_range * 0.2 + ] + upper_volume = SimReadySceneProcessor._convex_hull_volume(upper_points) + lower_volume = SimReadySceneProcessor._convex_hull_volume(lower_points) + + # Bottles usually have a smaller top (neck) than bottom; flip if necessary. + if upper_volume > lower_volume: + bottle_rotation = ( + Rotation.from_euler("x", 180.0, degrees=True).as_matrix() + @ bottle_rotation + ) + return bottle_rotation + + @staticmethod + def _convex_hull_volume(points: np.ndarray) -> float: + """Return the volume of a non-degenerate point set's convex hull.""" + if points.shape[0] < 4: + raise ValueError( + "Bottle standardization needs at least four points per end." + ) + try: + return float(ConvexHull(points).volume) + except QhullError as exc: + raise ValueError( + "Bottle standardization found a degenerate end volume." + ) from exc + + @staticmethod + def _three_floats(value: object, *, field_name: str) -> list[float]: + """Validate and convert a three-value layout field to floats.""" + if not isinstance(value, list) or len(value) != 3: + raise ValueError( + f"Coarse layout field {field_name} must contain three values." + ) + try: + return [float(item) for item in value] + except (TypeError, ValueError) as exc: + raise ValueError( + f"Coarse layout field {field_name} must contain numeric values." + ) from exc From bbf2a0692b7249590b0bab22a521afc886c69027 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Tue, 4 Aug 2026 11:31:38 +0800 Subject: [PATCH 27/41] Modified the scene segmentation: 1. Bold the asset outline 2. Avoid the 2d table aabb index place on the gray asset mask --- .../pipeline/scene_segmentation.py | 7 +- .../utils/scene_segmentation_utils.py | 140 ++++++++++++++++-- 2 files changed, 133 insertions(+), 14 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/pipeline/scene_segmentation.py b/embodichain/gen_sim/scene_engine/pipeline/scene_segmentation.py index 1505a970f..0963b9732 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/scene_segmentation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/scene_segmentation.py @@ -22,6 +22,8 @@ import shutil from typing import Any +from PIL import Image + 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 @@ -125,7 +127,7 @@ def segment_scene( 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( + table_validation_image_path, asset_union_mask = render_image_without_masks( image_path=resolved_image_path, mask_paths=asset_mask_paths, output_path=Path(debug_output_root) / "table_validation_base.png", @@ -134,6 +136,7 @@ def segment_scene( _segment_table( image_path=resolved_image_path, validation_image_path=table_validation_image_path, + label_avoid_mask=asset_union_mask, debug_output_root=debug_output_root, masks_output_root=masks_output_root, scene=scene, @@ -151,6 +154,7 @@ def segment_scene( def _segment_table( image_path: str | Path, validation_image_path: str | Path, + label_avoid_mask: Image.Image, debug_output_root: str | Path, masks_output_root: str | Path, scene: Scene, @@ -191,6 +195,7 @@ def _segment_table( candidates_image_path = render_numbered_mask_candidates( image_path=validation_image_path, candidates=candidates, + label_avoid_mask=label_avoid_mask, output_path=( Path(debug_output_root) / f"table_candidates_{prompt_label}.png" # Render with prompt label, for easily debug. 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 7c88d62a5..e0c3f772a 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 @@ -23,6 +23,8 @@ from PIL import Image, ImageChops, ImageDraw, ImageFilter, ImageFont +from embodichain.utils.logger import log_warning + @dataclass(frozen=True) class MaskCandidate: @@ -160,8 +162,8 @@ def render_image_without_masks( 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.""" +) -> tuple[Path, Image.Image]: + """Gray masked regions and return the image path with their combined mask.""" image = Image.open(image_path).convert("RGB") ignored_mask = Image.new("L", image.size, 0) for mask_path in mask_paths: @@ -174,7 +176,7 @@ def render_image_without_masks( 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 + return resolved_output_path, ignored_mask def render_numbered_mask_candidates( @@ -183,11 +185,13 @@ def render_numbered_mask_candidates( candidates: list[MaskCandidate], output_path: str | Path, mask_style: str = "fill", + label_avoid_mask: Image.Image | None = None, ) -> Path: """Overlay numbered mask candidates on their source image. Notice that: - mask_style can be either "fill" or "outline". - The label font and its background scale with the source image resolution. + - label_avoid_mask keeps labels outside known occluding regions. """ if mask_style not in {"fill", "outline"}: raise ValueError("mask_style must be 'fill' or 'outline'.") @@ -221,18 +225,51 @@ def render_numbered_mask_candidates( draw = ImageDraw.Draw(overlay) # Initialize a draw object. font = _load_label_font(image.size) + if label_avoid_mask is not None: + label_blocked_mask = label_avoid_mask.convert("L") + _require_image_size(label_blocked_mask, image.size) + else: + label_blocked_mask = None + label_occupied_mask = Image.new("L", image.size, 0) for candidate, mask in decoded_masks: bbox = mask.getbbox() if bbox is None: raise ValueError( f"Image Segmentation Server candidate {candidate.index} has an empty mask." ) + center = ( + ((bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2) + if label_blocked_mask is None + else _find_label_center_inside_bbox( + candidate_mask=mask, + candidate_bbox=bbox, + blocked_mask=ImageChops.lighter( + label_blocked_mask, label_occupied_mask + ), + label=str(candidate.index), + font=font, + ) + ) + if center is None: + # Keep every candidate selectable even when its AABB is entirely + # occluded. This is preferable to silently omitting its number. + log_warning( + "Could not place table-candidate label %s outside masked objects " + "while keeping it inside the candidate AABB; using the AABB center.", + candidate.index, + ) + center = ((bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2) + label_bounds = _number_label_bounds( + draw=draw, label=str(candidate.index), center=center, font=font + ) _draw_number_label( draw=draw, label=str(candidate.index), - center=((bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2), + center=center, font=font, ) + if label_blocked_mask is not None: + ImageDraw.Draw(label_occupied_mask).rectangle(label_bounds, fill=255) resolved_output_path = Path(output_path).expanduser().resolve() resolved_output_path.parent.mkdir(parents=True, exist_ok=True) @@ -248,9 +285,67 @@ def _require_image_size(mask: Image.Image, image_size: tuple[int, int]) -> None: ) +def _find_label_center_inside_bbox( + *, + candidate_mask: Image.Image, + candidate_bbox: tuple[int, int, int, int], + blocked_mask: Image.Image, + label: str, + font: ImageFont.ImageFont, +) -> tuple[float, float] | None: + """Find the nearest label centre that stays in a candidate AABB and avoids masks.""" + probe_draw = ImageDraw.Draw(Image.new("RGBA", candidate_mask.size)) + bounds_at_origin = _number_label_bounds( + draw=probe_draw, label=label, center=(0.0, 0.0), font=font + ) + minimum_x = candidate_bbox[0] - bounds_at_origin[0] + maximum_x = candidate_bbox[2] - 1 - bounds_at_origin[2] + minimum_y = candidate_bbox[1] - bounds_at_origin[1] + maximum_y = candidate_bbox[3] - 1 - bounds_at_origin[3] + if minimum_x > maximum_x or minimum_y > maximum_y: + return None + + # Expand each blocked region by the label's largest half-extent, so testing + # one candidate centre guarantees the complete label rectangle stays clear. + label_radius = max( + -bounds_at_origin[0], + bounds_at_origin[2], + -bounds_at_origin[1], + bounds_at_origin[3], + ) + blocked_centres = blocked_mask.filter(ImageFilter.MaxFilter(2 * label_radius + 1)) + bbox_center = ( + (candidate_bbox[0] + candidate_bbox[2]) / 2, + (candidate_bbox[1] + candidate_bbox[3]) / 2, + ) + for require_candidate_mask in (True, False): + for step in (max(1, round(min(candidate_mask.size) / 512)), 1): + best_center: tuple[float, float] | None = None + best_distance_squared = float("inf") + for y_coordinate in range(minimum_y, maximum_y + 1, step): + for x_coordinate in range(minimum_x, maximum_x + 1, step): + if blocked_centres.getpixel((x_coordinate, y_coordinate)): + continue + if require_candidate_mask and not candidate_mask.getpixel( + (x_coordinate, y_coordinate) + ): + continue + distance_squared = (x_coordinate - bbox_center[0]) ** 2 + ( + y_coordinate - bbox_center[1] + ) ** 2 + if distance_squared < best_distance_squared: + best_center = (float(x_coordinate), float(y_coordinate)) + best_distance_squared = distance_squared + if best_center is not None: + return best_center + return None + + def _mask_outer_outline(mask: Image.Image, image_size: tuple[int, int]) -> Image.Image: """Use dilation and subtraction to get the outer outline of a binary mask.""" - outline_width = max(1, round(min(image_size) / 400)) + # Asset-candidate images use outlines only; keep them visible at common + # image resolutions without obscuring the original object appearance. + outline_width = max(2, round(min(image_size) / 200)) dilated_mask = mask.filter(ImageFilter.MaxFilter(outline_width * 2 + 1)) return ImageChops.subtract(dilated_mask, mask) @@ -320,21 +415,40 @@ def _draw_number_label( font: ImageFont.ImageFont, ) -> None: """Draw a numbered label with red background and white text at the given center position.""" + label_bounds = _number_label_bounds( + draw=draw, label=label, center=center, font=font + ) label_box = draw.textbbox((0, 0), label, font=font) label_width = label_box[2] - label_box[0] label_height = label_box[3] - label_box[1] - 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, - ), + label_bounds, fill=(220, 0, 0, 255), outline=(255, 255, 255, 255), - width=max(1, padding // 3), + width=max(1, round(max(label_width, label_height) / 12)), ) draw.text((x, y), label, fill=(255, 255, 255, 255), font=font) + + +def _number_label_bounds( + *, + draw: ImageDraw.ImageDraw, + label: str, + center: tuple[float, float], + font: ImageFont.ImageFont, +) -> tuple[int, int, int, int]: + """Return the red label rectangle bounds for a label centre.""" + label_box = draw.textbbox((0, 0), label, font=font) + label_width = label_box[2] - label_box[0] + label_height = label_box[3] - label_box[1] + padding = max(4, round(max(label_width, label_height) / 4)) + x = center[0] - label_width / 2 + y = center[1] - label_height / 2 + return ( + round(x - padding), + round(y - padding), + round(x + label_width + padding), + round(y + label_height + padding), + ) From c45be975654f48b6d474ed791d7fc073f2066540 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Tue, 4 Aug 2026 11:57:47 +0800 Subject: [PATCH 28/41] Reformatted the scene export --- .../gen_sim/scene_engine/pipeline/generate.py | 7 +- .../scene_engine/pipeline/scene_export.py | 220 ---------------- .../pipeline/utils/scene_exporter.py | 249 ++++++++++++++++++ 3 files changed, 252 insertions(+), 224 deletions(-) delete mode 100644 embodichain/gen_sim/scene_engine/pipeline/scene_export.py create mode 100644 embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py diff --git a/embodichain/gen_sim/scene_engine/pipeline/generate.py b/embodichain/gen_sim/scene_engine/pipeline/generate.py index 92313f864..a92649a75 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.scene_export import export_scene +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_exporter import SceneExporter def generate_scene_from_image( @@ -110,12 +110,11 @@ def generate_scene_from_image( # 4. Scene Export log_info("Starting Scene Export") - export_scene( + scene_exporter = SceneExporter( scene=scene, output_root=resolved_output_root, - table_max_convex_hull_num=16, - asset_max_convex_hull_num=16, ) + scene_exporter.export() log_info("Completed Scene Export") return scene diff --git a/embodichain/gen_sim/scene_engine/pipeline/scene_export.py b/embodichain/gen_sim/scene_engine/pipeline/scene_export.py deleted file mode 100644 index 7593a3c95..000000000 --- a/embodichain/gen_sim/scene_engine/pipeline/scene_export.py +++ /dev/null @@ -1,220 +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 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( - *, - 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 a scene-only config and copy SimReady GLBs into ``mesh_assets``. - - Scene layouts are y-up. The simulator automatically converts each y-up GLB - to z-up, so this exporter copies each GLB unchanged and converts only its - world position and rotation for ``init_pos`` and ``init_rot``. ``body_scale`` - remains the original y-up scale associated with the GLB. This is not a - complete ``EmbodiedEnv``/``run-env`` configuration because a generated - scene does not determine a robot, its placement, or its control setup. - """ - if scene.table is None: - 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", - ) - 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() / "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("Scene export requires unique table and asset ids.") - - exported_entries = { - scene_object.id: _copy_scene_object_to_assets( - scene_object=scene_object, - mesh_assets_root=mesh_assets_root, - ) - for scene_object in scene_objects - } - scene_config = { - "format": "embodichain.scene-export/v1", - # This identifies the exported scene data only. It is deliberately not - # a Gymnasium environment ID because scene exports do not register or - # instantiate an EmbodiedEnv. - "scene_id": f"scene-engine-{int(time.time() * 1000)}", - "background": [ - _scene_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": [ - _scene_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 - ], - } - 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 scene_config_path - - -def _copy_scene_object_to_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 _scene_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 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") - - 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/utils/scene_exporter.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py new file mode 100644 index 000000000..3192a0a9e --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py @@ -0,0 +1,249 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from dataclasses import dataclass +import json +from 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 +from embodichain.utils.logger import log_info + +_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, +) + + +@dataclass(frozen=True) +class SceneExporterConfig: + """Collision-decomposition controls for scene export.""" + + table_max_convex_hull_num: int = _DEFAULT_MAX_CONVEX_HULL_NUM # Table hull limit. + asset_max_convex_hull_num: int = _DEFAULT_MAX_CONVEX_HULL_NUM # Asset hull limit. + + +class SceneExporter: + """Write one generated scene and its SimReady meshes as a scene export.""" + + def __init__( + self, + *, + scene: Scene, + output_root: str | Path, + config: SceneExporterConfig | None = None, + ) -> None: + self.scene = scene + self.output_root = Path(output_root).expanduser().resolve() + self.export_root = self.output_root / "scene_export" + self.scene_config_path: Path | None = None + self.config = config if config is not None else SceneExporterConfig() + self.table_max_convex_hull_num = _positive_int( + self.config.table_max_convex_hull_num, + field_name="table_max_convex_hull_num", + ) + self.asset_max_convex_hull_num = _positive_int( + self.config.asset_max_convex_hull_num, + field_name="asset_max_convex_hull_num", + ) + + def export(self) -> Path: + """Write a scene-only config and copy SimReady GLBs into ``mesh_assets``. + + Scene layouts are y-up. The simulator automatically converts each y-up + GLB to z-up, so this exporter copies each GLB unchanged and converts + only its world position and rotation for ``init_pos`` and ``init_rot``. + ``body_scale`` remains the original y-up scale associated with the GLB. + This is not a complete ``EmbodiedEnv``/``run-env`` configuration because + a generated scene does not determine a robot, its placement, or control. + """ + if self.scene.table is None: + raise ValueError("Cannot export a scene without a table.") + + mesh_assets_root = self.export_root / "mesh_assets" + mesh_assets_root.mkdir(parents=True, exist_ok=True) + scene_objects = [self.scene.table, *self.scene.assets] + object_ids = [scene_object.id for scene_object in scene_objects] + if len(set(object_ids)) != len(object_ids): + raise ValueError("Scene export requires unique table and asset ids.") + + exported_entries = { + scene_object.id: self._copy_scene_object_to_assets( + scene_object=scene_object, + mesh_assets_root=mesh_assets_root, + ) + for scene_object in scene_objects + } + scene_config = { + "format": "embodichain.scene-export/v1", + # This identifies the exported scene data only. It is deliberately not + # a Gymnasium environment ID because scene exports do not register or + # instantiate an EmbodiedEnv. + "scene_id": f"scene-engine-{int(time.time() * 1000)}", + "background": [ + self._scene_object_config( + scene_object=self.scene.table, + asset_relative_path=exported_entries[self.scene.table.id], + body_type="kinematic", + attrs=_TABLE_PHYSICS_ATTRS, + max_convex_hull_num=self.table_max_convex_hull_num, + ) + ], + "rigid_object": [ + self._scene_object_config( + scene_object=asset, + asset_relative_path=exported_entries[asset.id], + body_type="dynamic", + attrs=_ASSET_PHYSICS_ATTRS, + max_convex_hull_num=self.asset_max_convex_hull_num, + ) + for asset in self.scene.assets + ], + } + self.scene_config_path = self.export_root / "scene_config.json" + self.scene_config_path.write_text( + json.dumps(scene_config, indent=2, ensure_ascii=False) + "\n", + encoding="utf-8", + ) + log_info(f"Exported scene config: {self.scene_config_path}") + return self.scene_config_path + + @staticmethod + def _copy_scene_object_to_assets( + *, + scene_object: 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( + "SimReady GLB for scene object " + f"{object_id!r} not found: {source_glb_path}" + ) + destination_glb_path = mesh_assets_root / object_id / f"{object_id}.glb" + destination_glb_path.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(source_glb_path, destination_glb_path) + return destination_glb_path.relative_to(mesh_assets_root.parent).as_posix() + + @staticmethod + def _scene_object_config( + *, + scene_object: 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 scene-only object config from a final y-up object.""" + pos_y_up = SceneExporter._scene_vector(scene_object, "pos") + rot_y_up = SceneExporter._scene_vector(scene_object, "rot") + scale_y_up = SceneExporter._scene_vector(scene_object, "scale") + + 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, + } + + @staticmethod + 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 " + f"{field_name!r} vector." + ) + vector = [float(value) for value in values] + if not np.all(np.isfinite(vector)): + raise ValueError( + f"Scene object {scene_object.id!r} has non-finite " + f"{field_name!r}." + ) + return vector + + +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 + From 2436a375c6d55905061f7952676ed18f2fe50369 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Tue, 4 Aug 2026 14:41:38 +0800 Subject: [PATCH 29/41] Update the datastructure + Reformatted the pipeline --- .../gen_sim/scene_engine/core/asset.py | 51 -- .../gen_sim/scene_engine/core/scene.py | 27 +- .../gen_sim/scene_engine/core/scene_object.py | 85 +++ .../gen_sim/scene_engine/core/table.py | 51 -- .../gen_sim/scene_engine/pipeline/generate.py | 32 +- .../scene_engine/pipeline/scene_generation.py | 143 ++--- .../pipeline/scene_segmentation.py | 484 ----------------- .../pipeline/scene_understanding.py | 505 +++++++++++++++++- .../pipeline/utils/assets_gravity_settler.py | 71 ++- ...n_utils.py => image_segmentation_utils.py} | 0 .../pipeline/utils/scene_exporter.py | 84 +-- .../pipeline/utils/scene_generation_utils.py | 53 -- .../utils/simready_scene_processor.py | 53 +- 13 files changed, 743 insertions(+), 896 deletions(-) delete mode 100644 embodichain/gen_sim/scene_engine/core/asset.py create mode 100644 embodichain/gen_sim/scene_engine/core/scene_object.py delete mode 100644 embodichain/gen_sim/scene_engine/core/table.py delete mode 100644 embodichain/gen_sim/scene_engine/pipeline/scene_segmentation.py rename embodichain/gen_sim/scene_engine/pipeline/utils/{scene_segmentation_utils.py => image_segmentation_utils.py} (100%) diff --git a/embodichain/gen_sim/scene_engine/core/asset.py b/embodichain/gen_sim/scene_engine/core/asset.py deleted file mode 100644 index 81306d329..000000000 --- a/embodichain/gen_sim/scene_engine/core/asset.py +++ /dev/null @@ -1,51 +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 - -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 index ed41aa67a..79b2e4f2b 100644 --- a/embodichain/gen_sim/scene_engine/core/scene.py +++ b/embodichain/gen_sim/scene_engine/core/scene.py @@ -18,19 +18,28 @@ 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 +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject @dataclass class Scene: - """A scene containing a table and zero or more assets.""" + """A scene containing one table object and zero or more asset objects.""" - table: Table | None = None - assets: list[Asset] = field(default_factory=list) + objects: list[SceneObject] = field(default_factory=list) + + @property + def table(self) -> SceneObject | None: + """Return the sole table object, or ``None`` before understanding.""" + tables = [scene_object for scene_object in self.objects if scene_object.kind == "table"] + if len(tables) > 1: + raise ValueError("A scene may contain only one table object.") + return tables[0] if tables else None + + @property + def assets(self) -> list[SceneObject]: + """Return movable asset objects in their scene order.""" + return [scene_object for scene_object in self.objects if scene_object.kind == "asset"] def to_dict(self) -> dict[str, object]: - return { - "table": self.table.to_dict() if self.table is not None else None, - "assets": [asset.to_dict() for asset in self.assets], - } + """Serialize the canonical object collection for debugging artifacts.""" + return {"objects": [scene_object.to_dict() for scene_object in self.objects]} diff --git a/embodichain/gen_sim/scene_engine/core/scene_object.py b/embodichain/gen_sim/scene_engine/core/scene_object.py new file mode 100644 index 000000000..d9a837e87 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/core/scene_object.py @@ -0,0 +1,85 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Literal + + +@dataclass +class ObjectPhysics: + """Physics and collision settings shared by settling and scene export.""" + + body_type: Literal["dynamic", "kinematic"] # Runtime behaviour in simulation. + attrs: dict[str, float | int] # Rigid-body material and contact attributes. + max_convex_hull_num: int # Collision-decomposition hull budget. + + def __post_init__(self) -> None: + """Validate physics settings before a later stage consumes them.""" + if self.body_type not in {"dynamic", "kinematic"}: + raise ValueError("body_type must be 'dynamic' or 'kinematic'.") + if self.max_convex_hull_num <= 0: + raise ValueError("max_convex_hull_num must be positive.") + if not self.attrs: + raise ValueError("attrs must contain at least one physics attribute.") + if not all( + isinstance(name, str) and isinstance(value, (float, int)) + for name, value in self.attrs.items() + ): + raise ValueError("attrs must map strings to numeric physics values.") + + def to_dict(self) -> dict[str, object]: + """Serialize the physics settings for scene debugging artifacts.""" + return { + "body_type": self.body_type, + "attrs": self.attrs, + "max_convex_hull_num": self.max_convex_hull_num, + } + + +@dataclass +class SceneObject: + """One semantic object progressing through the Scene Engine pipeline.""" + + id: str # Stable scene-unique identifier. + kind: Literal["table", "asset"] # Table support body or movable scene asset. + category: str # Semantic category identified by scene understanding. + name: str # Human-readable visual name. + description: str # Detailed semantic and spatial description. + mask_path: str | None = None # Absolute path to the validated binary image mask. + simready_glb_path: str | None = None # Absolute path to the canonical SimReady GLB. + rot: list[float] | None = None # Final y-up Euler XYZ rotation in degrees. + pos: list[float] | None = None # Final y-up world position in metres. + scale: list[float] | None = None # Final y-up object scale. + physics: ObjectPhysics | None = None # Assigned when SimReady processing succeeds. + + def to_dict(self) -> dict[str, object]: + """Serialize this object and its currently available pipeline artifacts.""" + return { + "id": self.id, + "kind": self.kind, + "category": self.category, + "name": self.name, + "description": self.description, + "mask_path": self.mask_path, + "simready_glb_path": self.simready_glb_path, + "rot": self.rot, + "pos": self.pos, + "scale": self.scale, + "physics": self.physics.to_dict() if self.physics is not None else None, + } diff --git a/embodichain/gen_sim/scene_engine/core/table.py b/embodichain/gen_sim/scene_engine/core/table.py deleted file mode 100644 index bab0f94fb..000000000 --- a/embodichain/gen_sim/scene_engine/core/table.py +++ /dev/null @@ -1,51 +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 - -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/pipeline/generate.py b/embodichain/gen_sim/scene_engine/pipeline/generate.py index a92649a75..aa4f9bf5b 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generate.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generate.py @@ -22,10 +22,6 @@ 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, ) @@ -33,9 +29,6 @@ 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.utils.logger import log_info from embodichain.gen_sim.scene_engine.pipeline.scene_generation import ( @@ -67,29 +60,11 @@ def generate_scene_from_image( image_path=image_path, output_root=resolved_output_root, vlm_client=vlm_client, + image_segmentation_config_path=image_segmentation_config_path, ) log_info("Completed Scene Understanding") - # 2. 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 - ) - 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_info("Completed Scene Segmentation") - - # 3. Objects + Coarse Layout Generation + # 2. 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( @@ -101,14 +76,13 @@ def generate_scene_from_image( 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_info("Completed Objects + Coarse Layout Generation") - # 4. Scene Export + # 3. Scene Export log_info("Starting Scene Export") scene_exporter = SceneExporter( scene=scene, diff --git a/embodichain/gen_sim/scene_engine/pipeline/scene_generation.py b/embodichain/gen_sim/scene_engine/pipeline/scene_generation.py index b8174f7bd..d14ce364f 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/scene_generation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/scene_generation.py @@ -27,12 +27,8 @@ 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.core.scene_object import SceneObject from embodichain.gen_sim.scene_engine.pipeline.utils.assets_group_support_clamp import ( AssetsGroupSupportClamp, ) @@ -46,7 +42,6 @@ AssetsGravitySettler, ) from embodichain.gen_sim.scene_engine.pipeline.utils.scene_generation_utils import ( - export_baked_layout_object_glbs, layout_object_to_transform_matrix, load_glb_mesh, quaternion_wxyz_to_euler_xyz_degrees, @@ -68,7 +63,6 @@ def generate_scene_and_refine( output_root: str | Path, scene: Scene, *, - vlm_client: OpenAICompatibleVLM, geometry_generation_client: GeometryGenerationClient, ) -> Scene: @@ -98,18 +92,35 @@ def generate_scene_and_refine( 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, + # Simready all the assets(includes table). + # Treat table and assets seperately. + coarse_layout = _load_layout(coarse_geometry_output_root / "coarse_layout.json") + coarse_layout_by_id = { + layout_object["id"]: layout_object for layout_object in coarse_layout + } + simready_processor = SimReadySceneProcessor( scene=scene, - vlm_client=vlm_client, + coarse_layout_by_id=coarse_layout_by_id, + coarse_geometry_root=coarse_geometry_output_root, + simready_geometry_root=simready_geometry_output_root, + ) + simready_assets_layout = simready_processor.process_assets() + simready_table_layout = simready_processor.process_table() + # Concat then save the table info and the assets info in one JSON file. + simready_layout = [simready_table_layout, *simready_assets_layout] + (simready_geometry_output_root / "simready_layout.json").write_text( + json.dumps(simready_layout, indent=2, ensure_ascii=False) + "\n", + encoding="utf-8", + ) + + # Layout refinement will start with the table. + refined_table_layout, refined_assets_layout = _layout_refinement( + scene=scene, # Update this data structure internally. + simready_geometry_output_root=simready_geometry_output_root, # Contains simready assets and their current coarse layout JSON. + debug_output_root=debug_output_root, # Keep the table support surface info + optimized layout info (render with matplotlib) for debugging. ) # Write the Updated scene JSON for debugging. @@ -126,7 +137,6 @@ def _generate_coarse_results_from_masks( coarse_geometry_output_root: str | Path, scene: Scene, *, - vlm_client: OpenAICompatibleVLM, geometry_generation_client: GeometryGenerationClient, ) -> None: @@ -185,97 +195,6 @@ def _generate_coarse_results_from_masks( 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_processor = SimReadySceneProcessor( - scene=scene, - coarse_layout_by_id=coarse_layout_by_id, - coarse_geometry_root=coarse_geometry_output_root, - simready_geometry_root=simready_geometry_output_root, - ) - simready_assets_layout = simready_processor.process_assets() - simready_table_layout = simready_processor.process_table() - # Concat then save the table info and the assets info in one JSON file. - simready_layout = [simready_table_layout, *simready_assets_layout] - (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, @@ -306,7 +225,7 @@ def _update_scene_final_y_up_layout( def _copy_y_up_layout_to_scene_object( - scene_object: Table | Asset, + scene_object: SceneObject, layout_object: dict[str, object], ) -> None: """Copy one y-up layout object after validating its id and numeric vectors.""" @@ -335,7 +254,6 @@ 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 @@ -482,12 +400,19 @@ def _layout_refinement( # Notice that: we do not consider the assets like a bottle, which should be standing on the table but laid down # after the simulation. gravity_settler = AssetsGravitySettler( + scene=scene, table_layout=refined_table_layout, assets_layout=refined_assets_layout, geometry_root=simready_geometry_output_root, ) refined_assets_layout = gravity_settler.settle() + # Update the scene data structure with the final y-up layout values. + _update_scene_final_y_up_layout( + scene=scene, + table_layout=refined_table_layout, + assets_layout=refined_assets_layout, + ) return refined_table_layout, refined_assets_layout diff --git a/embodichain/gen_sim/scene_engine/pipeline/scene_segmentation.py b/embodichain/gen_sim/scene_engine/pipeline/scene_segmentation.py deleted file mode 100644 index 0963b9732..000000000 --- a/embodichain/gen_sim/scene_engine/pipeline/scene_segmentation.py +++ /dev/null @@ -1,484 +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 json -from pathlib import Path -import shutil -from typing import Any - -from PIL import Image - -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, asset_union_mask = 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, - label_avoid_mask=asset_union_mask, - debug_output_root=debug_output_root, - masks_output_root=masks_output_root, - scene=scene, - vlm_client=vlm_client, - image_segmentation_client=image_segmentation_client, - ) - # 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, - label_avoid_mask: Image.Image, - debug_output_root: str | Path, - masks_output_root: str | Path, - scene: Scene, - *, - vlm_client: OpenAICompatibleVLM, - image_segmentation_client: ImageSegmentationClient, -) -> None: - """Segment the table. (Now it only supports segment the complete tabletop)""" - if scene.table is None: - raise ValueError("Cannot segment a scene without a table.") - - table = scene.table - # Build the segmentation prompts for table. - for prompt_label, prompt in ( - ("name", table.name), - ("description", table.description), - ("table", "table"), - ("plane", "plane"), - ): - candidates = union_overlapping_mask_candidates( - build_mask_candidates( - image_segmentation_client.segment_single_object( - image_path=image_path, - prompt=prompt, - ) - ), - min_iou=0.8, # Union masks who have iou > 0.8 - ) - # If do not have candidate, then try segment the table with description, "table", "plane"... - # Notice that, this part could be extended with other segmentation prompt like - # a board, or newly-generated prompt from another VLM-calling etc. - if not candidates: - continue - - # Maybe the mask count = 1, but not correct; - # Maybe the mask count > 1; - # Thus, we need to validate with an VLM. - candidates_image_path = render_numbered_mask_candidates( - image_path=validation_image_path, - candidates=candidates, - label_avoid_mask=label_avoid_mask, - output_path=( - Path(debug_output_root) - / f"table_candidates_{prompt_label}.png" # Render with prompt label, for easily debug. - ), - ) - selected_mask_index = _validate_table_candidates_with_vlm( - table=table, - candidates=candidates, - candidates_image_path=candidates_image_path, - vlm_client=vlm_client, - ) - if selected_mask_index is None: - continue - - # Save result. - candidate = _candidate_by_index(candidates, selected_mask_index) - table.mask_path = str( - save_binary_mask( - candidate, - image_size=_image_size(image_path), - output_path=Path(masks_output_root) / "table_mask.png", - ) - ) - return - - raise ValueError("Unable to find a VLM-validated segmentation mask for the table.") - - -def _validate_table_candidates_with_vlm( - *, - table: 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 index 6e31770c3..82eaa2ffc 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/scene_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/scene_understanding.py @@ -21,13 +21,26 @@ from pathlib import Path import re import shutil +from typing import Any -from embodichain.gen_sim.scene_engine.core.asset import Asset +from PIL import Image + +from embodichain.gen_sim.scene_engine.clients.image_segmentation import ( + ImageSegmentationClient, +) from embodichain.gen_sim.scene_engine.core.scene import Scene -from embodichain.gen_sim.scene_engine.core.table import Table +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( OpenAICompatibleVLM, ) +from embodichain.gen_sim.scene_engine.pipeline.utils.image_segmentation_utils import ( + MaskCandidate, + build_mask_candidates, + render_image_without_masks, + render_numbered_mask_candidates, + save_binary_mask, + union_overlapping_mask_candidates, +) _SUPPORTED_IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png"} _CATEGORY_PATTERN = re.compile(r"^[a-z][a-z0-9_]*$") @@ -85,6 +98,50 @@ _USER_PROMPT = "Analyze the provided image and return only the required JSON object." +_TABLE_VALIDATION_SYSTEM_PROMPT = """You select the best table mask candidate. +The image contains table-mask candidates overlaid semi-transparently on the +scene. Gray regions are already-segmented non-table assets that were +intentionally removed for this validation; ignore them. Candidate numbers only +identify masks; do not treat the number or its background as scene content. + +Choose the candidate covering the main visible table. A table candidate is +acceptable when it covers the visible tabletop and/or legs, even if some edges +are incomplete, objects on the table occlude parts of it, or it slightly +overlaps those objects. Return null only when no candidate depicts the main +table. If there is one plausible candidate, select it rather than returning +null. + +Examples: +- Candidate 1 covers the tabletop and legs but misses a narrow edge: + {"selected_mask_index": 1} +- Candidate 1 is a cup and candidate 2 covers the main table: + {"selected_mask_index": 2} +- Every candidate is an object resting on the table, not the table itself: + {"selected_mask_index": null} + +Return JSON only, with exactly one key: selected_mask_index. Use a one-based +candidate index or null. Do not include Markdown or any other text.""" +_ASSET_ASSIGNMENT_SYSTEM_PROMPT = """You assign outlined mask candidates to a group of scene assets. +The image is the original scene with numbered candidate mask outlines. The +number labels identify candidates only; they are not scene content. Use the +provided category, name, and description of every asset to match each asset to +exactly one candidate. Descriptions can distinguish visually similar assets by +location. + +Extra candidate masks are normal and may be ignored. Never force a candidate +onto an asset. If any listed asset has no correct candidate, return +{"assignments": null}. + +Examples: +- Two listed paper cups match candidate 1 and candidate 3: + {"assignments": [{"asset_id": "paper_cup_001", "mask_index": 1}, {"asset_id": "paper_cup_002", "mask_index": 3}]} +- A listed asset is absent from every candidate: + {"assignments": null} + +Return JSON only, with exactly one key: assignments. It must be null or an +array of asset_id and mask_index objects. Do not include Markdown or any other +text.""" + def understand_scene( scene: Scene, @@ -92,12 +149,10 @@ def understand_scene( output_root: str | Path, *, vlm_client: OpenAICompatibleVLM, + image_segmentation_config_path: str | Path | None = None, 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. @@ -106,37 +161,75 @@ def understand_scene( shutil.rmtree(stage_output_root) stage_output_root.mkdir(parents=True, exist_ok=True) + _analyze_image_objects( # Update the scene data structure internally. + scene=scene, + image_path=resolved_image_path, + vlm_client=vlm_client, + json_max_attempts=json_max_attempts, + ) + + # Load the config and fail if the Image Segmentation Server is unavailable. + image_segmentation_client = ImageSegmentationClient.from_config( + image_segmentation_config_path + ) + try: + image_segmentation_client.check_health() # Error raising will happen internally. + _segment_scene( + image_path=resolved_image_path, + stage_output_root=stage_output_root, + scene=scene, + vlm_client=vlm_client, + image_segmentation_client=image_segmentation_client, + ) + finally: + image_segmentation_client.close() # Kill the session to avoid resource leaks. + + # Write the Updated scene JSON for debugging. + (stage_output_root / "scene.json").write_text( + json.dumps(scene.to_dict(), indent=2, ensure_ascii=False) + "\n", + encoding="utf-8", + ) + return scene + + +def _analyze_image_objects( + *, + scene: Scene, + image_path: str | Path, + vlm_client: OpenAICompatibleVLM, + json_max_attempts: int = 3, +) -> None: + """Analyze one image and update ``scene`` with validated semantic objects.""" + if json_max_attempts < 1: + raise ValueError("json_max_attempts must be at least 1.") + + resolved_image_path = _validate_image_path(image_path) last_validation_error: ValueError | None = None - for attempt in range(1, json_max_attempts + 1): + for _ in range(json_max_attempts): response_text = vlm_client.complete( image_path=resolved_image_path, system_prompt=_SYSTEM_PROMPT, user_prompt=_USER_PROMPT, ) try: - understood_scene = validate_scene_understanding_json(response_text) - scene.table = understood_scene.table - scene.assets = understood_scene.assets - validate_scene_understanding(scene) + analyzed_scene = _parse_image_object_analysis_response(response_text) + validate_scene_understanding(analyzed_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 + scene.objects = analyzed_scene.objects + return None assert last_validation_error is not None raise ValueError( - "VLM returned invalid scene-understanding JSON after " + "VLM returned invalid image-object analysis 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.""" +def _parse_image_object_analysis_response(response_text: str) -> Scene: + """Parse one VLM image-object analysis response into a semantic ``Scene``.""" json_text = _strip_json_code_fence(response_text) try: payload = json.loads(json_text) @@ -148,26 +241,28 @@ def validate_scene_understanding_json(response_text: str) -> Scene: id_counters: dict[str, int] = {} table_fields = _parse_scene_object_fields(payload["table"], field_name="table") - table = Table( + table = SceneObject( # id=_next_id(table_fields["category"], id_counters) # Use a fixed ID for the table. id="table", + kind="table", **table_fields, ) assets_value = payload["assets"] if not isinstance(assets_value, list): raise ValueError("VLM JSON key assets must be an array.") - assets: list[Asset] = [] + assets: list[SceneObject] = [] for index, asset in enumerate(assets_value): fields = _parse_scene_object_fields(asset, field_name=f"assets[{index}]") assets.append( - Asset( + SceneObject( id=_next_id(fields["category"], id_counters), + kind="asset", **fields, ) ) - return Scene(table=table, assets=assets) + return Scene(objects=[table, *assets]) def validate_scene_understanding(scene: Scene) -> None: @@ -252,3 +347,369 @@ def _next_id(category: str, counters: dict[str, int]) -> str: """Auto increment an ID for the same category, e.g. mug_001, mug_002, etc.""" counters[category] = counters.get(category, 0) + 1 return f"{category}_{counters[category]:03d}" + + +def _segment_scene( + *, + image_path: str | Path, + stage_output_root: str | Path, + scene: Scene, + vlm_client: OpenAICompatibleVLM, + image_segmentation_client: ImageSegmentationClient, +) -> None: + """Add validated table and asset mask paths to a semantic scene.""" + debug_output_root = ( + Path(stage_output_root) / "debug" + ) # Keeps the mask debug images. + masks_output_root = ( + Path(stage_output_root) / "masks" + ) # Keeps the validated masked images of each assets (include the table) + debug_output_root.mkdir() + masks_output_root.mkdir() + + # Segment the table and assets with VLM validation separately. + _segment_assets( + image_path=image_path, + debug_output_root=debug_output_root, + masks_output_root=masks_output_root, + scene=scene, + vlm_client=vlm_client, + image_segmentation_client=image_segmentation_client, + ) + # Prepare an image which do not contains any asset, for the VLM validation of the table + # segmentation more easily. + asset_mask_paths: list[str] = [] + for asset in scene.assets: + if asset.mask_path is None: + raise ValueError(f"Asset {asset.id!r} has no validated mask path.") + asset_mask_paths.append(asset.mask_path) + table_validation_image_path, asset_union_mask = render_image_without_masks( + image_path=image_path, + mask_paths=asset_mask_paths, + output_path=Path(debug_output_root) / "table_validation_base.png", + ) + # Segment the table. + _segment_table( + image_path=image_path, + validation_image_path=table_validation_image_path, + label_avoid_mask=asset_union_mask, + debug_output_root=debug_output_root, + masks_output_root=masks_output_root, + scene=scene, + vlm_client=vlm_client, + image_segmentation_client=image_segmentation_client, + ) + + +def _segment_table( + image_path: str | Path, + validation_image_path: str | Path, + label_avoid_mask: Image.Image, + debug_output_root: str | Path, + masks_output_root: str | Path, + scene: Scene, + *, + vlm_client: OpenAICompatibleVLM, + image_segmentation_client: ImageSegmentationClient, +) -> None: + """Segment the table. (Now it only supports segment the complete tabletop)""" + if scene.table is None: + raise ValueError("Cannot segment a scene without a table.") + + table = scene.table + # Build the segmentation prompts for table. + for prompt_label, prompt in ( + ("name", table.name), + ("description", table.description), + ("table", "table"), + ("plane", "plane"), + ): + candidates = union_overlapping_mask_candidates( + build_mask_candidates( + image_segmentation_client.segment_single_object( + image_path=image_path, + prompt=prompt, + ) + ), + min_iou=0.8, # Union masks who have iou > 0.8 + ) + # If do not have candidate, then try segment the table with description, "table", "plane"... + # Notice that, this part could be extended with other segmentation prompt like + # a board, or newly-generated prompt from another VLM-calling etc. + if not candidates: + continue + + # Maybe the mask count = 1, but not correct; + # Maybe the mask count > 1; + # Thus, we need to validate with an VLM. + candidates_image_path = render_numbered_mask_candidates( + image_path=validation_image_path, + candidates=candidates, + label_avoid_mask=label_avoid_mask, + output_path=( + Path(debug_output_root) + / f"table_candidates_{prompt_label}.png" # Render with prompt label, for easily debug. + ), + ) + selected_mask_index = _validate_table_candidates_with_vlm( + table=table, + candidates=candidates, + candidates_image_path=candidates_image_path, + vlm_client=vlm_client, + ) + if selected_mask_index is None: + continue + + # Save result. + candidate = _candidate_by_index(candidates, selected_mask_index) + table.mask_path = str( + save_binary_mask( + candidate, + image_size=_image_size(image_path), + output_path=Path(masks_output_root) / "table_mask.png", + ) + ) + return + + raise ValueError("Unable to find a VLM-validated segmentation mask for the table.") + + +def _validate_table_candidates_with_vlm( + *, + table: SceneObject, + candidates: list[MaskCandidate], + candidates_image_path: Path, + vlm_client: OpenAICompatibleVLM, + json_max_attempts: int = 3, +) -> int | None: + + if json_max_attempts < 1: + raise ValueError("json_max_attempts must be at least 1.") + + user_prompt = ( + "Table category: " + f"{table.category}\n" + f"Table name: {table.name}\n" + f"Table description: {table.description}\n" + f"Candidate indices range from 1 to {len(candidates)}." + ) + last_error: ValueError | None = None + for _ in range(json_max_attempts): + response_text = vlm_client.complete( + image_path=candidates_image_path, + system_prompt=_TABLE_VALIDATION_SYSTEM_PROMPT, + user_prompt=user_prompt, + ) + try: + return _parse_table_validation_response(response_text, candidates) + except ValueError as exc: + last_error = exc + + assert last_error is not None + raise ValueError( + "VLM returned invalid table-segmentation validation JSON after " + f"{json_max_attempts} attempts: {last_error}" + ) from last_error + + +def _parse_table_validation_response( + response_text: str, + candidates: list[MaskCandidate], +) -> int | None: + """Validate the strict VLM response schema for table candidate selection.""" + try: + payload = json.loads(_strip_json_code_fence(response_text)) + except json.JSONDecodeError as exc: + raise ValueError("VLM table validation response is not valid JSON.") from exc + if not isinstance(payload, dict) or set(payload) != {"selected_mask_index"}: + raise ValueError( + "VLM table validation JSON must contain only selected_mask_index." + ) + + selected_mask_index = payload["selected_mask_index"] + if selected_mask_index is None: + return None + if isinstance(selected_mask_index, bool) or not isinstance( + selected_mask_index, int + ): + raise ValueError("selected_mask_index must be an integer or null.") + _candidate_by_index(candidates, selected_mask_index) + return selected_mask_index + + +def _candidate_by_index( + candidates: list[MaskCandidate], + index: int, +) -> MaskCandidate: + for candidate in candidates: + if candidate.index == index: + return candidate + raise ValueError(f"VLM selected a nonexistent mask candidate: {index}.") + + +def _image_size(image_path: str | Path) -> tuple[int, int]: + from PIL import Image + + with Image.open(image_path) as image: + return image.size + + +def _segment_assets( + image_path: str | Path, + debug_output_root: str | Path, + masks_output_root: str | Path, + scene: Scene, + *, + vlm_client: OpenAICompatibleVLM, + image_segmentation_client: ImageSegmentationClient, +) -> None: + + # Group the assets by their categories. + assets_by_category: dict[str, list[SceneObject]] = {} + for asset in scene.assets: + assets_by_category.setdefault(asset.category, []).append(asset) + + image_size = _image_size(image_path) + for category, assets in assets_by_category.items(): + mask_rles: list[dict[str, Any]] = [] + # Use categories and names as segmentation prompt. + # Use category to segment first, then use each assets' name to segment. + prompts = [category, *dict.fromkeys(asset.name for asset in assets)] + for prompt in prompts: + mask_rles.extend( + image_segmentation_client.segment_single_object( + image_path=image_path, + prompt=prompt, + ) + ) + # Union duplicated mask candidates. + candidates = union_overlapping_mask_candidates( + build_mask_candidates(mask_rles), + min_iou=0.8, + ) + # If the number of candidate is less than the grouped assets, + # raise error directly. + if len(candidates) < len(assets): + raise ValueError( + f"Asset category {category!r} has {len(assets)} assets but only " + f"{len(candidates)} segmentation candidates." + ) + + candidates_image_path = render_numbered_mask_candidates( + image_path=image_path, + candidates=candidates, + output_path=Path(debug_output_root) / f"asset_candidates_{category}.png", + mask_style="outline", + ) + assignments = _validate_asset_candidates_with_vlm( + assets=assets, + candidates=candidates, + candidates_image_path=candidates_image_path, + vlm_client=vlm_client, + ) + if assignments is None: + raise ValueError( + f"VLM could not assign every {category!r} asset to a segmentation candidate." + ) + # Save results. + for asset in assets: + asset.mask_path = str( + save_binary_mask( + _candidate_by_index(candidates, assignments[asset.id]), + image_size=image_size, + output_path=Path(masks_output_root) / f"{asset.id}_mask.png", + ) + ) + + +def _validate_asset_candidates_with_vlm( + *, + assets: list[SceneObject], + candidates: list[MaskCandidate], + candidates_image_path: Path, + vlm_client: OpenAICompatibleVLM, + json_max_attempts: int = 3, +) -> dict[str, int] | None: + """Ask the VLM for a complete one-to-one asset-to-candidate assignment.""" + if json_max_attempts < 1: + raise ValueError("json_max_attempts must be at least 1.") + + assets_text = "\n".join( + "- " + f"id: {asset.id}; category: {asset.category}; name: {asset.name}; " + f"description: {asset.description}" + for asset in assets + ) + user_prompt = ( + "Asset group:\n" + f"{assets_text}\n\n" + f"Candidate indices range from 1 to {len(candidates)}." + ) + last_error: ValueError | None = None + for _ in range(json_max_attempts): + response_text = vlm_client.complete( + image_path=candidates_image_path, + system_prompt=_ASSET_ASSIGNMENT_SYSTEM_PROMPT, + user_prompt=user_prompt, + ) + try: + return _parse_asset_assignment_response(response_text, assets, candidates) + except ValueError as exc: + last_error = exc + + assert last_error is not None + raise ValueError( + "VLM returned invalid asset-segmentation assignment JSON after " + f"{json_max_attempts} attempts: {last_error}" + ) from last_error + + +def _parse_asset_assignment_response( + response_text: str, + assets: list[SceneObject], + candidates: list[MaskCandidate], +) -> dict[str, int] | None: + """Parse a strict complete assignment, or a valid missing-asset result.""" + try: + payload = json.loads(_strip_json_code_fence(response_text)) + except json.JSONDecodeError as exc: + raise ValueError("VLM asset assignment response is not valid JSON.") from exc + if not isinstance(payload, dict) or set(payload) != {"assignments"}: + raise ValueError("VLM asset assignment JSON must contain only assignments.") + + assignment_values = payload["assignments"] + if assignment_values is None: + return None + if not isinstance(assignment_values, list): + raise ValueError("assignments must be an array or null.") + + expected_asset_ids = {asset.id for asset in assets} + assignments: dict[str, int] = {} + assigned_mask_indices: set[int] = set() + for assignment in assignment_values: + if not isinstance(assignment, dict) or set(assignment) != { + "asset_id", + "mask_index", + }: + raise ValueError( + "Each assignment must contain only asset_id and mask_index." + ) + asset_id = assignment["asset_id"] + mask_index = assignment["mask_index"] + if not isinstance(asset_id, str) or not asset_id: + raise ValueError("assignment asset_id must be a non-empty string.") + if isinstance(mask_index, bool) or not isinstance(mask_index, int): + raise ValueError("assignment mask_index must be an integer.") + if asset_id in assignments: + raise ValueError(f"VLM assigned asset {asset_id!r} more than once.") + if mask_index in assigned_mask_indices: + raise ValueError( + f"VLM assigned candidate {mask_index} to more than one asset." + ) + _candidate_by_index(candidates, mask_index) + assignments[asset_id] = mask_index + assigned_mask_indices.add(mask_index) + + if set(assignments) != expected_asset_ids: + raise ValueError("VLM assignments must cover every asset in the group.") + return assignments diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/assets_gravity_settler.py b/embodichain/gen_sim/scene_engine/pipeline/utils/assets_gravity_settler.py index 3f68508cb..31e9e8443 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/assets_gravity_settler.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/assets_gravity_settler.py @@ -24,13 +24,18 @@ from scipy.spatial.transform import Rotation import trimesh +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_object import ( + ObjectPhysics, + SceneObject, +) from embodichain.gen_sim.scene_engine.pipeline.utils.scene_generation_utils import ( layout_object_to_transform_matrix, load_glb_mesh, transform_matrix_to_layout_object, ) from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.cfg import RigidObjectCfg +from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg from embodichain.lab.sim.shapes import MeshCfg from embodichain.utils.logger import log_info @@ -43,20 +48,21 @@ class AssetsGravitySettlerConfig: settle_steps: int = 300 # Fixed number of simulator steps to execute. physics_dt: float = 1.0 / 100.0 # Physics timestep in seconds. sim_device: str = "cpu" # Simulation device requested from EmbodiChain Lab. - max_convex_hull_num: int = 32 # VHACD hull budget for each collision mesh. class AssetsGravitySettler: - """Settle all assets together on one static table in a z-up simulation.""" + """Settle all assets together on one kinematic table in a z-up simulation.""" def __init__( self, *, + scene: Scene, table_layout: dict[str, object], assets_layout: list[dict[str, object]], geometry_root: str | Path, config: AssetsGravitySettlerConfig | None = None, ) -> None: + self.scene = scene self.table_layout = table_layout self.assets_layout = assets_layout self.geometry_root = Path(geometry_root).expanduser().resolve() @@ -69,8 +75,6 @@ def __init__( raise ValueError("Gravity-settle settle_steps must be positive.") if self.config.physics_dt <= 0.0: raise ValueError("Gravity-settle physics_dt must be positive.") - if self.config.max_convex_hull_num <= 0: - raise ValueError("Gravity-settle max_convex_hull_num must be positive.") def settle(self) -> list[dict[str, object]]: """Run gravity settling and return the resulting y-up asset layouts.""" @@ -81,12 +85,22 @@ def settle(self) -> list[dict[str, object]]: return self.settled_assets_layout table_id = self._require_layout_id(self.table_layout, name="Table") + table_object = self._require_scene_object(table_id, kind="table") asset_ids: set[str] = set() + asset_objects_by_id: dict[str, SceneObject] = {} for asset_layout in self.assets_layout: asset_id = self._require_layout_id(asset_layout, name="Asset") if asset_id in asset_ids: raise ValueError(f"Asset layouts contain duplicate id {asset_id!r}.") asset_ids.add(asset_id) + asset_objects_by_id[asset_id] = self._require_scene_object( + asset_id, kind="asset" + ) + expected_asset_ids = {asset.id for asset in self.scene.assets} + if asset_ids != expected_asset_ids: + raise ValueError( + "Gravity-settle layouts must contain exactly the scene asset ids." + ) y_up_to_z_up_matrix = np.eye(4) y_up_to_z_up_matrix[:3, :3] = np.array( @@ -127,8 +141,7 @@ def settle(self) -> list[dict[str, object]]: log_info( "Gravity settling started: " f"assets={len(prepared_assets)}, steps={self.config.settle_steps}, " - f"physics_dt={self.config.physics_dt:.4f} s, " - f"max_convex_hulls={self.config.max_convex_hull_num}." + f"physics_dt={self.config.physics_dt:.4f} s." ) sim = SimulationManager( SimulationManagerCfg( @@ -148,8 +161,9 @@ def settle(self) -> list[dict[str, object]]: self._simulation_euler_xyz_degrees(table_info["rigid_layout"]) ), body_scale=tuple(table_info["y_up_scale"]), - body_type="static", - max_convex_hull_num=self.config.max_convex_hull_num, + attrs=self._rigid_body_attrs(table_object.physics), + body_type=table_object.physics.body_type, + max_convex_hull_num=table_object.physics.max_convex_hull_num, acd_method="vhacd", ) ) @@ -166,8 +180,13 @@ def settle(self) -> list[dict[str, object]]: self._simulation_euler_xyz_degrees(rigid_layout) ), body_scale=tuple(asset_info["y_up_scale"]), - body_type="dynamic", - max_convex_hull_num=self.config.max_convex_hull_num, + attrs=self._rigid_body_attrs( + asset_objects_by_id[asset_id].physics + ), + body_type=asset_objects_by_id[asset_id].physics.body_type, + max_convex_hull_num=( + asset_objects_by_id[asset_id].physics.max_convex_hull_num + ), acd_method="vhacd", ) ) @@ -236,6 +255,36 @@ def _prepare_sim_body( ), } + def _require_scene_object(self, object_id: str, *, kind: str) -> SceneObject: + """Return one physics-ready scene object with the expected semantic kind.""" + matching_objects = [ + scene_object + for scene_object in self.scene.objects + if scene_object.id == object_id + ] + if len(matching_objects) != 1: + raise ValueError( + f"Gravity settling requires exactly one scene object {object_id!r}." + ) + scene_object = matching_objects[0] + if scene_object.kind != kind: + raise ValueError( + f"Scene object {object_id!r} must have kind {kind!r} before " + "gravity settling." + ) + if scene_object.physics is None: + raise ValueError( + f"Scene object {object_id!r} has no SimReady physics settings." + ) + return scene_object + + @staticmethod + def _rigid_body_attrs(physics: ObjectPhysics | None) -> RigidBodyAttributesCfg: + """Convert persisted SceneObject physics attributes into Lab config.""" + if physics is None: + raise ValueError("Gravity settling requires SimReady physics settings.") + return RigidBodyAttributesCfg(**physics.attrs) + @staticmethod def _mesh_to_z_up_world_for_aabb( *, diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_segmentation_utils.py b/embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py similarity index 100% rename from embodichain/gen_sim/scene_engine/pipeline/utils/scene_segmentation_utils.py rename to embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py index 3192a0a9e..cf1f55932 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py @@ -17,7 +17,6 @@ from __future__ import annotations -from dataclasses import dataclass import json from pathlib import Path import shutil @@ -26,27 +25,10 @@ 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.core.scene_object import SceneObject from embodichain.utils.logger import log_info -_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], @@ -57,14 +39,6 @@ ) -@dataclass(frozen=True) -class SceneExporterConfig: - """Collision-decomposition controls for scene export.""" - - table_max_convex_hull_num: int = _DEFAULT_MAX_CONVEX_HULL_NUM # Table hull limit. - asset_max_convex_hull_num: int = _DEFAULT_MAX_CONVEX_HULL_NUM # Asset hull limit. - - class SceneExporter: """Write one generated scene and its SimReady meshes as a scene export.""" @@ -73,21 +47,11 @@ def __init__( *, scene: Scene, output_root: str | Path, - config: SceneExporterConfig | None = None, ) -> None: self.scene = scene self.output_root = Path(output_root).expanduser().resolve() self.export_root = self.output_root / "scene_export" self.scene_config_path: Path | None = None - self.config = config if config is not None else SceneExporterConfig() - self.table_max_convex_hull_num = _positive_int( - self.config.table_max_convex_hull_num, - field_name="table_max_convex_hull_num", - ) - self.asset_max_convex_hull_num = _positive_int( - self.config.asset_max_convex_hull_num, - field_name="asset_max_convex_hull_num", - ) def export(self) -> Path: """Write a scene-only config and copy SimReady GLBs into ``mesh_assets``. @@ -104,7 +68,7 @@ def export(self) -> Path: mesh_assets_root = self.export_root / "mesh_assets" mesh_assets_root.mkdir(parents=True, exist_ok=True) - scene_objects = [self.scene.table, *self.scene.assets] + scene_objects = self.scene.objects object_ids = [scene_object.id for scene_object in scene_objects] if len(set(object_ids)) != len(object_ids): raise ValueError("Scene export requires unique table and asset ids.") @@ -126,18 +90,12 @@ def export(self) -> Path: self._scene_object_config( scene_object=self.scene.table, asset_relative_path=exported_entries[self.scene.table.id], - body_type="kinematic", - attrs=_TABLE_PHYSICS_ATTRS, - max_convex_hull_num=self.table_max_convex_hull_num, ) ], "rigid_object": [ self._scene_object_config( scene_object=asset, asset_relative_path=exported_entries[asset.id], - body_type="dynamic", - attrs=_ASSET_PHYSICS_ATTRS, - max_convex_hull_num=self.asset_max_convex_hull_num, ) for asset in self.scene.assets ], @@ -153,7 +111,7 @@ def export(self) -> Path: @staticmethod def _copy_scene_object_to_assets( *, - scene_object: Table | Asset, + scene_object: SceneObject, mesh_assets_root: Path, ) -> str: """Copy one referenced SimReady GLB and return its config-relative path.""" @@ -163,9 +121,7 @@ def _copy_scene_object_to_assets( 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." - ) + 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(): @@ -181,21 +137,20 @@ def _copy_scene_object_to_assets( @staticmethod def _scene_object_config( *, - scene_object: Table | Asset, + scene_object: SceneObject, asset_relative_path: str, - body_type: str, - attrs: dict[str, float | int], - max_convex_hull_num: int, ) -> dict[str, object]: """Build one z-up scene-only object config from a final y-up object.""" pos_y_up = SceneExporter._scene_vector(scene_object, "pos") rot_y_up = SceneExporter._scene_vector(scene_object, "rot") scale_y_up = SceneExporter._scene_vector(scene_object, "scale") + if scene_object.physics is None: + raise ValueError( + f"Scene object {scene_object.id!r} has no SimReady physics settings." + ) pos_z_up = _Y_UP_TO_Z_UP_ROTATION @ np.asarray(pos_y_up, dtype=float) - rotation_y_up = Rotation.from_euler( - "xyz", rot_y_up, degrees=True - ).as_matrix() + rotation_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 ) @@ -213,18 +168,18 @@ def _scene_object_config( "fpath": asset_relative_path, "compute_uv": False, }, - "attrs": attrs, - "body_type": body_type, + "attrs": scene_object.physics.attrs, + "body_type": scene_object.physics.body_type, "init_pos": pos_z_up.tolist(), "init_rot": rot_z_up.tolist(), # Do not permute this scale: it belongs to the original y-up GLB, # which SimulationManager itself converts to z-up. "body_scale": scale_y_up, - "max_convex_hull_num": max_convex_hull_num, + "max_convex_hull_num": scene_object.physics.max_convex_hull_num, } @staticmethod - def _scene_vector(scene_object: Table | Asset, field_name: str) -> list[float]: + def _scene_vector(scene_object: SceneObject, field_name: str) -> list[float]: """Read one finite final y-up layout vector from a scene object.""" values = getattr(scene_object, field_name) if not isinstance(values, list) or len(values) != 3: @@ -235,15 +190,6 @@ def _scene_vector(scene_object: Table | Asset, field_name: str) -> list[float]: vector = [float(value) for value in values] if not np.all(np.isfinite(vector)): raise ValueError( - f"Scene object {scene_object.id!r} has non-finite " - f"{field_name!r}." + f"Scene object {scene_object.id!r} has non-finite " f"{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/utils/scene_generation_utils.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_generation_utils.py index 49e5da138..2f648530a 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 @@ -100,59 +100,6 @@ def load_glb_mesh(glb_path: str | Path) -> trimesh.Trimesh: raise ValueError(f"GLB geometry is not a mesh: {resolved_glb_path}") -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 _three_floats(value: object, *, field_name: str) -> list[float]: # Validate whether the value is a list of three numeric values. diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_scene_processor.py b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_scene_processor.py index 33ee26a11..40b4a6f94 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_scene_processor.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_scene_processor.py @@ -28,8 +28,29 @@ import trimesh from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_object import ( + ObjectPhysics, + SceneObject, +) from embodichain.utils.logger import log_info +_TABLE_PHYSICS_ATTRS = { + "mass": 10.0, # Keep the table heavy if a simulator treats it as movable. + "static_friction": 0.95, # Resist lateral sliding at table contacts. + "dynamic_friction": 0.9, # Maintain high friction during sliding contacts. + "restitution": 0.01, # Prevent a table contact from producing visible bounce. +} +_ASSET_PHYSICS_ATTRS = { + "mass": 0.01, # Use a lightweight default for unconstrained generated assets. + "contact_offset": 0.003, # Start contact detection slightly before mesh contact. + "rest_offset": 0.001, # Keep a small stable separation after contact resolution. + "restitution": 0.01, # Prevent generated assets from bouncing on the table. + "max_depenetration_velocity": 10.0, # Cap corrective separation speed. + "min_position_iters": 32, # Use extra position iterations for stable contacts. + "min_velocity_iters": 8, # Use extra velocity iterations for stable contacts. +} +_FIXED_MAX_CONVEX_HULL_NUM = 16 # Shared VHACD hull budget for settling and export. + @dataclass(frozen=True) class SimReadySceneProcessorConfig: @@ -68,10 +89,7 @@ def process_table(self) -> dict[str, object]: """Process the required scene table and return its SimReady layout.""" if self.scene.table is None: raise ValueError("Cannot SimReady a scene without a table.") - self.simready_table_layout = self._process_object( - object_id=self.scene.table.id, - object_role="table", - ) + self.simready_table_layout = self._process_object(self.scene.table) return self.simready_table_layout def process_assets(self) -> list[dict[str, object]]: @@ -82,14 +100,14 @@ def process_assets(self) -> list[dict[str, object]]: if asset.id in asset_ids: raise ValueError(f"Scene assets contain duplicate id {asset.id!r}.") asset_ids.add(asset.id) - processed_assets.append( - self._process_object(object_id=asset.id, object_role="asset") - ) + processed_assets.append(self._process_object(asset)) self.simready_assets_layout = processed_assets return self.simready_assets_layout - def _process_object(self, *, object_id: str, object_role: str) -> dict[str, object]: + def _process_object(self, scene_object: SceneObject) -> dict[str, object]: """Canonicalize one coarse object and write its SimReady GLB.""" + object_id = scene_object.id + object_role = scene_object.kind if object_role not in {"table", "asset"}: raise ValueError(f"Unsupported SimReady object role {object_role!r}.") coarse_layout = self.coarse_layout_by_id.get(object_id) @@ -109,9 +127,28 @@ def _process_object(self, *, object_id: str, object_role: str) -> dict[str, obje raise FileNotFoundError( f"SimReady {object_role} geometry was not written: {output_path}" ) + scene_object.simready_glb_path = str(output_path) + scene_object.physics = self._fixed_physics_for_kind(object_role) log_info(f"Created SimReady {object_role}: {object_id!r}.") return {"id": object_id, **simready_transform} + @staticmethod + def _fixed_physics_for_kind(kind: str) -> ObjectPhysics: + """Create the fixed initial physics profile for one SimReady object.""" + if kind == "table": + return ObjectPhysics( + body_type="kinematic", + attrs=dict(_TABLE_PHYSICS_ATTRS), + max_convex_hull_num=_FIXED_MAX_CONVEX_HULL_NUM, + ) + if kind == "asset": + return ObjectPhysics( + body_type="dynamic", + attrs=dict(_ASSET_PHYSICS_ATTRS), + max_convex_hull_num=_FIXED_MAX_CONVEX_HULL_NUM, + ) + raise ValueError(f"Unsupported SceneObject kind {kind!r} for physics.") + def _canonicalize_object_mesh( self, *, From 5eb913425d462fb1b4c605b817893519bdc68c19 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Tue, 4 Aug 2026 15:29:37 +0800 Subject: [PATCH 30/41] Read .env for now --- embodichain/gen_sim/scene_engine/cli/start.py | 21 +--- .../clients/geometry_generation.py | 102 ++++++------------ .../clients/image_segmentation.py | 96 ++++++----------- .../scene_engine/configs/environment.py | 50 +++++++++ .../configs/scene_engine_config.json | 25 ----- .../gen_sim/scene_engine/core/scene.py | 12 ++- .../gen_sim/scene_engine/llms/load_config.py | 58 +++++----- .../llms/openai_compatible_client.py | 8 +- .../gen_sim/scene_engine/pipeline/generate.py | 13 +-- .../pipeline/scene_understanding.py | 7 +- 10 files changed, 156 insertions(+), 236 deletions(-) create mode 100644 embodichain/gen_sim/scene_engine/configs/environment.py delete mode 100644 embodichain/gen_sim/scene_engine/configs/scene_engine_config.json diff --git a/embodichain/gen_sim/scene_engine/cli/start.py b/embodichain/gen_sim/scene_engine/cli/start.py index 4052da257..637821a1a 100644 --- a/embodichain/gen_sim/scene_engine/cli/start.py +++ b/embodichain/gen_sim/scene_engine/cli/start.py @@ -28,10 +28,8 @@ 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.""" + """Generate one scene using the required ``gen_sim/.env`` settings.""" resolved_image_path = Path(image).expanduser().resolve() if not resolved_image_path.exists(): raise FileNotFoundError(f"Image input not found: {resolved_image_path}") @@ -48,12 +46,6 @@ def cli_scene_engine( 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. 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, ) print("Successfully completed!") @@ -75,18 +67,9 @@ def main(argv: Sequence[str] | None = None) -> None: required=True, help="Path to the output directory", ) - parser.add_argument( - "--config", - type=Path, - default=None, - help=( - "Optional Scene Engine JSON override. Without it, clients read the " - "packaged template and apply service environment-variable overrides." - ), - ) args = parser.parse_args(argv) - cli_scene_engine(args.image, args.output_root, config_path=args.config) + cli_scene_engine(args.image, args.output_root) if __name__ == "__main__": diff --git a/embodichain/gen_sim/scene_engine/clients/geometry_generation.py b/embodichain/gen_sim/scene_engine/clients/geometry_generation.py index d50a3a267..c84fe4c26 100644 --- a/embodichain/gen_sim/scene_engine/clients/geometry_generation.py +++ b/embodichain/gen_sim/scene_engine/clients/geometry_generation.py @@ -19,23 +19,15 @@ from contextlib import ExitStack import json -import os from pathlib import Path import time from typing import Any import requests -_DEFAULT_CONFIG_PATH = ( - Path(__file__).resolve().parents[1] / "configs" / "scene_engine_config.json" +from embodichain.gen_sim.scene_engine.configs.environment import ( + read_scene_engine_env_values, ) -_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: @@ -59,11 +51,9 @@ def __init__( 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 from_dotenv(cls) -> "GeometryGenerationClient": + """Create a client from its required ``gen_sim/.env`` settings.""" + return cls(**_load_dotenv_config()) def check_health(self) -> None: last_error: Exception | None = None @@ -104,9 +94,9 @@ def generate_objects( ) -> tuple[dict[str, Any], list[dict[str, Any]]]: """Generate objects through the geometry server's mask-list endpoint. - The SAM3D service represents both one-object and multi-object jobs as one + The service represents both one-object and multi-object jobs as one image plus a multipart ``masks`` list. The number of list items is the - only difference, so keeping one implementation prevents the two client + only difference, so keeping one implementation prevents the client paths from drifting apart. """ @@ -131,7 +121,7 @@ def generate_objects( ) resolved_object_masks.append((object_id, resolved_mask_path)) - # Send one multipart image + masks request, matching test_sam3d_client.py. + # Send one multipart image + masks request. response_data, response_objects = self._request_objects( image_path=resolved_image_path, object_masks=resolved_object_masks, @@ -226,7 +216,7 @@ def _request_objects( ) 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.""" + """Poll a queued geometry-generation job until it returns its result.""" if not isinstance(response_data, dict): raise RuntimeError( "Geometry Generation Server response must be a JSON object." @@ -408,79 +398,51 @@ def _image_content_type(image_path: Path) -> str: 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() - 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.") - config = dict(config) - _apply_environment_overrides(config) - - required_keys = ( - "base_url", - "timeout_s", - "max_attempts", - "health_path", - "generate_objects_path", +def _load_dotenv_config() -> dict[str, Any]: + values = read_scene_engine_env_values( + "SCENE_ENGINE_GEOMETRY_GENERATION_BASE_URL", + "SCENE_ENGINE_GEOMETRY_GENERATION_TIMEOUT_S", + "SCENE_ENGINE_GEOMETRY_GENERATION_MAX_ATTEMPTS", + "SCENE_ENGINE_GEOMETRY_GENERATION_HEALTH_PATH", + "SCENE_ENGINE_GEOMETRY_GENERATION_OBJECTS_PATH", ) - 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"]) + timeout_s = int(values["SCENE_ENGINE_GEOMETRY_GENERATION_TIMEOUT_S"]) except (TypeError, ValueError) as exc: raise ValueError( - "Geometry Generation Server config timeout_s must be an integer." + "SCENE_ENGINE_GEOMETRY_GENERATION_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." + "SCENE_ENGINE_GEOMETRY_GENERATION_TIMEOUT_S must be at least 1." ) try: - max_attempts = int(config["max_attempts"]) + max_attempts = int(values["SCENE_ENGINE_GEOMETRY_GENERATION_MAX_ATTEMPTS"]) except (TypeError, ValueError) as exc: raise ValueError( - "Geometry Generation Server config max_attempts must be an integer." + "SCENE_ENGINE_GEOMETRY_GENERATION_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." + "SCENE_ENGINE_GEOMETRY_GENERATION_MAX_ATTEMPTS must be at least 1." ) string_keys = ( - "base_url", - "health_path", - "generate_objects_path", + "SCENE_ENGINE_GEOMETRY_GENERATION_BASE_URL", + "SCENE_ENGINE_GEOMETRY_GENERATION_HEALTH_PATH", + "SCENE_ENGINE_GEOMETRY_GENERATION_OBJECTS_PATH", ) for key in string_keys: - if not isinstance(config[key], str) or not config[key].strip(): - raise ValueError( - f"Geometry Generation Server config key {key} must be a non-empty string." - ) + if not values[key].strip(): + raise ValueError(f"Scene Engine .env key {key} must be a non-empty string.") return { - "base_url": config["base_url"].strip(), + "base_url": values["SCENE_ENGINE_GEOMETRY_GENERATION_BASE_URL"].strip(), "timeout_s": timeout_s, "max_attempts": max_attempts, - "health_path": config["health_path"].strip(), - "generate_objects_path": config["generate_objects_path"].strip(), + "health_path": values["SCENE_ENGINE_GEOMETRY_GENERATION_HEALTH_PATH"].strip(), + "generate_objects_path": values[ + "SCENE_ENGINE_GEOMETRY_GENERATION_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 d3ded7218..2c56af91a 100644 --- a/embodichain/gen_sim/scene_engine/clients/image_segmentation.py +++ b/embodichain/gen_sim/scene_engine/clients/image_segmentation.py @@ -17,23 +17,14 @@ from __future__ import annotations -import json -import os from pathlib import Path from typing import Any import requests -_DEFAULT_CONFIG_PATH = ( - Path(__file__).resolve().parents[1] / "configs" / "scene_engine_config.json" +from embodichain.gen_sim.scene_engine.configs.environment import ( + read_scene_engine_env_values, ) -_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: @@ -56,12 +47,9 @@ def __init__( 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 from_dotenv(cls) -> "ImageSegmentationClient": + """Create a client from its required ``gen_sim/.env`` settings.""" + return cls(**_load_dotenv_config()) def check_health(self) -> None: last_error: requests.RequestException | None = None @@ -144,80 +132,56 @@ 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.") - config = dict(config) - _apply_environment_overrides(config) - - required_keys = ( - "base_url", - "timeout_s", - "max_attempts", - "health_path", - "segment_single_object_path", +def _load_dotenv_config() -> dict[str, Any]: + values = read_scene_engine_env_values( + "SCENE_ENGINE_IMAGE_SEGMENTATION_BASE_URL", + "SCENE_ENGINE_IMAGE_SEGMENTATION_TIMEOUT_S", + "SCENE_ENGINE_IMAGE_SEGMENTATION_MAX_ATTEMPTS", + "SCENE_ENGINE_IMAGE_SEGMENTATION_HEALTH_PATH", + "SCENE_ENGINE_IMAGE_SEGMENTATION_SINGLE_OBJECT_PATH", ) - 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"]) + timeout_s = int(values["SCENE_ENGINE_IMAGE_SEGMENTATION_TIMEOUT_S"]) except (TypeError, ValueError) as exc: raise ValueError( - "Image Segmentation Server config timeout_s must be an integer." + "SCENE_ENGINE_IMAGE_SEGMENTATION_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." + "SCENE_ENGINE_IMAGE_SEGMENTATION_TIMEOUT_S must be at least 1." ) try: - max_attempts = int(config["max_attempts"]) + max_attempts = int(values["SCENE_ENGINE_IMAGE_SEGMENTATION_MAX_ATTEMPTS"]) except (TypeError, ValueError) as exc: raise ValueError( - "Image Segmentation Server config max_attempts must be an integer." + "SCENE_ENGINE_IMAGE_SEGMENTATION_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." + "SCENE_ENGINE_IMAGE_SEGMENTATION_MAX_ATTEMPTS must be at least 1." ) - string_keys = ("base_url", "health_path", "segment_single_object_path") + string_keys = ( + "SCENE_ENGINE_IMAGE_SEGMENTATION_BASE_URL", + "SCENE_ENGINE_IMAGE_SEGMENTATION_HEALTH_PATH", + "SCENE_ENGINE_IMAGE_SEGMENTATION_SINGLE_OBJECT_PATH", + ) for key in string_keys: - if not isinstance(config[key], str) or not config[key].strip(): - raise ValueError( - f"Image Segmentation Server config key {key} must be a non-empty string." - ) + if not values[key].strip(): + raise ValueError(f"Scene Engine .env key {key} must be a non-empty string.") return { - "base_url": config["base_url"].strip(), + "base_url": values["SCENE_ENGINE_IMAGE_SEGMENTATION_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(), + "health_path": values["SCENE_ENGINE_IMAGE_SEGMENTATION_HEALTH_PATH"].strip(), + "segment_single_object_path": values[ + "SCENE_ENGINE_IMAGE_SEGMENTATION_SINGLE_OBJECT_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 - - 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/embodichain/gen_sim/scene_engine/configs/environment.py b/embodichain/gen_sim/scene_engine/configs/environment.py new file mode 100644 index 000000000..9b686d02f --- /dev/null +++ b/embodichain/gen_sim/scene_engine/configs/environment.py @@ -0,0 +1,50 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from pathlib import Path + +_SCENE_ENGINE_ENV_PATH = Path(__file__).resolve().parents[2] / ".env" + + +def read_scene_engine_env_values(*keys: str) -> dict[str, str]: + """Read only the requested Scene Engine settings from ``gen_sim/.env``.""" + if not _SCENE_ENGINE_ENV_PATH.is_file(): + raise FileNotFoundError( + f"Scene Engine .env file not found: {_SCENE_ENGINE_ENV_PATH}" + ) + + requested_keys = set(keys) + values: dict[str, str] = {} + for raw_line in _SCENE_ENGINE_ENV_PATH.read_text(encoding="utf-8").splitlines(): + line = raw_line.strip() + if not line or line.startswith("#") or "=" not in line: + continue + key, raw_value = line.split("=", maxsplit=1) + key = key.strip() + if key not in requested_keys: + continue + value = raw_value.strip() + if len(value) >= 2 and value[0] == value[-1] and value[0] in {'"', "'"}: + value = value[1:-1] + values[key] = value + + missing_keys = [key for key in keys if key not in values] + if missing_keys: + raise ValueError(f"Missing required Scene Engine .env keys: {missing_keys}") + return values diff --git a/embodichain/gen_sim/scene_engine/configs/scene_engine_config.json b/embodichain/gen_sim/scene_engine/configs/scene_engine_config.json deleted file mode 100644 index 642901ab3..000000000 --- a/embodichain/gen_sim/scene_engine/configs/scene_engine_config.json +++ /dev/null @@ -1,25 +0,0 @@ -{ - "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_objects_path": "/generate_multiple_objects" - } -} diff --git a/embodichain/gen_sim/scene_engine/core/scene.py b/embodichain/gen_sim/scene_engine/core/scene.py index 79b2e4f2b..07d937e41 100644 --- a/embodichain/gen_sim/scene_engine/core/scene.py +++ b/embodichain/gen_sim/scene_engine/core/scene.py @@ -30,7 +30,11 @@ class Scene: @property def table(self) -> SceneObject | None: """Return the sole table object, or ``None`` before understanding.""" - tables = [scene_object for scene_object in self.objects if scene_object.kind == "table"] + tables = [ + scene_object + for scene_object in self.objects + if scene_object.kind == "table" + ] if len(tables) > 1: raise ValueError("A scene may contain only one table object.") return tables[0] if tables else None @@ -38,7 +42,11 @@ def table(self) -> SceneObject | None: @property def assets(self) -> list[SceneObject]: """Return movable asset objects in their scene order.""" - return [scene_object for scene_object in self.objects if scene_object.kind == "asset"] + return [ + scene_object + for scene_object in self.objects + if scene_object.kind == "asset" + ] def to_dict(self) -> dict[str, object]: """Serialize the canonical object collection for debugging artifacts.""" diff --git a/embodichain/gen_sim/scene_engine/llms/load_config.py b/embodichain/gen_sim/scene_engine/llms/load_config.py index f2a786399..8a2af22d5 100644 --- a/embodichain/gen_sim/scene_engine/llms/load_config.py +++ b/embodichain/gen_sim/scene_engine/llms/load_config.py @@ -19,12 +19,10 @@ 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" +from embodichain.gen_sim.scene_engine.configs.environment import ( + read_scene_engine_env_values, ) @@ -39,55 +37,47 @@ class LLMConfig: 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}") - +def load_llm_config() -> LLMConfig: + """Load the required OpenAI-compatible LLM settings from ``gen_sim/.env``.""" + values = read_scene_engine_env_values( + "OPENAI_API_KEY", + "OPENAI_MODEL", + "OPENAI_BASE_URL", + "SCENE_ENGINE_OPENAI_DEFAULT_QUERY", + "OPENAI_MAX_ATTEMPTS", + ) try: - raw_config = json.loads(resolved_config_path.read_text(encoding="utf-8")) + default_query = json.loads(values["SCENE_ENGINE_OPENAI_DEFAULT_QUERY"]) except json.JSONDecodeError as exc: raise ValueError( - f"LLM config is not valid JSON: {resolved_config_path}" + "SCENE_ENGINE_OPENAI_DEFAULT_QUERY must contain a JSON object." ) 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.") + raise ValueError("SCENE_ENGINE_OPENAI_DEFAULT_QUERY must be a JSON object.") missing = [ key for key, value in { - "api_key": api_key, - "model": model, - "base_url": base_url, + "OPENAI_API_KEY": values["OPENAI_API_KEY"], + "OPENAI_MODEL": values["OPENAI_MODEL"], + "OPENAI_BASE_URL": values["OPENAI_BASE_URL"], }.items() - if not isinstance(value, str) or not value.strip() + if not value.strip() ] if missing: raise ValueError(f"Missing required LLM config keys: {missing}") try: - parsed_max_attempts = int(max_attempts) + parsed_max_attempts = int(values["OPENAI_MAX_ATTEMPTS"]) except (TypeError, ValueError) as exc: - raise ValueError("LLM config key max_attempts must be an integer.") from exc + raise ValueError("OPENAI_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.") + raise ValueError("OPENAI_MAX_ATTEMPTS must be at least 1.") return LLMConfig( - api_key=api_key.strip(), - model=model.strip(), - base_url=base_url.rstrip("/"), + api_key=values["OPENAI_API_KEY"].strip(), + model=values["OPENAI_MODEL"].strip(), + base_url=values["OPENAI_BASE_URL"].rstrip("/"), default_query=default_query, max_attempts=parsed_max_attempts, ) diff --git a/embodichain/gen_sim/scene_engine/llms/openai_compatible_client.py b/embodichain/gen_sim/scene_engine/llms/openai_compatible_client.py index 0b7cf3786..f83e316aa 100644 --- a/embodichain/gen_sim/scene_engine/llms/openai_compatible_client.py +++ b/embodichain/gen_sim/scene_engine/llms/openai_compatible_client.py @@ -36,11 +36,9 @@ 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 from_dotenv(cls) -> "OpenAICompatibleVLM": + """Create a client from its required ``gen_sim/.env`` settings.""" + return cls(load_llm_config()) def complete( self, diff --git a/embodichain/gen_sim/scene_engine/pipeline/generate.py b/embodichain/gen_sim/scene_engine/pipeline/generate.py index aa4f9bf5b..773a897bf 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generate.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generate.py @@ -40,17 +40,13 @@ 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) + vlm_client = OpenAICompatibleVLM.from_dotenv() scene = Scene() # 1. Scene Understanding @@ -60,16 +56,13 @@ def generate_scene_from_image( image_path=image_path, output_root=resolved_output_root, vlm_client=vlm_client, - image_segmentation_config_path=image_segmentation_config_path, ) log_info("Completed Scene Understanding") # 2. 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 - ) + # Load .env settings and fail if the Geometry Generation Server is unavailable. + geometry_generation_client = GeometryGenerationClient.from_dotenv() try: geometry_generation_client.check_health() # Error raising will happen internally. scene = generate_scene_and_refine( diff --git a/embodichain/gen_sim/scene_engine/pipeline/scene_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/scene_understanding.py index 82eaa2ffc..f91c665e6 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/scene_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/scene_understanding.py @@ -149,7 +149,6 @@ def understand_scene( output_root: str | Path, *, vlm_client: OpenAICompatibleVLM, - image_segmentation_config_path: str | Path | None = None, json_max_attempts: int = 3, ) -> Scene: @@ -168,10 +167,8 @@ def understand_scene( json_max_attempts=json_max_attempts, ) - # Load the config and fail if the Image Segmentation Server is unavailable. - image_segmentation_client = ImageSegmentationClient.from_config( - image_segmentation_config_path - ) + # Load .env settings and fail if the Image Segmentation Server is unavailable. + image_segmentation_client = ImageSegmentationClient.from_dotenv() try: image_segmentation_client.check_health() # Error raising will happen internally. _segment_scene( From 2c9477dabb0e6281a780d07a635d1f614da53e42 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Tue, 4 Aug 2026 15:32:11 +0800 Subject: [PATCH 31/41] Deleted all docs and all tests --- docs/source/features/generative_sim/index.rst | 1 - .../features/generative_sim/scene_engine.md | 156 --------------- docs/source/guides/cli.md | 56 ------ tests/gen_sim/scene_engine/test_cli.py | 149 -------------- tests/gen_sim/scene_engine/test_config.py | 126 ------------ tests/gen_sim/scene_engine/test_generate.py | 108 ----------- .../scene_engine/test_geometry_generation.py | 183 ------------------ .../scene_engine/test_image_segmentation.py | 111 ----------- tests/gen_sim/scene_engine/test_preview.py | 82 -------- .../gen_sim/scene_engine/test_scene_export.py | 84 -------- .../test_scene_generation_utils.py | 102 ---------- 11 files changed, 1158 deletions(-) delete mode 100644 docs/source/features/generative_sim/scene_engine.md delete mode 100644 tests/gen_sim/scene_engine/test_cli.py delete mode 100644 tests/gen_sim/scene_engine/test_config.py delete mode 100644 tests/gen_sim/scene_engine/test_generate.py delete mode 100644 tests/gen_sim/scene_engine/test_geometry_generation.py delete mode 100644 tests/gen_sim/scene_engine/test_image_segmentation.py delete mode 100644 tests/gen_sim/scene_engine/test_preview.py delete mode 100644 tests/gen_sim/scene_engine/test_scene_export.py delete mode 100644 tests/gen_sim/scene_engine/test_scene_generation_utils.py diff --git a/docs/source/features/generative_sim/index.rst b/docs/source/features/generative_sim/index.rst index 09d041571..1f7c759f7 100644 --- a/docs/source/features/generative_sim/index.rst +++ b/docs/source/features/generative_sim/index.rst @@ -7,4 +7,3 @@ 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 deleted file mode 100644 index 80b4d6764..000000000 --- a/docs/source/features/generative_sim/scene_engine.md +++ /dev/null @@ -1,156 +0,0 @@ -# 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. - -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 -{ - "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": "/predict" - }, - "geometry_generation": { - "base_url": "http://geometry-host:port", - "timeout_s": 600, - "max_attempts": 3, - "health_path": "/health", - "generate_objects_path": "/generate_multiple_objects" - } -} -``` - -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 - -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 b4b5786c9..f4cfb4ce0 100644 --- a/docs/source/guides/cli.md +++ b/docs/source/guides/cli.md @@ -59,62 +59,6 @@ The generated output contains the canonical source mesh under ``asset_source/``, --- -## Scene Engine - -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 \ - --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 template | Optional complete Scene Engine JSON override. Without it, supply the documented service environment variables; the packaged JSON is only a template. | - -``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. diff --git a/tests/gen_sim/scene_engine/test_cli.py b/tests/gen_sim/scene_engine/test_cli.py deleted file mode 100644 index 35446f951..000000000 --- a/tests/gen_sim/scene_engine/test_cli.py +++ /dev/null @@ -1,149 +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 - -from pathlib import Path - -import pytest - -from embodichain.gen_sim.scene_engine.cli import preview, 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_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] = {} - - 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"), - } - - -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/gen_sim/scene_engine/test_config.py b/tests/gen_sim/scene_engine/test_config.py deleted file mode 100644 index cef3704de..000000000 --- a/tests/gen_sim/scene_engine/test_config.py +++ /dev/null @@ -1,126 +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 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_generate.py b/tests/gen_sim/scene_engine/test_generate.py deleted file mode 100644 index 85a71d642..000000000 --- a/tests/gen_sim/scene_engine/test_generate.py +++ /dev/null @@ -1,108 +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 - -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 deleted file mode 100644 index 2b48eded4..000000000 --- a/tests/gen_sim/scene_engine/test_geometry_generation.py +++ /dev/null @@ -1,183 +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 - -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"]) - - -@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: - 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 deleted file mode 100644 index b4438f1c0..000000000 --- a/tests/gen_sim/scene_engine/test_image_segmentation.py +++ /dev/null @@ -1,111 +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 - -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=" ") - - -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() diff --git a/tests/gen_sim/scene_engine/test_preview.py b/tests/gen_sim/scene_engine/test_preview.py deleted file mode 100644 index 2501c4c23..000000000 --- a/tests/gen_sim/scene_engine/test_preview.py +++ /dev/null @@ -1,82 +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 - -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", - ) diff --git a/tests/gen_sim/scene_engine/test_scene_export.py b/tests/gen_sim/scene_engine/test_scene_export.py deleted file mode 100644 index c78bbbfbe..000000000 --- a/tests/gen_sim/scene_engine/test_scene_export.py +++ /dev/null @@ -1,84 +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 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 deleted file mode 100644 index af768eb22..000000000 --- a/tests/gen_sim/scene_engine/test_scene_generation_utils.py +++ /dev/null @@ -1,102 +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 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, - ) From 2d64a80f3886c207ed02bc9a7673b24ecef5597b Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Tue, 4 Aug 2026 15:34:28 +0800 Subject: [PATCH 32/41] Deleted json config in setup.py --- setup.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/setup.py b/setup.py index 17f1f055c..d3bf8fb98 100644 --- a/setup.py +++ b/setup.py @@ -120,9 +120,6 @@ 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, From 82acb6548932da871106f477a63b42d8fb07d974 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Tue, 4 Aug 2026 15:40:06 +0800 Subject: [PATCH 33/41] Added md + rst in docs/ --- docs/source/features/generative_sim/index.rst | 1 + .../features/generative_sim/scene_engine.md | 83 +++++++++++++++++++ 2 files changed, 84 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..03e4b3fa4 --- /dev/null +++ b/docs/source/features/generative_sim/scene_engine.md @@ -0,0 +1,83 @@ +# Scene Engine + +Scene Engine reconstructs a table-top scene from one image. It identifies the +table and visible objects, segments their masks, generates simulation-ready +meshes, refines the object layout on the table, and exports a scene that can be +loaded by EmbodiChain. + +## Quick Start + +Install EmbodiChain with the `gensim` extra first; see +[Installation](../../quick_start/install.md#optional-generative-simulation-gensim). + +Configure the required services in `embodichain/gen_sim/.env`, then generate a +scene: + +```bash +embodichain scene-engine \ + --image /path/to/scene.png \ + --output_root /path/to/scene_output +``` + +The same command is available through the package entry point: + +```bash +python -m embodichain scene-engine \ + --image /path/to/scene.png \ + --output_root /path/to/scene_output +``` + +## Configuration + +Scene Engine reads the LLM, segmentation, and geometry-generation settings +from `embodichain/gen_sim/.env`: + +```bash +OPENAI_API_KEY="your-api-key" +OPENAI_MODEL="your-model" +OPENAI_BASE_URL="https://api.openai.com/v1" +SCENE_ENGINE_OPENAI_DEFAULT_QUERY="{}" +OPENAI_MAX_ATTEMPTS=3 + +SCENE_ENGINE_IMAGE_SEGMENTATION_BASE_URL="http://host:port" +SCENE_ENGINE_IMAGE_SEGMENTATION_TIMEOUT_S=30 +SCENE_ENGINE_IMAGE_SEGMENTATION_MAX_ATTEMPTS=3 +SCENE_ENGINE_IMAGE_SEGMENTATION_HEALTH_PATH="/health" +SCENE_ENGINE_IMAGE_SEGMENTATION_SINGLE_OBJECT_PATH="/predict" + +SCENE_ENGINE_GEOMETRY_GENERATION_BASE_URL="http://host:port" +SCENE_ENGINE_GEOMETRY_GENERATION_TIMEOUT_S=600 +SCENE_ENGINE_GEOMETRY_GENERATION_MAX_ATTEMPTS=3 +SCENE_ENGINE_GEOMETRY_GENERATION_HEALTH_PATH="/health" +SCENE_ENGINE_GEOMETRY_GENERATION_OBJECTS_PATH="/generate_multiple_objects" +``` + +## Processing Flow + +- **Scene understanding**: analyzes the image and segments the table and visible objects. +- **Scene generation**: generates meshes, prepares SimReady geometry, detects the table support surface, and refines the table-top layout. +- **Scene export**: copies the final GLBs and writes a portable z-up scene export. + +## Output and Preview + +The important final outputs are: + +```text +scene_output/ +|-- scene_understanding/ # Object analysis, masks, and stage JSON +|-- scene_generation/ # Generated, SimReady, and layout-debug artifacts +`-- scene_export/ + |-- mesh_assets/ # Final GLBs + `-- scene_config.json # Exported scene description +``` + +Validate the export without opening a window: + +```bash +embodichain preview-scene \ + --output_root /path/to/scene_output \ + --headless +``` + +For an interactive preview, omit `--headless`. Add `--viser` to publish the +scene through Viser. From 8979f687f04210d09f757f088f37af33aa8cf85b Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Tue, 4 Aug 2026 15:56:10 +0800 Subject: [PATCH 34/41] Modified help, description and epilog --- embodichain/__main__.py | 2 +- embodichain/gen_sim/scene_engine/cli/start.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/embodichain/__main__.py b/embodichain/__main__.py index a841c9096..975e7649d 100644 --- a/embodichain/__main__.py +++ b/embodichain/__main__.py @@ -54,7 +54,7 @@ class Command: Command( name="scene-engine", target="embodichain.gen_sim.scene_engine.cli.start:main", - help="Generate a scene export from an input image.", + help="Generate a scene export from an input image using gen_sim/.env.", ), Command( name="preview-scene", diff --git a/embodichain/gen_sim/scene_engine/cli/start.py b/embodichain/gen_sim/scene_engine/cli/start.py index 637821a1a..59454b09f 100644 --- a/embodichain/gen_sim/scene_engine/cli/start.py +++ b/embodichain/gen_sim/scene_engine/cli/start.py @@ -53,7 +53,8 @@ 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="Generate a Scene Engine export from one input image.", + epilog="Service settings are read from embodichain/gen_sim/.env.", ) parser.add_argument( "--image", From 4f4102b569cf3640d4297eb421dc11b201a88682 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Tue, 4 Aug 2026 16:03:13 +0800 Subject: [PATCH 35/41] Added test files --- tests/gen_sim/scene_engine/test_clients.py | 272 ++++++++++++++++++ tests/gen_sim/scene_engine/test_config.py | 100 +++++++ .../scene_engine/test_group_table_aligner.py | 69 +++++ .../test_scene_core_and_export.py | 143 +++++++++ .../scene_engine/test_scene_understanding.py | 85 ++++++ .../scene_engine/test_support_and_layout.py | 176 ++++++++++++ 6 files changed, 845 insertions(+) create mode 100644 tests/gen_sim/scene_engine/test_clients.py create mode 100644 tests/gen_sim/scene_engine/test_config.py create mode 100644 tests/gen_sim/scene_engine/test_group_table_aligner.py create mode 100644 tests/gen_sim/scene_engine/test_scene_core_and_export.py create mode 100644 tests/gen_sim/scene_engine/test_scene_understanding.py create mode 100644 tests/gen_sim/scene_engine/test_support_and_layout.py diff --git a/tests/gen_sim/scene_engine/test_clients.py b/tests/gen_sim/scene_engine/test_clients.py new file mode 100644 index 000000000..59859eb84 --- /dev/null +++ b/tests/gen_sim/scene_engine/test_clients.py @@ -0,0 +1,272 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import pytest + +from embodichain.gen_sim.scene_engine.clients import geometry_generation +from embodichain.gen_sim.scene_engine.clients import image_segmentation +from embodichain.gen_sim.scene_engine.llms import load_config + + +class _Response: + """Minimal successful HTTP response used by client unit tests.""" + + def __init__(self, payload: object, *, content: bytes = b"") -> None: + self._payload = payload + self.content = content + + def raise_for_status(self) -> None: + return None + + def json(self) -> object: + return self._payload + + +class _Session: + """Capture HTTP calls without contacting an external service.""" + + def __init__(self, *, get_payload: object, post_payload: object | None = None) -> None: + self.get_payload = get_payload + self.post_payload = post_payload + self.get_calls: list[tuple[str, int]] = [] + self.post_call: dict[str, object] | None = None + + def get(self, url: str, *, timeout: int) -> _Response: + self.get_calls.append((url, timeout)) + return _Response(self.get_payload) + + def post(self, url: str, **kwargs: object) -> _Response: + self.post_call = {"url": url, **kwargs} + return _Response(self.post_payload) + + def close(self) -> None: + return None + + +def test_clients_load_their_required_dotenv_values( + monkeypatch: pytest.MonkeyPatch, +) -> None: + geometry_values = { + "SCENE_ENGINE_GEOMETRY_GENERATION_BASE_URL": "http://geometry/", + "SCENE_ENGINE_GEOMETRY_GENERATION_TIMEOUT_S": "60", + "SCENE_ENGINE_GEOMETRY_GENERATION_MAX_ATTEMPTS": "2", + "SCENE_ENGINE_GEOMETRY_GENERATION_HEALTH_PATH": "/health", + "SCENE_ENGINE_GEOMETRY_GENERATION_OBJECTS_PATH": "/objects", + } + segmentation_values = { + "SCENE_ENGINE_IMAGE_SEGMENTATION_BASE_URL": "http://segment/", + "SCENE_ENGINE_IMAGE_SEGMENTATION_TIMEOUT_S": "30", + "SCENE_ENGINE_IMAGE_SEGMENTATION_MAX_ATTEMPTS": "2", + "SCENE_ENGINE_IMAGE_SEGMENTATION_HEALTH_PATH": "/health", + "SCENE_ENGINE_IMAGE_SEGMENTATION_SINGLE_OBJECT_PATH": "/predict", + } + llm_values = { + "OPENAI_API_KEY": "test-key", + "OPENAI_MODEL": "test-model", + "OPENAI_BASE_URL": "http://llm/v1/", + "SCENE_ENGINE_OPENAI_DEFAULT_QUERY": '{"api-version": "1"}', + "OPENAI_MAX_ATTEMPTS": "2", + } + monkeypatch.setattr( + geometry_generation, "read_scene_engine_env_values", lambda *_: geometry_values + ) + monkeypatch.setattr( + image_segmentation, + "read_scene_engine_env_values", + lambda *_: segmentation_values, + ) + monkeypatch.setattr( + load_config, "read_scene_engine_env_values", lambda *_: llm_values + ) + + geometry_client = geometry_generation.GeometryGenerationClient.from_dotenv() + segmentation_client = image_segmentation.ImageSegmentationClient.from_dotenv() + llm_client_config = load_config.load_llm_config() + + assert geometry_client._base_url == "http://geometry" + assert geometry_client._generate_objects_path == "/objects" + assert segmentation_client._base_url == "http://segment" + assert segmentation_client._segment_single_object_path == "/predict" + assert llm_client_config.default_query == {"api-version": "1"} + assert llm_client_config.base_url == "http://llm/v1" + + +def test_geometry_dotenv_config_rejects_invalid_timeout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + values = { + "SCENE_ENGINE_GEOMETRY_GENERATION_BASE_URL": "http://geometry", + "SCENE_ENGINE_GEOMETRY_GENERATION_TIMEOUT_S": "0", + "SCENE_ENGINE_GEOMETRY_GENERATION_MAX_ATTEMPTS": "1", + "SCENE_ENGINE_GEOMETRY_GENERATION_HEALTH_PATH": "/health", + "SCENE_ENGINE_GEOMETRY_GENERATION_OBJECTS_PATH": "/objects", + } + monkeypatch.setattr( + geometry_generation, "read_scene_engine_env_values", lambda *_: values + ) + + with pytest.raises(ValueError, match="TIMEOUT_S must be at least 1"): + geometry_generation.GeometryGenerationClient.from_dotenv() + + +def test_service_health_checks_use_the_configured_health_path() -> None: + geometry_session = _Session(get_payload={"ok": True}) + geometry_client = geometry_generation.GeometryGenerationClient( + base_url="http://geometry", + timeout_s=60, + max_attempts=1, + health_path="/health", + generate_objects_path="/objects", + session=geometry_session, + ) + segmentation_session = _Session(get_payload={"ok": True}) + segmentation_client = image_segmentation.ImageSegmentationClient( + base_url="http://segment", + timeout_s=30, + max_attempts=1, + health_path="/health", + segment_single_object_path="/predict", + session=segmentation_session, + ) + + geometry_client.check_health() + segmentation_client.check_health() + + assert geometry_session.get_calls == [("http://geometry/health", 10)] + assert segmentation_session.get_calls == [("http://segment/health", 30)] + + +def test_segmentation_client_posts_prompt_and_returns_rle_masks(tmp_path: Path) -> None: + image_path = tmp_path / "scene.png" + image_path.write_bytes(b"png") + rle_mask = {"counts": [0, 1], "size": [1, 1]} + session = _Session( + get_payload={"ok": True}, + post_payload={"result": {"masks": [rle_mask]}}, + ) + client = image_segmentation.ImageSegmentationClient( + base_url="http://segment", + timeout_s=30, + max_attempts=1, + health_path="/health", + segment_single_object_path="/predict", + session=session, + ) + + assert client.segment_single_object(image_path=image_path, prompt="table") == [ + rle_mask + ] + assert session.post_call is not None + assert session.post_call["url"] == "http://segment/predict" + assert session.post_call["data"] == {"prompt": "table"} + + +def test_segmentation_client_accepts_instance_mask_response() -> None: + rle_mask = {"counts": [0, 1], "size": [1, 1]} + + masks = image_segmentation._extract_rle_masks( + {"data": {"instances": [{"mask_rle": rle_mask}]}} + ) + + assert masks == [rle_mask] + + +def test_geometry_response_requires_matching_ordered_objects() -> None: + response: dict[str, Any] = { + "ok": True, + "result": { + "objects": [ + { + "name": "table_001", + "mesh": "/results/table.glb", + "rotation_quaternion_wxyz": [1, 0, 0, 0], + "translation": [0, 1, 2], + "scale": [1, 1, 1], + } + ] + }, + } + + objects = geometry_generation._parse_objects_response( + response, + object_ids=["table_001"], + ) + + assert objects == [ + { + "mesh": "/results/table.glb", + "rotation_quaternion_wxyz": [1.0, 0.0, 0.0, 0.0], + "translation": [0.0, 1.0, 2.0], + "scale": [1.0, 1.0, 1.0], + } + ] + + +def test_geometry_client_posts_masks_and_downloads_glbs(tmp_path: Path) -> None: + class GeometrySession(_Session): + def post(self, url: str, **kwargs: object) -> _Response: + self.post_call = {"url": url, **kwargs} + return _Response( + { + "ok": True, + "result": { + "objects": [ + { + "name": "cup", + "mesh": "/results/cup.glb", + "rotation_quaternion_wxyz": [1, 0, 0, 0], + "translation": [0, 0, 0], + "scale": [1, 1, 1], + } + ] + }, + } + ) + + def get(self, url: str, *, timeout: int) -> _Response: + self.get_calls.append((url, timeout)) + return _Response({}, content=b"glTF-mesh") + + image_path = tmp_path / "scene.png" + mask_path = tmp_path / "cup.png" + image_path.write_bytes(b"png") + mask_path.write_bytes(b"png") + session = GeometrySession(get_payload={"ok": True}) + client = geometry_generation.GeometryGenerationClient( + base_url="http://geometry", + timeout_s=30, + max_attempts=1, + health_path="/health", + generate_objects_path="/objects", + session=session, + ) + + _, objects = client.generate_objects( + image_path=image_path, + object_masks=[("cup", mask_path)], + output_root=tmp_path / "output", + ) + + assert objects[0]["mesh"] == "/results/cup.glb" + assert session.post_call is not None + assert session.post_call["url"] == "http://geometry/objects" + assert (tmp_path / "output/cup.glb").read_bytes() == b"glTF-mesh" diff --git a/tests/gen_sim/scene_engine/test_config.py b/tests/gen_sim/scene_engine/test_config.py new file mode 100644 index 000000000..01914e144 --- /dev/null +++ b/tests/gen_sim/scene_engine/test_config.py @@ -0,0 +1,100 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from embodichain.gen_sim.scene_engine.cli import start +from embodichain.gen_sim.scene_engine.configs import environment + + +def test_read_scene_engine_env_values_reads_requested_keys( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + env_path = tmp_path / ".env" + env_path.write_text('OPENAI_MODEL="test-model"\nUNRELATED_VALUE=ignored\n') + monkeypatch.setattr(environment, "_SCENE_ENGINE_ENV_PATH", env_path) + + assert environment.read_scene_engine_env_values("OPENAI_MODEL") == { + "OPENAI_MODEL": "test-model" + } + + +def test_read_scene_engine_env_values_reports_missing_keys( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + env_path = tmp_path / ".env" + env_path.write_text("OPENAI_MODEL=test-model\n") + monkeypatch.setattr(environment, "_SCENE_ENGINE_ENV_PATH", env_path) + + with pytest.raises(ValueError, match="OPENAI_API_KEY"): + environment.read_scene_engine_env_values("OPENAI_MODEL", "OPENAI_API_KEY") + + +def test_scene_engine_help_exposes_only_runtime_arguments( + capsys: pytest.CaptureFixture[str], +) -> None: + with pytest.raises(SystemExit) as exc_info: + start.main(["--help"]) + + assert exc_info.value.code == 0 + output = capsys.readouterr().out + assert "--image" in output + assert "--output_root" in output + assert "gen_sim/.env" in output + assert "--config" not in output + + +def test_scene_engine_cli_forwards_validated_paths( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + image_path = tmp_path / "scene.png" + image_path.write_bytes(b"png") + captured: dict[str, Path] = {} + + def generate_scene(*, image_path: Path, output_root: Path) -> None: + captured["image_path"] = image_path + captured["output_root"] = output_root + + monkeypatch.setattr(start, "generate_scene_from_image", generate_scene) + output_root = tmp_path / "output" + + start.cli_scene_engine(image_path, output_root) + + assert captured == { + "image_path": image_path.resolve(), + "output_root": output_root.resolve(), + } + + +@pytest.mark.parametrize("image_name", ["missing.png", "scene.gif"]) +def test_scene_engine_cli_rejects_invalid_image_inputs( + tmp_path: Path, + image_name: str, +) -> None: + image_path = tmp_path / image_name + if image_path.suffix == ".gif": + image_path.write_bytes(b"gif") + + with pytest.raises((FileNotFoundError, ValueError)): + start.cli_scene_engine(image_path, tmp_path / "output") diff --git a/tests/gen_sim/scene_engine/test_group_table_aligner.py b/tests/gen_sim/scene_engine/test_group_table_aligner.py new file mode 100644 index 000000000..f34bbbeeb --- /dev/null +++ b/tests/gen_sim/scene_engine/test_group_table_aligner.py @@ -0,0 +1,69 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from pathlib import Path + +import pytest +import trimesh + +from embodichain.gen_sim.scene_engine.pipeline.utils.assets_group_table_aligner import ( + AssetsGroupTableAligner, + AssetsGroupTableAlignerConfig, +) + + +def _layout(object_id: str, y: float) -> dict[str, object]: + return { + "id": object_id, + "pos": [0.0, y, 0.0], + "rot": [0.0, 0.0, 0.0], + "scale": [1.0, 1.0, 1.0], + } + + +def test_group_table_aligner_preserves_relative_vertical_offsets(tmp_path: Path) -> None: + trimesh.creation.box(extents=(2.0, 1.0, 2.0)).export(tmp_path / "table.glb") + trimesh.creation.box(extents=(0.5, 1.0, 0.5)).export(tmp_path / "first.glb") + trimesh.creation.box(extents=(0.5, 1.0, 0.5)).export(tmp_path / "second.glb") + assets_layout = [_layout("first", 0.0), _layout("second", 0.3)] + + _, aligned_assets = AssetsGroupTableAligner( + table_layout=_layout("table", 0.0), + assets_layout=assets_layout, + geometry_root=tmp_path, + config=AssetsGroupTableAlignerConfig(clearance_m=0.1), + ).align() + + assert aligned_assets[0]["pos"][1] > assets_layout[0]["pos"][1] # type: ignore[index] + assert aligned_assets[1]["pos"][1] - aligned_assets[0]["pos"][1] == pytest.approx( # type: ignore[index] + 0.3 + ) + + +def test_group_table_aligner_returns_empty_assets_without_mesh_loading(tmp_path: Path) -> None: + table_layout = _layout("table", 0.0) + + aligned_table, aligned_assets = AssetsGroupTableAligner( + table_layout=table_layout, + assets_layout=[], + geometry_root=tmp_path, + ).align() + + assert aligned_table is table_layout + assert aligned_assets == [] diff --git a/tests/gen_sim/scene_engine/test_scene_core_and_export.py b/tests/gen_sim/scene_engine/test_scene_core_and_export.py new file mode 100644 index 000000000..14bc08a12 --- /dev/null +++ b/tests/gen_sim/scene_engine/test_scene_core_and_export.py @@ -0,0 +1,143 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import pytest + +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_object import ObjectPhysics, SceneObject +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_exporter import SceneExporter + + +def _scene_object( + *, + object_id: str, + kind: str, + glb_path: Path | None = None, + physics: ObjectPhysics | None = None, +) -> SceneObject: + return SceneObject( + id=object_id, + kind=kind, # type: ignore[arg-type] + category=kind, + name=object_id, + description=f"{kind} object", + simready_glb_path=str(glb_path) if glb_path is not None else None, + rot=[0.0, 0.0, 0.0], + pos=[1.0, 2.0, 3.0], + scale=[1.0, 2.0, 3.0], + physics=physics, + ) + + +def _physics(body_type: str) -> ObjectPhysics: + return ObjectPhysics( + body_type=body_type, # type: ignore[arg-type] + attrs={"mass": 1.0, "static_friction": 0.8}, + max_convex_hull_num=16, + ) + + +def test_scene_returns_one_table_and_ordered_assets() -> None: + table = _scene_object(object_id="table", kind="table") + asset = _scene_object(object_id="cup", kind="asset") + scene = Scene(objects=[table, asset]) + + assert scene.table is table + assert scene.assets == [asset] + assert scene.to_dict()["objects"][0]["id"] == "table" # type: ignore[index] + + +def test_scene_rejects_multiple_tables() -> None: + scene = Scene( + objects=[ + _scene_object(object_id="table_001", kind="table"), + _scene_object(object_id="table_002", kind="table"), + ] + ) + + with pytest.raises(ValueError, match="only one table"): + _ = scene.table + + +@pytest.mark.parametrize( + ("body_type", "attrs", "hulls"), + [ + ("static", {"mass": 1.0}, 1), + ("dynamic", {}, 1), + ("dynamic", {"mass": 1.0}, 0), + ], +) +def test_object_physics_rejects_invalid_values( + body_type: str, + attrs: dict[str, float], + hulls: int, +) -> None: + with pytest.raises(ValueError): + _physics = ObjectPhysics( # noqa: F841 + body_type=body_type, # type: ignore[arg-type] + attrs=attrs, + max_convex_hull_num=hulls, + ) + + +def test_scene_export_copies_meshes_and_converts_y_up_pose(tmp_path: Path) -> None: + table_glb = tmp_path / "table.glb" + asset_glb = tmp_path / "cup.glb" + table_glb.write_bytes(b"glTF-table") + asset_glb.write_bytes(b"glTF-cup") + table = _scene_object( + object_id="table", + kind="table", + glb_path=table_glb, + physics=_physics("kinematic"), + ) + asset = _scene_object( + object_id="cup", + kind="asset", + glb_path=asset_glb, + physics=_physics("dynamic"), + ) + + export_path = SceneExporter( + scene=Scene(objects=[table, asset]), + output_root=tmp_path / "output", + ).export() + exported = json.loads(export_path.read_text(encoding="utf-8")) + + assert (export_path.parent / "mesh_assets/table/table.glb").read_bytes() == b"glTF-table" + assert (export_path.parent / "mesh_assets/cup/cup.glb").read_bytes() == b"glTF-cup" + entry = exported["rigid_object"][0] + assert entry["uid"] == "cup" + assert entry["body_type"] == "dynamic" + assert entry["init_pos"] == [1.0, -3.0, 2.0] + assert entry["body_scale"] == [1.0, 2.0, 3.0] + assert np.allclose(entry["init_rot"], [0.0, 0.0, 0.0]) + + +def test_scene_export_requires_final_physics(tmp_path: Path) -> None: + glb_path = tmp_path / "table.glb" + glb_path.write_bytes(b"glTF") + table = _scene_object(object_id="table", kind="table", glb_path=glb_path) + + with pytest.raises(ValueError, match="no SimReady physics"): + SceneExporter(scene=Scene(objects=[table]), output_root=tmp_path).export() diff --git a/tests/gen_sim/scene_engine/test_scene_understanding.py b/tests/gen_sim/scene_engine/test_scene_understanding.py new file mode 100644 index 000000000..1fdb51f7b --- /dev/null +++ b/tests/gen_sim/scene_engine/test_scene_understanding.py @@ -0,0 +1,85 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.pipeline import scene_understanding + + +def _response(*, asset_name: str = "cup") -> str: + return json.dumps( + { + "table": { + "category": "dining_table", + "name": "wooden table", + "description": "A rectangular wooden table.", + }, + "assets": [ + { + "category": "cup", + "name": asset_name, + "description": "A small ceramic cup.", + } + ], + } + ) + + +def test_image_object_analysis_parses_code_fence_and_assigns_stable_ids() -> None: + scene = scene_understanding._parse_image_object_analysis_response( + f"```json\n{_response()}\n```" + ) + + assert scene.table is not None + assert scene.table.id == "table" + assert [asset.id for asset in scene.assets] == ["cup_001"] + + +def test_image_object_analysis_rejects_location_words_in_object_names() -> None: + with pytest.raises(ValueError, match="must not contain location"): + scene_understanding._parse_image_object_analysis_response( + _response(asset_name="left cup") + ) + + +def test_image_object_analysis_retries_then_updates_scene(tmp_path: Path) -> None: + class VLM: + def __init__(self) -> None: + self.responses = ["not-json", _response()] + + def complete(self, **_: object) -> str: + return self.responses.pop(0) + + image_path = tmp_path / "scene.png" + image_path.write_bytes(b"png") + scene = Scene() + + scene_understanding._analyze_image_objects( + scene=scene, + image_path=image_path, + vlm_client=VLM(), # type: ignore[arg-type] + json_max_attempts=2, + ) + + assert scene.table is not None + assert [asset.id for asset in scene.assets] == ["cup_001"] diff --git a/tests/gen_sim/scene_engine/test_support_and_layout.py b/tests/gen_sim/scene_engine/test_support_and_layout.py new file mode 100644 index 000000000..59057c1b2 --- /dev/null +++ b/tests/gen_sim/scene_engine/test_support_and_layout.py @@ -0,0 +1,176 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +import numpy as np +import pytest +from shapely.geometry import Point, Polygon +import trimesh + +from embodichain.gen_sim.scene_engine.pipeline.utils.assets_group_layout_optimizer import ( + AssetsSupportLayoutOptimizer, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.assets_group_support_clamp import ( + AssetsGroupSupportClamp, + AssetsGroupSupportClampConfig, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.table_support_surface import ( + TableSupportSurfaceDetector, +) + + +def _aabb(minimum_x: float, minimum_y: float, maximum_x: float, maximum_y: float) -> np.ndarray: + return np.array( + [ + [minimum_x, minimum_y], + [maximum_x, minimum_y], + [maximum_x, maximum_y], + [minimum_x, maximum_y], + ], + dtype=float, + ) + + +def _layout(object_id: str, x: float, y: float) -> dict[str, object]: + return {"id": object_id, "pos": [x, 0.0, -y]} + + +def _top_mesh(vertices_xy: list[tuple[float, float]], faces: list[list[int]], z: float) -> trimesh.Trimesh: + return trimesh.Trimesh( + vertices=np.array([[x, y, z] for x, y in vertices_xy], dtype=float), + faces=np.array(faces, dtype=int), + process=False, + ) + + +def test_support_detector_preserves_an_l_shaped_support_contour() -> None: + mesh = _top_mesh( + [(0, 0), (2, 0), (2, 1), (1, 1), (1, 2), (0, 2)], + [[0, 1, 3], [1, 2, 3], [0, 3, 5], [3, 4, 5]], + z=1.0, + ) + + region = TableSupportSurfaceDetector(table_world_mesh=mesh).detect() + + assert region.top_z == pytest.approx(1.0) + assert region.support_polygon.area == pytest.approx(3.0) + assert region.support_polygon.covers(Point(0.5, 1.5)) + assert not region.support_polygon.covers(Point(1.5, 1.5)) + + +def test_support_detector_prefers_main_tabletop_over_small_higher_piece() -> None: + main = _top_mesh( + [(0, 0), (2, 0), (2, 2), (0, 2)], [[0, 1, 2], [0, 2, 3]], z=1.0 + ) + decoration = _top_mesh( + [(0.25, 0.25), (0.75, 0.25), (0.75, 0.75), (0.25, 0.75)], + [[0, 1, 2], [0, 2, 3]], + z=1.2, + ) + mesh = trimesh.util.concatenate([main, decoration]) + + region = TableSupportSurfaceDetector(table_world_mesh=mesh).detect() + + assert region.top_z == pytest.approx(1.0) + assert region.support_polygon.area == pytest.approx(4.0) + + +def test_group_clamp_preserves_relative_layout_while_moving_inside_support() -> None: + support = Polygon([(0, 0), (4, 0), (4, 4), (0, 4)]) + aabbs = { + "first": _aabb(-0.5, 1.0, 0.5, 2.0), + "second": _aabb(1.0, 1.0, 2.0, 2.0), + } + layouts = [_layout("first", 0.0, 1.5), _layout("second", 1.5, 1.5)] + + refined = AssetsGroupSupportClamp( + support_region=support, + assets_aabb_2d_z_up_world_corners_by_id=aabbs, + assets_layout=layouts, + config=AssetsGroupSupportClampConfig(grid_resolution_m=0.05), + ).clamp() + + first, second = refined + assert first["pos"][0] > layouts[0]["pos"][0] # type: ignore[index] + assert second["pos"][0] - first["pos"][0] == pytest.approx(1.5) # type: ignore[index] + assert second["pos"][2] - first["pos"][2] == pytest.approx(0.0) # type: ignore[index] + + +def test_group_clamp_returns_unchanged_layout_when_already_contained() -> None: + layout = _layout("cup", 1.5, 1.5) + refined = AssetsGroupSupportClamp( + support_region=Polygon([(0, 0), (3, 0), (3, 3), (0, 3)]), + assets_aabb_2d_z_up_world_corners_by_id={"cup": _aabb(1.0, 1.0, 2.0, 2.0)}, + assets_layout=[layout], + ).clamp() + + assert refined == [layout] + + +def test_group_clamp_reports_infeasible_oversized_group() -> None: + clamp = AssetsGroupSupportClamp( + support_region=Polygon([(0, 0), (2, 0), (2, 2), (0, 2)]), + assets_aabb_2d_z_up_world_corners_by_id={"large": _aabb(0, 0, 3, 3)}, + assets_layout=[_layout("large", 1.5, 1.5)], + config=AssetsGroupSupportClampConfig(grid_resolution_m=0.1), + ) + + with pytest.raises(ValueError, match="cannot be placed"): + clamp.clamp() + + +def test_layout_optimizer_resolves_a_simple_pair_overlap() -> None: + optimizer = AssetsSupportLayoutOptimizer( + support_region=Polygon([(0, 0), (5, 0), (5, 5), (0, 5)]), + assets_aabb_2d_z_up_world_corners_by_id={ + "first": _aabb(1.0, 1.0, 2.0, 2.0), + "second": _aabb(1.5, 1.0, 2.5, 2.0), + }, + assets_layout=[_layout("first", 1.5, 1.5), _layout("second", 2.0, 1.5)], + ) + + refined = optimizer.optimize() + + refined_offsets = np.array( + [ + [refined[index]["pos"][0] - optimizer.assets_layout[index]["pos"][0], # type: ignore[index] + optimizer.assets_layout[index]["pos"][2] - refined[index]["pos"][2]] # type: ignore[index] + for index in range(2) + ] + ) + base_aabbs = np.stack( + [ + optimizer.assets_aabb_2d_z_up_world_corners_by_id["first"], + optimizer.assets_aabb_2d_z_up_world_corners_by_id["second"], + ] + ) + assert not optimizer._overlaps(base_aabbs, refined_offsets) + + +def test_layout_optimizer_rejects_unresolvable_overlap() -> None: + optimizer = AssetsSupportLayoutOptimizer( + support_region=Polygon([(0, 0), (2, 0), (2, 2), (0, 2)]), + assets_aabb_2d_z_up_world_corners_by_id={ + "first": _aabb(0, 0, 2, 2), + "second": _aabb(0, 0, 2, 2), + }, + assets_layout=[_layout("first", 1.0, 1.0), _layout("second", 1.0, 1.0)], + ) + + with pytest.raises(ValueError, match="cannot be resolved"): + optimizer.optimize() From d71033126b69ec11f27197e90a54595ab48f96f9 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Tue, 4 Aug 2026 16:12:07 +0800 Subject: [PATCH 36/41] Added extra dependencies in gen_sim: scene_engine + Added CI --- .github/workflows/main.yml | 5 +++++ pyproject.toml | 20 +++++++++++++++++++- 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index b85deaea9..3ad17a9bc 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -147,6 +147,11 @@ jobs: --extra-index-url https://download.blender.org/pypi/ pip install pytest-xdist + - name: Verify Scene Engine gensim installation + run: | + python -c "import matplotlib, numpy, open3d, requests, scipy, shapely, trimesh; from PIL import Image; import embodichain.gen_sim.scene_engine.pipeline.generate" + pytest tests/gen_sim/scene_engine -q + - name: Run default tests run: | echo "Default test suite (GPU-marked tests are skipped)" diff --git a/pyproject.toml b/pyproject.toml index d1daf53a1..2a127ccdd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,9 +52,27 @@ dependencies = [ ] [project.optional-dependencies] +scene-engine = [ + "requests", + "Pillow", + "numpy", + "scipy", + "shapely", + "trimesh", + "open3d", + "matplotlib" +] gensim = [ "bpy", - "pyrender==0.1.45" + "pyrender==0.1.45", + "requests", + "Pillow", + "numpy", + "scipy", + "shapely", + "trimesh", + "open3d", + "matplotlib" ] # cuRobo V2 is distributed from its source repository and provides separate # dependency sets for CUDA 12 and CUDA 13. Keep it optional so CPU-only and From 2298e541c54bcd73c582eedd48401dd32c6f32f5 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Tue, 4 Aug 2026 17:27:05 +0800 Subject: [PATCH 37/41] Fixed CI lint --- tests/gen_sim/scene_engine/test_clients.py | 4 +++- .../scene_engine/test_group_table_aligner.py | 8 ++++++-- .../scene_engine/test_scene_core_and_export.py | 9 +++++++-- .../scene_engine/test_support_and_layout.py | 18 +++++++++++------- 4 files changed, 27 insertions(+), 12 deletions(-) diff --git a/tests/gen_sim/scene_engine/test_clients.py b/tests/gen_sim/scene_engine/test_clients.py index 59859eb84..513f0c5c0 100644 --- a/tests/gen_sim/scene_engine/test_clients.py +++ b/tests/gen_sim/scene_engine/test_clients.py @@ -44,7 +44,9 @@ def json(self) -> object: class _Session: """Capture HTTP calls without contacting an external service.""" - def __init__(self, *, get_payload: object, post_payload: object | None = None) -> None: + def __init__( + self, *, get_payload: object, post_payload: object | None = None + ) -> None: self.get_payload = get_payload self.post_payload = post_payload self.get_calls: list[tuple[str, int]] = [] diff --git a/tests/gen_sim/scene_engine/test_group_table_aligner.py b/tests/gen_sim/scene_engine/test_group_table_aligner.py index f34bbbeeb..13f5a95ee 100644 --- a/tests/gen_sim/scene_engine/test_group_table_aligner.py +++ b/tests/gen_sim/scene_engine/test_group_table_aligner.py @@ -37,7 +37,9 @@ def _layout(object_id: str, y: float) -> dict[str, object]: } -def test_group_table_aligner_preserves_relative_vertical_offsets(tmp_path: Path) -> None: +def test_group_table_aligner_preserves_relative_vertical_offsets( + tmp_path: Path, +) -> None: trimesh.creation.box(extents=(2.0, 1.0, 2.0)).export(tmp_path / "table.glb") trimesh.creation.box(extents=(0.5, 1.0, 0.5)).export(tmp_path / "first.glb") trimesh.creation.box(extents=(0.5, 1.0, 0.5)).export(tmp_path / "second.glb") @@ -56,7 +58,9 @@ def test_group_table_aligner_preserves_relative_vertical_offsets(tmp_path: Path) ) -def test_group_table_aligner_returns_empty_assets_without_mesh_loading(tmp_path: Path) -> None: +def test_group_table_aligner_returns_empty_assets_without_mesh_loading( + tmp_path: Path, +) -> None: table_layout = _layout("table", 0.0) aligned_table, aligned_assets = AssetsGroupTableAligner( diff --git a/tests/gen_sim/scene_engine/test_scene_core_and_export.py b/tests/gen_sim/scene_engine/test_scene_core_and_export.py index 14bc08a12..d016d7aa7 100644 --- a/tests/gen_sim/scene_engine/test_scene_core_and_export.py +++ b/tests/gen_sim/scene_engine/test_scene_core_and_export.py @@ -24,7 +24,10 @@ import pytest from embodichain.gen_sim.scene_engine.core.scene import Scene -from embodichain.gen_sim.scene_engine.core.scene_object import ObjectPhysics, SceneObject +from embodichain.gen_sim.scene_engine.core.scene_object import ( + ObjectPhysics, + SceneObject, +) from embodichain.gen_sim.scene_engine.pipeline.utils.scene_exporter import SceneExporter @@ -124,7 +127,9 @@ def test_scene_export_copies_meshes_and_converts_y_up_pose(tmp_path: Path) -> No ).export() exported = json.loads(export_path.read_text(encoding="utf-8")) - assert (export_path.parent / "mesh_assets/table/table.glb").read_bytes() == b"glTF-table" + assert ( + export_path.parent / "mesh_assets/table/table.glb" + ).read_bytes() == b"glTF-table" assert (export_path.parent / "mesh_assets/cup/cup.glb").read_bytes() == b"glTF-cup" entry = exported["rigid_object"][0] assert entry["uid"] == "cup" diff --git a/tests/gen_sim/scene_engine/test_support_and_layout.py b/tests/gen_sim/scene_engine/test_support_and_layout.py index 59057c1b2..98ed30491 100644 --- a/tests/gen_sim/scene_engine/test_support_and_layout.py +++ b/tests/gen_sim/scene_engine/test_support_and_layout.py @@ -34,7 +34,9 @@ ) -def _aabb(minimum_x: float, minimum_y: float, maximum_x: float, maximum_y: float) -> np.ndarray: +def _aabb( + minimum_x: float, minimum_y: float, maximum_x: float, maximum_y: float +) -> np.ndarray: return np.array( [ [minimum_x, minimum_y], @@ -50,7 +52,9 @@ def _layout(object_id: str, x: float, y: float) -> dict[str, object]: return {"id": object_id, "pos": [x, 0.0, -y]} -def _top_mesh(vertices_xy: list[tuple[float, float]], faces: list[list[int]], z: float) -> trimesh.Trimesh: +def _top_mesh( + vertices_xy: list[tuple[float, float]], faces: list[list[int]], z: float +) -> trimesh.Trimesh: return trimesh.Trimesh( vertices=np.array([[x, y, z] for x, y in vertices_xy], dtype=float), faces=np.array(faces, dtype=int), @@ -74,9 +78,7 @@ def test_support_detector_preserves_an_l_shaped_support_contour() -> None: def test_support_detector_prefers_main_tabletop_over_small_higher_piece() -> None: - main = _top_mesh( - [(0, 0), (2, 0), (2, 2), (0, 2)], [[0, 1, 2], [0, 2, 3]], z=1.0 - ) + main = _top_mesh([(0, 0), (2, 0), (2, 2), (0, 2)], [[0, 1, 2], [0, 2, 3]], z=1.0) decoration = _top_mesh( [(0.25, 0.25), (0.75, 0.25), (0.75, 0.75), (0.25, 0.75)], [[0, 1, 2], [0, 2, 3]], @@ -148,8 +150,10 @@ def test_layout_optimizer_resolves_a_simple_pair_overlap() -> None: refined_offsets = np.array( [ - [refined[index]["pos"][0] - optimizer.assets_layout[index]["pos"][0], # type: ignore[index] - optimizer.assets_layout[index]["pos"][2] - refined[index]["pos"][2]] # type: ignore[index] + [ + refined[index]["pos"][0] - optimizer.assets_layout[index]["pos"][0], # type: ignore[index] + optimizer.assets_layout[index]["pos"][2] - refined[index]["pos"][2], + ] # type: ignore[index] for index in range(2) ] ) From 067d40e514b104c3c33dc6bce0e9dc6f974f8b2b Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Tue, 4 Aug 2026 18:05:08 +0800 Subject: [PATCH 38/41] Fixed scene-engine test file: make it have different name with the simready pipeline ones --- .../scene_engine/{test_config.py => test_scene_engine_config.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename tests/gen_sim/scene_engine/{test_config.py => test_scene_engine_config.py} (100%) diff --git a/tests/gen_sim/scene_engine/test_config.py b/tests/gen_sim/scene_engine/test_scene_engine_config.py similarity index 100% rename from tests/gen_sim/scene_engine/test_config.py rename to tests/gen_sim/scene_engine/test_scene_engine_config.py From 546be9b791709ce1d953a9f6eba20f08b84d1228 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Tue, 4 Aug 2026 18:31:43 +0800 Subject: [PATCH 39/41] Fixed dependencies in toml: scene_engine is under the management of gen_sim --- .../scene_engine/cli/task-1k-expanded.md | 516 ++++++++++++++++++ pyproject.toml | 12 +- 2 files changed, 517 insertions(+), 11 deletions(-) create mode 100644 embodichain/gen_sim/scene_engine/cli/task-1k-expanded.md diff --git a/embodichain/gen_sim/scene_engine/cli/task-1k-expanded.md b/embodichain/gen_sim/scene_engine/cli/task-1k-expanded.md new file mode 100644 index 000000000..0984d1b7d --- /dev/null +++ b/embodichain/gen_sim/scene_engine/cli/task-1k-expanded.md @@ -0,0 +1,516 @@ +# Task 1K 扩增方案（Scene 服务 Action 版） + +> 版本：v0.4 +> 日期：2026-08-04 +> 依据：1k+ tasks.md、task-1k-revision-recommendations.md、taskspec-semantic-protocol.md + +## 0. 核心目标 + +### 0.1 先认识几个词 + +| 词 | 含义 | 例子 | +|---|---|---| +| Scene | 机器人要操作的具体世界 | 一个关闭的抽屉、一个可抓取物体和一个把手 | +| Scene 模式 | 一类具有相同操作条件的 Scene | “关闭抽屉”；“打开抽屉且有可放置区域” | +| Action | 一个动作或一条动作序列 | E9 打开抽屉；E9 → E1 → E10 | +| Action 编排 | 按前置条件和动作结果连接 Easy Action | E9 打开后，E1 才能放入，最后 E10 关闭 | +| Scene Agent | 生成、恢复或编辑 Scene | 从抽屉图片恢复抽屉、把手和内部空间 | +| Action Agent | 在 Scene 中执行和编排 Action | 在打开的抽屉中执行 E1 放入物体 | +| TaskSpec | 规定 Scene 初始条件、goal 和允许 Action 的任务说明 | goal 是 contain(drawer, object) 且 drawer closed | +| Episode | 一次具体执行或扰动后的执行记录 | 同一任务中注入一次抓取滑移后的执行 | +| 任务扩增 | TaskSpec 的目标、条件或因果结构改变，形成有意义的新任务 | “放到桌面”改为“放入抽屉并关闭” | +| 数据扩大 | TaskSpec 不变，只增加 Scene、轨迹或 Episode | 同一任务换桌面、机器人或成功轨迹 | + +### 0.2 哪些算任务扩增 + +| 变化 | 分类 | 例子 | +|---|---|---| +| 改变最终 goal | 任务扩增 | on(object, table) → contain(drawer, object) | +| 增加数量、姿态、顺序或恢复要求 | 任务扩增 | 一个物体 → all；place → open → place → close | +| Scene 出现新状态，但 TaskSpec 和 goal 不变 | 数据扩大 | 同一个放置任务使用打开的不同抽屉场景 | +| 只更换物体、桌面、机器人或动作轨迹 | 数据扩大 | 杯子 A → 杯子 B；UR5 → Franka；轨迹 A → B | + +判断标准： + +```text +TaskSpec 的目标、必要条件或因果结构改变 → 任务扩增 +TaskSpec 不变，只增加 Scene、Action 轨迹或 Episode → 数据扩大 +``` + +Task 1K 的重点不是制造很多不同的物体或桌面，而是： + +``` +Scene Agent 提供能执行原子动作的场景条件 + ↓ +Action Agent 使用 Easy 原子动作 + ↓ +组合成更多有意义的 TaskSpec +``` + +Scene 是 Action 的基础；Action 是改变 Scene 的方式；TaskSpec 负责把二者组合成一个任务。 + +### 0.3 方案先看 + +Task 1K 按以下顺序扩增： + +```text +P0 固定 E1–E12，以及每个 Easy 需要的 Scene 条件 +P1 为每个 Easy 设计多种可执行 Scene 模式 +P2 用一个 Easy 生成最简单的 TaskSpec +P3 用前一个 Action 的结果连接下一个 Easy，形成短动作链 +P4 改变 goal、数量、姿态、顺序或铰接状态，形成有意义的新任务 +P5 检查 Scene、Action 和 goal，统计任务和数据数量 +``` + +核心不是把同一个任务换成很多物体，而是： + +```text +一个 Easy +→ 多种可执行 Scene +→ 多种 goal +→ 合法的 Easy Action 组合 +→ 更多语义不同的 TaskSpec +``` + +## 1. 最小流程 + +Task 生成有两条入口，最后都通过 TaskSpec 连接 Scene 和 Action。 + +### 1.1 TaskSpec 驱动 + +适合从 Easy 动作组合出新的任务： + +``` +组合 Easy Action 和目标 + ↓ +写出 TaskSpec：目标、动作和必要条件 + ↓ +Scene Agent 根据 TaskSpec 生成或恢复 Scene + ↓ +Action Agent 执行动作链 + ↓ +检查最终 goal +``` + +例如： + +``` +目标：把物体放入抽屉并关闭抽屉 +动作：E9 → E1 → E10 +TaskSpec：要求有关闭的抽屉、可抓取物体，最终 contain 且 closed +Scene Agent：生成满足这些条件的 Scene +Action Agent：执行 E9 → E1 → E10 +``` + +### 1.2 Scene 驱动 + +适合从图像、纯文字或编辑后的 Scene 中发现可以生成的任务： + +``` +图像 / 纯文字 / Scene 编辑输入 + ↓ +Scene Agent 生成或恢复 Scene + ↓ +TaskSpec 根据 Scene 中可用的物体、关系和状态，选择 goal 和 Easy Action + ↓ +Action Agent 执行 TaskSpec 规定的动作 + ↓ +检查最终 goal +``` + +例如，Scene Agent 从图像中发现一个关闭的抽屉、一个可抓取物体和一个把手，TaskSpec 可以选择： + +``` +goal：打开抽屉 → E9 +goal：打开并放入物体 → E9 → E1 +goal：打开、放入并关闭 → E9 → E1 → E10 +``` + +两条入口的区别是： + +``` +TaskSpec 驱动：先决定要做什么，再生成 Scene +Scene 驱动：先知道有什么，再决定能做什么 +``` + +## 2. 三个核心部分 + +### 2.1 Scene Agent + +Scene Agent 负责根据 Action 的需要生成、恢复或编辑 Scene。 + +它需要提供： + +- 需要被操作的物体； +- 物体的可抓取、可放置或可交互状态； +- 物体之间的关系； +- 桌面和容器的位置； +- 铰接式物体的开关状态； +- 下一步 Action 所需的目标区域或对象。 + +Scene Agent 当前支持的主要关系： + +| 关系 | 含义 | +|---|---| +| left、right、center | 物体相对于桌面的区域 | +| left_front、left_back | 桌面的左前、左后区域 | +| right_front、right_back | 桌面的右前、右后区域 | +| left_center、right_center | 桌面的左中、右中区域 | +| on | 物体在桌面或支撑面上 | +| stack_on | 一个刚体叠放在另一个刚体上 | +| contain | 刚体或铰接式容器包含刚体物体 | + +Scene Agent 可以生成以下场景： + +``` +关闭的抽屉 +打开的抽屉 +打开的抽屉 + 可放入的刚体物体 +打开的烤箱托盘 + 可放入的物体 +多个可堆叠物体 +有指定左右、前后或桌面区域关系的物体 +``` + +### 2.2 Action Agent + +Action Agent 只使用 1k+ tasks.md 中的 Easy 原子动作： + +| ID | Easy Action | 作用 | +|---|---|---| +| E1 | PickUp + 旧 Place | 抓取、移动到目标位姿并释放 | +| E2 | MoveHeldObject—upright | 将倒下物体扶正 | +| E3 | MoveHeldObject—horizontal | 将物体水平摆正 | +| E4 | MoveHeldObject—pour | 将容器移动到目标上方并倾倒 | +| E5 | Handover—vertical | 以竖直姿态交接物体 | +| E6 | Handover—horizontal | 以水平姿态交接物体 | +| E7 | CoordinatedPickUp | 双臂共同抓取盘子或托盘 | +| E8 | AssemblePlace | 将手机摆在手机支架上 | +| E9 | PullArticulatedPart | 沿关节轴拉开抽屉或烤箱托盘 | +| E10 | PushArticulatedPart | 沿关节轴推闭抽屉或烤箱托盘 | +| E11 | TurnKnob | 将旋钮转到目标角度或档位 | +| E12 | PressButton | 沿按钮法向按压并触发 | + +Action Agent 有两层能力： + +```text +原子能力：E1–E12 +编排能力：选择、排序、连接和执行多个 Easy Action +``` + +编排能力不是新的原子技能。它只负责判断：前一个 Action 的结果，是否满足下一个 Action 的前置条件。 + +### 2.3 TaskSpec：连接 Scene 和 Action + +TaskSpec 是二者之间的共同任务说明，不是第三个独立执行 Agent。 + +TaskSpec 至少说明： + +```yaml +TaskSpec: + scene: # Scene 的初始条件 + goal: # 最终要达到的条件 + constraints: # 执行限制和不变量 + actions: # 需要或允许的 Easy Action +``` + +它的连接关系是： + +```text +TaskSpec.scene 和 TaskSpec.constraints +→ Scene Agent 生成满足条件的 Scene + +TaskSpec.goal 和 TaskSpec.actions +→ Action Agent 选择并编排 Easy Action + +Scene + Action 执行结果 +→ 检查 TaskSpec.goal 是否成立 +``` + +因此，TaskSpec 是任务扩增的主要位置：改变 TaskSpec 的目标、必要条件或动作因果关系，才可能产生有意义的新任务。 + +## 3. Scene 如何服务原子 Action + +每个 Easy Action 都对应一类 Scene 条件。 + +| Easy Action | Scene 需要提供的条件 | 动作结果 | +|---|---|---| +| E1 | 可抓取物体、目标位置、支撑面或容器 | 物体被放到目标位置 | +| E2 | 倒下且可抓取的物体、稳定的 upright 姿态 | 物体被扶正 | +| E3 | 可抓取的长物体、合法 horizontal 姿态 | 物体被水平摆正 | +| E4 | 可倾倒容器、接收目标、对齐空间 | 物料被倒入目标 | +| E5 | 适合竖直交接的物体和两个可达区域 | 竖直交接完成 | +| E6 | 适合水平交接的物体和两个可达区域 | 水平交接完成 | +| E7 | 可由双臂共同抓取的盘子或托盘 | 双臂抓取完成 | +| E8 | 物体和匹配的支架或 socket | 物体完成装配放置 | +| E9 | 关闭的抽屉、柜门或烤箱托盘，以及可操作部件 | 铰接部件打开 | +| E10 | 已打开的铰接部件 | 铰接部件关闭 | +| E11 | 可旋转旋钮和目标档位 | 旋钮达到目标档位 | +| E12 | 可按压按钮和触发位置 | 按钮被触发 | + +扩增的基本问题是： + +``` +为了让某个 Easy Action 成立，Scene 还可以提供哪些不同的状态、关系和目标？ +``` + +## 4. 按规则进行底向上扩增 + +### 4.1 基本组合规则 + +TaskSpec 先规定 goal 和允许的 Action；Scene Agent 再提供满足这些 Action 前置条件的 Scene。若 Scene 已由图像或文字输入得到，TaskSpec 则从已有 Scene 中选择可执行的 goal 和 Action。 + +无论从哪条入口开始，每一步都必须满足： + +``` +pre(Action) ⊆ 当前 Scene 状态 +``` + +执行后更新 Scene： + +``` +下一个 Scene 状态 += 当前 Scene 状态 + Action 的结果 +``` + +然后继续选择下一个可以执行的 Easy Action。 + +``` +Scene_0 +→ Easy Action_1 +→ Scene_1 +→ Easy Action_2 +→ Scene_2 +→ ... +→ goal +``` + +第一阶段限制动作链长度： + +``` +长度 1：一个 Easy Action +长度 2：两个有明确前后关系的 Easy Action +长度 3：三个以内的线性组合 +``` + +### 4.2 实际扩增步骤 + +TaskSpec 驱动时： + +``` +Step 1 组合 Easy Action，确定 goal +Step 2 写出 TaskSpec 的 scene、goal 和 constraints +Step 3 Scene Agent 生成满足 TaskSpec 的 Scene +Step 4 Action Agent 执行动作链 +Step 5 检查最终 goal +``` + +Scene 驱动时： + +``` +Step 1 Scene Agent 从图像、文字或编辑指令得到 Scene +Step 2 根据 Scene 中的物体、关系和状态列出可用 Easy Action +Step 3 TaskSpec 选择 goal 和 Action 组合 +Step 4 Action Agent 执行动作链 +Step 5 检查最终 goal +``` + +### 4.3 组合例子 + +**例 1：扶正并放置** + +``` +Scene_0：倒下的罐头 + 桌面目标区域 +E2：罐头 → upright +E1：罐头 → 目标区域 +goal：upright(can) ∧ on(can, target_region) +``` + +动作链是 E2 → E1。只更换罐头、桌面或轨迹，不增加新的任务。 + +**例 2：打开抽屉、放入物体、关闭抽屉** + +``` +Scene_0：关闭的抽屉 + 抽屉把手 + 可抓取物体 +E9：打开抽屉 +E1：把物体放入抽屉 +E10：关闭抽屉 +goal：contain(drawer, object) ∧ closed(drawer) +``` + +Scene Agent 提供“铰接式容器可以放置物体”的场景能力；Action Agent 使用 E9、E1、E10 完成任务。 + +**例 3：姿态和区域同时变化** + +``` +Scene_0：倒下的红色罐头 + 杯子 + 桌面左右区域 +E2：红色罐头 → upright +E1：红色罐头 → 杯子左侧 +goal：upright(red_can) ∧ left(red_can, cup) +``` + +把 left 改成 right 会改变目标关系，可以形成新的任务；只改变罐头模型或桌面材质，不形成新任务。 + +## 5. 哪些组合可以形成新任务 + +任务扩增必须同时满足三件事： + +```text +1. TaskSpec 的 goal、必要条件或因果结构发生了变化 +2. Scene 能提供满足新要求的条件或状态 +3. Action 原子能力或动作编排能够完成新要求 +``` + +Action 可以保持同一个原子动作；只要 TaskSpec 的目标或必要条件改变，并且 Scene 和 Action 能够支持它，就可以形成新任务。只有 Scene 变化或只有 Action 轨迹变化，不足以称为任务扩增。 + +只保留以下规则： + +``` +Scene 提供不同的 Action 前置条件 +→ 可以生成不同的 Action 组合 + +Action 结果支持新的后续 Action +→ 可以增加动作链长度和因果关系 + +最终 goal、物体数量、必要顺序、姿态要求或必要能力改变 +→ 可以形成新的 TaskSpec +``` + +以下变化不作为新任务来源： + +``` +只更换物体模型 +只更换桌面外观 +只更换机器人 +只更换动作轨迹 +只改写语言 +``` + +这些变化只能称为数据量扩大： + +```text +Scene 数据更多，但 TaskSpec 不变 +Action 轨迹更多，但 TaskSpec 不变 +Episode 更多，但 TaskSpec 不变 +``` + +Scene 的变化只有在改变任务所需的状态、关系或目标，并被 TaskSpec 和 Action 编排使用时，才会参与新的任务组合。 + +## 6. 底向上扩增公式 + +设： + +``` +E = 12 个 Easy 原子 Action +S_e = 能支持原子 Action e 的 Scene 模式数 +G(e,s) = Scene 模式 s 下可以定义的语义 goal 数 +L(e,s,g) = 能完成 goal 的合法 Action 链数量 +Y = Scene、Action 和 goal 检查通过率 +``` + +候选任务数估算为： + +``` +C_task_bottomup + ≈ Σe∈E Σs∈S_e G(e,s) × L(e,s,g) +``` + +通过检查的任务数估算为： + +``` +N_task_bottomup + ≈ C_task_bottomup × Y +``` + +这里的 L 只统计由 E1–E12 组成的合法动作链。不同轨迹不增加任务数量。 + +### 6.1 简单估算例子 + +如果平均每个 Easy Action 有： + +``` +S = 4 个 Scene 模式 +G = 3 个语义 goal +L = 2 条合法 Action 链 +Y = 0.8 的通过率 +``` + +则： + +``` +C_task_bottomup + ≈ 12 × 4 × 3 × 2 + = 288 + +N_task_bottomup + ≈ 288 × 0.8 + ≈ 230 +``` + +要达到 1,000 个任务，需要增加不同的 Scene 状态、语义 goal 和合法 Action 组合，而不是只增加物体或轨迹。 + +### 6.2 铰接式容器估算例子 + +假设 Scene Agent 可以提供 4 种铰接式 Scene 模式： + +``` +1. 抽屉关闭 +2. 抽屉打开 +3. 抽屉打开且有可放置区域 +4. 抽屉打开且已有物体 +``` + +每种模式平均支持 3 个不同 goal、2 条合法 Action 链： + +``` +C_hinged + ≈ 4 × 3 × 2 + = 24 个候选任务 + +N_hinged + ≈ 24 × Y +``` + +可形成的任务包括： + +``` +打开抽屉 +打开并放入物体 +打开、放入并关闭 +打开、取出并恢复关闭 +``` + +差异来自 goal 和必要动作顺序，而不是抽屉材质或物体模型。 + +## 7. 执行顺序 + +``` +P0 固定 E1–E12 和 Scene 条件表 +P1 每个 Easy 先设计 3–4 种 Scene 模式 +P2 每个 Scene 模式先生成单动作任务 +P3 用 Action 结果连接长度为 2 和 3 的动作链 +P4 改变 goal、数量、顺序、姿态和铰接状态 +P5 检查 Scene、Action 和 goal +P6 统计任务、Scene、Action 和 Episode +``` + +第一轮可以用下面的预算估算： + +```text +12 个 Easy +× 每个 4 个 Scene 模式 +× 每个模式 3 个语义 goal +× 平均 2 条合法 Action 链 +× 0.8 通过率 +≈ 230 个候选任务 +``` + +之后继续增加 Scene 模式、goal 变化和合法 Action 组合，直到达到 Task 1K 的目标。 + +核心原则： + +``` +Scene Agent 负责提供 Action 能力成立所需的世界； +Action Agent 负责在这个世界中组合 Easy 原子动作； +TaskSpec 记录组合后的目标和必要条件； +只有目标或因果结构发生变化，才算新的任务。 +``` diff --git a/pyproject.toml b/pyproject.toml index 2a127ccdd..75964127d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,16 +52,6 @@ dependencies = [ ] [project.optional-dependencies] -scene-engine = [ - "requests", - "Pillow", - "numpy", - "scipy", - "shapely", - "trimesh", - "open3d", - "matplotlib" -] gensim = [ "bpy", "pyrender==0.1.45", @@ -72,7 +62,7 @@ gensim = [ "shapely", "trimesh", "open3d", - "matplotlib" + "matplotlib", ] # cuRobo V2 is distributed from its source repository and provides separate # dependency sets for CUDA 12 and CUDA 13. Keep it optional so CPU-only and From 27f94b9bbd5be838b54fa40eadffa7ff5ae1a943 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Tue, 4 Aug 2026 21:48:15 +0800 Subject: [PATCH 40/41] Fixed scene export bug + Deleted md in cli/ + Added test for the modified scene export --- .../scene_engine/cli/task-1k-expanded.md | 516 ------------------ .../pipeline/utils/scene_exporter.py | 6 +- .../test_scene_core_and_export.py | 23 + 3 files changed, 28 insertions(+), 517 deletions(-) delete mode 100644 embodichain/gen_sim/scene_engine/cli/task-1k-expanded.md diff --git a/embodichain/gen_sim/scene_engine/cli/task-1k-expanded.md b/embodichain/gen_sim/scene_engine/cli/task-1k-expanded.md deleted file mode 100644 index 0984d1b7d..000000000 --- a/embodichain/gen_sim/scene_engine/cli/task-1k-expanded.md +++ /dev/null @@ -1,516 +0,0 @@ -# Task 1K 扩增方案（Scene 服务 Action 版） - -> 版本：v0.4 -> 日期：2026-08-04 -> 依据：1k+ tasks.md、task-1k-revision-recommendations.md、taskspec-semantic-protocol.md - -## 0. 核心目标 - -### 0.1 先认识几个词 - -| 词 | 含义 | 例子 | -|---|---|---| -| Scene | 机器人要操作的具体世界 | 一个关闭的抽屉、一个可抓取物体和一个把手 | -| Scene 模式 | 一类具有相同操作条件的 Scene | “关闭抽屉”；“打开抽屉且有可放置区域” | -| Action | 一个动作或一条动作序列 | E9 打开抽屉；E9 → E1 → E10 | -| Action 编排 | 按前置条件和动作结果连接 Easy Action | E9 打开后，E1 才能放入，最后 E10 关闭 | -| Scene Agent | 生成、恢复或编辑 Scene | 从抽屉图片恢复抽屉、把手和内部空间 | -| Action Agent | 在 Scene 中执行和编排 Action | 在打开的抽屉中执行 E1 放入物体 | -| TaskSpec | 规定 Scene 初始条件、goal 和允许 Action 的任务说明 | goal 是 contain(drawer, object) 且 drawer closed | -| Episode | 一次具体执行或扰动后的执行记录 | 同一任务中注入一次抓取滑移后的执行 | -| 任务扩增 | TaskSpec 的目标、条件或因果结构改变，形成有意义的新任务 | “放到桌面”改为“放入抽屉并关闭” | -| 数据扩大 | TaskSpec 不变，只增加 Scene、轨迹或 Episode | 同一任务换桌面、机器人或成功轨迹 | - -### 0.2 哪些算任务扩增 - -| 变化 | 分类 | 例子 | -|---|---|---| -| 改变最终 goal | 任务扩增 | on(object, table) → contain(drawer, object) | -| 增加数量、姿态、顺序或恢复要求 | 任务扩增 | 一个物体 → all；place → open → place → close | -| Scene 出现新状态，但 TaskSpec 和 goal 不变 | 数据扩大 | 同一个放置任务使用打开的不同抽屉场景 | -| 只更换物体、桌面、机器人或动作轨迹 | 数据扩大 | 杯子 A → 杯子 B；UR5 → Franka；轨迹 A → B | - -判断标准： - -```text -TaskSpec 的目标、必要条件或因果结构改变 → 任务扩增 -TaskSpec 不变，只增加 Scene、Action 轨迹或 Episode → 数据扩大 -``` - -Task 1K 的重点不是制造很多不同的物体或桌面，而是： - -``` -Scene Agent 提供能执行原子动作的场景条件 - ↓ -Action Agent 使用 Easy 原子动作 - ↓ -组合成更多有意义的 TaskSpec -``` - -Scene 是 Action 的基础；Action 是改变 Scene 的方式；TaskSpec 负责把二者组合成一个任务。 - -### 0.3 方案先看 - -Task 1K 按以下顺序扩增： - -```text -P0 固定 E1–E12，以及每个 Easy 需要的 Scene 条件 -P1 为每个 Easy 设计多种可执行 Scene 模式 -P2 用一个 Easy 生成最简单的 TaskSpec -P3 用前一个 Action 的结果连接下一个 Easy，形成短动作链 -P4 改变 goal、数量、姿态、顺序或铰接状态，形成有意义的新任务 -P5 检查 Scene、Action 和 goal，统计任务和数据数量 -``` - -核心不是把同一个任务换成很多物体，而是： - -```text -一个 Easy -→ 多种可执行 Scene -→ 多种 goal -→ 合法的 Easy Action 组合 -→ 更多语义不同的 TaskSpec -``` - -## 1. 最小流程 - -Task 生成有两条入口，最后都通过 TaskSpec 连接 Scene 和 Action。 - -### 1.1 TaskSpec 驱动 - -适合从 Easy 动作组合出新的任务： - -``` -组合 Easy Action 和目标 - ↓ -写出 TaskSpec：目标、动作和必要条件 - ↓ -Scene Agent 根据 TaskSpec 生成或恢复 Scene - ↓ -Action Agent 执行动作链 - ↓ -检查最终 goal -``` - -例如： - -``` -目标：把物体放入抽屉并关闭抽屉 -动作：E9 → E1 → E10 -TaskSpec：要求有关闭的抽屉、可抓取物体，最终 contain 且 closed -Scene Agent：生成满足这些条件的 Scene -Action Agent：执行 E9 → E1 → E10 -``` - -### 1.2 Scene 驱动 - -适合从图像、纯文字或编辑后的 Scene 中发现可以生成的任务： - -``` -图像 / 纯文字 / Scene 编辑输入 - ↓ -Scene Agent 生成或恢复 Scene - ↓ -TaskSpec 根据 Scene 中可用的物体、关系和状态，选择 goal 和 Easy Action - ↓ -Action Agent 执行 TaskSpec 规定的动作 - ↓ -检查最终 goal -``` - -例如，Scene Agent 从图像中发现一个关闭的抽屉、一个可抓取物体和一个把手，TaskSpec 可以选择： - -``` -goal：打开抽屉 → E9 -goal：打开并放入物体 → E9 → E1 -goal：打开、放入并关闭 → E9 → E1 → E10 -``` - -两条入口的区别是： - -``` -TaskSpec 驱动：先决定要做什么，再生成 Scene -Scene 驱动：先知道有什么，再决定能做什么 -``` - -## 2. 三个核心部分 - -### 2.1 Scene Agent - -Scene Agent 负责根据 Action 的需要生成、恢复或编辑 Scene。 - -它需要提供： - -- 需要被操作的物体； -- 物体的可抓取、可放置或可交互状态； -- 物体之间的关系； -- 桌面和容器的位置； -- 铰接式物体的开关状态； -- 下一步 Action 所需的目标区域或对象。 - -Scene Agent 当前支持的主要关系： - -| 关系 | 含义 | -|---|---| -| left、right、center | 物体相对于桌面的区域 | -| left_front、left_back | 桌面的左前、左后区域 | -| right_front、right_back | 桌面的右前、右后区域 | -| left_center、right_center | 桌面的左中、右中区域 | -| on | 物体在桌面或支撑面上 | -| stack_on | 一个刚体叠放在另一个刚体上 | -| contain | 刚体或铰接式容器包含刚体物体 | - -Scene Agent 可以生成以下场景： - -``` -关闭的抽屉 -打开的抽屉 -打开的抽屉 + 可放入的刚体物体 -打开的烤箱托盘 + 可放入的物体 -多个可堆叠物体 -有指定左右、前后或桌面区域关系的物体 -``` - -### 2.2 Action Agent - -Action Agent 只使用 1k+ tasks.md 中的 Easy 原子动作： - -| ID | Easy Action | 作用 | -|---|---|---| -| E1 | PickUp + 旧 Place | 抓取、移动到目标位姿并释放 | -| E2 | MoveHeldObject—upright | 将倒下物体扶正 | -| E3 | MoveHeldObject—horizontal | 将物体水平摆正 | -| E4 | MoveHeldObject—pour | 将容器移动到目标上方并倾倒 | -| E5 | Handover—vertical | 以竖直姿态交接物体 | -| E6 | Handover—horizontal | 以水平姿态交接物体 | -| E7 | CoordinatedPickUp | 双臂共同抓取盘子或托盘 | -| E8 | AssemblePlace | 将手机摆在手机支架上 | -| E9 | PullArticulatedPart | 沿关节轴拉开抽屉或烤箱托盘 | -| E10 | PushArticulatedPart | 沿关节轴推闭抽屉或烤箱托盘 | -| E11 | TurnKnob | 将旋钮转到目标角度或档位 | -| E12 | PressButton | 沿按钮法向按压并触发 | - -Action Agent 有两层能力： - -```text -原子能力：E1–E12 -编排能力：选择、排序、连接和执行多个 Easy Action -``` - -编排能力不是新的原子技能。它只负责判断：前一个 Action 的结果，是否满足下一个 Action 的前置条件。 - -### 2.3 TaskSpec：连接 Scene 和 Action - -TaskSpec 是二者之间的共同任务说明，不是第三个独立执行 Agent。 - -TaskSpec 至少说明： - -```yaml -TaskSpec: - scene: # Scene 的初始条件 - goal: # 最终要达到的条件 - constraints: # 执行限制和不变量 - actions: # 需要或允许的 Easy Action -``` - -它的连接关系是： - -```text -TaskSpec.scene 和 TaskSpec.constraints -→ Scene Agent 生成满足条件的 Scene - -TaskSpec.goal 和 TaskSpec.actions -→ Action Agent 选择并编排 Easy Action - -Scene + Action 执行结果 -→ 检查 TaskSpec.goal 是否成立 -``` - -因此，TaskSpec 是任务扩增的主要位置：改变 TaskSpec 的目标、必要条件或动作因果关系，才可能产生有意义的新任务。 - -## 3. Scene 如何服务原子 Action - -每个 Easy Action 都对应一类 Scene 条件。 - -| Easy Action | Scene 需要提供的条件 | 动作结果 | -|---|---|---| -| E1 | 可抓取物体、目标位置、支撑面或容器 | 物体被放到目标位置 | -| E2 | 倒下且可抓取的物体、稳定的 upright 姿态 | 物体被扶正 | -| E3 | 可抓取的长物体、合法 horizontal 姿态 | 物体被水平摆正 | -| E4 | 可倾倒容器、接收目标、对齐空间 | 物料被倒入目标 | -| E5 | 适合竖直交接的物体和两个可达区域 | 竖直交接完成 | -| E6 | 适合水平交接的物体和两个可达区域 | 水平交接完成 | -| E7 | 可由双臂共同抓取的盘子或托盘 | 双臂抓取完成 | -| E8 | 物体和匹配的支架或 socket | 物体完成装配放置 | -| E9 | 关闭的抽屉、柜门或烤箱托盘，以及可操作部件 | 铰接部件打开 | -| E10 | 已打开的铰接部件 | 铰接部件关闭 | -| E11 | 可旋转旋钮和目标档位 | 旋钮达到目标档位 | -| E12 | 可按压按钮和触发位置 | 按钮被触发 | - -扩增的基本问题是： - -``` -为了让某个 Easy Action 成立，Scene 还可以提供哪些不同的状态、关系和目标？ -``` - -## 4. 按规则进行底向上扩增 - -### 4.1 基本组合规则 - -TaskSpec 先规定 goal 和允许的 Action；Scene Agent 再提供满足这些 Action 前置条件的 Scene。若 Scene 已由图像或文字输入得到，TaskSpec 则从已有 Scene 中选择可执行的 goal 和 Action。 - -无论从哪条入口开始，每一步都必须满足： - -``` -pre(Action) ⊆ 当前 Scene 状态 -``` - -执行后更新 Scene： - -``` -下一个 Scene 状态 -= 当前 Scene 状态 + Action 的结果 -``` - -然后继续选择下一个可以执行的 Easy Action。 - -``` -Scene_0 -→ Easy Action_1 -→ Scene_1 -→ Easy Action_2 -→ Scene_2 -→ ... -→ goal -``` - -第一阶段限制动作链长度： - -``` -长度 1：一个 Easy Action -长度 2：两个有明确前后关系的 Easy Action -长度 3：三个以内的线性组合 -``` - -### 4.2 实际扩增步骤 - -TaskSpec 驱动时： - -``` -Step 1 组合 Easy Action，确定 goal -Step 2 写出 TaskSpec 的 scene、goal 和 constraints -Step 3 Scene Agent 生成满足 TaskSpec 的 Scene -Step 4 Action Agent 执行动作链 -Step 5 检查最终 goal -``` - -Scene 驱动时： - -``` -Step 1 Scene Agent 从图像、文字或编辑指令得到 Scene -Step 2 根据 Scene 中的物体、关系和状态列出可用 Easy Action -Step 3 TaskSpec 选择 goal 和 Action 组合 -Step 4 Action Agent 执行动作链 -Step 5 检查最终 goal -``` - -### 4.3 组合例子 - -**例 1：扶正并放置** - -``` -Scene_0：倒下的罐头 + 桌面目标区域 -E2：罐头 → upright -E1：罐头 → 目标区域 -goal：upright(can) ∧ on(can, target_region) -``` - -动作链是 E2 → E1。只更换罐头、桌面或轨迹，不增加新的任务。 - -**例 2：打开抽屉、放入物体、关闭抽屉** - -``` -Scene_0：关闭的抽屉 + 抽屉把手 + 可抓取物体 -E9：打开抽屉 -E1：把物体放入抽屉 -E10：关闭抽屉 -goal：contain(drawer, object) ∧ closed(drawer) -``` - -Scene Agent 提供“铰接式容器可以放置物体”的场景能力；Action Agent 使用 E9、E1、E10 完成任务。 - -**例 3：姿态和区域同时变化** - -``` -Scene_0：倒下的红色罐头 + 杯子 + 桌面左右区域 -E2：红色罐头 → upright -E1：红色罐头 → 杯子左侧 -goal：upright(red_can) ∧ left(red_can, cup) -``` - -把 left 改成 right 会改变目标关系，可以形成新的任务；只改变罐头模型或桌面材质，不形成新任务。 - -## 5. 哪些组合可以形成新任务 - -任务扩增必须同时满足三件事： - -```text -1. TaskSpec 的 goal、必要条件或因果结构发生了变化 -2. Scene 能提供满足新要求的条件或状态 -3. Action 原子能力或动作编排能够完成新要求 -``` - -Action 可以保持同一个原子动作；只要 TaskSpec 的目标或必要条件改变，并且 Scene 和 Action 能够支持它，就可以形成新任务。只有 Scene 变化或只有 Action 轨迹变化，不足以称为任务扩增。 - -只保留以下规则： - -``` -Scene 提供不同的 Action 前置条件 -→ 可以生成不同的 Action 组合 - -Action 结果支持新的后续 Action -→ 可以增加动作链长度和因果关系 - -最终 goal、物体数量、必要顺序、姿态要求或必要能力改变 -→ 可以形成新的 TaskSpec -``` - -以下变化不作为新任务来源： - -``` -只更换物体模型 -只更换桌面外观 -只更换机器人 -只更换动作轨迹 -只改写语言 -``` - -这些变化只能称为数据量扩大： - -```text -Scene 数据更多，但 TaskSpec 不变 -Action 轨迹更多，但 TaskSpec 不变 -Episode 更多，但 TaskSpec 不变 -``` - -Scene 的变化只有在改变任务所需的状态、关系或目标，并被 TaskSpec 和 Action 编排使用时，才会参与新的任务组合。 - -## 6. 底向上扩增公式 - -设： - -``` -E = 12 个 Easy 原子 Action -S_e = 能支持原子 Action e 的 Scene 模式数 -G(e,s) = Scene 模式 s 下可以定义的语义 goal 数 -L(e,s,g) = 能完成 goal 的合法 Action 链数量 -Y = Scene、Action 和 goal 检查通过率 -``` - -候选任务数估算为： - -``` -C_task_bottomup - ≈ Σe∈E Σs∈S_e G(e,s) × L(e,s,g) -``` - -通过检查的任务数估算为： - -``` -N_task_bottomup - ≈ C_task_bottomup × Y -``` - -这里的 L 只统计由 E1–E12 组成的合法动作链。不同轨迹不增加任务数量。 - -### 6.1 简单估算例子 - -如果平均每个 Easy Action 有： - -``` -S = 4 个 Scene 模式 -G = 3 个语义 goal -L = 2 条合法 Action 链 -Y = 0.8 的通过率 -``` - -则： - -``` -C_task_bottomup - ≈ 12 × 4 × 3 × 2 - = 288 - -N_task_bottomup - ≈ 288 × 0.8 - ≈ 230 -``` - -要达到 1,000 个任务，需要增加不同的 Scene 状态、语义 goal 和合法 Action 组合，而不是只增加物体或轨迹。 - -### 6.2 铰接式容器估算例子 - -假设 Scene Agent 可以提供 4 种铰接式 Scene 模式： - -``` -1. 抽屉关闭 -2. 抽屉打开 -3. 抽屉打开且有可放置区域 -4. 抽屉打开且已有物体 -``` - -每种模式平均支持 3 个不同 goal、2 条合法 Action 链： - -``` -C_hinged - ≈ 4 × 3 × 2 - = 24 个候选任务 - -N_hinged - ≈ 24 × Y -``` - -可形成的任务包括： - -``` -打开抽屉 -打开并放入物体 -打开、放入并关闭 -打开、取出并恢复关闭 -``` - -差异来自 goal 和必要动作顺序，而不是抽屉材质或物体模型。 - -## 7. 执行顺序 - -``` -P0 固定 E1–E12 和 Scene 条件表 -P1 每个 Easy 先设计 3–4 种 Scene 模式 -P2 每个 Scene 模式先生成单动作任务 -P3 用 Action 结果连接长度为 2 和 3 的动作链 -P4 改变 goal、数量、顺序、姿态和铰接状态 -P5 检查 Scene、Action 和 goal -P6 统计任务、Scene、Action 和 Episode -``` - -第一轮可以用下面的预算估算： - -```text -12 个 Easy -× 每个 4 个 Scene 模式 -× 每个模式 3 个语义 goal -× 平均 2 条合法 Action 链 -× 0.8 通过率 -≈ 230 个候选任务 -``` - -之后继续增加 Scene 模式、goal 变化和合法 Action 组合，直到达到 Task 1K 的目标。 - -核心原则： - -``` -Scene Agent 负责提供 Action 能力成立所需的世界； -Action Agent 负责在这个世界中组合 Easy 原子动作； -TaskSpec 记录组合后的目标和必要条件； -只有目标或因果结构发生变化，才算新的任务。 -``` diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py index cf1f55932..fb66c30c4 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py @@ -116,7 +116,11 @@ def _copy_scene_object_to_assets( ) -> 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 {"", ".", ".."}: + if ( + Path(object_id).name != object_id + or "\\" in object_id + or object_id in {"", ".", ".."} + ): raise ValueError( f"Scene object id is not safe for a GLB filename: {object_id!r}" ) diff --git a/tests/gen_sim/scene_engine/test_scene_core_and_export.py b/tests/gen_sim/scene_engine/test_scene_core_and_export.py index d016d7aa7..56c12af1c 100644 --- a/tests/gen_sim/scene_engine/test_scene_core_and_export.py +++ b/tests/gen_sim/scene_engine/test_scene_core_and_export.py @@ -146,3 +146,26 @@ def test_scene_export_requires_final_physics(tmp_path: Path) -> None: with pytest.raises(ValueError, match="no SimReady physics"): SceneExporter(scene=Scene(objects=[table]), output_root=tmp_path).export() + + +def test_scene_export_rejects_backslash_in_object_id(tmp_path: Path) -> None: + glb_path = tmp_path / "table.glb" + glb_path.write_bytes(b"glTF") + table = _scene_object( + object_id="table", + kind="table", + glb_path=glb_path, + physics=_physics("kinematic"), + ) + unsafe_asset = _scene_object( + object_id=r"..\evil", + kind="asset", + glb_path=glb_path, + physics=_physics("dynamic"), + ) + + with pytest.raises(ValueError, match="not safe for a GLB filename"): + SceneExporter( + scene=Scene(objects=[table, unsafe_asset]), + output_root=tmp_path / "output", + ).export() From a5390af45483a3f2e97a15116d4829e6ffd58ab5 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Wed, 5 Aug 2026 14:23:20 +0800 Subject: [PATCH 41/41] fix: remove dexsim-provided gensim dependencies --- pyproject.toml | 4 ---- 1 file changed, 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 75964127d..db51c3041 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -57,11 +57,7 @@ gensim = [ "pyrender==0.1.45", "requests", "Pillow", - "numpy", "scipy", - "shapely", - "trimesh", - "open3d", "matplotlib", ] # cuRobo V2 is distributed from its source repository and provides separate