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/46] 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/46] 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/46] 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/46] 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/46] 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/46] 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/46] 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/46] 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/46] 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/46] 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/46] 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/46] 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/46] 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/46] 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/46] 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/46] 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/46] 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/46] 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/46] 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/46] 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/46] 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/46] 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/46] 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 b61ef2fc3cce4e74be23361c7f92d3213053ad66 Mon Sep 17 00:00:00 2001 From: PengXuanchao Date: Mon, 3 Aug 2026 16:34:54 +0800 Subject: [PATCH 24/46] Add Gradio UI with managed Viser previews --- .gitignore | 6 + .../gen_sim/gradio_ui/app_articraft.py | 884 +++++ .../gen_sim/gradio_ui/app_asset_engine.py | 283 ++ embodichain/gen_sim/gradio_ui/app_commands.py | 173 + embodichain/gen_sim/gradio_ui/app_config.py | 303 ++ embodichain/gen_sim/gradio_ui/app_media.py | 724 ++++ .../gen_sim/gradio_ui/app_processes.py | 193 + embodichain/gen_sim/gradio_ui/app_services.py | 28 + embodichain/gen_sim/gradio_ui/app_state.py | 187 + embodichain/gen_sim/gradio_ui/app_ui.py | 569 +++ .../gen_sim/gradio_ui/app_workflows.py | 3336 +++++++++++++++++ .../gen_sim/gradio_ui/assets/dexforce.png | Bin 0 -> 25276 bytes embodichain/gen_sim/gradio_ui/gradio_app.py | 85 + .../gradio_visualization_architecture.md | 276 ++ embodichain/gen_sim/gradio_ui/random_input.py | 540 +++ 15 files changed, 7587 insertions(+) create mode 100644 embodichain/gen_sim/gradio_ui/app_articraft.py create mode 100644 embodichain/gen_sim/gradio_ui/app_asset_engine.py create mode 100644 embodichain/gen_sim/gradio_ui/app_commands.py create mode 100644 embodichain/gen_sim/gradio_ui/app_config.py create mode 100644 embodichain/gen_sim/gradio_ui/app_media.py create mode 100644 embodichain/gen_sim/gradio_ui/app_processes.py create mode 100644 embodichain/gen_sim/gradio_ui/app_services.py create mode 100644 embodichain/gen_sim/gradio_ui/app_state.py create mode 100644 embodichain/gen_sim/gradio_ui/app_ui.py create mode 100644 embodichain/gen_sim/gradio_ui/app_workflows.py create mode 100644 embodichain/gen_sim/gradio_ui/assets/dexforce.png create mode 100644 embodichain/gen_sim/gradio_ui/gradio_app.py create mode 100644 embodichain/gen_sim/gradio_ui/gradio_visualization_architecture.md create mode 100644 embodichain/gen_sim/gradio_ui/random_input.py diff --git a/.gitignore b/.gitignore index 69061763a..e61027a6b 100644 --- a/.gitignore +++ b/.gitignore @@ -208,3 +208,9 @@ scripts/benchmark/rl/reports/* .worktrees/ # Local gym project workspace /gym_project/ +.debug_engine/ + +# Local Gradio UI dependencies, generated Articraft records, and bytecode +/embodichain/gen_sim/gradio_ui/.articraft/ +/embodichain/gen_sim/gradio_ui/.debug_engine/ +/embodichain/gen_sim/gradio_ui/__pycache__/ diff --git a/embodichain/gen_sim/gradio_ui/app_articraft.py b/embodichain/gen_sim/gradio_ui/app_articraft.py new file mode 100644 index 000000000..85d8544a0 --- /dev/null +++ b/embodichain/gen_sim/gradio_ui/app_articraft.py @@ -0,0 +1,884 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Codex-backed Articraft generation for the Asset engine. + +The integration uses Articraft's external-agent workflow: Articraft owns +record creation and validation while Codex authors the generated model. All +mutable run data is kept under ``ARTICRAFT_OUTPUT_ROOT``. +""" + +from __future__ import annotations + +import atexit +import os +import queue +import json +import shutil +import html +import signal +import socket +import subprocess +import sys +import threading +import time +import uuid +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +import gradio as gr + +from app_config import ( + ARTICRAFT_CONDA_ENV, + ARTICRAFT_OUTPUT_ROOT, + ARTICRAFT_REPOSITORY_URL, + ARTICRAFT_ROOT, + ARTICRAFT_VISER_PORT, +) +from app_processes import read_process_output, start_pipeline, terminate_process_group + +__all__ = [ + "build_articraft_panel", + "configure_articraft_environment", + "generate_articraft_asset", + "stop_articraft_viser_preview", +] + +_VISER_START_TIMEOUT_SECONDS = 15.0 +_VISER_STOP_TIMEOUT_SECONDS = 5.0 + + +def _command_path(name: str) -> str | None: + """Resolve commands even when Gradio did not inherit an interactive PATH.""" + configured = os.environ.get(f"{name.upper()}_EXE") + return configured or shutil.which(name) + + +def _conda_path() -> str | None: + configured = os.environ.get("CONDA_EXE") + if configured and Path(configured).is_file(): + return configured + return _command_path("conda") + + +def _conda_command(*args: str) -> list[str]: + conda = _conda_path() + if not conda: + raise RuntimeError("Conda was not found. Set CONDA_EXE before starting Gradio.") + return [conda, "run", "--no-capture-output", "-n", ARTICRAFT_CONDA_ENV, *args] + + +def _articraft_cli_command(*args: str) -> list[str]: + """Run the CLI from the checked-out source without installing it with pip.""" + return _conda_command("python", "-m", "cli.main", *args) + + +def _articraft_conda_environment_exists() -> bool: + """Check only for the named Conda environment, not package installation.""" + conda = _conda_path() + if not conda: + return False + try: + result = subprocess.run( + [conda, "env", "list", "--json"], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + timeout=30, + check=False, + ) + if result.returncode: + return False + environments = json.loads(result.stdout or "{}").get("envs", []) + return any(Path(path).name == ARTICRAFT_CONDA_ENV for path in environments) + except (OSError, json.JSONDecodeError, TypeError): + return False + + +def _run_check( + command: list[str], *, timeout: int = 45 +) -> subprocess.CompletedProcess[str]: + return subprocess.run( + command, + cwd=ARTICRAFT_ROOT, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + timeout=timeout, + check=False, + ) + + +def _short_output( + result: subprocess.CompletedProcess[str], *, limit: int = 1800 +) -> str: + output = (result.stdout or "").strip() + return output[-limit:] if len(output) > limit else (output or "(no output)") + + +def _check_requirements() -> tuple[list[str], list[str], str | None]: + """Return diagnostics and the Codex executable, without creating an asset.""" + errors: list[str] = [] + details: list[str] = [] + if not ( + ARTICRAFT_ROOT.is_dir() + and (ARTICRAFT_ROOT / ".git").exists() + and (ARTICRAFT_ROOT / "pyproject.toml").is_file() + ): + errors.append(f".articraft checkout is not ready: {ARTICRAFT_ROOT}") + if not _conda_path(): + errors.append("Conda is not on PATH. Set CONDA_EXE to the conda executable.") + elif not _articraft_conda_environment_exists(): + errors.append(f"Conda environment not found: {ARTICRAFT_CONDA_ENV}") + else: + details.append(f"Conda environment: {ARTICRAFT_CONDA_ENV}") + + codex = _command_path("codex") + if not codex: + errors.append("Codex CLI is not on PATH. Install it or set CODEX_EXE.") + elif not errors: + try: + result = _run_check([codex, "--version"]) + if result.returncode: + errors.append(f"Codex CLI check failed: {_short_output(result)}") + else: + details.append(f"Codex: {_short_output(result, limit=120)}") + except Exception as exc: + errors.append(f"Codex CLI check failed: {exc}") + + if not errors: + details.append(f".articraft checkout: {ARTICRAFT_ROOT}") + return details, errors, codex + + +def _prepare_articraft_checkout() -> tuple[bool, str]: + """Clone the configured checkout when absent, without overwriting a directory.""" + if ARTICRAFT_ROOT.exists(): + if (ARTICRAFT_ROOT / ".git").exists() and ( + ARTICRAFT_ROOT / "pyproject.toml" + ).is_file(): + return True, f".articraft checkout: {ARTICRAFT_ROOT}" + return ( + False, + f"{ARTICRAFT_ROOT} exists but is not an Articraft Git checkout; it was left untouched.", + ) + + git = _command_path("git") + if not git: + return False, "Git is not on PATH, so .articraft cannot be cloned." + try: + ARTICRAFT_ROOT.parent.mkdir(parents=True, exist_ok=True) + clone = subprocess.run( + [git, "clone", ARTICRAFT_REPOSITORY_URL, str(ARTICRAFT_ROOT)], + cwd=ARTICRAFT_ROOT.parent, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + timeout=300, + check=False, + ) + except Exception as exc: + return False, f"Unable to clone Articraft: {exc}" + if clone.returncode: + return False, f"Articraft clone failed: {_short_output(clone, limit=3000)}" + return True, f"Cloned .articraft from {ARTICRAFT_REPOSITORY_URL}" + + +def configure_articraft_environment() -> str: + """Clone the checkout, then verify the Conda environment and Codex.""" + checkout_ready, checkout_message = _prepare_articraft_checkout() + if not checkout_ready: + return "**Articulation is not ready.**\n\n- " + checkout_message + try: + for directory in ( + ARTICRAFT_OUTPUT_ROOT, + ARTICRAFT_OUTPUT_ROOT / "runs", + ARTICRAFT_OUTPUT_ROOT / "exports", + ): + directory.mkdir(parents=True, exist_ok=True) + except Exception as exc: + return f"**Unable to prepare the shared Articulation output folder:** `{exc}`" + details, errors, _ = _check_requirements() + if errors: + return "**Articulation is not ready.**\n\n" + "\n".join( + f"- {error}" for error in errors + ) + details.insert(0, checkout_message) + details.extend( + ( + f"Shared output: `{ARTICRAFT_OUTPUT_ROOT}`", + "Generation runs the `.articraft` checkout directly with `conda run`; no `pip install -e .` is required.", + ) + ) + return "**Articulation is ready.**\n\n" + "\n".join( + f"- {detail}" for detail in details + ) + + +def _record_id() -> str: + timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S") + # Articraft validates external IDs against the required ``rec_`` prefix. + return f"rec_ui_articraft_{timestamp}_{uuid.uuid4().hex[:8]}" + + +def _copy_reference_image(value: Any, run_root: Path) -> Path | None: + if not value: + return None + source = Path(str(value)) + if not source.is_file(): + raise ValueError( + "The reference image is no longer available; please upload it again." + ) + suffix = source.suffix.lower() or ".png" + if suffix not in {".png", ".jpg", ".jpeg", ".webp"}: + raise ValueError("Reference image must be PNG, JPG, JPEG, or WEBP.") + target = run_root / f"reference{suffix}" + shutil.copy2(source, target) + return target + + +def _active_model_path(record_dir: Path) -> Path: + candidates = sorted(record_dir.glob("revisions/*/model.py")) + if len(candidates) != 1: + raise FileNotFoundError( + f"Expected one active model.py in {record_dir}, found {len(candidates)}." + ) + return candidates[0] + + +def _make_result_bundle(record_id: str) -> tuple[Path, Path]: + materialized = ( + ARTICRAFT_OUTPUT_ROOT / "data" / "cache" / "record_materialization" / record_id + ) + if not (materialized / "model.urdf").is_file(): + raise FileNotFoundError( + "Articraft completed without a compiled model.urdf output." + ) + exports_root = ARTICRAFT_OUTPUT_ROOT / "exports" + exports_root.mkdir(parents=True, exist_ok=True) + archive = Path( + shutil.make_archive( + (exports_root / record_id).as_posix(), + "zip", + root_dir=materialized, + ) + ) + return materialized, archive + + +def _articraft_viser_iframe(record_id: str) -> str: + """Embed the Articulation Viser service through the Gradio page hostname.""" + srcdoc = ( + "" + ) + escaped_record_id = html.escape(record_id) + return ( + "
Viser articulation preview: " + f"{escaped_record_id}" + f"" + "
" + ) + + +class _ArticraftViserPreview: + """Own the single Articraft Viser process and its dedicated TCP port.""" + + def __init__(self, port: int) -> None: + self._port = port + self._lock = threading.Lock() + self._process: subprocess.Popen[str] | None = None + + def start(self, urdf_path: Path, record_id: str) -> str: + """Replace the active preview with a verified preview of one URDF.""" + if not urdf_path.is_file(): + raise FileNotFoundError(f"Compiled URDF is missing: {urdf_path}") + + with self._lock: + self._stop_managed_process() + self._clear_stale_listener() + process = start_pipeline(self._command(urdf_path)) + if not self._wait_until_owned(process): + terminate_process_group(process) + raise RuntimeError("New Articraft Viser preview did not bind its port.") + self._process = process + return _articraft_viser_iframe(record_id) + + def stop(self) -> None: + """Stop the preview process, if this panel started one.""" + with self._lock: + self._stop_managed_process() + + def _command(self, urdf_path: Path) -> list[str]: + return [ + sys.executable, + str(Path(__file__).with_name("app_media.py")), + "--asset_path", + str(urdf_path), + "--asset_type", + "articulation", + "--headless", + "--viser", + "--viser-host", + "0.0.0.0", + "--viser-port", + str(self._port), + ] + + def _stop_managed_process(self) -> None: + if self._process is not None: + terminate_process_group(self._process) + self._process = None + + def _clear_stale_listener(self) -> None: + if self._port_is_available(): + return + listener_pids = self._listener_pids() + if listener_pids is None: + raise RuntimeError( + "Cannot identify the process using the Articraft Viser port." + ) + if not listener_pids: + raise RuntimeError( + f"Port {self._port} is unavailable without a visible listener." + ) + + self._signal_listeners(listener_pids, signal.SIGTERM, "stop") + if self._wait_for_port_release(): + return + + remaining_pids = self._listener_pids() + if remaining_pids is None: + raise RuntimeError( + f"Cannot identify the stale Viser service on port {self._port}." + ) + self._signal_listeners(remaining_pids, signal.SIGKILL, "force-stop") + if not self._wait_for_port_release(): + raise RuntimeError( + f"The stale Viser service is still listening on port {self._port}." + ) + + def _signal_listeners( + self, listener_pids: set[int], signal_value: int, action: str + ) -> None: + for pid in listener_pids: + try: + os.kill(pid, signal_value) + except ProcessLookupError: + continue + except PermissionError as exc: + raise RuntimeError( + f"Cannot {action} Viser process {pid} using port {self._port}." + ) from exc + + def _wait_for_port_release(self) -> bool: + deadline = time.monotonic() + _VISER_STOP_TIMEOUT_SECONDS + while time.monotonic() < deadline: + if self._port_is_available(): + return True + time.sleep(0.1) + return self._port_is_available() + + def _wait_until_owned(self, process: subprocess.Popen[str]) -> bool: + deadline = time.monotonic() + _VISER_START_TIMEOUT_SECONDS + while time.monotonic() < deadline: + if process.poll() is not None: + return False + try: + with socket.create_connection(("127.0.0.1", self._port), timeout=0.2): + listener_pids = self._listener_pids() + if listener_pids is not None and process.pid in listener_pids: + return True + except OSError: + pass + time.sleep(0.25) + return False + + def _port_is_available(self) -> bool: + probe = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + try: + probe.bind(("0.0.0.0", self._port)) + except OSError: + return False + finally: + probe.close() + return True + + def _listener_pids(self) -> set[int] | None: + for command in self._listener_commands(): + try: + result = subprocess.run( + command, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + timeout=5, + check=False, + ) + except OSError: + continue + if result.returncode in (0, 1): + return {int(pid) for pid in result.stdout.split() if pid.isdecimal()} + return None + + def _listener_commands(self) -> list[list[str]]: + commands: list[list[str]] = [] + if lsof := _command_path("lsof"): + commands.append([lsof, "-nP", f"-iTCP:{self._port}", "-sTCP:LISTEN", "-t"]) + if fuser := _command_path("fuser"): + commands.append([fuser, "-n", "tcp", str(self._port)]) + return commands + + +_articraft_viser_preview = _ArticraftViserPreview(ARTICRAFT_VISER_PORT) + + +def stop_articraft_viser_preview() -> None: + """Stop the Viser subprocess currently owned by the Articraft panel. + + The preview runs independently from Gradio so it can be embedded through an + iframe. Expose its cleanup explicitly so application shutdown can release + the dedicated port instead of leaving an orphaned Viser server behind. + """ + _articraft_viser_preview.stop() + + +atexit.register(stop_articraft_viser_preview) + + +def _start_articraft_viser_preview(materialized: Path, record_id: str) -> str: + """Load the compiled URDF as an articulation and expose it through Viser.""" + return _articraft_viser_preview.start(materialized / "model.urdf", record_id) + + +def _external_check_is_unsupported(result: subprocess.CompletedProcess[str]) -> bool: + """Recognize the older Articraft CLI, which has no ``external check``.""" + output = (result.stdout or "").lower() + return "invalid choice: 'check'" in output and "external" in output + + +def _compile_report_failures(record_id: str) -> list[str]: + """Read blocking QC/test signals from the older CLI's compile report.""" + report_path = ( + ARTICRAFT_OUTPUT_ROOT + / "data" + / "cache" + / "record_materialization" + / record_id + / "compile_report.json" + ) + try: + report = json.loads(report_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + return [f"Compile report is unavailable: {report_path}"] + bundle = report.get("signal_bundle") if isinstance(report, dict) else None + signals = bundle.get("signals") if isinstance(bundle, dict) else None + if not isinstance(signals, list): + return ["Compile report contains no validation signals."] + failures: list[str] = [] + for signal in signals: + if not isinstance(signal, dict): + continue + if signal.get("severity") == "failure" or signal.get("blocking") is True: + failures.append( + str( + signal.get("summary") + or signal.get("code") + or "Unnamed validation failure" + ) + ) + return failures + + +def _build_codex_prompt( + *, + prompt: str, + record_id: str, + record_dir: Path, + model_path: Path, + reference_image: Path | None, +) -> str: + image_note = ( + f"A reference image is attached and also copied at {reference_image}. Use it as visual reference." + if reference_image + else "No reference image was supplied." + ) + return f"""You are the Codex external author for one Articraft articulated 3D asset. + +User request: +{prompt} + +{image_note} + +The Articraft source repository is {ARTICRAFT_ROOT}. The shared UI output/storage root is +{ARTICRAFT_OUTPUT_ROOT}. Articraft has already created this external workbench record: +record_id={record_id} +record_dir={record_dir} +active_model={model_path} + +Codex itself is launched from the Gradio environment, not the Articraft Conda environment. +For every Articraft CLI invocation, use this command prefix: + +{_conda_path()} run --no-capture-output -n {ARTICRAFT_CONDA_ENV} python -m cli.main + +Follow EXTERNAL_AGENT_DATA.md exactly. Read the design and link-naming guidance it references, +then use relevant SDK docs/examples. Edit only the active model.py for this record. Do not create +record folders or metadata manually, edit unrelated records, commit/push, or promote this +workbench record to the dataset. + +Create a realistic mechanically meaningful articulated object matching the request. Use semantic +parts, visible plausible joints, appropriate materials, and prompt-specific run_tests(). Iterate +until this succeeds: + +{_conda_path()} run --no-capture-output -n {ARTICRAFT_CONDA_ENV} python -m cli.main external --repo-root {ARTICRAFT_OUTPUT_ROOT} check {record_id} + +Then run: + +{_conda_path()} run --no-capture-output -n {ARTICRAFT_CONDA_ENV} python -m cli.main external --repo-root {ARTICRAFT_OUTPUT_ROOT} finalize {record_id} + +The Gradio app packages the compiled URDF and meshes after you finish. In your final response, +briefly state the articulation mechanisms and validation result.""" + + +def generate_articraft_asset(prompt_value: str, image_value: Any): + """Initialize a record, let Codex author it, and expose one result bundle.""" + prompt = (prompt_value or "").strip() + if not prompt: + yield None, "", "**Input error:** enter a description of the articulated object.", "", "" + return + + details, errors, codex = _check_requirements() + if errors or not codex: + message = ( + "\n".join(f"- {error}" for error in errors) or "Codex CLI is unavailable." + ) + yield None, "", f"**Articulation is not ready.**\n\n{message}", "", "" + return + + record_id = _record_id() + run_root = ARTICRAFT_OUTPUT_ROOT / "runs" / record_id + record_dir = ARTICRAFT_OUTPUT_ROOT / "data" / "records" / record_id + log_lines = [*details, f"Shared output: {ARTICRAFT_OUTPUT_ROOT}"] + try: + run_root.mkdir(parents=True, exist_ok=False) + reference_image = _copy_reference_image(image_value, run_root) + init_command = _articraft_cli_command( + "external", + "--repo-root", + str(ARTICRAFT_OUTPUT_ROOT), + "init", + "--agent", + "codex", + "--record-id", + record_id, + prompt, + ) + log_lines.append("$ " + " ".join(init_command[:-1]) + " ") + initialized = _run_check(init_command, timeout=90) + log_lines.append(_short_output(initialized, limit=4000)) + if initialized.returncode: + yield None, "", "**Articraft record initialization failed.**", "\n".join( + log_lines + ), "" + return + model_path = _active_model_path(record_dir) + except Exception as exc: + yield None, "", f"**Setup failed:** {exc}", "\n".join(log_lines), "" + return + + final_message = run_root / "codex_final_message.txt" + codex_command = [ + codex, + "exec", + "--sandbox", + "workspace-write", + "--color", + "never", + "-C", + str(ARTICRAFT_ROOT), + "--add-dir", + str(ARTICRAFT_OUTPUT_ROOT), + "--output-last-message", + str(final_message), + ] + if reference_image: + codex_command.extend(["--image", str(reference_image)]) + codex_command.append( + _build_codex_prompt( + prompt=prompt, + record_id=record_id, + record_dir=record_dir, + model_path=model_path, + reference_image=reference_image, + ) + ) + log_lines.append("$ codex exec --sandbox workspace-write …") + yield None, record_dir.as_posix(), "**Codex is generating and validating the Articraft model…**", "\n".join( + log_lines + ), "" + + try: + process = subprocess.Popen( + codex_command, + cwd=ARTICRAFT_ROOT, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + start_new_session=True, + env=os.environ.copy(), + ) + except Exception as exc: + yield None, record_dir.as_posix(), f"**Codex could not start:** {exc}", "\n".join( + log_lines + ), "" + return + + output_queue: queue.Queue[str] = queue.Queue() + reader = threading.Thread( + target=read_process_output, args=(process, output_queue), daemon=True + ) + reader.start() + while process.poll() is None: + try: + while True: + log_lines.append(output_queue.get_nowait()) + except queue.Empty: + pass + yield None, record_dir.as_posix(), "**Codex is generating and validating the Articraft model…**", "\n".join( + log_lines[-240:] + ), "" + time.sleep(0.75) + reader.join(timeout=2) + try: + while True: + log_lines.append(output_queue.get_nowait()) + except queue.Empty: + pass + + if final_message.is_file(): + final_text = final_message.read_text(encoding="utf-8", errors="replace").strip() + if final_text: + log_lines.append("\nCodex final response:\n" + final_text) + if process.returncode: + yield None, record_dir.as_posix(), f"**Codex generation failed** (exit code {process.returncode}).", "\n".join( + log_lines[-300:] + ), "" + return + + # Do not rely solely on Codex's final message: independently run the + # external validation and finalize gates before exposing an output bundle. + check_command = _articraft_cli_command( + "external", + "--repo-root", + str(ARTICRAFT_OUTPUT_ROOT), + "check", + record_id, + ) + log_lines.append("$ " + " ".join(check_command)) + yield ( + None, + record_dir.as_posix(), + "**Codex finished. Articraft is running the final validation gate…**", + "\n".join(log_lines[-300:]), + "", + ) + try: + checked = _run_check(check_command, timeout=300) + log_lines.append(_short_output(checked, limit=5000)) + except Exception as exc: + yield ( + None, + record_dir.as_posix(), + f"**Final Articraft validation could not run:** {exc}", + "\n".join(log_lines[-300:]), + "", + ) + return + if checked.returncode: + if not _external_check_is_unsupported(checked): + yield ( + None, + record_dir.as_posix(), + "**Articraft validation failed; no output bundle was published.**", + "\n".join(log_lines[-300:]), + "", + ) + return + # The older CLI reports external init/finalize/categories only. Its + # equivalent strict model validation is the top-level compile command. + compile_command = _articraft_cli_command( + "compile", + "--repo-root", + str(ARTICRAFT_OUTPUT_ROOT), + "--target", + "full", + "--validate", + "--strict-geom-qc", + record_id, + ) + log_lines.append( + "external check is unavailable; falling back to compile --validate." + ) + log_lines.append("$ " + " ".join(compile_command)) + yield ( + None, + record_dir.as_posix(), + "**Using this Articraft version's compile validation gate…**", + "\n".join(log_lines[-300:]), + "", + ) + try: + compiled = _run_check(compile_command, timeout=300) + log_lines.append(_short_output(compiled, limit=5000)) + except Exception as exc: + yield ( + None, + record_dir.as_posix(), + f"**Fallback Articraft validation could not run:** {exc}", + "\n".join(log_lines[-300:]), + "", + ) + return + if compiled.returncode: + yield ( + None, + record_dir.as_posix(), + "**Articraft validation failed; no output bundle was published.**", + "\n".join(log_lines[-300:]), + "", + ) + return + failures = _compile_report_failures(record_id) + if failures: + log_lines.append("Blocking compile-report failures: " + "; ".join(failures)) + yield ( + None, + record_dir.as_posix(), + "**Articraft validation found blocking model defects; no output bundle was published.**", + "\n".join(log_lines[-300:]), + "", + ) + return + + finalize_command = _articraft_cli_command( + "external", + "--repo-root", + str(ARTICRAFT_OUTPUT_ROOT), + "finalize", + record_id, + ) + log_lines.append("$ " + " ".join(finalize_command)) + try: + finalized = _run_check(finalize_command, timeout=300) + log_lines.append(_short_output(finalized, limit=5000)) + except Exception as exc: + yield ( + None, + record_dir.as_posix(), + f"**Articraft finalization could not run:** {exc}", + "\n".join(log_lines[-300:]), + "", + ) + return + if finalized.returncode: + yield ( + None, + record_dir.as_posix(), + "**Articraft finalization failed; no output bundle was published.**", + "\n".join(log_lines[-300:]), + "", + ) + return + + try: + materialized, archive = _make_result_bundle(record_id) + status = ( + "**Articraft generation completed and passed the Codex validation workflow.**\n\n" + f"- Record: `{record_dir}`\n- Compiled output: `{materialized}`\n- Downloadable bundle: `{archive}`" + ) + try: + preview_html = _start_articraft_viser_preview(materialized, record_id) + status += "\n- Interactive Viser preview: ready" + except Exception as exc: + preview_html = "" + status += f"\n- Interactive Viser preview could not start: `{exc}`" + log_lines.append(f"Viser preview failed: {exc}") + yield archive.as_posix(), record_dir.as_posix(), status, "\n".join( + log_lines[-300:] + ), preview_html + except Exception as exc: + yield None, record_dir.as_posix(), f"**Codex finished, but result packaging failed:** {exc}", "\n".join( + log_lines[-300:] + ), "" + + +def build_articraft_panel() -> None: + """Render the Articraft tab inside the Asset engine.""" + gr.Markdown( + "### Articulation\n" + "Generate an articulated object from text and an optional reference image. Codex writes and validates the Articraft model; only submit trusted requests." + ) + with gr.Row(): + configure_button = gr.Button("Configure Articulation & check Codex") + generate_button = gr.Button("Generate articulation", variant="primary") + environment_status = gr.Markdown("**Environment:** not checked.") + with gr.Row(): + prompt = gr.Textbox( + label="Articulated object description", + lines=5, + placeholder="e.g. A countertop toaster oven with a hinged door and rotating temperature knob.", + ) + image = gr.Image( + label="Optional reference image", + type="filepath", + image_mode="RGB", + sources=["upload"], + ) + with gr.Row(): + output_file = gr.File( + label="Compiled Articulation result bundle (.zip)", interactive=False + ) + record_folder = gr.Textbox( + label="Articulation record folder", interactive=False + ) + articulation_preview = gr.HTML( + "
" + "The interactive Viser articulation preview will appear here after generation." + "
" + ) + generation_status = gr.Markdown("**Status:** waiting for a description.") + generation_log = gr.Textbox( + label="Codex / Articraft log", lines=14, interactive=False + ) + + configure_button.click( + configure_articraft_environment, outputs=[environment_status], queue=False + ) + generate_button.click( + generate_articraft_asset, + inputs=[prompt, image], + outputs=[ + output_file, + record_folder, + generation_status, + generation_log, + articulation_preview, + ], + ) diff --git a/embodichain/gen_sim/gradio_ui/app_asset_engine.py b/embodichain/gen_sim/gradio_ui/app_asset_engine.py new file mode 100644 index 000000000..3c021ec90 --- /dev/null +++ b/embodichain/gen_sim/gradio_ui/app_asset_engine.py @@ -0,0 +1,283 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Standalone SimReady asset-engine workflow used by Debug mode. + +The upstream SimReady CLI works on a directory, while Gradio uploads files. +This adapter creates an isolated directory for every run, keeps material +sidecars together with the mesh, and exposes GLB previews before and after +processing. It deliberately has no DexSim dependency. +""" + +from __future__ import annotations + +import queue +import shutil +import threading +import uuid +from pathlib import Path +from typing import Any, Iterable + +import gradio as gr +import trimesh + +from app_articraft import build_articraft_panel +from app_config import ( + DEBUG_ASSET_ENGINE_ROOT, + EMBODICHAIN_ROOT, + SIMREADY_MESH_SUFFIXES, +) +from app_processes import read_process_output, start_pipeline + + +def _as_paths(value: Any) -> list[Path]: + if value is None: + return [] + values: Iterable[Any] = value if isinstance(value, (list, tuple)) else [value] + paths: list[Path] = [] + for item in values: + if isinstance(item, str): + paths.append(Path(item)) + elif isinstance(item, dict) and item.get("path"): + paths.append(Path(item["path"])) + return [path for path in paths if path.is_file()] + + +def _mesh_path(paths: Iterable[Path]) -> Path: + meshes = [path for path in paths if path.suffix.lower() in SIMREADY_MESH_SUFFIXES] + if not meshes: + supported = ", ".join(sorted(SIMREADY_MESH_SUFFIXES)) + raise ValueError( + f"Upload one mesh file ({supported}) and optional material files." + ) + return meshes[0] + + +def _safe_copy_uploads(upload_paths: list[Path], destination: Path) -> Path: + destination.mkdir(parents=True, exist_ok=False) + copied: list[Path] = [] + for index, source in enumerate(upload_paths): + # Upload file names are untrusted. Keep only their basename and avoid + # collisions without ever interpreting a supplied relative path. + name = source.name or f"upload_{index}" + target = destination / name + if target.exists(): + target = destination / f"{target.stem}_{index}{target.suffix}" + shutil.copy2(source, target) + copied.append(target) + return _mesh_path(copied) + + +def _export_preview(mesh_path: Path, destination: Path) -> Path: + """Convert every supported mesh type to GLB for one consistent viewer.""" + loaded = trimesh.load(mesh_path, force="scene", process=False) + if isinstance(loaded, trimesh.Trimesh): + scene = trimesh.Scene(loaded) + elif isinstance(loaded, trimesh.Scene): + scene = loaded + else: + raise ValueError(f"Unsupported mesh payload: {type(loaded)!r}") + if not scene.geometry: + raise ValueError("The uploaded asset contains no renderable geometry.") + destination.parent.mkdir(parents=True, exist_ok=True) + scene.export(destination) + return destination + + +def prepare_asset_input_preview(upload_value: Any): + """Validate an upload and return a normalized GLB preview without running SimReady.""" + try: + source = _mesh_path(_as_paths(upload_value)) + preview = DEBUG_ASSET_ENGINE_ROOT / "previews" / f"{uuid.uuid4().hex}.glb" + _export_preview(source, preview) + return ( + preview.as_posix(), + "**Asset input ready.** Review the model, then run SimReady.", + ) + except Exception as exc: + return None, f"**Input error:** {exc}" + + +def _find_simready_output(output_root: Path) -> Path: + candidates = sorted( + output_root.rglob("asset_simready.glb"), + key=lambda path: path.stat().st_mtime_ns, + reverse=True, + ) + if not candidates: + candidates = sorted( + output_root.rglob("asset_simready.obj"), + key=lambda path: path.stat().st_mtime_ns, + reverse=True, + ) + if not candidates: + raise FileNotFoundError( + "SimReady completed without asset_simready.glb or asset_simready.obj." + ) + return candidates[0] + + +def run_simready_asset(upload_value: Any, category: str): + """Run one upstream SimReady job and stream concise subprocess progress.""" + category = (category or "").strip() + if not category: + yield None, None, None, "**Input error:** enter an asset category.", "" + return + try: + uploads = _as_paths(upload_value) + _mesh_path(uploads) + run_root = DEBUG_ASSET_ENGINE_ROOT / "runs" / uuid.uuid4().hex + input_dir = run_root / "input" + output_root = run_root / "output" + source_mesh = _safe_copy_uploads(uploads, input_dir) + input_preview = _export_preview(source_mesh, run_root / "input_preview.glb") + except Exception as exc: + yield None, None, None, f"**Input error:** {exc}", "" + return + + command = [ + __import__("sys").executable, + "-m", + "embodichain.gen_sim.simready_pipeline.cli.start", + "--input_dir", + str(input_dir), + "--output_root", + str(output_root), + "--category", + category, + ] + log_lines = ["$ " + " ".join(command)] + yield input_preview.as_posix(), None, None, "**SimReady is running…**", "\n".join( + log_lines + ) + + try: + process = start_pipeline(command) + except Exception as exc: + yield input_preview.as_posix(), None, None, f"**Pipeline start failed:** {exc}", "\n".join( + log_lines + ) + return + + output_queue: queue.Queue[str] = queue.Queue() + reader = threading.Thread( + target=read_process_output, args=(process, output_queue), daemon=True + ) + reader.start() + while process.poll() is None: + try: + while True: + log_lines.append(output_queue.get_nowait()) + except queue.Empty: + pass + # Keep the browser responsive while the Blender/LLM stages run. + yield input_preview.as_posix(), None, None, "**SimReady is running…**", "\n".join( + log_lines[-160:] + ) + __import__("time").sleep(0.5) + reader.join(timeout=1) + try: + while True: + log_lines.append(output_queue.get_nowait()) + except queue.Empty: + pass + + if process.returncode != 0: + yield input_preview.as_posix(), None, None, f"**SimReady failed** (exit code {process.returncode}).", "\n".join( + log_lines[-220:] + ) + return + try: + result = _find_simready_output(output_root) + preview = ( + result + if result.suffix.lower() == ".glb" + else _export_preview(result, run_root / "output_preview.glb") + ) + yield input_preview.as_posix(), preview.as_posix(), result.as_posix(), "**SimReady completed.**", "\n".join( + log_lines[-220:] + ) + except Exception as exc: + yield input_preview.as_posix(), None, None, f"**Output error:** {exc}", "\n".join( + log_lines[-220:] + ) + + +def build_asset_engine_panel() -> dict[str, Any]: + """Create the Debug Asset-engine panel and return its event endpoints.""" + with gr.Column(visible=True) as panel: + gr.Markdown( + "## Asset engine\nConvert an existing mesh with SimReady, or generate a new articulated asset through Articraft and Codex. DexSim is not started in this engine." + ) + with gr.Tabs(): + with gr.Tab("SimReady"): + with gr.Row(): + uploads = gr.File( + label="3D asset and optional material files", + file_count="multiple", + type="filepath", + file_types=[ + ".glb", + ".gltf", + ".obj", + ".ply", + ".stl", + ".mtl", + ".png", + ".jpg", + ".jpeg", + ".webp", + ".bin", + ], + ) + category = gr.Textbox( + label="Asset category", + value="rigid_object", + placeholder="e.g. cup, chair, bottle", + ) + with gr.Row(): + input_model = gr.Model3D( + label="Input asset preview", + height=440, + clear_color=(0.94, 0.94, 0.94, 1.0), + ) + output_model = gr.Model3D( + label="SimReady asset preview", + height=440, + clear_color=(0.94, 0.94, 0.94, 1.0), + ) + with gr.Row(): + run_button = gr.Button("Run SimReady", variant="primary") + output_file = gr.File( + label="SimReady asset output", interactive=False + ) + status = gr.Markdown("**Status:** waiting for an asset.") + log = gr.Textbox(label="Pipeline log", lines=10, interactive=False) + with gr.Tab("Articulation"): + build_articraft_panel() + + uploads.change( + prepare_asset_input_preview, + inputs=[uploads], + outputs=[input_model, status], + queue=False, + ) + run_button.click( + run_simready_asset, + inputs=[uploads, category], + outputs=[input_model, output_model, output_file, status, log], + ) + return {"panel": panel} diff --git a/embodichain/gen_sim/gradio_ui/app_commands.py b/embodichain/gen_sim/gradio_ui/app_commands.py new file mode 100644 index 000000000..7dc1f4a1c --- /dev/null +++ b/embodichain/gen_sim/gradio_ui/app_commands.py @@ -0,0 +1,173 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""CLI command builders for EmbodiChain pipelines.""" + +from __future__ import annotations + +import sys +from typing import Protocol + +from app_config import ( + COMMANDS, + ROBOT_PROFILE_FRANKA, + ROBOT_PROFILE_UR5, + ROBOT_PROFILE_UR10, + SCENE_ID, +) + + +class ScenePathsLike(Protocol): + scene_id: str + image_path: object + fast_gym_config: object + agent_config: object + + +def robot_profile_cli_value(robot_profile: str | None) -> str | None: + return { + ROBOT_PROFILE_FRANKA: "franka", + ROBOT_PROFILE_UR5: "dual_ur5", + ROBOT_PROFILE_UR10: "dual_ur10", + }.get(robot_profile) + + +def _pipeline_paths(paths: ScenePathsLike) -> tuple[str, str]: + return ( + f"gym_project/{paths.scene_id}", + f"gym_project/action_agent_pipeline/configs/{paths.scene_id}", + ) + + +def build_initial_pipeline_command( + task_text: str, + paths: ScenePathsLike, + prompt2scene_prompt: str = "", + robot_profile: str | None = None, + load_template_material: bool = False, +) -> list[str]: + prompt_root, config_dir = _pipeline_paths(paths) + command = [ + sys.executable, + "-m", + COMMANDS["pipeline"]["module"], + "--image", + str(paths.image_path.resolve()), + "--prompt2scene-output-root", + prompt_root, + "--config-output-dir", + config_dir, + "--task_name", + SCENE_ID, + "--task_description", + task_text, + *COMMANDS["pipeline"]["base_args"], + ] + if profile := robot_profile_cli_value(robot_profile): + command.extend(["--robot-profile", profile]) + if prompt2scene_prompt.strip(): + command.extend(["--prompt2scene-prompt", prompt2scene_prompt.strip()]) + if load_template_material: + command.append("--load-template-material") + return command + + +def build_scene_edit_pipeline_command( + task_text: str, + env_text: str, + paths: ScenePathsLike, + robot_profile: str | None = None, + load_template_material: bool = False, +) -> list[str]: + prompt_root, config_dir = _pipeline_paths(paths) + command = [ + sys.executable, + "-m", + COMMANDS["pipeline"]["module"], + "--prompt2scene-output-root", + prompt_root, + "--prompt2scene-prompt", + env_text, + "--config-output-dir", + config_dir, + "--task_name", + SCENE_ID, + "--task_description", + task_text, + *COMMANDS["pipeline"]["base_args"], + ] + if profile := robot_profile_cli_value(robot_profile): + command.extend(["--robot-profile", profile]) + if load_template_material: + command.append("--load-template-material") + return command + + +def build_config_command_for_paths( + task_text: str, + paths: ScenePathsLike, + robot_profile: str | None = None, + load_template_material: bool = False, +) -> list[str]: + _, config_dir = _pipeline_paths(paths) + command = [ + sys.executable, + "-m", + COMMANDS["config"]["module"], + "--gym_project", + f"gym_project/{paths.scene_id}/gym_export", + "--output_dir", + config_dir, + "--task_name", + SCENE_ID, + "--task_description", + task_text, + *COMMANDS["config"]["base_args"], + ] + if profile := robot_profile_cli_value(robot_profile): + command.extend(["--robot-profile", profile]) + if load_template_material: + command.append("--load-template-material") + return command + + +def build_run_agent_command( + paths: ScenePathsLike, + *, + parallel_env: bool = False, + robot_profile: str | None = None, + supports_robot_profile: bool = False, +) -> list[str]: + agent = COMMANDS["agent"] + command = [ + sys.executable, + "-m", + agent["module"], + "--task_name", + SCENE_ID, + "--gym_config", + str(paths.fast_gym_config), + "--agent_config", + str(paths.agent_config), + *agent["base_args"], + "--num_envs", + agent["parallel_num_envs"] if parallel_env else agent["single_num_envs"], + ] + if parallel_env: + command.extend(agent["parallel_args"]) + if supports_robot_profile and (profile := robot_profile_cli_value(robot_profile)): + command.extend(["--robot-profile", profile]) + return command diff --git a/embodichain/gen_sim/gradio_ui/app_config.py b/embodichain/gen_sim/gradio_ui/app_config.py new file mode 100644 index 000000000..7de206cf2 --- /dev/null +++ b/embodichain/gen_sim/gradio_ui/app_config.py @@ -0,0 +1,303 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Central configuration for the Gradio application. + +Keep deployment-specific paths, UI copy, and CLI command definitions here so +application modules do not embed environment-specific values. +""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any + +PROXY_ENV_KEYS = ( + "HTTP_PROXY", + "HTTPS_PROXY", + "ALL_PROXY", + "FTP_PROXY", + "http_proxy", + "https_proxy", + "all_proxy", + "ftp_proxy", +) +DIRECT_NO_PROXY_VALUE = "*" + +# SimReady uses an OpenAI-compatible multimodal endpoint. Configure these +# values here for a local deployment, or provide the matching SIMREADY_* env +# vars before launch. Keep the API key out of commits; an empty value leaves +# any inherited OPENAI_* variables and SimReady's own JSON configuration intact. +SIMREADY_OPENAI_API_KEY = os.environ.get("SIMREADY_OPENAI_API_KEY", "") +SIMREADY_OPENAI_MODEL = os.environ.get("SIMREADY_OPENAI_MODEL", "") +SIMREADY_OPENAI_BASE_URL = os.environ.get("SIMREADY_OPENAI_BASE_URL", "") + + +def configure_direct_network_env(env: Any = None) -> None: + """Disable proxy inheritance for local pipeline and Gradio processes.""" + if env is None: + env = os.environ + for key in PROXY_ENV_KEYS: + env.pop(key, None) + env["NO_PROXY"] = DIRECT_NO_PROXY_VALUE + env["no_proxy"] = DIRECT_NO_PROXY_VALUE + env.setdefault("GRADIO_ANALYTICS_ENABLED", "False") + + +def configure_simready_llm_env(env: Any = None) -> None: + """Map app-level SimReady settings to the upstream CLI's environment.""" + if env is None: + env = os.environ + configured_values = { + "OPENAI_API_KEY": SIMREADY_OPENAI_API_KEY, + "OPENAI_MODEL": SIMREADY_OPENAI_MODEL, + "OPENAI_BASE_URL": SIMREADY_OPENAI_BASE_URL, + } + for key, value in configured_values.items(): + if value: + env[key] = value + + +APP_ROOT = Path(__file__).resolve().parent +EMBODICHAIN_ROOT = Path( + os.environ.get("EMBODICHAIN_ROOT", "/home/dex/桌面/EmbodiChain") +).expanduser() +ASSETS_DIR = APP_ROOT / "assets" +DEXFORCE_LOGO = ASSETS_DIR / "dexforce.png" +INTERACT_RANDOM_PREVIEW_DIR = APP_ROOT / ".gradio_previews" +DEBUG_ENGINE_ROOT = APP_ROOT / ".debug_engine" +DEBUG_ASSET_ENGINE_ROOT = DEBUG_ENGINE_ROOT / "assets" +ARTICRAFT_ROOT = Path( + os.environ.get("ARTICRAFT_ROOT", str(APP_ROOT / ".articraft")) +).expanduser() +ARTICRAFT_REPOSITORY_URL = os.environ.get( + "ARTICRAFT_REPOSITORY_URL", "https://github.com/mattzh72/articraft.git" +) +ARTICRAFT_CONDA_ENV = os.environ.get("ARTICRAFT_CONDA_ENV", "articraft") +# Keep every Articraft record, copied reference image, log, and downloadable +# result bundle under one app-owned directory rather than the source checkout. +ARTICRAFT_OUTPUT_ROOT = Path( + os.environ.get("ARTICRAFT_OUTPUT_ROOT", str(DEBUG_ENGINE_ROOT / "articraft")) +).expanduser() +DEBUG_SCENE_ENGINE_ROOT = DEBUG_ENGINE_ROOT / "scenes" +SCENE_ENGINE_CONFIG = ( + EMBODICHAIN_ROOT / "embodichain" / "gen_sim" / "scene_engine_config.json" +) +SCENE_ENGINE_VISER_PORT = int(os.environ.get("SCENE_ENGINE_VISER_PORT", "8080")) +# Articulation previews run as a separate Viser process from scene previews, +# so they need their own externally configurable port. +ARTICRAFT_VISER_PORT = int(os.environ.get("ARTICRAFT_VISER_PORT", "8081")) +SCENE_ID = "current" + +GYM_PROJECT_ROOT = EMBODICHAIN_ROOT / "gym_project" +ACTION_AGENT_ROOT = GYM_PROJECT_ROOT / "action_agent_pipeline" +IMAGE_DIR = ACTION_AGENT_ROOT / "images" +AUTO_LOG_DIR = ACTION_AGENT_ROOT / "auto_logs" +IMAGE_PATH = IMAGE_DIR / f"{SCENE_ID}.png" +PROMPT2SCENE_ROOT = GYM_PROJECT_ROOT / SCENE_ID +CONFIG_DIR = ACTION_AGENT_ROOT / "configs" / SCENE_ID +FAST_GYM_CONFIG = CONFIG_DIR / "fast_gym_config.json" +OUTPUTS_DIR = EMBODICHAIN_ROOT / "outputs" +CURRENT_GYM_EXPORT_DIR = PROMPT2SCENE_ROOT / "gym_export" +CURRENT_GYM_EXPORT_CONFIG = CURRENT_GYM_EXPORT_DIR / "gym_config.json" +GRADIO_SCENE_DIR = CONFIG_DIR / "gradio_scene" +GRADIO_SCENE_GLB = GRADIO_SCENE_DIR / "scene_current.glb" +GRADIO_INITIAL_SCENE_GLB = GRADIO_SCENE_DIR / "initial_scene.glb" +GRADIO_OBJECT_PREVIEW_GLB = GRADIO_SCENE_DIR / "object_preview.glb" +SCENE_MANIFEST = GRADIO_SCENE_DIR / "scene_manifest.json" +PENDING_PREFIX = "_gradio_pending_" +REPLACED_PREFIX = "_gradio_replaced_" +GRADIO_SCENE_TRANSFORM_POLICY = "dexsim_gltf_y_up_to_sim_z_up_v1" + +PROCESS_STOP_TIMEOUT_S = 8.0 +TEXT_REWRITE_SUFFIXES = {".json", ".jsonl", ".txt", ".yaml", ".yml", ".md", ".csv"} +VIDEO_SUFFIXES = {".mp4", ".avi", ".mov", ".mkv", ".webm"} +LEROBOT_PREVIEW_DIR = OUTPUTS_DIR / "lerobot_previews" +COMBINED_PREVIEW_DIR = OUTPUTS_DIR / "combined_previews" +LEROBOT_PREVIEW_MAX_FRAMES = 360 +COMBINED_VIDEO_FPS = 25 + +TOP_MODE_AUTO = "auto" +TOP_MODE_INTERACT = "interact" +TOP_MODE_PARALLEL_ENV = "parallel_env" +APP_MODE_DEMO = "demo" +APP_MODE_DEBUG = "debug" +DEBUG_ENGINE_ASSET = "asset_engine" +DEBUG_ENGINE_SCENE = "scene_engine" +DEBUG_ENGINE_ACTION = "action_engine" +DEBUG_ENGINES = ( + (DEBUG_ENGINE_ASSET, "Asset_engine"), + (DEBUG_ENGINE_SCENE, "Scene_engine"), + (DEBUG_ENGINE_ACTION, "Action_engine"), +) + +# SimReady accepts one mesh plus optional material/texture sidecar files. The +# File component deliberately permits the sidecars so OBJ/GLTF uploads retain +# their appearance during both preview and processing. +SIMREADY_MESH_SUFFIXES = {".glb", ".gltf", ".obj", ".ply", ".stl"} + +LANGUAGE_EN = "en" +LANGUAGE_ZH = "zh" +BUTTON_LABELS = { + LANGUAGE_EN: { + "auto": "Auto", + "interact": "Interact", + "parallel_env": "Parallel Simulation", + "rerun_simulation": "Run Task", + "generate": "Generate", + "start": "Start", + "random_input": "Random Task", + "random_scene_input": "Random Scene", + "reset": "Reset", + "stop": "Stop", + "language": "中文", + }, + LANGUAGE_ZH: { + "auto": "自动", + "interact": "交互", + "parallel_env": "并行仿真", + "rerun_simulation": "运行任务", + "generate": "生成", + "start": "开始", + "random_input": "随机任务", + "random_scene_input": "随机场景", + "reset": "重置", + "stop": "停止", + "language": "English", + }, +} +UI_TEXT = { + LANGUAGE_EN: { + "heading": "# Generative Simulation User Interface", + "instruction": "Upload one image, enter one task, then EmbodiChain will generate simulation data what you want.", + "robot": "Robot", + "input_image": "Input image", + "task_description": "Task description", + "task_placeholder": "Put the middle bottle on the book", + "scene_description": "Scene description", + "scene_placeholder": "Optional: describe how to edit the current scene", + "scene_mode": "Generation mode", + "scene_mode_initial": "Initial generation", + "scene_mode_edit": "Edit current scene", + "scene_mode_task_only": "Change task only", + "single_video_preview": "LeRobot Data Preview", + "parallel_video_preview": "Parallel Env Data Preview", + "current_task": "Current task", + "progress": "Progress", + "initial_preview": "Initial scene preview", + "edited_preview": "Edited scene preview", + "object_preview": "Generated object GLBs preview", + }, + LANGUAGE_ZH: { + "heading": "# 生成式仿真用户界面", + "instruction": "上传一张图片，输入一个任务，EmbodiChain 将生成所需的仿真数据。", + "robot": "机器人", + "input_image": "输入图像", + "task_description": "任务描述", + "task_placeholder": "把中间的水瓶放到书上", + "scene_description": "场景描述", + "scene_placeholder": "可选：描述如何编辑当前场景", + "scene_mode": "生成模式", + "scene_mode_initial": "初始生成", + "scene_mode_edit": "编辑当前场景", + "scene_mode_task_only": "仅修改任务", + "single_video_preview": "LeRobot 数据预览", + "parallel_video_preview": "并行环境数据预览", + "current_task": "当前任务", + "progress": "进度", + "initial_preview": "初始场景预览", + "edited_preview": "编辑后场景预览", + "object_preview": "生成对象 GLB 预览", + }, +} + +PIPELINE_MODE_INITIAL = "initial" +PIPELINE_MODE_EDIT = "edit" +PIPELINE_MODE_TASK_ONLY = "task_only" +SCENE_MODE_INITIAL = "initial" +SCENE_MODE_EDIT = "edit" +SCENE_MODE_TASK_ONLY = "task_only" +ROBOT_PROFILE_FRANKA = "Franka" +ROBOT_PROFILE_UR5 = "UR5" +ROBOT_PROFILE_UR10 = "UR10" +ROBOT_PROFILES = [ROBOT_PROFILE_FRANKA, ROBOT_PROFILE_UR5, ROBOT_PROFILE_UR10] +DEFAULT_ROBOT_PROFILE = ROBOT_PROFILE_UR5 +RUN_LOG_MODE_AUTO = "auto" +RUN_LOG_MODE_INTERACT = "interact" + +# Command modules and immutable argument defaults. Dynamic values are added by +# command builders in app_commands.py. +COMMANDS = { + "pipeline": { + "module": "embodichain.gen_sim.action_agent_pipeline.cli.run_agent_pipeline", + "base_args": ( + "--use-prompt2scene", + "--overwrite-config", + "--regenerate", + "--skip-run-agent", + ), + }, + "config": { + "module": "embodichain.gen_sim.action_agent_pipeline.cli.generate_action_agent_config", + "base_args": ("--overwrite",), + }, + "agent": { + "module": "embodichain.gen_sim.action_agent_pipeline.cli.run_agent", + "help_args": ("--help",), + "base_args": ("--regenerate", "--renderer", "fast-rt"), + "parallel_args": ("--arena_space", "2.2", "--filter_dataset_saving"), + "parallel_num_envs": "9", + "single_num_envs": "1", + }, + # Scene Engine is dispatched by EmbodiChain's registered top-level CLI. + # The scene_engine package itself has no __main__.py in this checkout. + "scene_engine": { + "module": "embodichain", + "base_args": ("scene-engine",), + "preview_script": "embodichain/gen_sim/scene_engine/cli/preview.py", + }, +} + +PHASE_DEFINITIONS = { + "idle": (0, "Idle"), + "received": (5, "Input received"), + "started": (10, "Local pipeline started"), + "scene_intake": (20, "Scene understanding"), + "relations": (35, "Segmentation and spatial relations"), + "asset_generation": (55, "3D asset generation"), + "gym_export": (70, "Scene export"), + "config": (82, "Action config generated"), + "preview": (90, "3D preview loaded"), + "complete": (100, "Complete"), + "failed": (100, "Failed"), +} +TIMING_PHASE_LABELS = { + "relations": "Segmentation / spatial relations", + "asset_generation": "Object generation", + "gym_export": "Scene generation / export", + "action_graph_execution": "Action graph execution", +} +TIMING_PHASE_ORDER = ( + "relations", + "asset_generation", + "gym_export", + "action_graph_execution", +) + +SERVER_NAME = os.environ.get("GRADIO_SERVER_NAME", "0.0.0.0") +SERVER_PORT = int(os.environ.get("GRADIO_SERVER_PORT", "7860")) +DEFAULT_CONCURRENCY_LIMIT = 1 diff --git a/embodichain/gen_sim/gradio_ui/app_media.py b/embodichain/gen_sim/gradio_ui/app_media.py new file mode 100644 index 000000000..41271c0f5 --- /dev/null +++ b/embodichain/gen_sim/gradio_ui/app_media.py @@ -0,0 +1,724 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Run-log archival and video/dataset preview generation.""" + +from __future__ import annotations + +import argparse +import json +import math +import shutil +import subprocess +from collections.abc import Sequence +from datetime import datetime +from pathlib import Path +from typing import Any + +import numpy as np +from PIL import Image, ImageDraw + +from app_config import * # noqa: F403 - media paths and limits are configuration. +from app_state import format_timing_lines, runtime, runtime_lock, snapshot_timing_locked + + +def archive_run_log( + *, + mode: str, + task_description: str = "", + scene_description: str = "", + outcome: str, + audience_video: Path | None = None, +) -> Path | None: + with runtime_lock: + run_logs = list(runtime.log_lines) + status_text = runtime.status + last_error = runtime.last_error + runtime_video = runtime.video_path + timing_durations, simulation_duration = snapshot_timing_locked() + + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f") + run_dir = make_next_log_archive_dir() + video_paths, video_errors = archive_audience_video( + run_dir, + audience_video or runtime_video, + ) + log_path = run_dir / "log.md" + content = [ + f"mode: {mode}", + "", + f"Timestamp: {timestamp}", + f"Outcome: {outcome}", + "", + "## Task description", + "", + task_description or "", + "", + "## Scene description", + "", + scene_description or "", + "", + "## Status", + "", + status_text or "", + ] + if last_error: + content.extend(["", "## Last error", "", last_error]) + if video_paths: + content.extend( + [ + "", + "## Archived audience video", + "", + *[path.as_posix() for path in video_paths], + ] + ) + if video_errors: + content.extend(["", "## Video archive errors", "", *video_errors]) + content.extend( + [ + "", + "## Logs", + "", + "```text", + "\n".join(run_logs) if run_logs else "(no logs)", + "```", + "", + "## Timing", + "", + *format_timing_lines(timing_durations, simulation_duration), + "", + ] + ) + + try: + run_dir.mkdir(parents=True, exist_ok=True) + log_path.write_text("\n".join(content), encoding="utf-8") + except Exception as exc: + with runtime_lock: + runtime.log_lines.append(f"Failed to archive run log: {exc}") + return None + return log_path + + +def make_next_log_archive_dir() -> Path: + AUTO_LOG_DIR.mkdir(parents=True, exist_ok=True) + existing_indices = [ + int(path.name) + for path in AUTO_LOG_DIR.iterdir() + if path.is_dir() and path.name.isdigit() + ] + next_index = (max(existing_indices) + 1) if existing_indices else 1 + while True: + candidate = AUTO_LOG_DIR / f"{next_index:04d}" + if not candidate.exists(): + try: + candidate.mkdir(parents=True, exist_ok=False) + return candidate + except FileExistsError: + pass + next_index += 1 + + +def archive_audience_video( + run_dir: Path, + video_path: Path | None, +) -> tuple[list[Path], list[str]]: + copied_paths: list[Path] = [] + errors: list[str] = [] + if video_path is None: + return copied_paths, errors + if not video_path.is_file(): + return copied_paths, [f"Audience video not found: {video_path}"] + destination = run_dir / "audience_video" / video_path.name + try: + destination.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(video_path, destination) + except Exception as exc: + errors.append(f"Failed to archive audience video {video_path}: {exc}") + return copied_paths, errors + copied_paths.append(destination.relative_to(run_dir)) + return copied_paths, errors + + +def collect_output_videos() -> list[Path]: + if not OUTPUTS_DIR.is_dir(): + return [] + return sorted( + path + for path in OUTPUTS_DIR.rglob("*") + if path.is_file() and path.suffix.lower() in VIDEO_SUFFIXES + ) + + +def collect_audience_output_videos() -> list[Path]: + videos = collect_output_videos() + audience_videos = [ + path + for path in videos + if "audience" in path.relative_to(OUTPUTS_DIR).as_posix().lower() + ] + if audience_videos: + return audience_videos + return [ + path + for path in videos + if "audience" in path.relative_to(OUTPUTS_DIR).as_posix().lower() + ] + + +def latest_audience_output_video(min_mtime_ns: int | None = None) -> Path | None: + latest_path: Path | None = None + latest_mtime = -1 + for path in collect_audience_output_videos(): + try: + mtime = path.stat().st_mtime_ns + except OSError: + continue + if min_mtime_ns is not None and mtime < min_mtime_ns: + continue + if mtime > latest_mtime: + latest_path = path + latest_mtime = mtime + return latest_path + + +def configured_lerobot_roots() -> list[Path]: + roots: list[Path] = [] + env_root = os.environ.get("EMBODICHAIN_DATASET_ROOT") + if env_root: + roots.append(Path(env_root).expanduser()) + roots.append(Path("~/.cache/embodichain_datasets").expanduser()) + + config_roots = read_lerobot_save_paths(CURRENT_PATHS.fast_gym_config) + roots.extend(config_roots) + + normalized: list[Path] = [] + seen: set[Path] = set() + for root in roots: + root = root.expanduser() + if not root.is_absolute(): + root = EMBODICHAIN_ROOT / root + try: + resolved = root.resolve() + except OSError: + resolved = root + if resolved in seen: + continue + seen.add(resolved) + normalized.append(root) + return normalized + + +def read_lerobot_save_paths(config_path: Path) -> list[Path]: + if not config_path.is_file(): + return [] + try: + config = json.loads(config_path.read_text(encoding="utf-8")) + except Exception: + return [] + + paths: list[Path] = [] + + def visit(value: Any, key_path: tuple[str, ...] = ()) -> None: + if isinstance(value, dict): + if key_path[-2:] == ("lerobot", "params") and isinstance( + value.get("save_path"), str + ): + paths.append(Path(value["save_path"])) + for key, child in value.items(): + visit(child, (*key_path, str(key))) + elif isinstance(value, list): + for item in value: + visit(item, key_path) + + visit(config) + return paths + + +def collect_lerobot_datasets() -> list[Path]: + datasets: list[Path] = [] + for root in configured_lerobot_roots(): + if not root.is_dir(): + continue + try: + candidates = list(root.iterdir()) + except OSError: + continue + for candidate in candidates: + if not candidate.is_dir(): + continue + if (candidate / "meta" / "info.json").is_file() or ( + candidate / "data" + ).is_dir(): + datasets.append(candidate) + return datasets + + +def latest_lerobot_dataset(min_mtime_ns: int | None = None) -> Path | None: + latest_path: Path | None = None + latest_mtime = -1 + for dataset_path in collect_lerobot_datasets(): + if not lerobot_dataset_has_frames(dataset_path): + continue + mtime = latest_lerobot_dataset_mtime_ns(dataset_path) + if min_mtime_ns is not None and mtime < min_mtime_ns: + continue + if mtime > latest_mtime: + latest_path = dataset_path + latest_mtime = mtime + return latest_path + + +def lerobot_dataset_has_frames(dataset_path: Path) -> bool: + data_dir = dataset_path / "data" + return data_dir.is_dir() and any(data_dir.rglob("*.parquet")) + + +def latest_lerobot_dataset_mtime_ns(dataset_path: Path) -> int: + latest_mtime = -1 + for child in dataset_path.rglob("*"): + if not child.is_file(): + continue + try: + latest_mtime = max(latest_mtime, child.stat().st_mtime_ns) + except OSError: + continue + if latest_mtime >= 0: + return latest_mtime + try: + return dataset_path.stat().st_mtime_ns + except OSError: + return -1 + + +def build_lerobot_preview_video(dataset_path: Path) -> Path | None: + parquet_paths = sorted((dataset_path / "data").rglob("*.parquet")) + if not parquet_paths: + return None + + latest_source_mtime = max( + latest_lerobot_dataset_mtime_ns(dataset_path), + *(path.stat().st_mtime_ns for path in parquet_paths), + ) + output_path = LEROBOT_PREVIEW_DIR / f"{dataset_path.name}_data_preview.mp4" + if output_path.is_file() and output_path.stat().st_mtime_ns >= latest_source_mtime: + return output_path + + try: + import imageio.v2 as imageio + import pandas as pd + except Exception as exc: + with runtime_lock: + runtime.log_lines.append( + f"LeRobot preview skipped; missing dependency: {exc}" + ) + return None + + try: + data_frame = pd.concat( + [pd.read_parquet(path) for path in parquet_paths], + ignore_index=True, + ) + except Exception as exc: + with runtime_lock: + runtime.log_lines.append(f"LeRobot preview skipped; read failed: {exc}") + return None + + if data_frame.empty: + return None + + try: + fps = read_lerobot_fps(dataset_path) or 25 + fps = max(1, min(int(round(fps)), 30)) + frames = render_lerobot_data_frames(data_frame, dataset_path.name) + if not frames: + return None + output_path.parent.mkdir(parents=True, exist_ok=True) + with imageio.get_writer(output_path, fps=fps, codec="libx264") as writer: + for frame in frames: + writer.append_data(frame) + except Exception as exc: + with runtime_lock: + runtime.log_lines.append(f"LeRobot preview skipped; render failed: {exc}") + return None + + return output_path + + +def video_duration_seconds(video_path: Path) -> float | None: + command = [ + "ffprobe", + "-v", + "error", + "-show_entries", + "format=duration", + "-of", + "default=noprint_wrappers=1:nokey=1", + str(video_path), + ] + try: + result = subprocess.run( + command, + capture_output=True, + text=True, + check=False, + timeout=15, + ) + duration = float(result.stdout.strip()) + except (OSError, ValueError, subprocess.TimeoutExpired): + return None + return duration if duration > 0 else None + + +def build_single_env_combined_video( + audience_video: Path | None, + lerobot_video: Path | None, +) -> Path | None: + """Create a synchronized side-by-side simulation and LeRobot video.""" + if ( + audience_video is None + or lerobot_video is None + or not audience_video.is_file() + or not lerobot_video.is_file() + ): + return None + + audience_duration = video_duration_seconds(audience_video) + lerobot_duration = video_duration_seconds(lerobot_video) + if audience_duration is None or lerobot_duration is None: + return None + + latest_source_mtime = max( + audience_video.stat().st_mtime_ns, + lerobot_video.stat().st_mtime_ns, + ) + output_path = ( + COMBINED_PREVIEW_DIR + / f"{safe_filename_part(audience_video.stem)}_with_lerobot.mp4" + ) + if output_path.is_file() and output_path.stat().st_mtime_ns >= latest_source_mtime: + return output_path + + lerobot_time_scale = audience_duration / lerobot_duration + filter_graph = ( + f"[0:v]fps={COMBINED_VIDEO_FPS},scale=960:540:force_original_aspect_ratio=decrease," + "pad=960:540:(ow-iw)/2:(oh-ih)/2:color=black,setsar=1," + "setpts=PTS-STARTPTS[sim];" + f"[1:v]fps={COMBINED_VIDEO_FPS},scale=960:540:force_original_aspect_ratio=decrease," + "pad=960:540:(ow-iw)/2:(oh-ih)/2:color=black,setsar=1," + f"setpts=(PTS-STARTPTS)*{lerobot_time_scale:.9f}[data];" + "[sim][data]hstack=inputs=2:shortest=1,format=yuv420p[video]" + ) + command = [ + "ffmpeg", + "-y", + "-i", + str(audience_video), + "-i", + str(lerobot_video), + "-filter_complex", + filter_graph, + "-map", + "[video]", + "-an", + "-c:v", + "libx264", + "-preset", + "veryfast", + "-crf", + "23", + "-movflags", + "+faststart", + str(output_path), + ] + try: + output_path.parent.mkdir(parents=True, exist_ok=True) + result = subprocess.run( + command, + capture_output=True, + text=True, + check=False, + timeout=180, + ) + if result.returncode == 0 and output_path.is_file(): + return output_path + with runtime_lock: + runtime.log_lines.append( + "Combined video skipped: " + + ( + result.stderr.strip().splitlines()[-1] + if result.stderr + else "ffmpeg failed" + ) + ) + except (OSError, subprocess.TimeoutExpired) as exc: + with runtime_lock: + runtime.log_lines.append(f"Combined video skipped: {exc}") + return None + + +def read_lerobot_fps(dataset_path: Path) -> int | None: + info_path = dataset_path / "meta" / "info.json" + if not info_path.is_file(): + return None + try: + info = json.loads(info_path.read_text(encoding="utf-8")) + except Exception: + return None + fps = info.get("fps") + if isinstance(fps, (int, float)): + return int(fps) + return None + + +def render_lerobot_data_frames(data_frame: Any, dataset_name: str) -> list[np.ndarray]: + total_rows = len(data_frame) + frame_indices = np.linspace( + 0, + total_rows - 1, + num=min(total_rows, LEROBOT_PREVIEW_MAX_FRAMES), + dtype=int, + ) + state = series_to_matrix(data_frame.get("observation.state")) + action = series_to_matrix(data_frame.get("action")) + qvel = series_to_matrix(data_frame.get("observation.qvel")) + timestamps = numeric_column(data_frame, "timestamp", total_rows) + + frames: list[np.ndarray] = [] + for row_index in frame_indices: + image = Image.new("RGB", (960, 544), (247, 248, 250)) + draw = ImageDraw.Draw(image) + draw_lerobot_header( + draw, + dataset_name=dataset_name, + row_index=int(row_index), + total_rows=total_rows, + timestamp=float(timestamps[row_index]) if len(timestamps) else None, + ) + draw_signal_panel( + draw, (40, 96, 920, 220), state, row_index, "observation.state" + ) + draw_signal_panel(draw, (40, 244, 920, 368), action, row_index, "action") + draw_bar_panel(draw, (40, 392, 920, 506), qvel, row_index, "observation.qvel") + frames.append(np.asarray(image)) + return frames + + +def series_to_matrix(series: Any, max_dims: int = 12) -> np.ndarray: + if series is None: + return np.empty((0, 0), dtype=float) + rows: list[np.ndarray] = [] + for value in series: + array = np.asarray(value, dtype=float).reshape(-1) + if array.size: + rows.append(array[:max_dims]) + if not rows: + return np.empty((0, 0), dtype=float) + width = max(row.size for row in rows) + matrix = np.full((len(rows), width), np.nan, dtype=float) + for index, row in enumerate(rows): + matrix[index, : row.size] = row + return matrix + + +def numeric_column(data_frame: Any, column: str, fallback_length: int) -> np.ndarray: + if column not in data_frame: + return np.arange(fallback_length, dtype=float) + try: + values = np.asarray(data_frame[column], dtype=float) + except Exception: + values = np.arange(fallback_length, dtype=float) + return values + + +def draw_lerobot_header( + draw: ImageDraw.ImageDraw, + *, + dataset_name: str, + row_index: int, + total_rows: int, + timestamp: float | None, +) -> None: + draw.text((40, 28), "LeRobot dataset preview", fill=(17, 24, 39)) + short_name = dataset_name if len(dataset_name) <= 78 else f"{dataset_name[:75]}..." + draw.text((40, 54), short_name, fill=(75, 85, 99)) + progress = 0 if total_rows <= 1 else row_index / (total_rows - 1) + draw.text((750, 28), f"frame {row_index + 1}/{total_rows}", fill=(17, 24, 39)) + if timestamp is not None: + draw.text((750, 54), f"t = {timestamp:.2f}s", fill=(75, 85, 99)) + draw.rectangle((40, 78, 920, 82), fill=(224, 231, 239)) + draw.rectangle((40, 78, int(40 + 880 * progress), 82), fill=(37, 99, 235)) + + +def draw_signal_panel( + draw: ImageDraw.ImageDraw, + box: tuple[int, int, int, int], + matrix: np.ndarray, + row_index: int, + title: str, +) -> None: + x0, y0, x1, y1 = box + draw.rounded_rectangle(box, radius=8, fill=(255, 255, 255), outline=(209, 213, 219)) + draw.text((x0 + 14, y0 + 10), title, fill=(17, 24, 39)) + if matrix.size == 0: + draw.text((x0 + 14, y0 + 48), "No numeric data", fill=(107, 114, 128)) + return + plot_box = (x0 + 14, y0 + 36, x1 - 14, y1 - 16) + draw_timeseries(draw, plot_box, matrix, row_index) + + +def draw_bar_panel( + draw: ImageDraw.ImageDraw, + box: tuple[int, int, int, int], + matrix: np.ndarray, + row_index: int, + title: str, +) -> None: + x0, y0, x1, y1 = box + draw.rounded_rectangle(box, radius=8, fill=(255, 255, 255), outline=(209, 213, 219)) + draw.text((x0 + 14, y0 + 10), title, fill=(17, 24, 39)) + if matrix.size == 0 or row_index >= len(matrix): + draw.text((x0 + 14, y0 + 48), "No numeric data", fill=(107, 114, 128)) + return + values = matrix[row_index] + finite = values[np.isfinite(values)] + if finite.size == 0: + return + max_abs = max(float(np.nanmax(np.abs(finite))), 1e-6) + base_y = y1 - 30 + left = x0 + 18 + available_width = x1 - x0 - 36 + bar_count = min(len(values), 12) + bar_gap = 8 + bar_width = max(8, (available_width - bar_gap * (bar_count - 1)) // bar_count) + for index in range(bar_count): + value = values[index] + if not np.isfinite(value): + continue + x = left + index * (bar_width + bar_gap) + height = int((abs(float(value)) / max_abs) * 58) + color = (22, 163, 74) if value >= 0 else (220, 38, 38) + y_top = base_y - height + draw.rectangle((x, y_top, x + bar_width, base_y), fill=color) + draw.text((x, base_y + 5), str(index), fill=(107, 114, 128)) + + +def draw_timeseries( + draw: ImageDraw.ImageDraw, + box: tuple[int, int, int, int], + matrix: np.ndarray, + row_index: int, +) -> None: + x0, y0, x1, y1 = box + draw.rectangle(box, outline=(229, 231, 235)) + sample_count = min(len(matrix), LEROBOT_PREVIEW_MAX_FRAMES) + if sample_count <= 1: + return + sampled = matrix[ + np.linspace(0, len(matrix) - 1, num=sample_count, dtype=int), + : min(matrix.shape[1], 8), + ] + finite = sampled[np.isfinite(sampled)] + if finite.size == 0: + return + minimum = float(np.nanmin(finite)) + maximum = float(np.nanmax(finite)) + if math.isclose(minimum, maximum): + minimum -= 1.0 + maximum += 1.0 + palette = [ + (37, 99, 235), + (5, 150, 105), + (217, 119, 6), + (220, 38, 38), + (124, 58, 237), + (8, 145, 178), + (79, 70, 229), + (202, 138, 4), + ] + + def point(sample_index: int, value: float) -> tuple[int, int]: + x = int(x0 + (x1 - x0) * sample_index / (sample_count - 1)) + y = int(y1 - (y1 - y0) * (value - minimum) / (maximum - minimum)) + return x, y + + for dim in range(sampled.shape[1]): + points = [ + point(index, float(value)) + for index, value in enumerate(sampled[:, dim]) + if np.isfinite(value) + ] + if len(points) >= 2: + draw.line(points, fill=palette[dim % len(palette)], width=2) + + cursor_x = int(x0 + (x1 - x0) * row_index / max(len(matrix) - 1, 1)) + draw.line((cursor_x, y0, cursor_x, y1), fill=(17, 24, 39), width=2) + + +def safe_filename_part(value: str) -> str: + safe = "".join( + char if char.isalnum() or char in {"-", "_"} else "_" for char in value.strip() + ) + return safe.strip("_")[:80] + + +def run_articraft_viser_preview(args: argparse.Namespace) -> None: + """Load an Articraft URDF and publish its initial scene topology to Viser. + + The generic asset-preview command starts Viser lazily. For an asset loaded + before that first capture, explicitly marking the topology dirty and + capturing once more ensures the initial scene is sent to the browser. + + Args: + args: Parsed preview-asset command-line arguments. + """ + from embodichain.lab.scripts import preview_asset + from embodichain.lab.sim.sim_manager import SimulationManager + from embodichain.utils.logger import log_info + + sim = SimulationManager(preview_asset.build_sim_cfg(args)) + try: + if args.env_map: + log_info(f"Setting environment map: {args.env_map} ...", color="green") + sim.set_indirect_lighting(args.env_map) + + assets = preview_asset.load_assets(sim, args) + log_info(f"Loaded {len(assets)} asset(s) successfully.", color="green") + if args.viser: + sim.start_visualization() + sim.notify_visualization_topology_changed() + sim.capture_visualization_safely(force=True) + preview_asset._run_preview_mode(sim, assets, args) + finally: + log_info("Destroying simulation ...", color="green") + sim.destroy() + + +def articraft_viser_preview_cli(argv: Sequence[str] | None = None) -> None: + """Run the Articraft-aware variant of the generic preview-asset CLI. + + Args: + argv: Arguments excluding the program name, or ``None`` for ``sys.argv``. + """ + from embodichain.lab.scripts import preview_asset + + parser = preview_asset._create_parser() + run_articraft_viser_preview(parser.parse_args(argv)) + + +if __name__ == "__main__": + articraft_viser_preview_cli() diff --git a/embodichain/gen_sim/gradio_ui/app_processes.py b/embodichain/gen_sim/gradio_ui/app_processes.py new file mode 100644 index 000000000..993ac076c --- /dev/null +++ b/embodichain/gen_sim/gradio_ui/app_processes.py @@ -0,0 +1,193 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Pipeline subprocess execution and progress detection.""" + +from __future__ import annotations + +import os +import queue +import signal +import subprocess +import sys +import time +from pathlib import Path + +from app_config import * # noqa: F403 - process settings are central configuration. +from app_state import PHASES + +_RUN_AGENT_SUPPORTS_ROBOT_PROFILE: bool | None = None + + +def run_agent_cli_supports_robot_profile() -> bool: + global _RUN_AGENT_SUPPORTS_ROBOT_PROFILE + if _RUN_AGENT_SUPPORTS_ROBOT_PROFILE is not None: + return _RUN_AGENT_SUPPORTS_ROBOT_PROFILE + try: + result = subprocess.run( + [ + sys.executable, + "-m", + COMMANDS["agent"]["module"], + *COMMANDS["agent"]["help_args"], + ], + cwd=EMBODICHAIN_ROOT, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + env=build_pipeline_env(), + timeout=20, + ) + help_text = (result.stdout or "").lower() + _RUN_AGENT_SUPPORTS_ROBOT_PROFILE = "--robot-profile" in help_text + except Exception: + _RUN_AGENT_SUPPORTS_ROBOT_PROFILE = False + return _RUN_AGENT_SUPPORTS_ROBOT_PROFILE + + +def build_run_agent_command( + paths: ScenePaths, *, parallel_env: bool = False, robot_profile: str | None = None +) -> list[str]: + from app_commands import build_run_agent_command as build_command + + return build_command( + paths, + parallel_env=parallel_env, + robot_profile=robot_profile, + supports_robot_profile=run_agent_cli_supports_robot_profile(), + ) + + +def start_pipeline(command: list[str]) -> subprocess.Popen[str]: + env = build_pipeline_env() + env["PYTHONUNBUFFERED"] = "1" + return subprocess.Popen( + command, + cwd=EMBODICHAIN_ROOT, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + start_new_session=True, + env=env, + ) + + +def build_pipeline_env() -> dict[str, str]: + env = os.environ.copy() + configure_direct_network_env(env) + configure_simready_llm_env(env) + return env + + +def terminate_process_group(process: subprocess.Popen[str]) -> None: + if process.poll() is not None: + return + try: + os.killpg(process.pid, signal.SIGTERM) + except ProcessLookupError: + return + except Exception: + process.terminate() + + deadline = time.monotonic() + PROCESS_STOP_TIMEOUT_S + while time.monotonic() < deadline: + if process.poll() is not None: + return + time.sleep(0.2) + + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + return + except Exception: + process.kill() + + +def detect_phase_from_files(current_key: str, paths: ScenePaths) -> str: + candidates = [ + ("scene_intake", paths.prompt_root / "scene_intake" / "result.json"), + ("relations", paths.prompt_root / "image_segments" / "result.json"), + ( + "relations", + paths.prompt_root / "image_spatial_relations" / "result.json", + ), + ("gym_export", paths.prompt_root / "gym_export" / "gym_config.json"), + ("config", paths.fast_gym_config), + ("preview", paths.gradio_scene_glb), + ] + best_key = current_key + best_progress = PHASES.get(best_key, PHASES["idle"]).progress + + if any(paths.prompt_root.glob("unified_scene_gen/**/*.glb")): + best_key, best_progress = _choose_later_phase( + best_key, + best_progress, + "asset_generation", + ) + for phase_key, marker in candidates: + if marker.exists(): + best_key, best_progress = _choose_later_phase( + best_key, + best_progress, + phase_key, + ) + return best_key + + +def _choose_later_phase( + current_key: str, + current_progress: int, + candidate_key: str, +) -> tuple[str, int]: + candidate_progress = PHASES[candidate_key].progress + if candidate_progress > current_progress: + return candidate_key, candidate_progress + return current_key, current_progress + + +def update_phase_from_log(line: str, current_key: str) -> str: + text = line.lower() + mapping = [ + ("scene_intake", "scene_intake"), + ("image_segments", "relations"), + ("image_spatial_relations", "relations"), + ("unified_scene_gen", "asset_generation"), + ("glb", "asset_generation"), + ("gym_export", "gym_export"), + ("generated gym config", "config"), + ("fast_gym_config", "config"), + ] + best_key = current_key + best_progress = PHASES.get(best_key, PHASES["idle"]).progress + for needle, phase_key in mapping: + if needle in text: + best_key, best_progress = _choose_later_phase( + best_key, + best_progress, + phase_key, + ) + return best_key + + +def read_process_output( + process: subprocess.Popen[str], + output_queue: queue.Queue[str], +) -> None: + if process.stdout is None: + return + for line in process.stdout: + output_queue.put(line.rstrip()) diff --git a/embodichain/gen_sim/gradio_ui/app_services.py b/embodichain/gen_sim/gradio_ui/app_services.py new file mode 100644 index 000000000..36b7fccdf --- /dev/null +++ b/embodichain/gen_sim/gradio_ui/app_services.py @@ -0,0 +1,28 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Compatibility facade for application services. + +The executable workflow lives in :mod:`app_workflows`; the Gradio view lives +in :mod:`app_ui`. Keep this module small so existing imports remain valid +while callers move to the focused modules. +""" + +from __future__ import annotations + +from app_ui import build_demo + +__all__ = ["build_demo"] diff --git a/embodichain/gen_sim/gradio_ui/app_state.py b/embodichain/gen_sim/gradio_ui/app_state.py new file mode 100644 index 000000000..260469af1 --- /dev/null +++ b/embodichain/gen_sim/gradio_ui/app_state.py @@ -0,0 +1,187 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Shared, thread-safe runtime state and timing helpers.""" + +from __future__ import annotations + +import subprocess +import threading +import time +import uuid +from collections import deque +from dataclasses import dataclass, field +from pathlib import Path + +from app_config import ( + DEFAULT_ROBOT_PROFILE, + LANGUAGE_EN, + PHASE_DEFINITIONS, + SCENE_MODE_INITIAL, + TIMING_PHASE_LABELS, + TIMING_PHASE_ORDER, +) + + +@dataclass(frozen=True) +class Phase: + progress: int + label: str + + +PHASES = {key: Phase(*value) for key, value in PHASE_DEFINITIONS.items()} + + +@dataclass +class RuntimeState: + is_busy: bool = False + run_token: str = field(default_factory=lambda: uuid.uuid4().hex) + auto_loop_active: bool = False + auto_loop_token: str | None = None + auto_round: int = 0 + auto_scene_mode: str = SCENE_MODE_INITIAL + auto_parallel_env: bool = False + auto_robot_profile: str = DEFAULT_ROBOT_PROFILE + language: str = LANGUAGE_EN + process: subprocess.Popen[str] | None = None + sim_process: subprocess.Popen[str] | None = None + scene_preview_process: subprocess.Popen[str] | None = None + sim_started: bool = False + sim_finished: bool = False + sim_returncode: int | None = None + phase_key: str = "idle" + status: str = "Idle." + task_text: str = "" + input_task_text: str = "" + input_scene_text: str = "" + image_path: Path | None = None + video_path: Path | None = None + last_sent_video_signature: tuple[str, int] | None = None + lerobot_video_path: Path | None = None + lerobot_dataset_path: Path | None = None + submitted_input_revision: int = 0 + object_model_path: Path | None = None + scene_model_path: Path | None = None + edited_scene_model_path: Path | None = None + last_error: str | None = None + log_lines: deque[str] = field(default_factory=deque) + timing_started_ns: int | None = None + current_timing_phase_key: str | None = None + current_timing_phase_started_ns: int | None = None + phase_durations_ns: dict[str, int] = field(default_factory=dict) + simulation_started_monotonic_ns: int | None = None + simulation_duration_ns: int | None = None + + +runtime = RuntimeState() +runtime_lock = threading.Lock() + + +def clear_run_timing_locked() -> None: + runtime.timing_started_ns = None + runtime.current_timing_phase_key = None + runtime.current_timing_phase_started_ns = None + runtime.phase_durations_ns.clear() + runtime.simulation_started_monotonic_ns = None + runtime.simulation_duration_ns = None + + +def start_run_timing_locked(phase_key: str) -> None: + now_ns = time.monotonic_ns() + runtime.timing_started_ns = now_ns + runtime.current_timing_phase_key = phase_key + runtime.current_timing_phase_started_ns = now_ns + runtime.phase_durations_ns.clear() + runtime.simulation_started_monotonic_ns = None + runtime.simulation_duration_ns = None + + +def record_phase_transition_locked(new_phase_key: str) -> None: + current_key = runtime.current_timing_phase_key + current_started_ns = runtime.current_timing_phase_started_ns + now_ns = time.monotonic_ns() + if current_key is None or current_started_ns is None: + runtime.timing_started_ns = runtime.timing_started_ns or now_ns + runtime.current_timing_phase_key = new_phase_key + runtime.current_timing_phase_started_ns = now_ns + return + if new_phase_key == current_key: + return + runtime.phase_durations_ns[current_key] = runtime.phase_durations_ns.get( + current_key, 0 + ) + max(0, now_ns - current_started_ns) + runtime.current_timing_phase_key = new_phase_key + runtime.current_timing_phase_started_ns = now_ns + + +def set_runtime_phase_locked(new_phase_key: str) -> None: + record_phase_transition_locked(new_phase_key) + runtime.phase_key = new_phase_key + + +def record_simulation_started_locked() -> None: + runtime.simulation_started_monotonic_ns = time.monotonic_ns() + runtime.simulation_duration_ns = None + + +def record_simulation_finished_locked() -> None: + started_ns = runtime.simulation_started_monotonic_ns + if started_ns is not None: + runtime.simulation_duration_ns = max(0, time.monotonic_ns() - started_ns) + runtime.simulation_started_monotonic_ns = None + + +def snapshot_timing_locked() -> tuple[dict[str, int], int | None]: + durations = dict(runtime.phase_durations_ns) + current_key = runtime.current_timing_phase_key + current_started_ns = runtime.current_timing_phase_started_ns + if ( + current_key + and current_started_ns is not None + and current_key not in {"complete", "failed", "idle"} + ): + durations[current_key] = durations.get(current_key, 0) + max( + 0, time.monotonic_ns() - current_started_ns + ) + simulation_duration_ns = runtime.simulation_duration_ns + if ( + simulation_duration_ns is None + and runtime.simulation_started_monotonic_ns is not None + ): + simulation_duration_ns = max( + 0, time.monotonic_ns() - runtime.simulation_started_monotonic_ns + ) + return durations, simulation_duration_ns + + +def format_duration_ns(duration_ns: int) -> str: + seconds = duration_ns / 1_000_000_000 + if seconds < 60: + return f"{seconds:.2f}s" + minutes = int(seconds // 60) + return f"{minutes}m {seconds - minutes * 60:05.2f}s" + + +def format_timing_lines( + phase_durations_ns: dict[str, int], simulation_duration_ns: int | None +) -> list[str]: + timing_values = dict(phase_durations_ns) + if simulation_duration_ns is not None: + timing_values["action_graph_execution"] = simulation_duration_ns + return [ + f"- {TIMING_PHASE_LABELS[key]}: {format_duration_ns(value) if (value := timing_values.get(key)) is not None else 'skipped'}" + for key in TIMING_PHASE_ORDER + ] diff --git a/embodichain/gen_sim/gradio_ui/app_ui.py b/embodichain/gen_sim/gradio_ui/app_ui.py new file mode 100644 index 000000000..8e1a10b4b --- /dev/null +++ b/embodichain/gen_sim/gradio_ui/app_ui.py @@ -0,0 +1,569 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Gradio layout and event bindings. + +The workflow layer supplies all callbacks; this module only owns presentation +and wires components to those callbacks. +""" + +from __future__ import annotations + +from app_workflows import * # noqa: F401,F403 - callbacks/constants form the UI contract. +from app_asset_engine import build_asset_engine_panel + + +def select_application_mode(selected_mode: str | None): + """Switch between the full product UI and the focused engine UI.""" + is_debug = selected_mode == APP_MODE_DEBUG + debug_css = ( + "" if is_debug else "" + ) + return ( + gr.update(variant="secondary" if is_debug else "primary"), + gr.update(variant="primary" if is_debug else "secondary"), + gr.update(visible=is_debug), + gr.update(value=debug_css), + APP_MODE_DEBUG if is_debug else APP_MODE_DEMO, + gr.update(visible=is_debug), + ) + + +def select_debug_engine(selected_engine: str): + """Expose an explicit active state without starting any pipeline.""" + button_updates = tuple( + gr.update(variant="primary" if engine == selected_engine else "secondary") + for engine, _ in DEBUG_ENGINES + ) + return ( + *button_updates, + gr.update(visible=selected_engine == DEBUG_ENGINE_ASSET), + gr.update(visible=selected_engine == DEBUG_ENGINE_SCENE), + gr.update(visible=selected_engine == DEBUG_ENGINE_ACTION), + ) + + +def action_engine_snapshot(): + """Adapt the shared runtime snapshot to the five Action-engine widgets.""" + video, task, progress, status, initial, edited, _objects = ui_snapshot() + return video, task, progress, status, initial or edited + + +def run_action_engine_panel(task_text: str, robot_profile: str | None): + run_action_engine_from_current(task_text, robot_profile) + return action_engine_snapshot() + + +def build_demo() -> gr.Blocks: + with gr.Blocks(title="EmbodiChain Gradio") as demo: + app_mode = gr.State(APP_MODE_DEMO) + run_mode = gr.State(TOP_MODE_INTERACT) + action_mode = gr.State(None) + language = gr.State(LANGUAGE_EN) + last_seen_input_revision = gr.State(0) + interact_prebuilt_scene_dir = gr.State(None) + mode_style = gr.HTML(value="", visible=True) + with gr.Row(): + demo_mode_button = gr.Button("Demo", variant="primary") + debug_mode_button = gr.Button("Debug", variant="secondary") + with gr.Row(visible=False) as debug_controls: + asset_engine_button = gr.Button("Asset_engine", variant="primary") + scene_engine_button = gr.Button("Scene_engine", variant="secondary") + action_engine_button = gr.Button("Action_engine", variant="secondary") + with gr.Column(visible=False) as debug_engine_area: + asset_engine = build_asset_engine_panel() + with gr.Column(visible=False) as scene_engine_panel: + gr.Markdown( + "## Scene engine\n" + "Upload one image to generate a Scene Engine export. " + "The resulting Viser page is shown below." + ) + with gr.Row(): + with gr.Column(scale=1): + debug_scene_image = gr.Image( + label=UI_TEXT[LANGUAGE_EN]["input_image"], + sources=["upload", "webcam"], + type="filepath", + format="png", + height=300, + ) + debug_scene_run = gr.Button("Generate scene", variant="primary") + with gr.Column(scale=2): + debug_scene_progress = gr.Slider( + 0, + 100, + value=0, + step=1, + label=UI_TEXT[LANGUAGE_EN]["progress"], + interactive=False, + ) + debug_scene_status = gr.Markdown(format_status("Idle.")) + debug_scene_output = gr.Textbox( + label="Scene output directory (hash-named)", + interactive=False, + ) + debug_scene_preview = gr.HTML( + "
" + "The Viser preview will appear here after generation." + "
" + ) + with gr.Column(visible=False) as action_engine_panel: + gr.Markdown( + "## Action engine\nUses the Gym scene produced by Scene engine (not merely a rendered GLB), then generates the action config and launches DexSim. This retains collisions, poses and physics metadata required by simulation." + ) + with gr.Row(): + with gr.Column(scale=1): + debug_action_task = gr.Textbox( + label="Task description", + placeholder="e.g. Put the bottle on the table", + ) + debug_action_robot = gr.Radio( + choices=ROBOT_PROFILES, + value=DEFAULT_ROBOT_PROFILE, + label=UI_TEXT[LANGUAGE_EN]["robot"], + ) + debug_action_load = gr.Button("Load current scene") + debug_action_run = gr.Button("Run DexSim", variant="primary") + with gr.Column(scale=2): + debug_action_scene = gr.Model3D( + label="Input Gym scene preview", + height=420, + clear_color=(0.94, 0.94, 0.94, 1.0), + ) + debug_action_video = gr.Video( + label=UI_TEXT[LANGUAGE_EN]["single_video_preview"], + height=320, + autoplay=True, + loop=True, + ) + debug_action_current_task = gr.Textbox( + label=UI_TEXT[LANGUAGE_EN]["current_task"], + interactive=False, + ) + debug_action_progress = gr.Slider( + 0, + 100, + value=0, + step=1, + label=UI_TEXT[LANGUAGE_EN]["progress"], + interactive=False, + ) + debug_action_status = gr.Markdown( + format_status("Load or generate a scene first.") + ) + debug_action_refresh_timer = gr.Timer(2.0) + with gr.Row(equal_height=True, elem_classes="demo-only"): + if DEXFORCE_LOGO.is_file(): + gr.Image( + value=str(DEXFORCE_LOGO), + show_label=False, + container=False, + height=58, + width=183, + ) + heading = gr.Markdown(UI_TEXT[LANGUAGE_EN]["heading"]) + with gr.Row(elem_classes="demo-only"): + auto_button = gr.Button("Auto", variant="secondary") + interact_button = gr.Button("Interact", variant="primary") + parallel_env_button = gr.Button("Parallel Simulation", variant="secondary") + language_button = gr.Button("中文", variant="secondary") + with gr.Row(elem_classes="demo-only"): + with gr.Column(scale=4): + instruction = gr.HTML( + "
" + "Upload one image, enter one task, then EmbodiChain " + " will generate what you want." + "
" + ) + with gr.Column(scale=1): + robot_profile = gr.Radio( + choices=ROBOT_PROFILES, + value=DEFAULT_ROBOT_PROFILE, + label=UI_TEXT[LANGUAGE_EN]["robot"], + ) + + with gr.Row(elem_classes="demo-only"): + with gr.Column(scale=1): + image_input = gr.Image( + label=UI_TEXT[LANGUAGE_EN]["input_image"], + sources=["upload", "webcam"], + type="filepath", + format="png", + height=320, + ) + with gr.Row(): + with gr.Column(): + task_input = gr.Textbox( + label=UI_TEXT[LANGUAGE_EN]["task_description"], + placeholder=UI_TEXT[LANGUAGE_EN]["task_placeholder"], + lines=1, + ) + random_task_input_button = gr.Button("Random Task") + with gr.Column(): + env_input = gr.Textbox( + label=UI_TEXT[LANGUAGE_EN]["scene_description"], + placeholder=UI_TEXT[LANGUAGE_EN]["scene_placeholder"], + lines=1, + ) + random_scene_input_button = gr.Button("Random Scene") + scene_mode = gr.Radio( + choices=scene_mode_choices(LANGUAGE_EN), + value=SCENE_MODE_INITIAL, + label=UI_TEXT[LANGUAGE_EN]["scene_mode"], + ) + with gr.Row(): + generate_button = gr.Button("Generate", variant="primary") + rerun_simulation_button = gr.Button("Run Task", variant="secondary") + reset_button = gr.Button("Reset", variant="stop") + with gr.Column(scale=2): + current_image = gr.Video( + label=UI_TEXT[LANGUAGE_EN]["single_video_preview"], + height=420, + elem_id="embodichain-video-preview", + autoplay=True, + loop=True, + ) + current_task = gr.Textbox( + label=UI_TEXT[LANGUAGE_EN]["current_task"], + interactive=False, + lines=2, + ) + + progress = gr.Slider( + minimum=0, + maximum=100, + value=0, + step=1, + label=UI_TEXT[LANGUAGE_EN]["progress"], + interactive=False, + elem_classes="demo-only", + ) + status = gr.Markdown(format_status("Idle."), elem_classes="demo-only") + with gr.Row(elem_classes="demo-only"): + model = gr.Model3D( + label=UI_TEXT[LANGUAGE_EN]["initial_preview"], + height=520, + clear_color=(0.94, 0.94, 0.94, 1.0), + ) + edited_model = gr.Model3D( + label=UI_TEXT[LANGUAGE_EN]["edited_preview"], + height=520, + clear_color=(0.94, 0.94, 0.94, 1.0), + ) + object_model = gr.Model3D( + label=UI_TEXT[LANGUAGE_EN]["object_preview"], + height=360, + clear_color=(0.94, 0.94, 0.94, 1.0), + elem_classes="demo-only", + ) + + refresh_timer = gr.Timer(2.0) + top_mode_outputs = [ + auto_button, + interact_button, + parallel_env_button, + generate_button, + rerun_simulation_button, + random_task_input_button, + random_scene_input_button, + reset_button, + current_image, + run_mode, + action_mode, + ] + demo_mode_button.click( + select_application_mode, + inputs=[gr.State(APP_MODE_DEMO)], + outputs=[ + demo_mode_button, + debug_mode_button, + debug_controls, + mode_style, + app_mode, + debug_engine_area, + ], + queue=False, + ) + debug_mode_button.click( + select_application_mode, + inputs=[gr.State(APP_MODE_DEBUG)], + outputs=[ + demo_mode_button, + debug_mode_button, + debug_controls, + mode_style, + app_mode, + debug_engine_area, + ], + queue=False, + ) + for engine, button in zip( + (engine for engine, _ in DEBUG_ENGINES), + ( + asset_engine_button, + scene_engine_button, + action_engine_button, + ), + ): + button.click( + select_debug_engine, + inputs=[gr.State(engine)], + outputs=[ + asset_engine_button, + scene_engine_button, + action_engine_button, + asset_engine["panel"], + scene_engine_panel, + action_engine_panel, + ], + queue=False, + ) + debug_scene_run.click( + run_scene_engine, + inputs=[debug_scene_image], + outputs=[ + debug_scene_progress, + debug_scene_status, + debug_scene_output, + debug_scene_preview, + ], + ) + debug_action_load.click( + action_engine_snapshot, + outputs=[ + debug_action_video, + debug_action_current_task, + debug_action_progress, + debug_action_status, + debug_action_scene, + ], + queue=False, + ) + debug_action_run.click( + run_action_engine_panel, + inputs=[debug_action_task, debug_action_robot], + outputs=[ + debug_action_video, + debug_action_current_task, + debug_action_progress, + debug_action_status, + debug_action_scene, + ], + ) + debug_action_refresh_timer.tick( + action_engine_snapshot, + outputs=[ + debug_action_video, + debug_action_current_task, + debug_action_progress, + debug_action_status, + debug_action_scene, + ], + queue=False, + ) + auto_button.click( + select_top_mode, + inputs=[ + gr.State(TOP_MODE_AUTO), + gr.State(None), + run_mode, + action_mode, + language, + ], + outputs=top_mode_outputs, + queue=False, + ) + interact_button.click( + select_top_mode, + inputs=[ + gr.State(TOP_MODE_INTERACT), + gr.State(None), + run_mode, + action_mode, + language, + ], + outputs=top_mode_outputs, + queue=False, + ) + parallel_env_button.click( + select_top_mode, + inputs=[ + gr.State(None), + gr.State(TOP_MODE_PARALLEL_ENV), + run_mode, + action_mode, + language, + ], + outputs=top_mode_outputs, + queue=False, + ) + language_button.click( + toggle_language, + inputs=[language, run_mode, action_mode], + outputs=[ + auto_button, + interact_button, + parallel_env_button, + generate_button, + rerun_simulation_button, + random_task_input_button, + random_scene_input_button, + reset_button, + language_button, + heading, + instruction, + robot_profile, + image_input, + task_input, + env_input, + scene_mode, + current_image, + current_task, + progress, + model, + edited_model, + object_model, + language, + ], + queue=False, + ) + generate_button.click( + run_generate_for_top_mode, + inputs=[ + run_mode, + action_mode, + scene_mode, + robot_profile, + image_input, + task_input, + env_input, + interact_prebuilt_scene_dir, + language, + ], + outputs=[ + image_input, + task_input, + env_input, + current_image, + current_task, + progress, + status, + model, + edited_model, + object_model, + ], + ) + random_task_input_button.click( + randomize_interact_task_input, + inputs=[run_mode, language], + outputs=[ + image_input, + task_input, + env_input, + scene_mode, + interact_prebuilt_scene_dir, + model, + edited_model, + object_model, + ], + queue=False, + ) + random_scene_input_button.click( + randomize_interact_scene_input, + inputs=[run_mode, language], + outputs=[env_input], + queue=False, + ) + rerun_simulation_button.click( + rerun_current_simulation, + inputs=[ + run_mode, + action_mode, + robot_profile, + ], + outputs=[ + image_input, + task_input, + env_input, + current_image, + current_task, + progress, + status, + model, + edited_model, + object_model, + ], + queue=False, + ) + image_input.upload( + clear_interact_prebuilt_scene, + outputs=[interact_prebuilt_scene_dir], + queue=False, + ) + scene_mode.change( + scene_mode_input_updates, + inputs=[scene_mode], + outputs=[task_input, env_input], + queue=False, + ) + reset_button.click( + clear_interact_prebuilt_scene, + outputs=[interact_prebuilt_scene_dir], + queue=False, + ) + reset_button.click( + run_reset_or_stop, + inputs=[run_mode], + outputs=[ + image_input, + task_input, + env_input, + current_image, + current_task, + progress, + status, + model, + edited_model, + object_model, + ], + queue=False, + ) + refresh_timer.tick( + synced_ui_snapshot, + inputs=[run_mode, action_mode, last_seen_input_revision], + outputs=[ + image_input, + task_input, + env_input, + current_image, + current_task, + progress, + status, + model, + edited_model, + object_model, + rerun_simulation_button, + last_seen_input_revision, + scene_mode, + robot_profile, + parallel_env_button, + action_mode, + ], + queue=False, + ) + return demo diff --git a/embodichain/gen_sim/gradio_ui/app_workflows.py b/embodichain/gen_sim/gradio_ui/app_workflows.py new file mode 100644 index 000000000..412443ae4 --- /dev/null +++ b/embodichain/gen_sim/gradio_ui/app_workflows.py @@ -0,0 +1,3336 @@ +# ---------------------------------------------------------------------------- +# 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 hashlib +import html +import io +import json +import importlib.util +import math +import os +import queue +import shutil +import signal +import socket +import subprocess +import sys +import threading +import time +import uuid +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path +from typing import Any, Iterable + +from random_input import IMAGE_DIR as AUTO_BASE_IMAGE_DIR +from random_input import ( + auto_image_directories, + available_auto_task_indices, + generate_auto_scene_description, + generate_auto_text_input, + get_prebuilt_scene_dir, + parse_task_id, +) +from app_config import * # noqa: F403 - services intentionally consume central config. +from app_processes import ( + build_pipeline_env, + build_run_agent_command, + detect_phase_from_files, + read_process_output, + run_agent_cli_supports_robot_profile, + start_pipeline, + terminate_process_group, + update_phase_from_log, +) +from app_media import * # noqa: F403 - workflow consumes media service helpers. + +configure_direct_network_env() + +import gradio as gr +import numpy as np +import trimesh +from PIL import Image, ImageDraw, ImageOps + +_RUN_AGENT_SUPPORTS_ROBOT_PROFILE: bool | None = None +VIDEO_SYNC_JS = r""" +() => { + const audienceRootId = "embodichain-audience-video"; + const lerobotRootId = "embodichain-lerobot-video"; + let syncing = false; + + function findVideo(rootId) { + const root = document.getElementById(rootId); + return root ? root.querySelector("video") : null; + } + + function sourceLoaded(video) { + return Boolean(video && (video.currentSrc || video.src)); + } + + function copyTime(source, target) { + if (!sourceLoaded(source) || !sourceLoaded(target)) { + return; + } + const sourceTime = source.currentTime || 0; + if (!Number.isFinite(sourceTime)) { + return; + } + const duration = Number.isFinite(target.duration) ? target.duration : sourceTime; + const targetTime = Math.min(sourceTime, duration); + if (Math.abs((target.currentTime || 0) - targetTime) > 0.35) { + try { + target.currentTime = targetTime; + } catch (_) { + // Some browsers reject seeking before metadata is fully available. + } + } + } + + function syncPlayback(sourceRootId, targetRootId, shouldPlay) { + if (syncing) { + return; + } + const source = findVideo(sourceRootId); + const target = findVideo(targetRootId); + if (!sourceLoaded(source) || !sourceLoaded(target)) { + return; + } + + syncing = true; + copyTime(source, target); + + const release = () => { + window.setTimeout(() => { + syncing = false; + }, 0); + }; + + if (shouldPlay) { + const result = target.play(); + if (result && typeof result.finally === "function") { + result.catch(() => {}).finally(release); + } else { + release(); + } + } else { + target.pause(); + release(); + } + } + + function bindOne(rootId, peerRootId) { + const video = findVideo(rootId); + if (!sourceLoaded(video) || video.dataset.embodichainSyncBound === "true") { + return; + } + video.dataset.embodichainSyncBound = "true"; + video.addEventListener("play", () => syncPlayback(rootId, peerRootId, true)); + video.addEventListener("pause", () => syncPlayback(rootId, peerRootId, false)); + } + + function bindVideos() { + bindOne(audienceRootId, lerobotRootId); + bindOne(lerobotRootId, audienceRootId); + } + + bindVideos(); + window.setInterval(bindVideos, 1000); + const observer = new MutationObserver(bindVideos); + observer.observe(document.body, { childList: true, subtree: true }); +} +""" + + +# Runtime ownership lives in app_state; this module only orchestrates it. +from app_state import ( + PHASES, + Phase, + RuntimeState, + clear_run_timing_locked, + format_duration_ns, + format_timing_lines, + record_phase_transition_locked, + record_simulation_finished_locked, + record_simulation_started_locked, + runtime, + runtime_lock, + set_runtime_phase_locked, + snapshot_timing_locked, + start_run_timing_locked, +) + + +@dataclass(frozen=True) +class ScenePaths: + scene_id: str + image_path: Path + prompt_root: Path + config_dir: Path + + @property + def fast_gym_config(self) -> Path: + return self.config_dir / "fast_gym_config.json" + + @property + def agent_config(self) -> Path: + return self.config_dir / "agent_config.json" + + @property + def gradio_scene_dir(self) -> Path: + return self.config_dir / "gradio_scene" + + @property + def gradio_scene_glb(self) -> Path: + return self.gradio_scene_dir / "scene_current.glb" + + @property + def gradio_object_preview_glb(self) -> Path: + return self.gradio_scene_dir / "object_preview.glb" + + @property + def scene_manifest(self) -> Path: + return self.gradio_scene_dir / "scene_manifest.json" + + @property + def object_preview_manifest(self) -> Path: + return self.gradio_scene_dir / "object_preview_manifest.json" + + +CURRENT_PATHS = ScenePaths( + scene_id=SCENE_ID, + image_path=IMAGE_PATH, + prompt_root=PROMPT2SCENE_ROOT, + config_dir=CONFIG_DIR, +) + + +def make_stage_paths(run_token: str) -> ScenePaths: + scene_id = f"{PENDING_PREFIX}{run_token[:12]}" + return ScenePaths( + scene_id=scene_id, + image_path=IMAGE_DIR / f"{scene_id}.png", + prompt_root=GYM_PROJECT_ROOT / scene_id, + config_dir=ACTION_AGENT_ROOT / "configs" / scene_id, + ) + + +def make_replaced_paths(run_token: str) -> ScenePaths: + scene_id = f"{REPLACED_PREFIX}{run_token[:12]}" + return ScenePaths( + scene_id=scene_id, + image_path=IMAGE_DIR / f"{scene_id}.png", + prompt_root=GYM_PROJECT_ROOT / scene_id, + config_dir=ACTION_AGENT_ROOT / "configs" / scene_id, + ) + + +def save_input( + image_value: str | np.ndarray | Image.Image, + task_text: str, + image_path: Path, +) -> Path: + if image_value is None: + raise ValueError("Please upload an image first.") + if not task_text.strip(): + raise ValueError("Please enter a task description.") + + image_path.parent.mkdir(parents=True, exist_ok=True) + if isinstance(image_value, str): + image = Image.open(image_value) + elif isinstance(image_value, np.ndarray): + image = Image.fromarray(image_value) + elif isinstance(image_value, Image.Image): + image = image_value + else: + raise TypeError(f"Unsupported image input type: {type(image_value)!r}") + + image = ImageOps.exif_transpose(image).convert("RGB") + image.save(image_path, format="PNG") + return image_path + + +def reset_current_scene() -> list[str]: + process: subprocess.Popen[str] | None = None + sim_process: subprocess.Popen[str] | None = None + with runtime_lock: + runtime.run_token = uuid.uuid4().hex + runtime.auto_loop_active = False + runtime.auto_loop_token = None + runtime.auto_round = 0 + process = runtime.process + sim_process = runtime.sim_process + runtime.process = None + runtime.sim_process = None + runtime.sim_started = False + runtime.sim_finished = False + runtime.sim_returncode = None + runtime.is_busy = False + runtime.phase_key = "idle" + runtime.status = "Idle." + runtime.task_text = "" + runtime.input_task_text = "" + runtime.input_scene_text = "" + runtime.image_path = None + runtime.video_path = None + runtime.lerobot_video_path = None + runtime.lerobot_dataset_path = None + runtime.object_model_path = None + runtime.scene_model_path = None + runtime.edited_scene_model_path = None + runtime.last_error = None + runtime.log_lines.clear() + clear_run_timing_locked() + + if process is not None: + terminate_process_group(process) + if sim_process is not None: + terminate_process_group(sim_process) + + return cleanup_current_and_staging() + + +def cleanup_current_and_staging() -> list[str]: + errors: list[str] = [] + paths: list[Path] = [ + PROMPT2SCENE_ROOT, + CONFIG_DIR, + IMAGE_PATH, + *pending_artifact_paths(), + ] + for path in paths: + errors.extend(remove_path(path)) + errors.extend(cleanup_outputs_preserving_videos()) + return errors + + +def cleanup_auto_generated_artifacts(extra_image_path: Path | None = None) -> list[str]: + errors: list[str] = [] + paths: list[Path] = [ + PROMPT2SCENE_ROOT, + CONFIG_DIR, + IMAGE_PATH, + *pending_artifact_paths(), + ] + if extra_image_path is not None: + paths.append(extra_image_path) + + for path in paths: + if is_protected_auto_base_image(path): + continue + errors.extend(remove_path(path)) + + with runtime_lock: + runtime.image_path = None + runtime.input_task_text = "" + runtime.input_scene_text = "" + runtime.lerobot_video_path = None + runtime.lerobot_dataset_path = None + runtime.object_model_path = None + runtime.scene_model_path = None + runtime.edited_scene_model_path = None + errors.extend(cleanup_outputs_preserving_videos()) + return errors + + +def is_protected_auto_base_image(path: Path) -> bool: + try: + path.resolve().relative_to(AUTO_BASE_IMAGE_DIR.resolve()) + except ValueError: + return False + except FileNotFoundError: + return False + return True + + +def pending_artifact_paths() -> list[Path]: + paths: list[Path] = [] + for root in (GYM_PROJECT_ROOT, ACTION_AGENT_ROOT / "configs"): + if root.is_dir(): + paths.extend(root.glob(f"{PENDING_PREFIX}*")) + paths.extend(root.glob(f"{REPLACED_PREFIX}*")) + if IMAGE_DIR.is_dir(): + paths.extend(IMAGE_DIR.glob(f"{PENDING_PREFIX}*.png")) + paths.extend(IMAGE_DIR.glob(f"{REPLACED_PREFIX}*.png")) + return paths + + +def remove_path(path: Path) -> list[str]: + try: + if path.is_dir(): + shutil.rmtree(path) + else: + path.unlink(missing_ok=True) + except Exception as exc: + return [f"Failed to remove {path}: {exc}"] + return [] + + +def cleanup_outputs_preserving_videos() -> list[str]: + if not OUTPUTS_DIR.exists(): + return [] + if OUTPUTS_DIR.is_file(): + if OUTPUTS_DIR.suffix.lower() in VIDEO_SUFFIXES: + return [] + return remove_path(OUTPUTS_DIR) + + errors: list[str] = [] + for path in sorted( + OUTPUTS_DIR.rglob("*"), + key=lambda item: len(item.parts), + reverse=True, + ): + if path.is_file() and path.suffix.lower() not in VIDEO_SUFFIXES: + errors.extend(remove_path(path)) + + for path in sorted( + OUTPUTS_DIR.rglob("*"), + key=lambda item: len(item.parts), + reverse=True, + ): + if not path.is_dir(): + continue + try: + path.rmdir() + except OSError: + pass + except Exception as exc: + errors.append(f"Failed to remove empty output directory {path}: {exc}") + return errors + + +from app_commands import ( + build_config_command_for_paths, + build_initial_pipeline_command, + build_scene_edit_pipeline_command, + robot_profile_cli_value, +) + + +def build_edit_pipeline_command( + task_text: str, + env_text: str, + robot_profile: str | None = None, + load_template_material: bool = False, +) -> list[str]: + return build_scene_edit_pipeline_command( + task_text, env_text, CURRENT_PATHS, robot_profile, load_template_material + ) + + +def build_task_only_config_command( + task_text: str, + robot_profile: str | None = None, + load_template_material: bool = False, +) -> list[str]: + return build_config_command_for_paths( + task_text, CURRENT_PATHS, robot_profile, load_template_material + ) + + +def format_current_task(task_text: str, env_text: str = "") -> str: + return "\n".join( + part for part in ((task_text or "").strip(), (env_text or "").strip()) if part + ) + + +def build_gradio_scene_from_fast_config( + config_path: Path, + scene_dir: Path | None = None, +) -> Path: + config_dir = config_path.parent + if scene_dir is None: + scene_dir = config_dir / "gradio_scene" + scene_glb = scene_dir / "scene_current.glb" + scene_manifest = scene_dir / "scene_manifest.json" + with config_path.open("r", encoding="utf-8") as file: + config = json.load(file) + config_stat = config_path.stat() + + scene = trimesh.Scene() + manifest: dict[str, Any] = { + "source_config": os.path.relpath(config_path, scene_dir), + "source_config_size": config_stat.st_size, + "source_config_mtime_ns": config_stat.st_mtime_ns, + "transform_policy": GRADIO_SCENE_TRANSFORM_POLICY, + "objects": [], + } + + object_count = 0 + for role, obj in iter_scene_objects(config): + shape = obj.get("shape") if isinstance(obj, dict) else None + if not isinstance(shape, dict) or shape.get("shape_type") != "Mesh": + continue + raw_fpath = shape.get("fpath") + if not raw_fpath: + continue + mesh_path = resolve_mesh_path(config_dir, str(raw_fpath)) + if not mesh_path.is_file(): + raise FileNotFoundError( + f"Mesh file not found for {obj.get('uid')}: {mesh_path}" + ) + + transform = object_transform(obj) + frame_transform = gltf_to_sim_frame_transform(mesh_path) + if frame_transform is not None: + transform = transform @ frame_transform + add_mesh_to_scene(scene, mesh_path, transform, str(obj.get("uid", "object"))) + manifest["objects"].append( + { + "uid": obj.get("uid"), + "role": role, + "source_mesh": os.path.relpath(mesh_path, scene_dir), + "source_mesh_size": mesh_path.stat().st_size, + "source_mesh_mtime_ns": mesh_path.stat().st_mtime_ns, + "gltf_to_sim_frame": frame_transform is not None, + } + ) + object_count += 1 + + if object_count == 0: + raise ValueError(f"No mesh objects found in {config_path}") + + scene_dir.mkdir(parents=True, exist_ok=True) + scene.export(scene_glb) + scene_manifest.write_text( + json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", + encoding="utf-8", + ) + return scene_glb + + +def gradio_scene_is_current( + scene_glb: Path, + manifest_path: Path, + config_path: Path, +) -> bool: + if ( + not scene_glb.is_file() + or not manifest_path.is_file() + or not config_path.is_file() + ): + return False + try: + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + config = json.loads(config_path.read_text(encoding="utf-8")) + config_stat = config_path.stat() + except Exception: + return False + + if manifest.get("source_config") != os.path.relpath( + config_path, manifest_path.parent + ): + return False + if manifest.get("source_config_size") != config_stat.st_size: + return False + if manifest.get("source_config_mtime_ns") != config_stat.st_mtime_ns: + return False + if manifest.get("transform_policy") != GRADIO_SCENE_TRANSFORM_POLICY: + return False + + expected_objects = [] + try: + for role, obj in iter_scene_objects(config): + shape = obj.get("shape") if isinstance(obj, dict) else None + if not isinstance(shape, dict) or shape.get("shape_type") != "Mesh": + continue + raw_fpath = shape.get("fpath") + if not raw_fpath: + continue + mesh_path = resolve_mesh_path(config_path.parent, str(raw_fpath)) + mesh_stat = mesh_path.stat() + frame_transform = gltf_to_sim_frame_transform(mesh_path) + expected_objects.append( + { + "uid": obj.get("uid"), + "role": role, + "source_mesh": os.path.relpath(mesh_path, manifest_path.parent), + "source_mesh_size": mesh_stat.st_size, + "source_mesh_mtime_ns": mesh_stat.st_mtime_ns, + "gltf_to_sim_frame": frame_transform is not None, + } + ) + except OSError: + return False + return manifest.get("objects") == expected_objects + + +def collect_generated_object_glbs(paths: ScenePaths) -> list[Path]: + if not paths.prompt_root.is_dir(): + return [] + + glb_paths: list[Path] = [] + seen: set[Path] = set() + for glb_dir in sorted(paths.prompt_root.rglob("glb_gen")): + if not glb_dir.is_dir(): + continue + candidates = [ + path for path in glb_dir.rglob("*_simready.glb") if is_previewable_glb(path) + ] + if not candidates: + candidates = [ + path for path in glb_dir.rglob("*.glb") if is_previewable_glb(path) + ] + for path in sorted(candidates): + resolved = path.resolve() + if resolved in seen: + continue + seen.add(resolved) + glb_paths.append(path) + return glb_paths + + +def is_previewable_glb(path: Path) -> bool: + if not path.is_file() or path.name.startswith("."): + return False + return not any(part.startswith(".") for part in path.relative_to(path.anchor).parts) + + +def build_object_preview_scene( + glb_paths: list[Path], + scene_dir: Path, +) -> Path: + if not glb_paths: + raise ValueError("No generated object GLBs found") + + scene_dir.mkdir(parents=True, exist_ok=True) + preview_glb = scene_dir / "object_preview.glb" + preview_manifest = scene_dir / "object_preview_manifest.json" + scene = trimesh.Scene() + manifest: dict[str, Any] = {"objects": []} + + cursor = 0.0 + spacing = 0.35 + added_count = 0 + for object_index, mesh_path in enumerate(glb_paths): + meshes = load_mesh_geometries(mesh_path) + if not meshes: + continue + + bounds = combined_bounds(meshes) + extents = bounds[1] - bounds[0] + max_extent = float(max(extents.max(), 1e-6)) + scale = 1.0 / max_extent + scaled_width = max(float(extents[0]) * scale, 0.2) + placement_x = cursor + scaled_width / 2.0 + cursor += scaled_width + spacing + + transform = ( + trimesh.transformations.translation_matrix( + [ + placement_x, + 0.0, + 0.0, + ] + ) + @ trimesh.transformations.scale_matrix(scale) + @ trimesh.transformations.translation_matrix( + [ + -float((bounds[0][0] + bounds[1][0]) / 2.0), + -float((bounds[0][1] + bounds[1][1]) / 2.0), + -float(bounds[0][2]), + ] + ) + ) + + for mesh_index, mesh in enumerate(meshes): + mesh.apply_transform(transform) + name = f"object_{object_index}_{mesh_index}" + scene.add_geometry(mesh, node_name=name, geom_name=name) + added_count += 1 + + manifest["objects"].append( + { + "source_mesh": os.path.relpath(mesh_path, scene_dir), + "size": mesh_path.stat().st_size, + "mtime_ns": mesh_path.stat().st_mtime_ns, + } + ) + + if added_count == 0: + raise ValueError("No renderable meshes found in generated object GLBs") + + if cursor > spacing: + scene.apply_transform( + trimesh.transformations.translation_matrix( + [-(cursor - spacing) / 2.0, 0.0, 0.0] + ) + ) + scene.export(preview_glb) + preview_manifest.write_text( + json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", + encoding="utf-8", + ) + return preview_glb + + +def object_preview_is_current( + manifest_path: Path, + glb_paths: list[Path], +) -> bool: + if not manifest_path.is_file(): + return False + try: + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + except Exception: + return False + expected = [] + for path in glb_paths: + try: + stat = path.stat() + except OSError: + return False + expected.append( + { + "source_mesh": os.path.relpath(path, manifest_path.parent), + "size": stat.st_size, + "mtime_ns": stat.st_mtime_ns, + } + ) + return manifest.get("objects") == expected + + +def load_mesh_geometries(mesh_path: Path) -> list[trimesh.Trimesh]: + loaded = trimesh.load(mesh_path, force="scene", process=False) + gltf_to_sim_transform = gltf_to_sim_frame_transform(mesh_path) + if isinstance(loaded, trimesh.Trimesh): + mesh = loaded.copy() + if gltf_to_sim_transform is not None: + mesh.apply_transform(gltf_to_sim_transform) + return [mesh] + if isinstance(loaded, trimesh.Scene): + meshes: list[trimesh.Trimesh] = [] + for geometry in loaded.dump(concatenate=False): + if isinstance(geometry, trimesh.Trimesh): + mesh = geometry.copy() + if gltf_to_sim_transform is not None: + mesh.apply_transform(gltf_to_sim_transform) + meshes.append(mesh) + return meshes + raise TypeError(f"Unsupported mesh type for {mesh_path}: {type(loaded)!r}") + + +def gltf_to_sim_frame_transform(mesh_path: Path) -> np.ndarray | None: + if mesh_path.suffix.lower() not in {".glb", ".gltf"}: + return None + # Match DexSim's native GLTF Y-up to simulation Z-up conversion. + transform = np.eye(4) + transform[:3, :3] = np.array( + [ + [1.0, 0.0, 0.0], + [0.0, 0.0, -1.0], + [0.0, 1.0, 0.0], + ], + dtype=float, + ) + return transform + + +def combined_bounds(meshes: list[trimesh.Trimesh]) -> np.ndarray: + valid_bounds = [ + mesh.bounds + for mesh in meshes + if mesh.vertices is not None and len(mesh.vertices) > 0 + ] + if not valid_bounds: + raise ValueError("Mesh has no vertices") + bounds = np.asarray(valid_bounds, dtype=float) + return np.stack([bounds[:, 0, :].min(axis=0), bounds[:, 1, :].max(axis=0)]) + + +def iter_scene_objects(config: dict[str, Any]) -> Iterable[tuple[str, dict[str, Any]]]: + for role in ("background", "rigid_object"): + value = config.get(role, []) + if isinstance(value, dict): + value = [value] + if not isinstance(value, list): + continue + for obj in value: + if isinstance(obj, dict): + yield role, obj + + +def resolve_mesh_path(config_dir: Path, raw_fpath: str) -> Path: + mesh_path = Path(raw_fpath).expanduser() + if not mesh_path.is_absolute(): + mesh_path = config_dir / mesh_path + return mesh_path.resolve() + + +def object_transform(obj: dict[str, Any]) -> np.ndarray: + scale = vector3(obj.get("body_scale"), [1.0, 1.0, 1.0]) + + scale_matrix = np.eye(4) + scale_matrix[0, 0] = scale[0] + scale_matrix[1, 1] = scale[1] + scale_matrix[2, 2] = scale[2] + + init_local_pose = matrix4(obj.get("init_local_pose")) + if init_local_pose is not None: + return init_local_pose @ scale_matrix + + position = vector3(obj.get("init_pos"), [0.0, 0.0, 0.0]) + rotation_degrees = vector3(obj.get("init_rot"), [0.0, 0.0, 0.0]) + root_matrix = euler_xyz_degrees_matrix(rotation_degrees, position) + return root_matrix @ scale_matrix + + +def euler_xyz_degrees_matrix( + rotation_degrees: list[float], + position: list[float], +) -> np.ndarray: + rx, ry, rz = (math.radians(value) for value in rotation_degrees) + cx, sx = math.cos(rx), math.sin(rx) + cy, sy = math.cos(ry), math.sin(ry) + cz, sz = math.cos(rz), math.sin(rz) + + rot_x = np.array( + [ + [1.0, 0.0, 0.0, 0.0], + [0.0, cx, -sx, 0.0], + [0.0, sx, cx, 0.0], + [0.0, 0.0, 0.0, 1.0], + ], + dtype=float, + ) + rot_y = np.array( + [ + [cy, 0.0, sy, 0.0], + [0.0, 1.0, 0.0, 0.0], + [-sy, 0.0, cy, 0.0], + [0.0, 0.0, 0.0, 1.0], + ], + dtype=float, + ) + rot_z = np.array( + [ + [cz, -sz, 0.0, 0.0], + [sz, cz, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.0], + [0.0, 0.0, 0.0, 1.0], + ], + dtype=float, + ) + matrix = rot_x @ rot_y @ rot_z + matrix[:3, 3] = position + return matrix + + +def matrix4(value: Any) -> np.ndarray | None: + if not isinstance(value, (list, tuple)) or len(value) != 4: + return None + try: + matrix = np.asarray(value, dtype=float) + except (TypeError, ValueError): + return None + if matrix.shape != (4, 4) or not np.all(np.isfinite(matrix)): + return None + return matrix + + +def vector3(value: Any, default: list[float]) -> list[float]: + if not isinstance(value, (list, tuple)) or len(value) != 3: + return list(default) + return [float(value[0]), float(value[1]), float(value[2])] + + +def add_mesh_to_scene( + scene: trimesh.Scene, + mesh_path: Path, + transform: np.ndarray, + uid: str, +) -> None: + loaded = trimesh.load(mesh_path, force="scene", process=False) + if isinstance(loaded, trimesh.Trimesh): + loaded.apply_transform(transform) + scene.add_geometry(loaded, node_name=uid, geom_name=uid) + return + + if isinstance(loaded, trimesh.Scene): + loaded.apply_transform(transform) + for index, geometry in enumerate(loaded.dump(concatenate=False)): + if isinstance(geometry, trimesh.Trimesh): + scene.add_geometry( + geometry, + node_name=f"{uid}_{index}", + geom_name=f"{uid}_{index}", + ) + return + + raise TypeError(f"Unsupported mesh type for {mesh_path}: {type(loaded)!r}") + + +def promote_stage_to_current(stage: ScenePaths, run_token: str) -> list[str]: + backup = make_replaced_paths(run_token) + promotion_errors: list[str] = [] + cleanup_errors: list[str] = [] + + for required_path in (stage.prompt_root, stage.config_dir, stage.image_path): + if not required_path.exists(): + raise FileNotFoundError(f"Generated artifact missing: {required_path}") + + cleanup_errors.extend(remove_path(backup.prompt_root)) + cleanup_errors.extend(remove_path(backup.config_dir)) + cleanup_errors.extend(remove_path(backup.image_path)) + + moved_to_backup: list[tuple[Path, Path]] = [] + moved_to_current: list[tuple[Path, Path]] = [] + try: + move_if_exists(PROMPT2SCENE_ROOT, backup.prompt_root, moved_to_backup) + move_if_exists(CONFIG_DIR, backup.config_dir, moved_to_backup) + move_if_exists(IMAGE_PATH, backup.image_path, moved_to_backup) + + move_required(stage.prompt_root, PROMPT2SCENE_ROOT, moved_to_current) + move_required(stage.config_dir, CONFIG_DIR, moved_to_current) + move_required(stage.image_path, IMAGE_PATH, moved_to_current) + rewrite_promoted_paths(stage) + except Exception as exc: + promotion_errors.append(f"Failed to promote generated scene: {exc}") + restore_promoted_paths(moved_to_current, moved_to_backup, promotion_errors) + raise RuntimeError("\n".join(promotion_errors)) from exc + + cleanup_errors.extend(remove_path(backup.prompt_root)) + cleanup_errors.extend(remove_path(backup.config_dir)) + cleanup_errors.extend(remove_path(backup.image_path)) + return cleanup_errors + + +def move_if_exists(src: Path, dst: Path, moved: list[tuple[Path, Path]]) -> None: + if not src.exists(): + return + dst.parent.mkdir(parents=True, exist_ok=True) + src.rename(dst) + moved.append((src, dst)) + + +def move_required(src: Path, dst: Path, moved: list[tuple[Path, Path]]) -> None: + if not src.exists(): + raise FileNotFoundError(src) + dst.parent.mkdir(parents=True, exist_ok=True) + src.rename(dst) + moved.append((dst, src)) + + +def restore_promoted_paths( + moved_to_current: list[tuple[Path, Path]], + moved_to_backup: list[tuple[Path, Path]], + errors: list[str], +) -> None: + for current_path, original_stage_path in reversed(moved_to_current): + try: + if current_path.exists(): + original_stage_path.parent.mkdir(parents=True, exist_ok=True) + current_path.rename(original_stage_path) + except Exception as exc: + errors.append( + f"Failed to restore staging artifact {original_stage_path}: {exc}" + ) + + for original_current_path, backup_path in reversed(moved_to_backup): + try: + if backup_path.exists() and not original_current_path.exists(): + original_current_path.parent.mkdir(parents=True, exist_ok=True) + backup_path.rename(original_current_path) + except Exception as exc: + errors.append( + f"Failed to restore previous scene {original_current_path}: {exc}" + ) + + +def rewrite_promoted_paths(stage: ScenePaths) -> None: + replacements = [ + (str(stage.config_dir), str(CONFIG_DIR)), + (str(stage.prompt_root), str(PROMPT2SCENE_ROOT)), + (str(stage.image_path), str(IMAGE_PATH)), + (stage.scene_id, SCENE_ID), + ] + for root in (PROMPT2SCENE_ROOT, CONFIG_DIR): + if not root.is_dir(): + continue + for path in root.rglob("*"): + if not path.is_file() or path.suffix.lower() not in TEXT_REWRITE_SUFFIXES: + continue + text = path.read_text(encoding="utf-8") + new_text = text + for old, new in replacements: + new_text = new_text.replace(old, new) + if new_text != text: + path.write_text(new_text, encoding="utf-8") + + +def ensure_initial_scene_snapshot(*, overwrite: bool = False) -> Path: + if not GRADIO_SCENE_GLB.is_file(): + build_gradio_scene_from_fast_config(FAST_GYM_CONFIG, GRADIO_SCENE_DIR) + if overwrite or not GRADIO_INITIAL_SCENE_GLB.is_file(): + GRADIO_SCENE_DIR.mkdir(parents=True, exist_ok=True) + shutil.copy2(GRADIO_SCENE_GLB, GRADIO_INITIAL_SCENE_GLB) + return GRADIO_INITIAL_SCENE_GLB + + +def prepare_current_scene_for_edit() -> Path: + scene_state = PROMPT2SCENE_ROOT / "gym_export" / "scene_state" / "result.json" + if not scene_state.is_file(): + raise FileNotFoundError( + f"Current prompt2scene scene state not found: {scene_state}" + ) + if not FAST_GYM_CONFIG.is_file(): + raise FileNotFoundError(f"Current gym config not found: {FAST_GYM_CONFIG}") + + initial_scene_path = ensure_initial_scene_snapshot() + errors = remove_path(GRADIO_SCENE_GLB) + errors.extend(remove_path(SCENE_MANIFEST)) + if errors: + raise RuntimeError("\n".join(errors)) + return initial_scene_path + + +def prebuilt_scene_dir_for_image_value( + image_value: str | np.ndarray | Image.Image, +) -> Path | None: + if not isinstance(image_value, str): + return None + image_path = Path(image_value).expanduser() + task_index = parse_task_id(image_path.name) + if task_index is None: + return None + + try: + resolved_image = image_path.resolve() + except FileNotFoundError: + return None + filename = image_path.name + matches_auto_image = False + for image_dir in auto_image_directories(): + candidate = image_dir / filename + if not candidate.is_file(): + continue + try: + if candidate.resolve() == resolved_image: + matches_auto_image = True + break + except FileNotFoundError: + continue + if not matches_auto_image: + return None + + scene_dir = get_prebuilt_scene_dir(task_index) + return scene_dir if scene_dir.is_dir() else None + + +def copy_prebuilt_scene_to_stage(prebuilt_scene_dir: Path, stage: ScenePaths) -> None: + required_paths = [ + prebuilt_scene_dir / "gym_export" / "gym_config.json", + prebuilt_scene_dir / "gym_export" / "scene_state" / "result.json", + prebuilt_scene_dir / "gym_export" / "scene_state" / "unified_scene.json", + prebuilt_scene_dir / "gym_export" / "scene_state" / "unified_scene_gen.json", + ] + missing = [path for path in required_paths if not path.is_file()] + if missing: + missing_text = ", ".join(str(path) for path in missing) + raise FileNotFoundError(f"Prebuilt scene is incomplete: {missing_text}") + + cleanup_errors = [] + cleanup_errors.extend(remove_path(stage.prompt_root)) + cleanup_errors.extend(remove_path(stage.config_dir)) + if cleanup_errors: + raise RuntimeError("\n".join(cleanup_errors)) + stage.prompt_root.parent.mkdir(parents=True, exist_ok=True) + shutil.copytree(prebuilt_scene_dir, stage.prompt_root) + + +def build_interact_random_initial_preview(prebuilt_scene_dir: Path) -> Path: + scene_id = prebuilt_scene_dir.name + preview_dir = INTERACT_RANDOM_PREVIEW_DIR / scene_id + config_path = prebuilt_scene_dir / "gym_export" / "gym_config.json" + return build_gradio_scene_from_fast_config(config_path, preview_dir) + + +def current_scene_available_for_task_only() -> bool: + return CURRENT_GYM_EXPORT_CONFIG.is_file() + + +def rerun_simulation_is_available() -> bool: + return ( + CURRENT_PATHS.fast_gym_config.is_file() and CURRENT_PATHS.agent_config.is_file() + ) + + +def run_generate( + image_value: str | np.ndarray | Image.Image, + task_text: str, + env_text: str, + *, + force_initial: bool = False, + scene_mode: str = SCENE_MODE_INITIAL, + parallel_env: bool = False, + robot_profile: str | None = None, + load_template_material: bool = False, + run_log_mode: str = RUN_LOG_MODE_INTERACT, + prebuilt_scene_dir: Path | None = None, + launch_simulation: bool = True, +): + task_text = (task_text or "").strip() + env_text = (env_text or "").strip() + if force_initial or scene_mode == SCENE_MODE_INITIAL: + mode = PIPELINE_MODE_INITIAL + elif scene_mode == SCENE_MODE_EDIT: + mode = PIPELINE_MODE_EDIT + elif scene_mode == SCENE_MODE_TASK_ONLY: + mode = PIPELINE_MODE_TASK_ONLY + else: + raise ValueError(f"Unsupported scene mode: {scene_mode}") + requested_mode = mode + if requested_mode == PIPELINE_MODE_TASK_ONLY: + env_text = "" + resolved_prebuilt_scene_dir = ( + prebuilt_scene_dir or prebuilt_scene_dir_for_image_value(image_value) + ) + use_prebuilt_scene = resolved_prebuilt_scene_dir is not None + supervisor_mode = PIPELINE_MODE_INITIAL if use_prebuilt_scene else mode + old_sim_process: subprocess.Popen[str] | None = None + with runtime_lock: + if runtime.is_busy: + yield ui_snapshot(extra_status="A pipeline run is already in progress.") + return + old_sim_process = runtime.sim_process + runtime.sim_process = None + runtime.sim_started = False + runtime.sim_finished = False + runtime.sim_returncode = None + + if old_sim_process is not None: + terminate_process_group(old_sim_process) + + token = uuid.uuid4().hex + stage = ( + CURRENT_PATHS + if supervisor_mode in {PIPELINE_MODE_EDIT, PIPELINE_MODE_TASK_ONLY} + else make_stage_paths(token) + ) + initial_scene_path: Path | None = None + existing_object_preview_path = ( + GRADIO_OBJECT_PREVIEW_GLB + if supervisor_mode in {PIPELINE_MODE_EDIT, PIPELINE_MODE_TASK_ONLY} + and GRADIO_OBJECT_PREVIEW_GLB.is_file() + else None + ) + prebuilt_initial_scene_dir: Path | None = None + try: + if use_prebuilt_scene: + if not task_text: + raise ValueError("Please enter a task description.") + if requested_mode == PIPELINE_MODE_EDIT and not env_text: + raise ValueError("Please enter a scene description to edit.") + image_path = save_input(image_value, task_text, stage.image_path) + prebuilt_initial_scene_dir = resolved_prebuilt_scene_dir + copy_prebuilt_scene_to_stage(prebuilt_initial_scene_dir, stage) + initial_scene_path = build_interact_random_initial_preview( + prebuilt_initial_scene_dir + ) + elif mode == PIPELINE_MODE_EDIT: + if not task_text: + raise ValueError("Please enter a task description.") + if not env_text: + raise ValueError("Please enter a scene description to edit.") + initial_scene_path = prepare_current_scene_for_edit() + image_path = IMAGE_PATH if IMAGE_PATH.is_file() else None + elif mode == PIPELINE_MODE_TASK_ONLY: + if not task_text: + raise ValueError("Please enter a task description.") + if not CURRENT_GYM_EXPORT_CONFIG.is_file(): + raise FileNotFoundError( + f"Current gym export not found: {CURRENT_GYM_EXPORT_CONFIG}" + ) + if GRADIO_INITIAL_SCENE_GLB.is_file(): + initial_scene_path = GRADIO_INITIAL_SCENE_GLB + elif GRADIO_SCENE_GLB.is_file(): + initial_scene_path = GRADIO_SCENE_GLB + image_path = IMAGE_PATH if IMAGE_PATH.is_file() else None + else: + image_path = save_input(image_value, task_text, stage.image_path) + except Exception as exc: + with runtime_lock: + runtime.phase_key = "failed" + runtime.status = f"Input error: {exc}" + runtime.last_error = str(exc) + runtime.log_lines.clear() + clear_run_timing_locked() + runtime.log_lines.append(runtime.status) + if run_log_mode == RUN_LOG_MODE_INTERACT: + archive_run_log( + mode=RUN_LOG_MODE_INTERACT, + task_description=task_text, + scene_description=env_text, + outcome="input_error", + ) + yield ui_snapshot() + return + + should_edit_prebuilt_scene = ( + prebuilt_initial_scene_dir is not None + and requested_mode != PIPELINE_MODE_TASK_ONLY + and bool(env_text) + ) + if mode == PIPELINE_MODE_EDIT and prebuilt_initial_scene_dir is None: + command = build_edit_pipeline_command( + task_text, + env_text, + robot_profile, + load_template_material, + ) + elif mode == PIPELINE_MODE_TASK_ONLY and prebuilt_initial_scene_dir is None: + command = build_task_only_config_command( + task_text, + robot_profile, + load_template_material, + ) + elif should_edit_prebuilt_scene: + command = build_scene_edit_pipeline_command( + task_text, + env_text, + stage, + robot_profile, + load_template_material, + ) + elif prebuilt_initial_scene_dir is not None: + command = build_config_command_for_paths( + task_text, + stage, + robot_profile, + load_template_material, + ) + else: + command = build_initial_pipeline_command( + task_text, + stage, + env_text, + robot_profile, + load_template_material, + ) + display_task_text = format_current_task(task_text, env_text) + with runtime_lock: + runtime.run_token = token + runtime.is_busy = True + runtime.phase_key = "received" + if mode == PIPELINE_MODE_EDIT: + runtime.status = "Starting scene edit..." + elif should_edit_prebuilt_scene: + runtime.status = "Prebuilt scene loaded. Starting scene edit..." + elif mode == PIPELINE_MODE_TASK_ONLY: + runtime.status = "Current scene found. Regenerating action config only..." + else: + runtime.status = "Input saved. Starting local pipeline..." + runtime.task_text = display_task_text + runtime.input_task_text = task_text + runtime.input_scene_text = env_text + runtime.image_path = image_path + runtime.submitted_input_revision += 1 + runtime.lerobot_video_path = None + runtime.lerobot_dataset_path = None + runtime.object_model_path = existing_object_preview_path + runtime.scene_model_path = initial_scene_path + runtime.edited_scene_model_path = None + runtime.last_error = None + runtime.sim_started = False + runtime.sim_finished = False + runtime.sim_returncode = None + runtime.log_lines.clear() + clear_run_timing_locked() + runtime.log_lines.append("$ " + " ".join(command)) + yield ui_snapshot() + + try: + process = start_pipeline(command) + except Exception as exc: + with runtime_lock: + runtime.is_busy = False + runtime.process = None + runtime.phase_key = "failed" + runtime.status = f"Pipeline start failed: {exc}" + runtime.last_error = str(exc) + if run_log_mode == RUN_LOG_MODE_INTERACT: + archive_run_log( + mode=RUN_LOG_MODE_INTERACT, + task_description=task_text, + scene_description=env_text, + outcome="pipeline_start_failed", + ) + yield ui_snapshot() + return + + output_queue: queue.Queue[str] = queue.Queue() + reader = threading.Thread( + target=read_process_output, + args=(process, output_queue), + daemon=True, + ) + supervisor = threading.Thread( + target=supervise_pipeline, + args=( + token, + stage, + supervisor_mode, + process, + display_task_text, + task_text, + env_text, + output_queue, + reader, + parallel_env, + robot_profile, + run_log_mode, + initial_scene_path, + should_edit_prebuilt_scene, + launch_simulation, + ), + daemon=True, + ) + + with runtime_lock: + if runtime.run_token != token: + terminate_process_group(process) + return + runtime.process = process + start_run_timing_locked("started") + runtime.phase_key = "started" + runtime.status = "Local pipeline started." + reader.start() + supervisor.start() + yield ui_snapshot() + + while True: + with runtime_lock: + still_current = runtime.run_token == token + busy = runtime.is_busy + if not still_current or not busy: + break + time.sleep(1.0) + yield ui_snapshot() + yield ui_snapshot() + + +def start_auto_loop_state() -> str | None: + process: subprocess.Popen[str] | None = None + sim_process: subprocess.Popen[str] | None = None + with runtime_lock: + if runtime.auto_loop_active or runtime.is_busy: + runtime.status = "A pipeline run is already in progress." + return None + + available_tasks = available_auto_task_indices() + if not available_tasks: + image_dirs = ", ".join(str(path) for path in auto_image_directories()) + message = ( + "Auto cannot start: no task input images were found. " + "Add task1_0.png through task5_3.png to one of: " + f"{image_dirs}" + ) + runtime.phase_key = "failed" + runtime.status = message + runtime.last_error = message + runtime.log_lines.clear() + clear_run_timing_locked() + runtime.log_lines.append(message) + return None + + token = uuid.uuid4().hex + process = runtime.process + sim_process = runtime.sim_process + runtime.process = None + runtime.sim_process = None + runtime.sim_started = False + runtime.sim_finished = False + runtime.sim_returncode = None + runtime.auto_loop_active = True + runtime.auto_loop_token = token + runtime.auto_round = 0 + runtime.auto_scene_mode = SCENE_MODE_INITIAL + runtime.auto_parallel_env = False + runtime.auto_robot_profile = DEFAULT_ROBOT_PROFILE + runtime.phase_key = "received" + runtime.status = "Auto loop starting." + runtime.video_path = None + runtime.lerobot_video_path = None + runtime.lerobot_dataset_path = None + runtime.last_error = None + runtime.log_lines.clear() + clear_run_timing_locked() + + if process is not None: + terminate_process_group(process) + if sim_process is not None: + terminate_process_group(sim_process) + return token + + +def auto_loop_is_active(loop_token: str) -> bool: + with runtime_lock: + return runtime.auto_loop_active and runtime.auto_loop_token == loop_token + + +def finish_auto_loop(loop_token: str, status_text: str | None = None) -> None: + with runtime_lock: + if runtime.auto_loop_token != loop_token: + return + runtime.auto_loop_active = False + runtime.auto_loop_token = None + runtime.auto_round = 0 + runtime.auto_scene_mode = SCENE_MODE_INITIAL + runtime.auto_parallel_env = False + runtime.auto_robot_profile = DEFAULT_ROBOT_PROFILE + if status_text is not None: + runtime.status = status_text + + +def stop_auto_loop_if_running() -> bool: + process: subprocess.Popen[str] | None = None + sim_process: subprocess.Popen[str] | None = None + with runtime_lock: + if not runtime.auto_loop_active: + return False + runtime.run_token = uuid.uuid4().hex + runtime.auto_loop_active = False + runtime.auto_loop_token = None + runtime.auto_round = 0 + runtime.auto_scene_mode = SCENE_MODE_INITIAL + runtime.auto_parallel_env = False + runtime.auto_robot_profile = DEFAULT_ROBOT_PROFILE + process = runtime.process + sim_process = runtime.sim_process + runtime.process = None + runtime.sim_process = None + runtime.sim_started = False + runtime.sim_finished = False + runtime.sim_returncode = None + runtime.is_busy = False + runtime.phase_key = "idle" + runtime.status = "Stopped." + runtime.task_text = "" + runtime.input_task_text = "" + runtime.input_scene_text = "" + runtime.image_path = None + runtime.video_path = None + runtime.lerobot_video_path = None + runtime.lerobot_dataset_path = None + runtime.object_model_path = None + runtime.scene_model_path = None + runtime.edited_scene_model_path = None + runtime.last_error = None + runtime.log_lines.clear() + clear_run_timing_locked() + + if process is not None: + terminate_process_group(process) + if sim_process is not None: + terminate_process_group(sim_process) + return True + + +def wait_for_current_simulation_to_exit( + loop_token: str, + base_image: str, + auto_task: str, + auto_scene: str, +): + while auto_loop_is_active(loop_token): + with runtime_lock: + sim_running = runtime.sim_process is not None + if not sim_running: + break + time.sleep(1.0) + yield ( + base_image, + auto_task, + auto_scene, + *ui_snapshot(extra_status="Auto waiting for Dexsim to exit."), + ) + + +def run_generate_for_top_mode( + run_mode: str, + action_mode: str | None, + scene_mode: str, + robot_profile: str | None, + image_value: str | np.ndarray | Image.Image, + task_text: str, + env_text: str, + interact_prebuilt_scene_dir: str | None, + language: str | None, +): + parallel_env = action_mode == TOP_MODE_PARALLEL_ENV + if run_mode != TOP_MODE_AUTO: + selected_prebuilt_scene_dir = ( + Path(interact_prebuilt_scene_dir) if interact_prebuilt_scene_dir else None + ) + for snapshot in run_generate( + image_value, + task_text, + env_text, + force_initial=False, + scene_mode=scene_mode, + parallel_env=parallel_env, + robot_profile=robot_profile, + load_template_material=False, + run_log_mode=RUN_LOG_MODE_INTERACT, + prebuilt_scene_dir=selected_prebuilt_scene_dir, + ): + yield ( + gr.update(), + gr.update(), + gr.update(), + *snapshot, + ) + return + + loop_token = start_auto_loop_state() + if loop_token is None: + yield ( + gr.update(), + gr.update(), + gr.update(), + *ui_snapshot(), + ) + return + + with runtime_lock: + runtime.language = language or LANGUAGE_EN + + def set_auto_control_state( + scene_mode: str, + parallel_env: bool, + robot_profile: str | None, + ) -> None: + with runtime_lock: + runtime.auto_scene_mode = scene_mode + runtime.auto_parallel_env = parallel_env + runtime.auto_robot_profile = robot_profile or DEFAULT_ROBOT_PROFILE + + def run_auto_phase( + phase_name: str, + base_image: str, + task_text: str, + scene_text: str, + *, + scene_mode: str, + parallel_env: bool, + robot_profile: str | None, + force_initial: bool = False, + prebuilt_scene_dir: Path | None = None, + ): + for snapshot in run_generate( + base_image, + task_text, + scene_text, + force_initial=force_initial, + scene_mode=scene_mode, + parallel_env=parallel_env, + robot_profile=robot_profile, + load_template_material=False, + run_log_mode=RUN_LOG_MODE_AUTO, + prebuilt_scene_dir=prebuilt_scene_dir, + ): + yield ( + base_image, + task_text, + scene_text, + *snapshot, + ) + if not auto_loop_is_active(loop_token): + break + + if not auto_loop_is_active(loop_token): + archive_run_log( + mode=RUN_LOG_MODE_AUTO, + task_description=task_text, + scene_description=scene_text, + outcome="stopped", + ) + return "stopped" + + with runtime_lock: + pipeline_failed = runtime.phase_key == "failed" + pipeline_error = runtime.last_error + if pipeline_failed: + cleanup_auto_generated_artifacts() + if pipeline_error: + with runtime_lock: + runtime.last_error = pipeline_error + runtime.log_lines.append( + f"{phase_name} generation failed: {pipeline_error}" + ) + archive_run_log( + mode=RUN_LOG_MODE_AUTO, + task_description=task_text, + scene_description=scene_text, + outcome="pipeline_failed", + ) + return "pipeline_failed" + + for snapshot in wait_for_current_simulation_to_exit( + loop_token, + base_image, + task_text, + scene_text, + ): + yield snapshot + + if not auto_loop_is_active(loop_token): + archive_run_log( + mode=RUN_LOG_MODE_AUTO, + task_description=task_text, + scene_description=scene_text, + outcome="stopped", + ) + return "stopped" + + with runtime_lock: + simulation_completed = ( + runtime.sim_started + and runtime.sim_finished + and runtime.sim_process is None + ) + round_outcome = "completed" if simulation_completed else "simulation_failed" + archive_run_log( + mode=RUN_LOG_MODE_AUTO, + task_description=task_text, + scene_description=scene_text, + outcome=round_outcome, + ) + yield ( + base_image, + task_text, + scene_text, + *ui_snapshot(extra_status=f"{phase_name}: {round_outcome}."), + ) + return round_outcome + + def run_auto_parallel_simulation( + base_image: str, + task_text: str, + scene_text: str, + *, + robot_profile: str | None, + ): + with runtime_lock: + simulation_token = runtime.run_token + runtime.sim_started = False + runtime.sim_finished = False + runtime.sim_returncode = None + runtime.last_error = None + runtime.status = "Starting parallel simulation..." + clear_run_timing_locked() + runtime.log_lines.append("Auto phase: starting parallel simulation.") + + simulation_error = launch_current_simulation( + simulation_token, + parallel_env=True, + robot_profile=robot_profile, + run_log_mode=RUN_LOG_MODE_AUTO, + task_description=task_text, + scene_description=scene_text, + ) + if simulation_error is not None: + with runtime_lock: + runtime.phase_key = "failed" + runtime.status = ( + f"Parallel simulation launch failed: {simulation_error}" + ) + runtime.last_error = simulation_error + runtime.log_lines.append(runtime.status) + archive_run_log( + mode=RUN_LOG_MODE_AUTO, + task_description=task_text, + scene_description=scene_text, + outcome="simulation_launch_failed", + ) + yield ( + base_image, + task_text, + scene_text, + *ui_snapshot(), + ) + return "simulation_failed" + + yield ( + base_image, + task_text, + scene_text, + *ui_snapshot(extra_status="Parallel simulation started."), + ) + for snapshot in wait_for_current_simulation_to_exit( + loop_token, + base_image, + task_text, + scene_text, + ): + yield snapshot + + if not auto_loop_is_active(loop_token): + archive_run_log( + mode=RUN_LOG_MODE_AUTO, + task_description=task_text, + scene_description=scene_text, + outcome="stopped", + ) + return "stopped" + + with runtime_lock: + simulation_completed = ( + runtime.sim_started + and runtime.sim_finished + and runtime.sim_process is None + ) + round_outcome = "completed" if simulation_completed else "simulation_failed" + archive_run_log( + mode=RUN_LOG_MODE_AUTO, + task_description=task_text, + scene_description=scene_text, + outcome=round_outcome, + ) + yield ( + base_image, + task_text, + scene_text, + *ui_snapshot(extra_status=f"Parallel simulation: {round_outcome}."), + ) + return round_outcome + + while auto_loop_is_active(loop_token): + auto_task = "" + auto_scene = "" + task_label = "unknown" + with runtime_lock: + runtime.auto_round += 1 + auto_round = runtime.auto_round + runtime.status = f"Auto round {auto_round}: cleaning previous artifacts." + runtime.last_error = None + runtime.sim_started = False + runtime.sim_finished = False + runtime.sim_returncode = None + runtime.log_lines.clear() + clear_run_timing_locked() + runtime.log_lines.append(f"Auto round {auto_round} started.") + + cleanup_errors = cleanup_auto_generated_artifacts() + if cleanup_errors: + with runtime_lock: + runtime.log_lines.extend(cleanup_errors) + + if not auto_loop_is_active(loop_token): + break + + try: + with runtime_lock: + selected_language = runtime.language + auto_input = generate_auto_text_input( + language=selected_language, + include_scene=False, + ) + except Exception as exc: + if not auto_loop_is_active(loop_token): + break + with runtime_lock: + runtime.phase_key = "failed" + runtime.status = f"Auto text generation failed: {exc}" + runtime.last_error = str(exc) + clear_run_timing_locked() + runtime.log_lines.append(runtime.status) + yield ( + gr.update(), + gr.update(), + gr.update(), + *ui_snapshot(), + ) + archive_run_log( + mode=RUN_LOG_MODE_AUTO, + task_description=auto_task, + scene_description=auto_scene, + outcome="text_generation_failed", + ) + continue + + base_image = auto_input.base_image_path.as_posix() + auto_task = auto_input.task_description + task_label = f"task{auto_input.task_index[0]}_{auto_input.task_index[1]}" + with runtime_lock: + runtime.task_text = format_current_task(auto_task, auto_scene) + runtime.input_task_text = auto_task + runtime.input_scene_text = auto_scene + runtime.image_path = auto_input.base_image_path + runtime.video_path = None + runtime.lerobot_video_path = None + runtime.lerobot_dataset_path = None + runtime.phase_key = "received" + runtime.status = ( + f"Auto round {auto_round}: selected {task_label}. " + "Starting prompt2scene pipeline." + ) + runtime.last_error = None + runtime.log_lines.append( + f"Auto selected {task_label}: task={auto_task!r}, scene={auto_scene!r}" + ) + if auto_input.prebuilt_scene_dir is not None: + runtime.log_lines.append( + f"Auto prebuilt scene: {auto_input.prebuilt_scene_dir}" + ) + yield ( + base_image, + auto_task, + auto_scene, + *ui_snapshot(extra_status=f"Auto text generated: {task_label}."), + ) + + if not auto_loop_is_active(loop_token): + archive_run_log( + mode=RUN_LOG_MODE_AUTO, + task_description=auto_task, + scene_description=auto_scene, + outcome="stopped", + ) + break + + with runtime_lock: + selected_language = runtime.language + phase_results: list[str] = [] + + phase_results.append("stopped") + set_auto_control_state(SCENE_MODE_INITIAL, False, robot_profile) + phase_generator = run_auto_phase( + "Initial generation", + base_image, + auto_task, + auto_scene, + scene_mode=SCENE_MODE_INITIAL, + parallel_env=False, + robot_profile=robot_profile, + force_initial=True, + prebuilt_scene_dir=auto_input.prebuilt_scene_dir, + ) + try: + while True: + yield next(phase_generator) + except StopIteration as exc: + phase_results[0] = str(exc.value) + + if phase_results[0] == "stopped": + break + if phase_results[0] == "pipeline_failed": + continue + + try: + auto_edit_scene = generate_auto_scene_description( + task_index=auto_input.task_index, + language=selected_language, + ensure_scene=True, + ) + except Exception as exc: + with runtime_lock: + runtime.phase_key = "failed" + runtime.status = f"Auto scene description generation failed: {exc}" + runtime.last_error = str(exc) + clear_run_timing_locked() + runtime.log_lines.append(runtime.status) + yield ( + gr.update(), + gr.update(), + gr.update(), + *ui_snapshot(), + ) + archive_run_log( + mode=RUN_LOG_MODE_AUTO, + task_description=auto_task, + scene_description=auto_scene, + outcome="text_generation_failed", + ) + continue + + phase_results.append("stopped") + set_auto_control_state(SCENE_MODE_EDIT, False, robot_profile) + phase_generator = run_auto_phase( + "Scene edit", + base_image, + auto_task, + auto_edit_scene, + scene_mode=SCENE_MODE_EDIT, + parallel_env=False, + robot_profile=robot_profile, + force_initial=False, + ) + try: + while True: + yield next(phase_generator) + except StopIteration as exc: + phase_results[1] = str(exc.value) + + if phase_results[1] == "stopped": + break + if phase_results[1] == "pipeline_failed": + continue + + phase_results.append("stopped") + set_auto_control_state(SCENE_MODE_EDIT, True, robot_profile) + phase_generator = run_auto_parallel_simulation( + base_image, + auto_task, + auto_edit_scene, + robot_profile=robot_profile, + ) + try: + while True: + yield next(phase_generator) + except StopIteration as exc: + phase_results[2] = str(exc.value) + + if phase_results[2] == "stopped": + break + + phase_results.append("stopped") + set_auto_control_state(SCENE_MODE_TASK_ONLY, False, ROBOT_PROFILE_FRANKA) + phase_generator = run_auto_phase( + "Task-only Franka", + base_image, + auto_task, + "", + scene_mode=SCENE_MODE_TASK_ONLY, + parallel_env=False, + robot_profile=ROBOT_PROFILE_FRANKA, + force_initial=False, + ) + try: + while True: + yield next(phase_generator) + except StopIteration as exc: + phase_results[3] = str(exc.value) + + if phase_results[3] == "stopped": + break + if phase_results[3] == "pipeline_failed": + continue + + phase_results.append("stopped") + set_auto_control_state(SCENE_MODE_TASK_ONLY, True, ROBOT_PROFILE_FRANKA) + phase_generator = run_auto_parallel_simulation( + base_image, + auto_task, + "", + robot_profile=ROBOT_PROFILE_FRANKA, + ) + try: + while True: + yield next(phase_generator) + except StopIteration as exc: + phase_results[4] = str(exc.value) + + if phase_results[4] == "stopped": + break + + cleanup_errors = cleanup_auto_generated_artifacts() + if cleanup_errors: + with runtime_lock: + runtime.log_lines.extend(cleanup_errors) + + finish_auto_loop(loop_token) + + +def _scene_engine_phase_from_log(line: str, current_key: str) -> str: + """Map the standalone Scene Engine's stage names to the shared progress UI.""" + text = line.lower() + mapping = ( + ("scene understanding", "scene_intake"), + ("scene segmentation", "relations"), + ("coarse layout", "asset_generation"), + ("scene export", "gym_export"), + ) + current_progress = PHASES.get(current_key, PHASES["idle"]).progress + for needle, phase_key in mapping: + if needle in text and PHASES[phase_key].progress > current_progress: + return phase_key + return current_key + + +def _scene_engine_updates( + output_root: Path | None = None, + preview_html: str | None = None, +) -> tuple[int, str, str | None, str]: + with runtime_lock: + phase = PHASES.get(runtime.phase_key, PHASES["idle"]) + status = format_status( + runtime.status, + phase=phase, + busy=runtime.is_busy, + last_error=runtime.last_error, + ) + return ( + phase.progress, + status, + output_root.as_posix() if output_root is not None else None, + preview_html or "", + ) + + +def _prepare_scene_engine_input( + image_value: str | np.ndarray | Image.Image, +) -> tuple[str, Path, Path]: + """Normalize an uploaded image and store it under a stable content hash.""" + if image_value is None: + raise ValueError("Please upload an image first.") + if isinstance(image_value, str): + image = Image.open(image_value) + elif isinstance(image_value, np.ndarray): + image = Image.fromarray(image_value) + elif isinstance(image_value, Image.Image): + image = image_value + else: + raise TypeError(f"Unsupported image input type: {type(image_value)!r}") + + normalized = ImageOps.exif_transpose(image).convert("RGB") + image_bytes = io.BytesIO() + normalized.save(image_bytes, format="PNG") + scene_hash = hashlib.sha256(image_bytes.getvalue()).hexdigest()[:16] + output_root = DEBUG_SCENE_ENGINE_ROOT / scene_hash + output_root.mkdir(parents=True, exist_ok=True) + image_path = output_root / "input.png" + image_path.write_bytes(image_bytes.getvalue()) + return scene_hash, output_root, image_path + + +def _wait_for_viser(port: int, process: subprocess.Popen[str]) -> bool: + """Wait briefly for Viser's HTTP listener, without treating Ctrl-C as success.""" + deadline = time.monotonic() + 15.0 + while time.monotonic() < deadline: + if process.poll() is not None: + return False + try: + with socket.create_connection(("127.0.0.1", port), timeout=0.2): + return True + except OSError: + time.sleep(0.25) + return False + + +def _viser_iframe(port: int, scene_hash: str) -> str: + """Embed the Viser service using the same hostname as the Gradio page.""" + srcdoc = ( + "" + ) + return ( + f"
Viser preview: {html.escape(scene_hash)}" + f"" + "
" + ) + + +def run_scene_engine(image_value: str | np.ndarray | Image.Image): + """Generate one image-conditioned scene and expose its Viser preview.""" + output_root: Path | None = None + preview_html = "" + try: + scene_hash, output_root, image_path = _prepare_scene_engine_input(image_value) + if not SCENE_ENGINE_CONFIG.is_file(): + raise FileNotFoundError( + f"Scene Engine config not found: {SCENE_ENGINE_CONFIG}" + ) + except Exception as exc: + with runtime_lock: + set_runtime_phase_locked("failed") + runtime.status = f"Input error: {exc}" + runtime.last_error = str(exc) + yield _scene_engine_updates(output_root, preview_html) + return + + old_preview: subprocess.Popen[str] | None = None + busy_message: str | None = None + with runtime_lock: + if runtime.is_busy: + runtime.status = "Another pipeline is already running." + runtime.last_error = runtime.status + busy_message = runtime.status + else: + old_preview = runtime.scene_preview_process + runtime.scene_preview_process = None + token = uuid.uuid4().hex + runtime.run_token = token + runtime.is_busy = True + set_runtime_phase_locked("received") + runtime.status = ( + f"Image saved. Generating Scene Engine output {scene_hash}." + ) + runtime.last_error = None + runtime.image_path = image_path + runtime.log_lines.clear() + clear_run_timing_locked() + + if busy_message is not None: + yield _scene_engine_updates(output_root, preview_html) + return + + if old_preview is not None: + terminate_process_group(old_preview) + + command = [ + sys.executable, + "-m", + COMMANDS["scene_engine"]["module"], + *COMMANDS["scene_engine"]["base_args"], + "--image", + str(image_path), + "--output_root", + str(output_root), + "--config", + str(SCENE_ENGINE_CONFIG), + ] + with runtime_lock: + runtime.log_lines.append("$ " + " ".join(command)) + yield _scene_engine_updates(output_root, preview_html) + + try: + process = start_pipeline(command) + except Exception as exc: + with runtime_lock: + runtime.is_busy = False + set_runtime_phase_locked("failed") + runtime.status = f"Scene Engine start failed: {exc}" + runtime.last_error = str(exc) + yield _scene_engine_updates(output_root, preview_html) + return + + output_queue: queue.Queue[str] = queue.Queue() + reader = threading.Thread( + target=read_process_output, args=(process, output_queue), daemon=True + ) + with runtime_lock: + if runtime.run_token != token: + terminate_process_group(process) + return + runtime.process = process + start_run_timing_locked("started") + set_runtime_phase_locked("started") + runtime.status = "Scene Engine generation started." + reader.start() + + while process.poll() is None: + drained = drain_output_queue(output_queue) + with runtime_lock: + for line in drained: + runtime.log_lines.append(line) + set_runtime_phase_locked( + _scene_engine_phase_from_log(line, runtime.phase_key) + ) + if (output_root / "scene_export" / "scene_config.json").is_file(): + set_runtime_phase_locked("gym_export") + runtime.status = PHASES[runtime.phase_key].label + "." + yield _scene_engine_updates(output_root, preview_html) + time.sleep(0.5) + + reader.join(timeout=1.0) + with runtime_lock: + for line in drain_output_queue(output_queue): + runtime.log_lines.append(line) + set_runtime_phase_locked( + _scene_engine_phase_from_log(line, runtime.phase_key) + ) + runtime.process = None + + scene_export = output_root / "scene_export" / "scene_config.json" + if process.returncode != 0 or not scene_export.is_file(): + detail = ( + f"Scene Engine exited with code {process.returncode}." + if process.returncode != 0 + else f"Scene Engine did not create {scene_export}." + ) + with runtime_lock: + runtime.is_busy = False + set_runtime_phase_locked("failed") + runtime.status = detail + runtime.last_error = detail + yield _scene_engine_updates(output_root, preview_html) + return + + port = SCENE_ENGINE_VISER_PORT + preview_command = [ + sys.executable, + COMMANDS["scene_engine"]["preview_script"], + str(output_root), + "--viser", + "--viser-host", + "0.0.0.0", + "--viser-port", + str(port), + ] + try: + preview_process = start_pipeline(preview_command) + except Exception as exc: + with runtime_lock: + runtime.is_busy = False + set_runtime_phase_locked("failed") + runtime.status = f"Viser preview start failed: {exc}" + runtime.last_error = str(exc) + yield _scene_engine_updates(output_root, preview_html) + return + + with runtime_lock: + runtime.log_lines.append("$ " + " ".join(preview_command)) + set_runtime_phase_locked("preview") + runtime.status = "Starting Viser preview..." + yield _scene_engine_updates(output_root, preview_html) + + if not _wait_for_viser(port, preview_process): + terminate_process_group(preview_process) + with runtime_lock: + runtime.is_busy = False + set_runtime_phase_locked("failed") + runtime.status = "Viser preview did not start." + runtime.last_error = runtime.status + yield _scene_engine_updates(output_root, preview_html) + return + + preview_html = _viser_iframe(port, scene_hash) + with runtime_lock: + runtime.scene_preview_process = preview_process + runtime.is_busy = False + set_runtime_phase_locked("complete") + runtime.status = "Scene generated successfully. Viser preview is ready." + runtime.last_error = None + yield _scene_engine_updates(output_root, preview_html) + + +def run_action_engine_from_current(task_text: str, robot_profile: str | None): + """Launch DexSim for the Gym scene most recently generated by Scene engine.""" + task_text = (task_text or "").strip() + failure: str | None = None + with runtime_lock: + if not task_text: + runtime.status = "Enter a task description first." + runtime.last_error = "Task description is required." + failure = runtime.status + elif not rerun_simulation_is_available(): + runtime.status = "Generate a scene first." + runtime.last_error = "Current Gym scene/config is unavailable." + failure = runtime.status + elif ( + runtime.process is not None + or runtime.sim_process is not None + or runtime.is_busy + ): + runtime.status = "Another pipeline or simulation is already running." + runtime.last_error = "Busy." + failure = runtime.status + elif not action_agent_cli_is_available(): + runtime.status = ( + "Action-agent CLI is unavailable in this EmbodiChain environment." + ) + runtime.last_error = ( + "Missing embodichain.gen_sim.action_agent_pipeline.cli.run_agent" + ) + failure = runtime.status + else: + token = uuid.uuid4().hex + runtime.run_token = token + runtime.task_text = task_text + runtime.input_task_text = task_text + runtime.input_scene_text = "" + runtime.status = "Starting DexSim action simulation..." + runtime.last_error = None + runtime.log_lines.append(runtime.status) + + if failure: + return ui_snapshot() + + error = launch_current_simulation( + token, + robot_profile=robot_profile, + run_log_mode=RUN_LOG_MODE_INTERACT, + task_description=task_text, + ) + if error: + with runtime_lock: + runtime.status = error + runtime.last_error = error + return ui_snapshot() + + +def action_agent_cli_is_available() -> bool: + """Avoid spawning a subprocess when the optional action-agent package is absent.""" + try: + return importlib.util.find_spec(COMMANDS["agent"]["module"]) is not None + except (ImportError, ModuleNotFoundError): + return False + + +def supervise_pipeline( + token: str, + stage: ScenePaths, + mode: str, + process: subprocess.Popen[str], + display_task_text: str, + task_description: str, + scene_description: str, + output_queue: queue.Queue[str], + reader: threading.Thread, + parallel_env: bool, + robot_profile: str | None, + run_log_mode: str, + initial_scene_path: Path | None, + show_generated_scene_as_edit: bool, + launch_simulation: bool = True, +) -> None: + is_edit = mode == PIPELINE_MODE_EDIT + is_task_only = mode == PIPELINE_MODE_TASK_ONLY + scene_build_error: str | None = None + simulation_error: str | None = None + simulation_started = False + try: + while True: + with runtime_lock: + still_current = runtime.run_token == token + if not still_current: + terminate_process_group(process) + return + + drained = drain_output_queue(output_queue) + if drained: + with runtime_lock: + for line in drained: + runtime.log_lines.append(line) + set_runtime_phase_locked( + update_phase_from_log(line, runtime.phase_key) + ) + + with runtime_lock: + detected_key = detect_phase_from_files(runtime.phase_key, stage) + set_runtime_phase_locked(detected_key) + if detected_key in PHASES and runtime.phase_key != "failed": + runtime.status = PHASES[detected_key].label + "." + + glb_paths = collect_generated_object_glbs(stage) + if glb_paths and ( + not stage.gradio_object_preview_glb.is_file() + or not object_preview_is_current( + stage.object_preview_manifest, + glb_paths, + ) + ): + try: + object_preview_path = build_object_preview_scene( + glb_paths, + stage.gradio_scene_dir, + ) + with runtime_lock: + runtime.object_model_path = object_preview_path + set_runtime_phase_locked( + _choose_later_phase( + runtime.phase_key, + PHASES.get(runtime.phase_key, PHASES["idle"]).progress, + "asset_generation", + )[0] + ) + runtime.status = ( + f"Generated object GLB preview loaded " + f"({len(glb_paths)} files)." + ) + except Exception as exc: + with runtime_lock: + runtime.log_lines.append(f"Object preview pending: {exc}") + + if ( + not is_edit + and not is_task_only + and stage.fast_gym_config.is_file() + and not gradio_scene_is_current( + stage.gradio_scene_glb, + stage.scene_manifest, + stage.fast_gym_config, + ) + ): + try: + scene_path = build_gradio_scene_from_fast_config( + stage.fast_gym_config, + stage.gradio_scene_dir, + ) + scene_build_error = None + with runtime_lock: + if show_generated_scene_as_edit: + if initial_scene_path is not None: + runtime.scene_model_path = initial_scene_path + runtime.edited_scene_model_path = scene_path + else: + runtime.scene_model_path = scene_path + set_runtime_phase_locked("preview") + runtime.status = "3D preview loaded." + runtime.last_error = None + except Exception as exc: + scene_build_error = str(exc) + with runtime_lock: + runtime.log_lines.append( + f"3D preview error: {scene_build_error}" + ) + runtime.last_error = scene_build_error + + if process.poll() is not None: + break + time.sleep(0.5) + + reader.join(timeout=1.0) + drained = drain_output_queue(output_queue) + with runtime_lock: + for line in drained: + runtime.log_lines.append(line) + set_runtime_phase_locked(update_phase_from_log(line, runtime.phase_key)) + + glb_paths = collect_generated_object_glbs(stage) + if glb_paths and ( + not stage.gradio_object_preview_glb.is_file() + or not object_preview_is_current(stage.object_preview_manifest, glb_paths) + ): + try: + object_preview_path = build_object_preview_scene( + glb_paths, + stage.gradio_scene_dir, + ) + with runtime_lock: + runtime.object_model_path = object_preview_path + except Exception as exc: + with runtime_lock: + runtime.log_lines.append(f"Object preview skipped: {exc}") + + if is_task_only and process.returncode == 0 and stage.fast_gym_config.is_file(): + try: + scene_path = build_gradio_scene_from_fast_config( + stage.fast_gym_config, + stage.gradio_scene_dir, + ) + scene_build_error = None + if GRADIO_SCENE_GLB.is_file(): + GRADIO_SCENE_DIR.mkdir(parents=True, exist_ok=True) + shutil.copy2(GRADIO_SCENE_GLB, GRADIO_INITIAL_SCENE_GLB) + with runtime_lock: + runtime.scene_model_path = scene_path + runtime.edited_scene_model_path = None + set_runtime_phase_locked("preview") + runtime.status = "3D preview loaded." + runtime.last_error = None + except Exception as exc: + scene_build_error = str(exc) + elif ( + stage.fast_gym_config.is_file() + and (not is_edit or process.returncode == 0) + and not gradio_scene_is_current( + stage.gradio_scene_glb, + stage.scene_manifest, + stage.fast_gym_config, + ) + ): + try: + scene_path = build_gradio_scene_from_fast_config( + stage.fast_gym_config, + stage.gradio_scene_dir, + ) + scene_build_error = None + with runtime_lock: + if is_edit or show_generated_scene_as_edit: + runtime.edited_scene_model_path = scene_path + else: + runtime.scene_model_path = scene_path + set_runtime_phase_locked("preview") + runtime.status = "3D preview loaded." + runtime.last_error = None + except Exception as exc: + scene_build_error = str(exc) + + cleanup_errors: list[str] = [] + promotion_error: str | None = None + pipeline_output_ready = ( + stage.fast_gym_config.is_file() and stage.agent_config.is_file() + if is_task_only + else stage.fast_gym_config.is_file() + ) + missing_output_name = ( + f"{stage.fast_gym_config.name} and/or {stage.agent_config.name}" + if is_task_only + else FAST_GYM_CONFIG.name + ) + pipeline_succeeded = ( + process.returncode == 0 and pipeline_output_ready and not scene_build_error + ) + if pipeline_succeeded: + if is_edit: + with runtime_lock: + if runtime.run_token == token: + runtime.image_path = ( + IMAGE_PATH if IMAGE_PATH.is_file() else None + ) + if GRADIO_OBJECT_PREVIEW_GLB.is_file(): + runtime.object_model_path = GRADIO_OBJECT_PREVIEW_GLB + if GRADIO_INITIAL_SCENE_GLB.is_file(): + runtime.scene_model_path = GRADIO_INITIAL_SCENE_GLB + if GRADIO_SCENE_GLB.is_file(): + runtime.edited_scene_model_path = GRADIO_SCENE_GLB + if launch_simulation: + simulation_error = launch_current_simulation( + token, + parallel_env=parallel_env, + robot_profile=robot_profile, + run_log_mode=run_log_mode, + task_description=task_description, + scene_description=scene_description, + ) + simulation_started = simulation_error is None + elif is_task_only: + with runtime_lock: + if runtime.run_token == token: + runtime.image_path = ( + IMAGE_PATH if IMAGE_PATH.is_file() else None + ) + if GRADIO_OBJECT_PREVIEW_GLB.is_file(): + runtime.object_model_path = GRADIO_OBJECT_PREVIEW_GLB + if GRADIO_INITIAL_SCENE_GLB.is_file(): + runtime.scene_model_path = GRADIO_INITIAL_SCENE_GLB + elif GRADIO_SCENE_GLB.is_file(): + runtime.scene_model_path = GRADIO_SCENE_GLB + runtime.edited_scene_model_path = None + if launch_simulation: + simulation_error = launch_current_simulation( + token, + parallel_env=parallel_env, + robot_profile=robot_profile, + run_log_mode=run_log_mode, + task_description=task_description, + scene_description=scene_description, + ) + simulation_started = simulation_error is None + else: + try: + cleanup_errors = promote_stage_to_current(stage, token) + except Exception as exc: + promotion_error = str(exc) + else: + initial_scene_error: str | None = None + try: + if show_generated_scene_as_edit and initial_scene_path: + GRADIO_SCENE_DIR.mkdir(parents=True, exist_ok=True) + shutil.copy2(initial_scene_path, GRADIO_INITIAL_SCENE_GLB) + else: + ensure_initial_scene_snapshot(overwrite=True) + except Exception as exc: + initial_scene_error = str(exc) + with runtime_lock: + if runtime.run_token == token: + runtime.image_path = IMAGE_PATH + if GRADIO_OBJECT_PREVIEW_GLB.is_file(): + runtime.object_model_path = GRADIO_OBJECT_PREVIEW_GLB + if show_generated_scene_as_edit: + if GRADIO_INITIAL_SCENE_GLB.is_file(): + runtime.scene_model_path = GRADIO_INITIAL_SCENE_GLB + elif initial_scene_path is not None: + runtime.scene_model_path = initial_scene_path + if GRADIO_SCENE_GLB.is_file(): + runtime.edited_scene_model_path = GRADIO_SCENE_GLB + elif GRADIO_INITIAL_SCENE_GLB.is_file(): + runtime.scene_model_path = GRADIO_INITIAL_SCENE_GLB + elif GRADIO_SCENE_GLB.is_file(): + runtime.scene_model_path = GRADIO_SCENE_GLB + if not show_generated_scene_as_edit: + runtime.edited_scene_model_path = None + if initial_scene_error: + runtime.log_lines.append( + f"Initial scene snapshot skipped: {initial_scene_error}" + ) + if launch_simulation: + simulation_error = launch_current_simulation( + token, + parallel_env=parallel_env, + robot_profile=robot_profile, + run_log_mode=run_log_mode, + task_description=task_description, + scene_description=scene_description, + ) + simulation_started = simulation_error is None + + archive_after_status = False + archive_outcome = "completed" + with runtime_lock: + if runtime.run_token != token: + return + runtime.is_busy = False + runtime.process = None + if pipeline_succeeded and not promotion_error: + set_runtime_phase_locked("complete") + runtime.status = "Pipeline completed successfully." + runtime.task_text = display_task_text + runtime.image_path = IMAGE_PATH if IMAGE_PATH.is_file() else None + if GRADIO_OBJECT_PREVIEW_GLB.is_file(): + runtime.object_model_path = GRADIO_OBJECT_PREVIEW_GLB + if is_edit or show_generated_scene_as_edit: + if GRADIO_INITIAL_SCENE_GLB.is_file(): + runtime.scene_model_path = GRADIO_INITIAL_SCENE_GLB + elif initial_scene_path is not None: + runtime.scene_model_path = initial_scene_path + if GRADIO_SCENE_GLB.is_file(): + runtime.edited_scene_model_path = GRADIO_SCENE_GLB + elif GRADIO_INITIAL_SCENE_GLB.is_file(): + runtime.scene_model_path = GRADIO_INITIAL_SCENE_GLB + runtime.edited_scene_model_path = None + elif GRADIO_SCENE_GLB.is_file(): + runtime.scene_model_path = GRADIO_SCENE_GLB + runtime.edited_scene_model_path = None + if simulation_started: + runtime.status += "\nDexsim simulation launched." + if cleanup_errors: + runtime.status += "\nCleanup completed with errors; see Last error." + runtime.last_error = "\n".join(cleanup_errors) + if simulation_error: + runtime.status += ( + "\nDexsim launch failed; Gradio preview is still available." + ) + runtime.last_error = simulation_error + elif process.returncode == 0 and not pipeline_output_ready: + set_runtime_phase_locked("failed") + runtime.status = f"Pipeline ended without {missing_output_name}." + runtime.last_error = runtime.status + archive_outcome = "pipeline_output_missing" + elif scene_build_error: + set_runtime_phase_locked("failed") + runtime.status = f"3D preview failed: {scene_build_error}" + runtime.last_error = scene_build_error + archive_outcome = "preview_failed" + elif promotion_error: + set_runtime_phase_locked("failed") + runtime.status = f"Scene promotion failed: {promotion_error}" + runtime.last_error = promotion_error + archive_outcome = "promotion_failed" + else: + set_runtime_phase_locked("failed") + runtime.status = ( + f"Pipeline failed with return code {process.returncode}." + ) + runtime.last_error = runtime.status + archive_outcome = "pipeline_failed" + if pipeline_succeeded and not promotion_error: + archive_outcome = ( + "dexsim_launch_failed" if simulation_error else "completed" + ) + archive_after_status = ( + run_log_mode == RUN_LOG_MODE_INTERACT and not simulation_started + ) + if archive_after_status: + archive_run_log( + mode=RUN_LOG_MODE_INTERACT, + task_description=task_description or display_task_text, + scene_description=scene_description, + outcome=archive_outcome, + ) + except Exception as exc: + should_archive_exception = False + with runtime_lock: + if runtime.run_token == token: + runtime.is_busy = False + runtime.process = None + set_runtime_phase_locked("failed") + runtime.status = f"Pipeline supervision failed: {exc}" + runtime.last_error = str(exc) + runtime.log_lines.append(runtime.status) + should_archive_exception = run_log_mode == RUN_LOG_MODE_INTERACT + if should_archive_exception: + archive_run_log( + mode=RUN_LOG_MODE_INTERACT, + task_description=task_description or display_task_text, + scene_description=scene_description, + outcome="pipeline_supervision_failed", + ) + + +def launch_current_simulation( + token: str, + *, + parallel_env: bool = False, + robot_profile: str | None = None, + run_log_mode: str = RUN_LOG_MODE_INTERACT, + task_description: str = "", + scene_description: str = "", +) -> str | None: + if not CURRENT_PATHS.fast_gym_config.is_file(): + return f"Dexsim launch skipped; missing {CURRENT_PATHS.fast_gym_config}" + if not CURRENT_PATHS.agent_config.is_file(): + return f"Dexsim launch skipped; missing {CURRENT_PATHS.agent_config}" + + command = build_run_agent_command( + CURRENT_PATHS, + parallel_env=parallel_env, + robot_profile=robot_profile, + ) + started_at_ns = time.time_ns() + try: + process = start_pipeline(command) + except Exception as exc: + return f"Dexsim launch failed: {exc}" + + output_queue: queue.Queue[str] = queue.Queue() + reader = threading.Thread( + target=read_process_output, + args=(process, output_queue), + daemon=True, + ) + monitor = threading.Thread( + target=monitor_simulation, + args=( + token, + process, + output_queue, + reader, + started_at_ns, + run_log_mode, + task_description, + scene_description, + parallel_env, + ), + daemon=True, + ) + + with runtime_lock: + if runtime.run_token != token: + stale = True + else: + stale = False + runtime.sim_process = process + runtime.sim_started = True + runtime.sim_finished = False + runtime.sim_returncode = None + record_simulation_started_locked() + runtime.log_lines.append("$ " + " ".join(command)) + + if stale: + terminate_process_group(process) + return None + + reader.start() + monitor.start() + return None + + +def monitor_simulation( + token: str, + process: subprocess.Popen[str], + output_queue: queue.Queue[str], + reader: threading.Thread, + started_at_ns: int, + run_log_mode: str, + task_description: str, + scene_description: str, + parallel_env: bool, +) -> None: + while process.poll() is None: + append_simulation_logs(token, process, drain_output_queue(output_queue)) + time.sleep(0.5) + + reader.join(timeout=1.0) + append_simulation_logs(token, process, drain_output_queue(output_queue)) + latest_video = latest_audience_output_video(min_mtime_ns=started_at_ns) + latest_dataset = latest_lerobot_dataset(min_mtime_ns=started_at_ns) + lerobot_video = ( + build_lerobot_preview_video(latest_dataset) + if latest_dataset is not None + else None + ) + combined_video = ( + build_single_env_combined_video(latest_video, lerobot_video) + if not parallel_env + else None + ) + display_video = combined_video or latest_video + + should_archive = False + archive_outcome = "completed" + with runtime_lock: + if runtime.run_token != token or runtime.sim_process is not process: + return + record_simulation_finished_locked() + runtime.sim_process = None + runtime.sim_finished = True + runtime.sim_returncode = process.returncode + runtime.video_path = display_video + runtime.lerobot_dataset_path = latest_dataset + runtime.lerobot_video_path = ( + None if combined_video is not None else lerobot_video + ) + if process.returncode == 0: + runtime.status = ( + "Pipeline completed successfully.\nDexsim simulation finished." + ) + if latest_video is None: + runtime.log_lines.append("Audience video not found in outputs.") + if latest_dataset is None: + runtime.log_lines.append( + "LeRobot dataset with recorded frames not found." + ) + elif lerobot_video is None: + runtime.log_lines.append( + f"LeRobot dataset found, but preview was not generated: {latest_dataset}" + ) + elif combined_video is not None: + runtime.log_lines.append( + f"Single-env combined video created: {combined_video}" + ) + else: + runtime.status = ( + "Pipeline completed successfully.\n" + f"Dexsim simulation exited with return code {process.returncode}." + ) + runtime.log_lines.append( + f"Dexsim simulation exited with return code {process.returncode}." + ) + archive_outcome = "simulation_failed" + should_archive = run_log_mode == RUN_LOG_MODE_INTERACT + if should_archive: + archive_run_log( + mode=RUN_LOG_MODE_INTERACT, + task_description=task_description, + scene_description=scene_description, + outcome=archive_outcome, + audience_video=display_video, + ) + + +def append_simulation_logs( + token: str, + process: subprocess.Popen[str], + lines: list[str], +) -> None: + if not lines: + return + with runtime_lock: + if runtime.run_token != token or runtime.sim_process is not process: + return + for line in lines: + runtime.log_lines.append(line) + + +def drain_output_queue(output_queue: queue.Queue[str]) -> list[str]: + lines: list[str] = [] + while True: + try: + lines.append(output_queue.get_nowait()) + except queue.Empty: + return lines + + +def run_reset(): + cleanup_errors = reset_current_scene() + last_error = "\n".join(cleanup_errors) if cleanup_errors else None + status_text = ( + "Reset complete." + if not cleanup_errors + else "Reset completed, but some cleanup failed." + ) + return ( + None, + "", + "", + None, + "", + PHASES["idle"].progress, + format_status(status_text, last_error=last_error), + None, + None, + None, + ) + + +def stop_current_run_without_cleanup(): + process: subprocess.Popen[str] | None = None + sim_process: subprocess.Popen[str] | None = None + with runtime_lock: + runtime.run_token = uuid.uuid4().hex + runtime.auto_loop_active = False + runtime.auto_loop_token = None + runtime.auto_round = 0 + process = runtime.process + sim_process = runtime.sim_process + runtime.process = None + runtime.sim_process = None + runtime.sim_started = False + runtime.sim_finished = False + runtime.sim_returncode = None + runtime.is_busy = False + runtime.phase_key = "idle" + runtime.status = "Stopped." + runtime.task_text = "" + runtime.input_task_text = "" + runtime.input_scene_text = "" + runtime.image_path = None + runtime.video_path = None + runtime.lerobot_video_path = None + runtime.lerobot_dataset_path = None + runtime.object_model_path = None + runtime.scene_model_path = None + runtime.edited_scene_model_path = None + runtime.last_error = None + runtime.log_lines.clear() + + if process is not None: + terminate_process_group(process) + if sim_process is not None: + terminate_process_group(sim_process) + + return ( + None, + "", + "", + None, + "", + PHASES["idle"].progress, + format_status("Stopped."), + None, + None, + None, + ) + + +def run_reset_or_stop(run_mode: str): + if run_mode == TOP_MODE_AUTO: + return stop_current_run_without_cleanup() + return run_reset() + + +def rerun_current_simulation( + run_mode: str | None, + action_mode: str | None, + robot_profile: str | None, +): + def _rerun_outputs(): + return ( + gr.update(), + gr.update(), + gr.update(), + *ui_snapshot(), + ) + + if run_mode != TOP_MODE_INTERACT: + with runtime_lock: + runtime.status = "Rerun 3D is only available in Interact mode." + runtime.last_error = runtime.status + runtime.log_lines.append(runtime.status) + return _rerun_outputs() + + if not rerun_simulation_is_available(): + with runtime_lock: + runtime.status = ( + "Current simulation files are not available. Generate once first." + ) + runtime.last_error = runtime.status + runtime.log_lines.append(runtime.status) + return _rerun_outputs() + + token = uuid.uuid4().hex + with runtime_lock: + if runtime.process is not None: + runtime.status = "Another pipeline run is in progress. Stop it first." + runtime.last_error = runtime.status + runtime.log_lines.append(runtime.status) + return _rerun_outputs() + if runtime.sim_process is not None: + runtime.status = "Another Dexsim process is running. Stop it first." + runtime.last_error = runtime.status + runtime.log_lines.append(runtime.status) + return _rerun_outputs() + if runtime.is_busy: + runtime.status = "Another run is in progress. Stop it first." + runtime.last_error = runtime.status + runtime.log_lines.append(runtime.status) + return _rerun_outputs() + + runtime.run_token = token + runtime.sim_started = False + runtime.sim_finished = False + runtime.sim_returncode = None + runtime.last_error = None + clear_run_timing_locked() + runtime.log_lines.append("Starting Dexsim rerun (run_agent only).") + + simulation_error = launch_current_simulation( + runtime.run_token, + parallel_env=action_mode == TOP_MODE_PARALLEL_ENV, + robot_profile=robot_profile, + run_log_mode=RUN_LOG_MODE_INTERACT, + task_description=runtime.task_text, + scene_description=runtime.input_scene_text, + ) + if simulation_error is not None: + with runtime_lock: + runtime.status = f"Dexsim rerun launch failed: {simulation_error}" + runtime.last_error = simulation_error + runtime.log_lines.append(runtime.status) + + return _rerun_outputs() + + +def randomize_interact_task_input(run_mode: str | None, language: str | None): + """Fill the Interact form with one available template task.""" + if run_mode != TOP_MODE_INTERACT: + return ( + gr.update(), + gr.update(), + gr.update(), + gr.update(), + None, + gr.update(), + None, + None, + ) + auto_input = generate_auto_text_input( + language=language or LANGUAGE_EN, + include_scene=False, + ) + initial_preview = None + if auto_input.prebuilt_scene_dir is not None: + initial_preview = build_interact_random_initial_preview( + auto_input.prebuilt_scene_dir + ).as_posix() + return ( + auto_input.base_image_path.as_posix(), + gr.update(value=auto_input.task_description, interactive=True), + gr.update(), + SCENE_MODE_INITIAL, + ( + auto_input.prebuilt_scene_dir.as_posix() + if auto_input.prebuilt_scene_dir + else None + ), + initial_preview, + None, + None, + ) + + +def randomize_interact_scene_input(run_mode: str | None, language: str | None): + """Fill only the scene text in the Interact form.""" + if run_mode != TOP_MODE_INTERACT: + return gr.update() + scene_description = generate_auto_scene_description( + language=language or LANGUAGE_EN, + ensure_scene=True, + ) + return gr.update(value=scene_description, interactive=True) + + +def clear_interact_prebuilt_scene() -> None: + return None + + +def button_updates( + language: str | None, + run_mode: str | None, + action_mode: str | None, +) -> tuple[Any, Any, Any, Any, Any, Any, Any, Any]: + """Build localized labels while preserving the selected button variants.""" + labels = BUTTON_LABELS.get(language or LANGUAGE_EN, BUTTON_LABELS[LANGUAGE_EN]) + is_auto = run_mode == TOP_MODE_AUTO + is_interact = run_mode != TOP_MODE_AUTO + is_parallel_env = action_mode == TOP_MODE_PARALLEL_ENV + can_rerun = ( + run_mode == TOP_MODE_INTERACT + and rerun_simulation_is_available() + and not runtime.is_busy + and runtime.process is None + and runtime.sim_process is None + ) + return ( + gr.update( + value=labels["auto"], + variant="primary" if is_auto else "secondary", + ), + gr.update( + value=labels["interact"], + variant="primary" if is_interact else "secondary", + ), + gr.update( + value=labels["parallel_env"], + variant="primary" if is_parallel_env else "secondary", + interactive=not is_auto, + ), + gr.update(value=labels["start"] if is_auto else labels["generate"]), + gr.update( + value=labels["rerun_simulation"], + visible=is_interact, + interactive=can_rerun, + ), + gr.update(value=labels["random_input"], visible=is_interact), + gr.update(value=labels["random_scene_input"], visible=is_interact), + gr.update(value=labels["stop"] if is_auto else labels["reset"]), + ) + + +def auto_control_updates( + run_mode: str | None, + action_mode: str | None, +) -> tuple[Any, Any, Any, str | None]: + if run_mode != TOP_MODE_AUTO: + return ( + gr.update(interactive=True), + gr.update(interactive=True), + gr.update(), + action_mode, + ) + + with runtime_lock: + scene_mode = runtime.auto_scene_mode + parallel_env = runtime.auto_parallel_env + robot_profile = runtime.auto_robot_profile + labels = BUTTON_LABELS.get(runtime.language, BUTTON_LABELS[LANGUAGE_EN]) + return ( + gr.update(value=scene_mode, interactive=False), + gr.update(value=robot_profile, interactive=False), + gr.update( + value=labels["parallel_env"], + variant="primary" if parallel_env else "secondary", + interactive=False, + ), + TOP_MODE_PARALLEL_ENV if parallel_env else None, + ) + + +def video_preview_label(language: str | None, action_mode: str | None) -> str: + text = UI_TEXT.get(language or LANGUAGE_EN, UI_TEXT[LANGUAGE_EN]) + key = ( + "parallel_video_preview" + if action_mode == TOP_MODE_PARALLEL_ENV + else "single_video_preview" + ) + return text[key] + + +def scene_mode_choices(language: str | None) -> list[tuple[str, str]]: + text = UI_TEXT.get(language or LANGUAGE_EN, UI_TEXT[LANGUAGE_EN]) + return [ + (text["scene_mode_initial"], SCENE_MODE_INITIAL), + (text["scene_mode_edit"], SCENE_MODE_EDIT), + (text["scene_mode_task_only"], SCENE_MODE_TASK_ONLY), + ] + + +def scene_mode_input_updates(scene_mode: str | None) -> tuple[Any, Any]: + """Set field availability for the selected scene operation.""" + is_task_only = scene_mode == SCENE_MODE_TASK_ONLY + return ( + gr.update(interactive=True), + gr.update(interactive=not is_task_only), + ) + + +def localized_ui_updates( + language: str | None, + action_mode: str | None, +) -> tuple[Any, ...]: + """Return updates for every non-button, user-facing static UI string.""" + text = UI_TEXT.get(language or LANGUAGE_EN, UI_TEXT[LANGUAGE_EN]) + instruction_html = ( + "
{text['instruction']}
" + ) + return ( + gr.update(value=text["heading"]), + gr.update(value=instruction_html), + gr.update(label=text["robot"]), + gr.update(label=text["input_image"]), + gr.update( + label=text["task_description"], + placeholder=text["task_placeholder"], + ), + gr.update( + label=text["scene_description"], + placeholder=text["scene_placeholder"], + ), + gr.update( + label=text["scene_mode"], + choices=scene_mode_choices(language), + ), + gr.update(label=video_preview_label(language, action_mode)), + gr.update(label=text["current_task"]), + gr.update(label=text["progress"]), + gr.update(label=text["initial_preview"]), + gr.update(label=text["edited_preview"]), + gr.update(label=text["object_preview"]), + ) + + +def toggle_language( + language: str | None, + run_mode: str | None, + action_mode: str | None, +): + next_language = LANGUAGE_ZH if language != LANGUAGE_ZH else LANGUAGE_EN + with runtime_lock: + runtime.language = next_language + labels = BUTTON_LABELS[next_language] + return ( + *button_updates(next_language, run_mode, action_mode), + gr.update(value=labels["language"]), + *localized_ui_updates(next_language, action_mode), + next_language, + ) + + +def select_top_mode( + selected_run_mode: str | None, + selected_action_mode: str | None, + current_run_mode: str, + current_action_mode: str | None, + language: str | None, +): + run_mode = selected_run_mode or current_run_mode or TOP_MODE_INTERACT + action_mode = current_action_mode + if selected_action_mode == TOP_MODE_PARALLEL_ENV: + action_mode = ( + None if action_mode == TOP_MODE_PARALLEL_ENV else TOP_MODE_PARALLEL_ENV + ) + elif selected_action_mode: + action_mode = selected_action_mode + if ( + run_mode != current_run_mode + or action_mode != current_action_mode + or run_mode != TOP_MODE_AUTO + ): + stop_auto_loop_if_running() + return ( + *button_updates(language, run_mode, action_mode), + gr.update(label=video_preview_label(language, action_mode)), + run_mode, + action_mode, + ) + + +def ui_snapshot(extra_status: str | None = None): + with runtime_lock: + phase = PHASES.get(runtime.phase_key, PHASES["idle"]) + video_value = None + video_signature = None + if runtime.video_path and runtime.video_path.is_file(): + video_value = runtime.video_path.as_posix() + video_signature = (video_value, runtime.video_path.stat().st_mtime_ns) + if runtime.auto_loop_active: + video_update = video_value + elif video_signature != runtime.last_sent_video_signature: + runtime.last_sent_video_signature = video_signature + video_update = video_value + else: + video_update = gr.update() + object_model_value = ( + runtime.object_model_path.as_posix() + if runtime.object_model_path and runtime.object_model_path.is_file() + else None + ) + model_value = ( + runtime.scene_model_path.as_posix() + if runtime.scene_model_path and runtime.scene_model_path.is_file() + else None + ) + edited_model_value = ( + runtime.edited_scene_model_path.as_posix() + if runtime.edited_scene_model_path + and runtime.edited_scene_model_path.is_file() + else None + ) + task_text = runtime.task_text + status_text = runtime.status + if extra_status: + status_text = f"{status_text}\n{extra_status}" + busy = runtime.is_busy + last_error = runtime.last_error + return ( + video_update, + task_text, + phase.progress, + format_status( + status_text, + phase=phase, + busy=busy, + last_error=last_error, + ), + model_value, + edited_model_value, + object_model_value, + ) + + +def synced_ui_snapshot( + run_mode: str | None = None, + action_mode: str | None = None, + last_seen_input_revision: int | None = None, +): + sync_inputs = False + with runtime_lock: + submitted_input_revision = runtime.submitted_input_revision + sync_inputs = ( + runtime.auto_loop_active + or run_mode == TOP_MODE_AUTO + or submitted_input_revision != (last_seen_input_revision or 0) + ) + image_value = ( + runtime.image_path.as_posix() + if runtime.image_path and runtime.image_path.is_file() + else None + ) + input_task_text = runtime.input_task_text + input_scene_text = runtime.input_scene_text + can_rerun = ( + runtime.process is None + and runtime.sim_process is None + and not runtime.is_busy + and rerun_simulation_is_available() + ) + + if sync_inputs: + input_values = (image_value, input_task_text, input_scene_text) + else: + input_values = (gr.update(), gr.update(), gr.update()) + return ( + *input_values, + *ui_snapshot(), + gr.update( + visible=run_mode == TOP_MODE_INTERACT, + interactive=run_mode == TOP_MODE_INTERACT and can_rerun, + ), + submitted_input_revision, + *auto_control_updates(run_mode, action_mode), + ) + + +def format_status( + status_text: str, + *, + phase: Phase | None = None, + busy: bool = False, + last_error: str | None = None, +) -> str: + if phase is None: + phase = PHASES["idle"] + state = "running" if busy else "ready" + parts = [ + f"**State:** {state}", + f"**Phase:** {phase.progress}% - {phase.label}", + f"**Status:** {status_text}", + ] + if last_error: + escaped_error = last_error.replace("`", "'") + if "\n" in escaped_error: + parts.append(f"**Last error:**\n```text\n{escaped_error}\n```") + else: + parts.append(f"**Last error:** `{escaped_error}`") + return "\n\n".join(parts) diff --git a/embodichain/gen_sim/gradio_ui/assets/dexforce.png b/embodichain/gen_sim/gradio_ui/assets/dexforce.png new file mode 100644 index 0000000000000000000000000000000000000000..6a8e11b0037138c09e4d7efb3ad763573edc7dfe GIT binary patch literal 25276 zcmbSy1yqz@*YC_w(w$Nc-Q6Vu(kUek1JW>bmq>~r4I&^2D$OVi4I?QnD&3$o3`m2( zeeiwX|N6eWZmfG}EtkVH=Q(Geoxi=$J_&j{s)TrScpwmnP+d*g00hEF2Z6xQJJ`Sz z`?Xj&@E}(Ks8L`kHEMNy9wd`EBhz?Hu?6-ET*OK(Y#f za9fzGgDD0{8<&_&vQH1%#xer2ft!%m<9%^9l0swGHI+@L~DS49X5ZFmGqL zud}BI^X-hbcAkE|a!`QQe{I4Y{?D`?KL2JEz%qeATeyG_zu@hr{tmQ<{WA{k=k4}) zaeJ76gPVi9gNLsVFjnZFv2Z6(Ur!$=&;QNT|NQ*l766#5t^Lm$|4Ulj-Tzs_$5+K4 z*v7vZ^1mGIV;BT?5HN7?@$~bCIjHypWU}1u1}?4S?O^Nc>22ug>GmH{>HUXfW?^A| zF=kE^XAgVN03WXZ@`8i1t*?U|6p)%2pP(?Gu$ZBcsI;)Kw6Hjjpro{*;J=1yd)hlc z3Hq-?Mfn7U422}51%;(W{`*isY3yx%ZU3)>?P1bSJiXm*fekym+d4W3z&#wH%>S@a zTFKMR(;HYApibn!K37*#()0Fw;_L?e;A5crfLUEdNnA)uT%1psU+6En+S=0U9zMRd z9xw-WWjQFo9DZkKdub^#K`ALQ5lKEf5eGXy5h+mzJ}C)t2|h^)2|EW{QPC$3Vs`&| zU)dApcdG*bdEfs3@qJxyXTS(;-TuGlxi!sO-H}#v_5q|C^l#I7;I~E0s!#> zw+8#KeF*%|Tm1Xme=+X=36EeWxIN)t%stC2sF(KGJ=iC;u26rp8}5l-01e=TQ*)TODA7^x9})O_e`1 zdSh>IZA_KPVfd9;YM)uNFPFlfW#XpuwY+1rwsJ`#XD$WpjMfqcN+5C^)e(ipzzIq! zko^n-aA|8rogK-t52NuXnA4o*_JsmF>%I0S$03>LDa9o>`zW(Q17PSDS0*4 z9R(FN?XPa56WahZp(u{gZob)%&3YCPS=i}( zUum{gIRgj^HPYxWQS*%hD|C&RTH&^S@s!LgE9J8vG_Kg7ixVVx+~C0wun@!%Fz-iw z9x7jupj@t2q-V4FtyK1M=M!7L-_ZIYuW4;0Afa@C1AR0_@$qLR_6FVgLMU9*iqzT2 z$xFibAM21YZZ0nba(6a1Y?eD?=DY{_+!k7;3e=U$^;}(f@TvJrT3Q;FhHYs+djHbk z@A%`*7`#!G^Ol^Hv@k;}v3u2Nsb=bYAOM8xr(0=ZF7rSfu(Ix#T5E?vfC+qdZ0 znpl-Rm4QOV?gNYyCW2f;=I~p8cyexL48~NNnY?M~+ijempXYu2@!740uxVyBGFX&6 ztbF__wp25lsO4ME;IGE@f)Gh`>!E2!c-Ygig!p&?4R)0@Kg@mReIA>vbyEX2>qx)k z+307sU*2iwHD2zuGy(s{c`Zw{c6WDWkG?7$9T9{zIG+FhWf8ntDtpjPe$l`VHDCdF zb#_3yEaR(Vpj;9Q7{P&~=)Tg@y(PU-L`mnxFxJe>e5KP-mBvEE%~d%k7nwdy ztnBfRN0k;$-tWM}n%H7m#>CkJ2XYOHg#7@8OCR$gy-*eyf zn2A2h&uj%YA8Vmanz~X2n2BRS(37~#zEJtm&KUfrqi>ox+M+7TCCz}8)Z=*fDVJ=Q z?gLKed$U+c>9Dr&^tW%FO2l&AfG3~{6C^GA^eJ7HDa=w{w^r2CLpp^fwoi9oG%t zI{dbct2?fPJz5VQ*3XfWmgxh6`SzZ)@lGRF&W1#Wy`Nv#h4#HBi#oTDRaGyr@TuT8 zSEn+GTFRY*_R}97Dy`aHJt#um`ov$^)+4;YRhf$`(F1jb}-z)W?(;L%A#uRr)dd{Q&5vP zju~K#y0QnD>FnAo?-j1NMr99^-ja~}NbkRkhlPUlf}9RXm;6YpZ*cZCRS{}_GW2Jx z`}j-pqX;#H+kLp}E`gUGp=#v%$XXtN#p`2ziRTAAlxLCdtNk2>Kd&#*CeNSxUO?cs z9XF%4>c3E`gCy3Acjj^}8jY4+r>HIjU=h_3pZ(?kU`<&J+|SR?mj!~(?Ve{nJMu&2 z(bC3|2oy(T2lzQd$`xd3oV4%5<^$)e9SjS_^9aUWRKt7>i%}(R4$cnFYmjU-TV@A@ zyXg0%Ix6YT`^->EYKXv&9BLE|!Gi{JGo7Sg^{&XQkAWy>c zU0d6HSlwEnTZICOo40Hyq%zsnfJ+_pC&?+{SR!(L?iq{GUYsd(`aQbs=OxNIsvPQwA3u_co{@wA2 z51lq>!KCx>w9AZQ7-s?|QM1q&IhmK7oGdvT*@`vgnuA8C7`3$1&VCFbhO--mr}H^H z0j;Yxw&Tb&7x~-~PSc;#{$hvZ)&}WdGj@7NhtcLRn%n=rRWR<+irQkVM`1W*aal;n zZhmQm9%0u@?L2V=5Ev&a_1%*)rZHxdDSEJ`Tw*EEEEyeSvZRtkEvpf=pW+{$^pF+g zVOd2{{?PYK`}rrwNkVk=v;5BUq@Qj!XyNKvrk3RcS>%fwhWl+TwIvFGgUQ^q*_oT- zZ^!U*ADnnOQFX6$Vx<3ErzK2DZGPi9AlwZ>Nyn|K3mExOmbr!4R$!~@ zu?eKfHNDiXx_2I`7G`;XW~r*G`m?TZb1;2Bj2~W6)33QNJ2t|fWki>2O;BFSUh%0-6=Let7|yCO75{1BU%+;p$5B5b>=Uw46w?^X5b9`Sj@)WW~Dc*{KVBzXRTw(8PPtl%`Hk57;m8cyDa_xu2ZW5bt zO?YsAeblKLp{k0bnZ;|{lSp}h#a22+c&6Vi#1G!{C(*2OM$y@;`tfId}XriB?Alz`5p#^gzO!R^a{<2Yn=b6`v4 zh+342FExFi;Wexf*+<`ZPbNfY-)K=$P5X>L3)QFxl7H5gm1!riB)zoAY97j~{t01h z$rHHt~jkWKDIZP;S%kykz07R|v*| z@6LuU;vHnMS6<=RJxrdmk>RYOt1%R}man3)$E-Btx97!OltCqS=NkZ-7`|t+(=OsE zYTnQ>Yl!uAG%3-e~+4gnnBl9=w)T4WW^^AbO(~havD>ve#gT*Un(T z8Rs`I{L}V*dye;@xC^2O+AT`J*og)@;O5AxVxh1`bEgnO?|iF-r}RYCvaSax40wnm-~ zgvKsF1wS6)w3}!8oJW#0cn1_sDAeUQWCGH z;#|A)6Uek4w}&>3UZ+8w;do;y)vQh2)A8VasghwN0X_NBh*dh* z*q{-fsLRN0uCUU*(8~7vbi2uHzjz`mqFC$-wx3|Lmfh;|B1bJ=3guh9NU^&Qb=dJ^ zAbfm$fH;M7QOe((p<{OkING?+!5WJ*RCYZtS>?YYh)1Aa(svqU2;`%TY#0W3YN=@A zvZcaVj4Al}b{`0v9$)9xK7Iav)uUox&FU$`+Y4(_Ccw|vFXIhtudb=kbOP8bhBU2T z5WsS_7_Jip0U`z^TlquDSDyz`f0namOcMPx-bqUDUMI6jLo=6@Yxkrvo6LR~g@ns= z`t0=>v%U=^_24PYhLP<%04Q0!GUZ*0#U8`aHq9F#l8S3}BNYsX?)~=5!~6MgLo9c6 z{Wkx}6VLpfZd$@EO5V1w^Q|JSrqh^RqtvKBu#F+G@OBk3Py}iuumz5>pARj(e9XKR z7xO5}M3w$6{iR(@WI>-l(ax`Xu~7?3p7a3QT>A7)KR-_dGD@|8(f!fQy}c(uGFn$( z&qki1b`eUEL*{-Th4b@dXi2PIc2!USd;>AE9 zFTJ-4O@35PY-b+yU5{creCOSFtEm@i0N2FUENrM36)9kcSrE(lt`1dSKR*PJT?Gd} z{35MV{s{X_qgyCT@T}ZkWdB<{jQ9!BuVb;cXDbS#8k^y2b*HgB2nInGImQ$!b{ih8 zo*|>5x8sVk2CwnqCLOj6$g~Bk6~7Ml;FOv9$U$hbzVmOC-&F2*mdM5LXl(l*1D;TW z$Gvs|h2JMi!a4ZNA0X@Ev>4WobWBXBw*)*f>T~+f!*0TfY)$1&IDr)oSdJLzBRde3!;(T4PdEHnE8S#-^kh-7V zbr33<`P-kB;2n-D8qqqmtq{HYgYNFV*c6*}`5O}3B#m0y7JXk@V?KYqTmFyEBCNC| z&WP+K^DXZ&};Odx(0Tc zxQ>rtAY2DEz!H9G6D$t8bG}6^IF@d8a&tAq?hyuq<&_%j_ zOV;%S2DKDDT8mpI_Yc87OW*1|`SL|3YBIjP>JqPVJpOr}oAaS(W!0P?#ev!TTRpS{ z0SJ3d=4FwuM>I0|rQkE0_!KGg1`mhT-WUDYR4MC`DcSpf0BG`fI~i;3A^qxIl+EA; zy6+<4I2)Y28CD0Hv`hHap^<}ri+V-|!*or{e>YFz8Sgw7B=kb&wI}!zfIP%g!&s@G zL$Mniy##u-HB$@0eWXijJkCkqr{#LQmxp|-r1IqSF<7ii7Qtz7k3!R(!i2%z`1jKkG^aE?TK`Uz_;UAsq`{MoRsSV9;4 zw4_2C3Fm=&DoHjwf{sq4RGU;MjS1uu!@y4`r^kt&nW0g;6*zn*Z7J{OBz7oe;qChva z&|@rj#cz9-%ATTOWJ|?6d+!lOH3tq>9e4fP(VvLy4wg-et2AN>=5_4($9OPruuKJ) zTvyeXhqcX7RRYd}%slp_ z2%z%j_y2UJwG+D{P&*;6OD$6-uclcQMAwKJ$Y30{PqE}x2rj)(CBZ|cLy1>g&%&zE zJzx*uQ`gCAZSfB?BDov0<`L0tO-;3F8$7Q$FlrJQHC8MP;cZz(>9IBBRg=u9Q|$8@ z+>NOyHNM8N01y|P!ULa}(iLpX7ea-MUMETvDF5EkV@PgXcF9$Fe^iJfm4`bw!wgAO z(RmEr(fQ^c-gU6WQsRaWCpoIpA_|YL0lXxT%dycL4L4!tFfL!zH+td-5r?EdU*WC< z8L>%aIhkJu*Wy)2Pv1SWd4~8b-zp1mhxUi`gA(Jp4VlOc$?W{gqh4EbI2V#GR1ICD zIS&>k%8~mriNTxNN_j@?Pag0E5(v2HA%)rq$v9BNG5vF5!jc^&Ynqt>`+!lKMSN#h zSJ>g&a6gW*%mO%+bA_hx^YF<;lsutZ0{n`8hp@xIKHe(mP&ih8&__C&TARN|Dmrx6 zbZ68|xqK0ZiU`9=g}g-^p6dq$6R>xi3AQnr$NYheAwP^VoOA6ETHLAZ^rvI!iCoMB zawG;m)J0KSE`_n=i)lx-Y5%&Dk(YY~pF^}gfI=Zq`u!PQPMD6kexz?fJh#N0^|=zW zoy;pK9U;-9v9DiQ2x~HnVj?(bn-A_wS>eYPc=q0(8Cw7A;Uc)C`)=(PvbqvA(t?g1 zzdB=(IsMJ2D%5XZ_|Oe!P@#6T&k=^T3qWRqWdYNgx=#sHL-kL!I?Gw!$Fil%L|DqP z!vciHkPRNp*FG!?92fjgc@S1G=ZX7Ab}Z|HM|cODrbg?W*qVmcNwZcy(6n%gb0TIH zepLBfGn~?x^B_od^;qi7(Dn&xWLkkJ6et0I`HBQDVkTo|0Todpjx(Y+9m($ZW}E%3 zvUXhYzs=9@-YYj1wV4G!4;xAp%4z^H0b#q}bg^7AW@^;c{F%@R`_+qC-QK;2*{)dr zCw_t=3E`)I6tQ);>bbcLLiS3NCdiCjWcME}Dkj-P;J|qGTuD+>=Tq`a|4{B5k0QEv z;g!t#+`;ZZeGRGjyT^jR4@gAI%ldxetE+mJ)a7)5{(u{2 z)s;Q8Z8nOfdM>DxB;HtHh?wIFD?tsBt9Z}OJ4O{vDn1l7nFpZ;Mr^uGO6Z_q4oHT zOuhc9Nbg0VDm(O<<`7&5WDLzr7<^O6Ntf1pB9s;0=K-T$F9)Za`lDuMXQvy5l9S2B zTp?o(4GYgQ2~U?FFS(jtKKBj-HzwzYuw8r<7;AXb>wGVodNz>wpkR|S)pxtwpSSgU z0P|Iyoyr$+!rxPpL#ON*qc&GG&&MaT#~jeVIoH?XehcHkhHwC|`s+gM1npR`F1hks z(EzCHx>Vd$wmfypzSlc%}oJuXQJo1Keaz&X$G;7 zI-Q#lnZZXDE9KVa)WEcU@AS^l+{3D_k;gY5;@3Sk!}i!v%zhgzQg|UL-a$wX7vZLu z1;%W^*$kX&x^eb zr)+Q~hTQ47OEqq;q!oi5&K`NCPzB6za43eUOt+62v{#TY5N9t~0AHABeG(tx(h^5O8Q-201Wn|$z>P)DMJpv89DsJv)ZV39SW{ zAoHKPcgRszd<%x*PF_#y6MRBMJiXQBMat|9RYp>o7THz25UVO!EXE(K*ANky&tkZ4Mq%b>M8`B3rbP+oN}=eCpPer$X0dn-yy18&XE<!F(IEc&Nq=Hnr#ALz=Cj2~X=!QG29Nx{58MUYiHYNTu2Uy| zg|n77&u;9058Z3OrrQxSahbhmnwpjifW=I3dyG-&6^TbhMPkNkilM4-5YN}K8+;A+ z(it&WkJm!V?&BMxlh3?53fS^EJCujA>Qlt757DKKdL;LYKrv6Nv|?@5%+2$HwJyG6 ziy8#sI`jNmER~`kQJGpgng_URvBbpXY9K0dFUtctLCVcaCSj1$K;XOantA5c&}A| zv8BI<^{X{`u3TbRJ8q0OHS%>D8?Obvx|vl94b$Cmx8*o8r3h1MVC_W#^8 zMKDlbW4BT};k3u?;@ydDE0$c1+OA!ZwSLnm;sLF8$E*g)qh(&+p%PJZ{oo>Xhua{r z!8U+JT_C|Ca=}rx|5tC3_c=8{on#O2D>4TB~{X1Ec5(} zpl+eiAUsKy6`I9vbVp2{>~ILa|AeDpVc&3ESeY#0l*aQ4<=>7dp{Z>nV*kUnW7AS! z6Zxoo1^Mpf9BAnoqks5EgZx-ZvKgSFU1mg%AXc|dbpAsR_tJuK?A~7lANhfCyIumR zeMO)B=14CRFROx(r@^}a%9J0(m;?`FI9YC5P!}!TF4S6!9G}ug$kbljDs&oq%O_KU zLuieYcoSb1@iZp);JK>p<|b7^!-qbbV!2fusqa{DuIp4@MGfW)UE1+^v?qByQ~NEh z*W{j24ZfD0k^qX0F7sHXqivt;XhIlI?zMkA?F-Rn?-KlDEqq&6LfBy`l6)AoKd?%0 z-%VWMakF-gAkbvx?l*SpIVCrR81ISCQ^R;N9HkS=MQ<>z?o{?6HDCOMDJwLDi79!u z{gE{OZS69MUwaq9?*ehdl1Gs+>4mtX@V4>^UZ>=peve&1EW}^Q=m){oza39as}~w| zYSxpqEQ@dgegwH|!+;#_E!9RbefZO!9s-M@ln&EcxEZ>|J3J=UhS0~Y`Y0pi77PWZ z#*d3P9l{#-3*-{6LT^Aq%#Dtk96vq{fYCMS#kAA>Cx*fH+ z`tEq#XDT&aMMYMk8*g@=&;ozRPf%fscA#BkYd*P`f~%Ky>P2eJ8Bxid6gvs4nYds8 zgpnS?mb{j%`c|zpe{rPK@)jLg(L=KbZ%bB-wuWhFa`zBdAfh5(eLVF+^sfCv&V1-1 zJIYR5xb}N*r{Y?CCR7_5k7iJ+E95DY<$z=_ulcM-Pr1O#c%D9w%(jt9jJ06nQlv)y zg4;lw2)8S^$+Ah#xc4#PxSR94K7=Qp?q8J20L*>VZbaaW7c##F5AoE(5+Uf1JSbFd$e|fj}TDz4Y#FK|R6|)LRLwtqdQWKhe znH_f8`8?C~i(t@g`$kp&Ph33iA&)H8W39qNb%ZcsyeK`(V~^hK_SlmI&Rgvr znR@punrB@0MNgL9lHMMK!Vijyd8gg1@4k!#B5|o$Zl!VvXf+H65L}mh?YnuL$VnA? z7#F-pa8UtVvBNvQeKkdO&W;0<1GiGP1W`3y)*8L$sr}>_og0^rkv{M-YlY`=_7e~Z66ok8_^l1b2tQK1Owxn|YyC_jpgO4nGKRLzW3>7dRf6% zLa97xJG{P`u%V^rqJ^}L_#qH}?BVEMOJUr7rVOpD(wbMUx&g$a3OB7ag38}^cqASxzy7aAPu7)qAC@ZYqz$#a2 zYACNNY74WsYX=4cU=r~>wsBF0@6+Q?=;R71m~w%Z4dKRY+ZP>{VK#xiW0LJ!zs0;P z8J^wq4G8#T*Mnzm8(%W6h;n2|I&_G&uIGBcKN31uM-faf6R_SXt0A)}o1>dPG%;Kf zV>+8m4Zc*SS^v~{p6mN7Gvl#o31seBOT=BW!t`4ADgB`U5>?(GbjtSdW0WmXU&Rsc z2O?d^(vNL520%XqY}d@1@{@`I7bsk#?kRyyj3y(^+GU)9S^bJUbtdJDvpTHV+|}S; zM#|+G#_TGwNG5ppTf~G6UElkMb4NHl>Xw%GY0S)D4(Ca5$c2>#5hD%6UYDBoN0fzf zV8j6@1LQoC;;#j=g!Y7Y)_05?EO^4@uwUP?*p(q6RW}6DNArm<80@G3S{}Rh(Meg% zUNHe}e+OcoO8}Xwb0&9pFXOe(V()R|Yr$npAGeWMqQ_bNb7y-}A_G)WayC!a*G9}E z&z`wSSIIo;7>WahA}TNz@Lwk%XYF=-V9HzU^gx9h$5<3fos^Enk?JdJP<&JXqR^KU8y&#tOvsNC}Ns3l@cCdUqt&<8R*_ zE7;y#o4gQ>9|1hYu!HbstGXNC9f;EhVoso2&J^)L4yfo*M@mT2tZ!#n<*;JxB<2^% z5(PewmSoLuMpPZ%#W`QN!tI0BZ^Ndvt0WogXQ=EUNW632eR6*-*2|%8WOFi|vmoGz z)La}O6>DakKLhug)m&BK3RcSro z2g>l6G||T+uD*g3T~qm+)2PaYJ~tb>(p@vNRc0#nhFOHjr$X=ok*sp4_D_}WwzjAm zn~GL#HmggI%It%N*uN1?y^8q`z=n;zt(B(r0v3We*Juud<7$Z2hErua{R=rD zM=a!`g{q9A2KySq3tG@Uvg!ui^pRA30K5gjCXugGBwxRM%+<~#<74GXWA{WE$9{9_ z)`);Y$yz|OvEd4WOo8dE2X_$@b}htzVv$Y2TjLHz?Tq5?b84!pacFq?uoS9QxLPL5 zX}vTHY@G-X`=(LlPZuY;s1CFVWuUh*-0 zp3+}n7kC%hvMCC5cR;x&e%zt;t_%$?MeB9%B#Tw1tawmI{dkk7ZY^B)l$WfBZb@rV zF;EoDmTu5$87Y9l0eLIe2iYgHk^Nc@k%gw2y*LIF9%BI!`5lsZlf=)+Tdekn3bWXT zCrW%#k>Jf?8WegULlj7*$sWd@c5h#Ghy9Cdl(C+e39u=lw~)?GapP)0&>F46E5|p! zWGe1kn0xyS-5|jJl~)?5_0V3|jL1`|aEbkz>%|}9z-y(E>zOS72De?Az{uf-?+MxX=qr!Wo3mFf~MJVcZ1RbQwjgV6&8s zxg$@W#QP?Z2jxRu7n(h`WP%baG|~dJ!+h-+G|Xkd?hIO}$QGw(fWs!K96wYk%HjT? zny&Hf3}PNy_gx1c0u7g;aZots1O?EByX1#*oqM)AXNiK`miol@1&z;`34!_&)IhEP zGJ;cG^i1Re*cNN%Q`Cr7wROisIri`Owx}ivPS`Vzt)W-_tW#U|z^Q^PZ~#PFhJ6^t z1a7cX@z*G}H&8a{Pg6Io_l^O?q|Dcz?w|Qh@LHF3QpZ+o-+G6BPq^~c(xzUCYx7!Dlxfv&&DwxHQfL@!f;Ym-UGa}Al%c}b%jJ~eekzq+42i!N00NL4 zi8l=U+rJ`DvNd1(cZ~#1mSjVS`vKo zIe-K*MOBg)m7W;r!6Q@6esXgxY~xA~PNZ;#?h7p{#{Y=Ffn?k8$I^l*t(pGdMHD^q4Jum!R!qV-mD{KOw$k&I2W5M;G!$YxP}!=n`}%>@7(j z=3lODebg)5EU8yF! z6mEO88j5NW{N*LmE+)ROHX_kS)$EIvMJH}cwt8%!w|mBmA0Dg)BSsGUNrLxh7F2oQ zRcZZO9|_KY7KHLIKJ6GE93HVhtCRhdD#+XXCa2Mary|Wp?ry0sAqqc|8j^ycdq$|L zT(af%+-oNvS(RlYyH5o~^-@aU+=c|Ij|JAE%ZBwlVM{N#BF)*5iHH7qx}jZ_sYO(S zt4Vi?vu4-7YGLEaJbU?C7wBS8E-?@V`#_Tsf3;qR1a&4PH8W!t8fUPLBt@wOc!~$=o@8aXfy4ZzZT5;4n5`fsm!Iw_{md{?%&8nDS>hx+eJV&5G zvQChPD)evDMl?rkE~R>#+<*gw9!EVWRA?HLH_3TtBB$5ST}x|~k?P>0FY5Je7|K^C zhj=$t`TCWmv?g(Gm9wHA;=sB$pQA2oOT1dwO+%;d%2=;;YBD4pxUaN;5x8eoTCnG8 z4vDKzmiaRlVBT+Vz_&4aODm72w(8Rh3;uPPjK<+7eIiL==8sAPTO?^gK^t+a4`o)$zc`dwiC9C5J&udQ6~q$eX-}o^sg#SR8_^l;&g;~ z_U+MW{9Xg0y#zWo&(^#Vm%lKV_?KSAuMUC2HIL}` z!OY#C??a(F;|>_?NraETHBSd{_DJGMP6d#NZN&gi+of9H*};G{=oq!v9{uu_D9%}R zv=^R=JUR%+2W&yuigyO~Z*Qg$L112p&QTnlhb#)*vZ7p6tg~?}2#IvRN1(L0tmfPc zdE5sRx2aNiXDcUajoC^U$cH@ci=j;ZdnLoUb3wt1pq1!M(j)m5MN5neBD(1!wEXk{ zbDzc--0b*reH=F4!C^ORcIr=qfpQP)xCIMk!v2Of{XZuXr`EV-n(u~usutA@#bJg* zQ+XlpF!Hb3#lf28#V1fk4fE7dbCRyP%TxJK)2QPd%v%x%wCI?=>P;${np-hZ7EVdivTZ@{FBJQQpObeg9iTKMGtwOSw4tIBV zHQ3om8-R%8@Ec{os~!r`q%TAnGswvHvO&|>%{i+AXJyM-axD})R%C7+i2>^!n~R?W zwzhXIEiK!-%CLf#I@+W6op|8sfT{P};;z{7dRYveQ1*QwM5-v4&v&uO-*`@)hfc=8 z)BNw9OfUak-o;*6(9bS5wE9thi8-HX+Sk)FT{H|g+9axiQla$BB{Z}A?q$oKCKVI< z>gIjfZbt}dD(AKS${{{V0>PM)xz#asxkhxm5tknh81J+Qz8t%~Z83k&!ZrYNoGIug zi)p`+kI9SeAAiCCN1~bOx9kB+2A)N=ov+6==_o5}1Sru?{$Z2^2S%@&rs;u_agc<6 zVsCBv#oEIhA+|T~E7EVR@O7PwZQpoLyS0GgEctUvNFnN)mq%M>c2q<@>S|sgd0veZ z8RcKvL>*XxW+pSJiim@b|FqpXZZHf(%MpV`@*h%kc!B~5PY7>A4BcuUwpPJQ+b*B{ z`gY;RHtR;c?cMqnNn>GHjpG3%Y-hA^DWtX^)=5-J*3S<^-PX$G--)3XEo73|nnS%U zxE%iLpmm~S?Z0yY;-`_2{@FA_6ddfbB549Z4Sx9;wiMmBM!k|0k&hxZ<@)r?JI;vi zD=UL4?1N$5LX22e$@4#6lY3Q4`e9gidpsdhzYUtbb?RW)Lq$hbob%9e`a`4lj$-d$7Is+=4 zxs|LHQE#J+i2CN9<`=bMN@4`cy!4X{m)gS`>GC+D@GE)nmr4&25z&db%nYq-g(WRs zXNz443R)Xl;5GvX3cCZiWTb}GLU_jNLEe8&1wwZw{%&OoB!{yKO20e5*v8z$zD{I& zshcewc`M|YX5OKkw0r5eK=S|#0Z(zS8XKionSUMvFXYt5k$~yl3Ent>IPOX<`7a6c z=4y4ixb<8>nK5>7hp?W2zt=vM(>Laqi4w20>M{TnUZ?xy_ErchCEZC4R`w6rS&qyf|nZk=d7NI-3#_ zQn;!Kc0I4}h7qXTW>jT@PYh9Y>a@a&Y9S#ZrIlI($u3Pfz@4w1onRZ=XWNq$=dD0H zLlvXuK>533Yr5NzW|$)^Wb^GT5;!qOy3YbBG1OxS2f*I(J`LcS1j*P@qe9WfqzyiH z?q%7t!!}|jm*T!0&%g9P(HRC7ykGQk5=a$Nea)2)cpeHjCQfy73X@Tn|UuU3b!z(ek5fT#{U1~6UmA;{tS z_p0Nqw~xQ+R4KA1T9)k;7HI>}Qt{4?G)vcp{%m@GsMS@>-x1YmE&&{K+72nViVRb{FJ#HZrDaj%B1WY%N#f1SGh3Vbq^0X+-M97Ywin6a5ehRW;m zXCL0p02CoU=@vQ#8d)<5y3U+r7X$mki+S1?e@lA_%`NQ}0PXBA8nTN8>n520#UJ(C z3kfUZKN4nVpMogxusQ{gXi{${qX15F;{(N9&0_V3v$U`k|Ghlrt(pgtW7hWZ18+RD zRg*9LH&buHLA-<_sAYIa&!k~h+ys~y_-q0t3FX%QC%m~sCn0PDzr2ClD8OTbe+?%* z2Dgk%Exj-D+#n;$b`4)agt!laOB*Qlg59`i05up z{`mz({4KDR56T9&{1FEF=tu=V(oXsHCc7(Ur~wy3;|+;!d848fhN}J^;DKNK?4OaT z-B{ox5274#sVE z6Pz0rV*PzvdI=7Q{s&$i=mvaE_=nPTON2R<>;-sXRVJS13g49eke+raa-pvPWxq38 zFTWN9mX54KNiR9?_-&kHu7!aGS)zi{Y8gZ37l58iMf)fWL}8T_iZ$(WdS>R*^!b&W z9Hdg+X5rOs4oXIiDd=j3`tUA4z5n9kLc=WoHF#%F=8me>IC}C1yBZXZ<=ZPbd4Otw z(MV{^6V3^B`m|%c&b-eXwg%A!S~VXmWo#5tl;TX`n+kOvz5{i{E6{tkE6euVV+XSA zlP=NN*b(?-Y|}2Kpac@GMPxF*N2T)Df<=(yr`|S|>!#c)ROWtH#QRat)> zt0<*b;gm9~*RdbCMQaw($2t|%-`p(CmcHpj;6G7fARlU+h94RFz-L#nPNbd5n^KL` z1BZ%D(wYyrCr6SQG$I4i=|(PfO!@F;+VjHs_Sltrxyj}8 zDNs#~NXc~bTq>O6ZlXXFxNTRs(k_wh@cNNq`K|nDaJey9j54IV4fl<={)HGZ*j`i= z)hx}hu2aTks!!`4%*(M(P1NsB&zRyeCUkl`}OIXpKs-IQ;iu4jDQg?Ot#L_){p(U-z%J z;-8yk>wCB)$G*M2?d(KTGE`C0pI7Vhaz!T%*RagovJb;Ef#szlFHU&jSBvpIm8k*! zdOG4iL{3KxYFN2Rc}T4q(1^dFj+u7f$SZDL?x)u;m2W8(dH=+TZe_Ov-~nJ~3c>p1 zBP!!)i{=++1K){{h!R~bzpcyNHwXyi2YKTT8x0E^_R;pI(nCUf0!S#=l#8o$#BCJJ z$J=8zj8N+B(I5U6Q1qt>h6+-Pio!SCt2NkpIti8n7ZnXtU;Z4?Nyc8V>O36en|V|| zj4}dd+?ChZr64PA)F(z*3v#3xP38M}Q8=$ANL z6RxSKSJ=iiG>IxQ1$<|!ve=I<0HS>+`DIRb_pY0}hd z(umWv5T*Q;J(0=*qV*;@YnXMr?Gyq6K$+F91!x=SH@wMOP`ic9#l0_Z0q-i}l%gJa zaS>egKT5jpK&an8e&>#(bQ#%OheWoFvxIpLILlk#l#Gk3UmQ*f~Mr-;gJz(XCu+d zE6o+JfLO581sa&#Z*e~3MF7-F_DL-#+yb_VbnL~-27m?l3na@H zxDN1BHqI7=Fnj;dP@8pEhYGV+0#CZd?Vo%0WQlChPyoYsMU*ODRk0Y z<0!i4C>sw+32n*P>-(;|e%0*k2b<;n`g--Myh!MX&1c2q^!|XB+Q{V%`8%qydwdp* z1%m&1x31V8WP^-S;L!m*SFbn>f0`$%f&sn3__oyaZOtiD(5@QJi*S2;#B|)E3^DKF zJizCT-W3uvW-_;6gT99b!M3ld1gV%*|2X{8Krp;Ft0Iy6sk588oC^*7w*jZnR)sgu zrYgBZFslKvDNk-57o3#1q;CFih8`$+X(_}+v)pA;Q~}3IsOZOx4!aaNC#KY#&M%_` zpD?VqGWU)6o!sY(9DbK&PS>g9pJdUv=~-vqW7;PYf~;MI+}>A%CaI~P_9RBqV#@rA zVZqWh1#FPEWUCwM-N156+X?GvWd67o6WI(@8bjq;pK|4d6AEtcWQ}T}l(5>=-6=BL zUA^Pw*bax0QC;U*b)ChKA0(^zNBpTHaXS&j7v}LJgkn?Q#3%o3$*0%97U(@J%Rsfb z#1#tzL>V;I>XEBJBZ2`XgKoyKvq?P((2&BM{*4iCs`gmOFD)H;;~|LFvE^EJGkx2t zrt9KiiLl5Zsm%9ol-r3!KHLh*RQfdvy$^(IUnOpKup7`&+;)oYf@~D8JVtop;qT>v z6?swH<8ap>k>+zIV-DmDx^7?5$v{gdi*v*dSd23J`^kUFlZASR8;MZ@_LOWnqTSAV z>}L~jCbYpgX>4o$o1*XzC>u(u?`b)&=Oe_t+wV!_k=&wE^Os^m>%{zn*UqD`qb^?t zI%f~5P`6u8{>x9|dPpCfIMTST2d?Me$u$}=9|WljCP}O$ z+s0}A=Wm8bsMnNLH7#e_#S(gYdN65nqB-qa2Bjf(CfjOfxtVvkpM2OPgwyM{6yynl zZ|G{jO&0TgSC4eSbi{J}X-D0>c@vW=BXpdtb+Ju|L`+%oJ5L;uzcps~M^yu~@>Awj z*cml7RIXoOi*Y`cqAZ%7Lf`}82a};}{6+X0(fsTYP{<)scQQ=0w3rOQtuA=*h=wx& z!PwXdUn&fBmie*A3bS5Y#)&dV-avA3ais!ln{Y|KdIkMV%ReJVeo4ainR?QC^UZ6w z85wl}CCb!dI|nt2MRDTZPaa4^vDex5IjZaHGo6Os4NOmC^!4?-<06cr9C2JY-#^Ga zjQTbY4q(o2RqX^jQr6KBgy3#nnzE?1k(WnbF<`hl6GaP`FMh1hn-Tu2=$^j84@MkZ z0!Q94KI98wbfy04DT$;?QB%_N!g`1$q+&uv4d3O{#!6itl6Nm-6^&y8zH1dt%gE>fl5iav zt2Hz(Tu>-sf*v49u|2H+-F?4z`dPOajGT)Z|y3*HH$R9WjTp zCT#EvPO|@)#|T@mG}E^%zPQ35<+ElAI?k?HW~@Aqpue5%89CtsE!pIG9vw_xiV}%6SfQ$eX zp!V6Vz-ng3pBv}qpu>W$UMvxFT+a|>(#H#I3YXzgJjRm~Z6IvImr9vv zPK|`#gUoM>Cn~YK+-(|FR56Fht6)^jm+0@ORIK_Kmb3vE(kuc|n+K0@FQVGm^(Y1R zu63}KOVwWvjO~z=b?mk4h!?tZHFnZdJJ`!Ik!%rSM%Z;f0zyuj!bInaCMBK&zu+Cq z&}Y`n(ENQKT;A*!U2i1=FtCWH@@1s})AEyBfncvAhFK=C9~(kLX&m}8dw8F&Gjgl>Z+KqU?6d9;`!`q%=%dt&b(Kt+{Z5_QuLTphCt zhw{0=!-6@W{AZ;K+TqbsU;2a8$om=%wOT~e#@T|49?_;oLIW_j{mlh7#FdRsyVM%; zb2{d6K)|x2eOwZyv~S0K-O(D*ofWTd$w*CxTzT~ZyY~vQMl786+X4% z-!2eV1_qjJFXzU+?20rh;vOyh9&py{5lmuBs#J$F1NPeNyH{J>(V68b=1l#9oO4D#Y}&b#O%RKgB^ zmv@mxH&=MjTtlSUN<`~5h2QlJq`?4F0 zzf+GK?;4sGG^kUDh=IXRBb4NeON44;kLNiy2;u(*$ppWBAB=?3$70Wi@%XTNEp|t! zB$dpx&Kt_&Pp}gD|u||x*ng`a9({Rf044SH1b`-M)_RGsS z>-pvgZZHK(?(*eYj@xb&e4u89owb?by~bqxi!bWhn>45eHTMEK_PIsKfbqkn?`hn`&g`IefTLjv2|GL^FriaRu>0rW%6&6v!*p_~=5!-Zz>RWLM{wifhjfMHVDRNBk)B}gZ^C? zr&Br~`b6!t1-*ET&6?hCmnduljd4D2;$KXpi!ZxR{gRclWE{MjpKxXMY3 z(~E}l^=(U?l2(k!Z6slE$bOyzo5v}wYHWT=(!XSwIP~OE3Ce!r2K7MP)S`3+DC90L zxTQaocy<1j;VWo?QV)R4GOa_Nk-Mb(?pPm&O{r`KN%!52aV=*nFn+~+ zKj(F7zUN02K%RAeO~3=r7j$s0z_>@|1pWhQotu9GoJnyevhvchm(bSlrqzvH)3>g1 z_dd~2nee_E3J+2Up{0y4V1351TFLQQ>Qj>LtJ4uLu3iDA6lrGqEu;F3ibzeopZapJ zag?#VPl3G$>pX{kutM}(B1jbrE(O5?Q@=k4LMF2?M`5U|SJJ&L*ALb=UuVLOW!lF| z!=G5;e4$~&QZd!vR=8&`TwgG%MdO5|2`6UORqSoPPEjG!oIL+8WDL;CIDHXSCoIg4 zr+)AOu#E68czt21zCMmR)qYn}1zq<^2z-%2!Rotp6hx8Ks5rt84`f<6h&Jk7Il&`a z`-xBVfnORXNXB)CHJy}GsjQnYhVb*#Qpu#I`Tg|#lRj+V_CqP*;IQQfdu^eeP0B{i zV>x#tsvVuj8?aWOVPPM3{M2yYTETx|n{j#?9-jV^%Q4GSDl{=iSS>DsRw2HdSjhv2 z$_;B?UX^&&VJB1hv1eDqnhmiK1tYji5l~AIvL;>kj>w0G@3%|7GWAwFc@Jiku(xH> zA0=+9qi=0|QM)x|%R5)eN*}o zN0jbc75hoxm1f~8XT)ieVKYrCzVyUFHC24IbsMWflC6Nd`sOzKlH#jnw9)+WM5i^& zg9a*bl3N~f-(QMaG3EOy<+odotDHbgK4(lBSyE4P)HL=zHHwAoQIM09lN{fn*FTTO zT}&8G?h|9L+w9fwwT)P^+6ryDLOR5++Gc_Z8UQO`T+ADh3{(Q9?VIQPP2u~58FH<_ zS9X1+NPE5D9$ZT_V;UydY3Xr*w$=73fmvz z#)PEw&q|fHFNLm6hsj9HX|_Fj;~zWk#$=Rxm_MELYYv)jo4Mff=HFrMoH0L3esa*p zc}ixH!3t|z_!BLk01UMxzJs4O5%aEZ0WkD}ODMU%T~aMmrh3e`Dw*1kXEz7nivox$a)y-(RxzL3YtX%MJK6r}bL5Yg z?7X~WcXxN!X%h3-DFCBl@^_+7(!Gf%u+1<&Io&nLfkZzT zgrIjKD{&q987iU)Kn4R#^_hb))osAcEO zx_+($MUZw3@2|%g(Hp$gBwq&OY^0b?bbK0uMhqL#MbX~ov4^*UN8~+sEq~%7a z0Gcbs0(10BF2M8!%d?27Hr&6iy31$lUah#YY+I6P03?12?0M3v9oYgrx1*8l@L0|L z)Yr=17x7*6#(G=!WJ!P#Q(?<0uIv6u&GlG^@p{XR>S*8LGM8T%J_QC#bo%lW{!vjR zaJGqB0-5v!{!=9|g0RrY0#9xDgCe5*b;`zwTKk?$@T-W(D_0|V-)ILF-Y2J_cDvu( zKN=Nqr$Gya&@Jltd$t?=_r7J|<1I2&;*`8eX(`4*}% zFi3Ysics6+B;*cfaSwUQUW8YOpT}xCfu@N(&kwEZ5_14kwvBc}=6BW35IEt`1QGpr zH8m4It7Y<@;$A`S0+m+e>rj>yUA^J2XqYQB@RwirH*Q3o6gVN{vnE3OZ%O|$v7*;_ zi`;7^J6;Q6Ufz$zoF<$vL#al(jt``tFZf{GDdAGp1B%=|nFz|nCk|GqWbQ|+JPz5; zQ=8R^0Fsrmp)!1j@j#c+e}agWZC;FjXmbCwgY66Z`0YpJ9QIakZ<-2rvro!D2hW>Q zZE=#dv!@=)X`08dOiVBz$PG+_Hu{^USxu;MMO~gKz>glQQZQL+8Sk?KC;yQv#lmJj9iT^#zkOq0 zs^u_cT{v*sBs1L4f&>URDzdi7t)C}_;aXE#`7Si8RnYH!`;_h%NoD+WTc%-On$FL4 z0~=)h$eP;(b&*bB&n8gS|!e#5B%VLCp6h4ezo%mK3P5q(^*P_dX%cwi%Ik zNq5HAgTLU*QTr812%;_eb7pH`&PO!O#>d+;pRWH5n3fU$O+gmHgjYb<$E%G|Aq*NLP#VIp zg;sd;Kw~Z-w*)DFepm52%7xo238`;;{)y*|w`-0qBJe4B`doaWqewd|UKQ(Mh8YvbkI zkNdA|(MwPWcuKH~uQSKPaFTbZ9nXgi4%6nQ*>I^q`XIf@tU=XftI>17;v_iqUsKC?ZvT>4#bs+Qg)Rw1~yWCvW zBcI;$akb>L=W*yn-23294JZ% z!KOhGDaNX+AF_M9@OL2!cd)vsac4U(09?J+j^=C;2;=^Gu(hy|p+i~R2W)h6^b&r9 z>Zx_H8t*0-A!wns`Kbd_c1a0g#%C(J#%rlnIfkhiFq2`-plG;z-q{+M2*= z^MFEl&S!`ubVy-a=?Tg!7lNyijug2=d$#RyzKDT$EfC_(L8>6;|3VY2nA@m_?H95f zhx6}n2W&P*ama=eiNpgDinfMP16WauDbfPTeHueEi?j%1>JXsD3jrYcsOSYiIJn(O zP;Q@boAN#(%8O0`b)=o+jWWFt`&7PnV5$NOQZ_E3a8F&CRC>erw|+!F3OK*B)CVSZ zH>(#=-rz!NVp=jRlZqF{x~s+u9sNsR$O#yE8w(8s9yExeJ&g#yZeTc|ecBQX;RfFr zRs;p)O=TRfxbF0+Pf#v@KCW`qmYKEL29`Q@pTTjY%K;`#xKYA?OH&i(wbc7#w=_s~ o% None: + """Release UI-owned preview subprocesses without masking app shutdown.""" + try: + stop_articraft_viser_preview() + except Exception: + # Shutdown must not be blocked by an already-exited preview process. + pass + + +def _handle_shutdown_signal(signum: int, _frame: object) -> None: + """Terminate the Articraft Viser child before leaving the Gradio process.""" + _stop_child_previews() + if signum == signal.SIGINT: + raise KeyboardInterrupt + raise SystemExit(128 + signum) + + +def _install_shutdown_handlers() -> None: + """Install cleanup-aware handlers for the normal Gradio stop signals.""" + signal.signal(signal.SIGINT, _handle_shutdown_signal) + signal.signal(signal.SIGTERM, _handle_shutdown_signal) + + +def main() -> None: + if not EMBODICHAIN_ROOT.is_dir(): + raise FileNotFoundError(f"EmbodiChain root not found: {EMBODICHAIN_ROOT}") + demo = build_demo() + demo.queue(default_concurrency_limit=DEFAULT_CONCURRENCY_LIMIT) + _install_shutdown_handlers() + try: + demo.launch( + server_name=SERVER_NAME, + server_port=SERVER_PORT, + allowed_paths=[ + str(EMBODICHAIN_ROOT), + str(ASSETS_DIR), + str(DEBUG_ENGINE_ROOT), + ], + ) + finally: + _stop_child_previews() + + +if __name__ == "__main__": + main() diff --git a/embodichain/gen_sim/gradio_ui/gradio_visualization_architecture.md b/embodichain/gen_sim/gradio_ui/gradio_visualization_architecture.md new file mode 100644 index 000000000..6fd293aed --- /dev/null +++ b/embodichain/gen_sim/gradio_ui/gradio_visualization_architecture.md @@ -0,0 +1,276 @@ +# Gradio 可视化系统架构 + +本文档以当前代码为准，描述 Gradio Demo、Debug 下的三个引擎，以及它们与 EmbodiChain、SimReady、Articraft 和 DexSim 的边界。`gradio_app.py` 只负责启动；界面、资产工作流、场景工作流和进程管理分散在专用模块中。 + +## 架构总览 + +```text +gradio_app.py + │ 启动、队列、allowed_paths + ▼ +app_services.py（兼容门面） + ▼ +app_ui.py ───────────► app_asset_engine.py ───► SimReady CLI + │ 布局、模式和事件绑定 │ │ + │ │ └──────────► app_articraft.py ───► Articraft CLI + Codex CLI + ▼ +app_workflows.py ────► EmbodiChain Scene Engine CLI + Viser + │ └► prompt2scene / action-agent pipeline + DexSim + ├──────────────► app_commands.py 命令构造 + ├──────────────► app_processes.py 子进程、环境、日志和阶段检测 + ├──────────────► app_state.py 共享 RuntimeState、锁和计时 + ├──────────────► app_media.py 视频、数据集预览和日志归档 + └──────────────► app_config.py 路径、端口、文案和固定参数 +``` + +| 模块 | 职责 | +| --- | --- | +| `gradio_app.py` | 唯一启动入口；校验 `EMBODICHAIN_ROOT`，创建 Blocks，设置队列和本地文件访问路径。 | +| `app_ui.py` | Demo/Debug 布局、引擎面板切换和回调绑定；不实现 pipeline。 | +| `app_asset_engine.py` | SimReady 上传适配、输入/输出 GLB 预览、处理日志，以及 Asset engine 的 Articraft 标签页。 | +| `app_articraft.py` | Articraft checkout/环境检查、外部记录创建、Codex 生成与校验、URDF bundle 和 Viser 关节预览。 | +| `app_workflows.py` | Demo 的 prompt2scene/action-agent 工作流、独立 Scene Engine 工作流、GLB 预览、场景提升和 DexSim。 | +| `app_processes.py` | 子进程环境、进程组终止、stdout 读取、Demo pipeline 阶段检测。 | +| `app_state.py` | `RuntimeState`、互斥锁、进度阶段、运行 token 和耗时统计。 | +| `app_commands.py` | prompt2scene、动作配置和 `run_agent` 的参数构造。 | +| `app_media.py` | 观众视频、LeRobot 数据预览、组合视频和运行日志归档。 | +| `app_config.py` | 路径、环境变量、端口、UI 文案、引擎模式和 CLI 固定参数。 | + +## 启动、路径和网络环境 + +从本项目目录启动： + +```bash +conda run -n embodichain python gradio_app.py +``` + +| 变量 | 默认值 | 用途 | +| --- | --- | --- | +| `EMBODICHAIN_ROOT` | `/home/dex/桌面/EmbodiChain` | EmbodiChain 根目录。 | +| `GRADIO_SERVER_NAME` | `0.0.0.0` | Gradio 监听地址。 | +| `GRADIO_SERVER_PORT` | `7860` | Gradio 监听端口。 | +| `SCENE_ENGINE_VISER_PORT` | `8080` | 独立 Scene Engine 的 Viser 端口。 | +| `ARTICRAFT_VISER_PORT` | `8081` | Articraft 关节预览的 Viser 端口。 | +| `ARTICRAFT_ROOT` | `<项目>/.articraft` | Articraft checkout。 | +| `ARTICRAFT_CONDA_ENV` | `articraft` | 运行 Articraft CLI 的 Conda 环境。 | +| `ARTICRAFT_OUTPUT_ROOT` | `<项目>/.debug_engine/articraft` | Articraft 记录、运行日志和导出 bundle。 | + +`demo.launch()` 仅开放 EmbodiChain 根目录、`assets/` 和 `.debug_engine/` 给浏览器读取。pipeline 子进程由 `build_pipeline_env()` 创建环境：它清除代理变量、设置 `NO_PROXY=no_proxy=*`、关闭 Gradio analytics，并把非空的 SimReady 配置映射为 `OPENAI_*`。这不会改写启动 Gradio 的父进程环境。 + +## 页面与引擎 + +顶部的 `Demo` / `Debug` 只切换可见面板，不会启动任务；切换后共享运行状态保留。Debug 有三个按钮：`Asset_engine`、`Scene_engine`、`Action_engine`。它们的实际输入和产物并不完全相同： + +| Engine | 输入 | 预览/下载 | 实际产物 | 是否启动 DexSim | +| --- | --- | --- | --- | --- | +| Asset engine / SimReady | 一个网格、可选材质附件、类别 | 输入 GLB、SimReady GLB、原始输出下载 | `.debug_engine/assets/runs//` | 否 | +| Asset engine / Articulation | 文字、可选参考图 | URDF articulation 的 Viser、zip 下载 | `.debug_engine/articraft/` | 否 | +| Scene engine | 一张图片 | Scene Engine 的 Viser | `.debug_engine/scenes//` | 否 | +| Action engine | `current` Gym 场景、任务、机器人 | `current` 的 GLB 和 DexSim 视频 | EmbodiChain `gym_project/current` 与 `outputs/` | 是 | + +因此，Debug 的 Scene engine 是独立的图像条件场景生成器；它不会提升、复制或转换输出到 `gym_project/current`。Action engine 只消费 Demo/prompt2scene 工作流已经生成的 `current` Gym 场景。界面中的 “Scene engine” 文案表达的是所需场景类型，并不意味着独立 Scene Engine 输出已自动连到 Action engine。 + +## Demo：端到端 Gym 场景和 DexSim + +Demo 提供 `Auto`、`Interact`、`Parallel Simulation` 三种运行状态，以及图像、任务、场景描述、生成模式、机器人、随机输入、视频和 GLB 预览。它们与顶部的 Demo/Debug 模式无关。 + +`run_generate()` 是 Demo 的主入口。初始生成会在 staging 场景中运行 prompt2scene/action-agent pipeline，成功后才 promote 为固定的 `current`；随后默认启动 DexSim。编辑和仅改任务复用已有 `current`： + +```text +Initial generation + image + task + → _gradio_pending_ + → run_agent_pipeline --skip-run-agent + → fast_gym_config / agent_config / GLB previews + → promote 到 current + → run_agent（DexSim） + +Edit current scene + current + task + scene description + → 编辑 pipeline + → current + → run_agent（DexSim） + +Change task only + current + task + → generate_action_agent_config + → current + → run_agent（DexSim） +``` + +场景生成期间，工作流会从 `fast_gym_config.json` 构建场景 GLB，并将生成的对象 GLB 合并为对象预览。`launch_simulation=False` 是可用的工作流参数，但当前 Debug Scene panel 不调用这条 Demo 工作流；它调用独立的 `run_scene_engine()`。 + +正式场景固定在： + +```text +gym_project/current/ +gym_project/current/gym_export/ +gym_project/action_agent_pipeline/images/current.png +gym_project/action_agent_pipeline/configs/current/ + fast_gym_config.json + agent_config.json + gradio_scene/ + scene_current.glb + initial_scene.glb + object_preview.glb +``` + +初始生成使用 `_gradio_pending_` 路径。提升失败或 pipeline 失败时，已有 `current` 保持不变；成功提升后会重写 staging 中的路径引用。`Reset` 会清理当前场景和 staging 产物；`Stop` 通过进程组终止正在运行的 pipeline 或 DexSim。 + +## Asset engine + +### SimReady：单资产目录适配 + +SimReady CLI 接收目录，而 Gradio 接收上传文件。上传文件会复制到隔离目录，文件名只保留 basename，重名追加序号，避免上传路径或重名影响处理： + +```text +mesh + sidecar files + → .debug_engine/assets/runs//input/ + → trimesh 导出 input_preview.glb + → SimReady CLI + → output/**/asset_simready.glb（优先）或 asset_simready.obj + → GLB 预览 + 原始文件下载 +``` + +主网格支持 `.glb`、`.gltf`、`.obj`、`.ply`、`.stl`；可一并上传 `.mtl`、纹理和 `.bin` 等附件。执行命令为： + +```bash +python -m embodichain.gen_sim.simready_pipeline.cli.start \ + --input_dir \ + --output_root \ + --category +``` + +处理函数以 generator 持续返回最近的 stdout；完成时优先预览 `asset_simready.glb`，只有 OBJ 时再转为 GLB。此路径不依赖 DexSim。 + +### Articulation：Articraft + Codex + +Articulation 标签页根据文本和可选参考图生成一个可下载的 articulated asset。先点击环境检查：若 `ARTICRAFT_ROOT` 不存在，应用会 clone `ARTICRAFT_REPOSITORY_URL`；随后检查 Conda、指定的 Articraft 环境和 Codex CLI。该操作会创建 checkout 和 `.debug_engine/articraft/` 中的输出目录，现有的非 Articraft 目录不会被覆盖。 + +生成流程： + +```text +description + optional image + → Articraft external init（创建 rec_ui_articraft_* 记录） + → 启动 Codex CLI，仅授权编辑该记录的 active model.py + → Articraft external check + └─ 旧版 CLI 无 check 时：compile --validate --strict-geom-qc + compile_report + → Articraft external finalize + → materialized model.urdf + meshes + → exports/.zip + Viser articulation preview +``` + +产物、记录和参考图均在 `ARTICRAFT_OUTPUT_ROOT` 下，不能直接当作 Demo 的 Gym 场景或 SimReady 资产；若要进入后续仿真，需要另行定义并实现转换/导入流程。Articraft Viser 每次成功预览会终止旧的 Articraft 预览进程，再以 `0.0.0.0:` 启动新进程。 + +## 独立 Scene engine 和 Viser + +Scene engine 只接收图像。上传图像会先进行 EXIF 归正并转为 RGB PNG，以 PNG 字节的 SHA-256 前 16 位作为目录名；相同图像会复用同一目录： + +```text +image + → .debug_engine/scenes//input.png + → python -m embodichain scene-engine + --image + --output_root + --config /embodichain/gen_sim/scene_engine_config.json + → /scene_export/scene_config.json + → preview.py --viser --viser-host 0.0.0.0 --viser-port 8080 + → Gradio iframe +``` + +当 `scene_export/scene_config.json` 存在且生成进程返回成功时，应用才启动 Viser。iframe 使用 Gradio 页面当前的协议和主机名转向 Viser 端口，因此从其他设备访问时，浏览器必须能访问该端口。每次新 Scene Engine 任务开始前会终止旧的 Scene Viser 进程。输出目录会显示在 UI 中，便于检查 hash 命名的场景导出。 + +## Action engine：Gym 场景契约 + +Action engine 不接收裸 GLB。普通 GLB 只有渲染数据，而 DexSim 还需要碰撞、物理参数、初始位姿、资源相对路径和 action 配置。当前实现的前置条件是： + +```text +gym_project/current/gym_export/ +gym_project/action_agent_pipeline/configs/current/fast_gym_config.json +gym_project/action_agent_pipeline/configs/current/agent_config.json +``` + +点击 `Load current scene` 只读取共享状态快照。点击 `Run DexSim` 会先检查任务、`current` 的 Gym/action 配置、运行占用和可导入的 `embodichain.gen_sim.action_agent_pipeline.cli.run_agent`，再以当前配置调用 `run_agent`。它不会因为新的任务文本重建动作图；任务改变时应在 Demo 里使用 `Change task only`，或者实现显式的配置再生成步骤。 + +运行命令的核心参数为： + +```bash +python -m embodichain.gen_sim.action_agent_pipeline.cli.run_agent \ + --task_name current \ + --gym_config <.../fast_gym_config.json> \ + --agent_config <.../agent_config.json> \ + --regenerate --renderer fast-rt --num_envs <1|9> +``` + +并行模式额外传入 arena 和数据保存过滤参数。`--robot-profile` 仅在通过 `run_agent --help` 探测到该参数时加入。DexSim 完成后会寻找 audience 视频和 LeRobot 数据集；单环境可组合两种预览视频。 + +## 共享状态、并发和进度 + +Demo、独立 Scene engine 和 Action engine 共享 `RuntimeState` 与 `runtime_lock`，其中包含运行 token、pipeline/DexSim/Scene Viser 进程、输入、预览、日志、阶段和计时。运行 token 用于丢弃过期线程的更新。Articraft Viser 使用单独的锁和进程引用；SimReady 使用自己的同步 generator。 + +`demo.queue(default_concurrency_limit=1)` 将队列中的高成本回调串行化。Demo 的 `Timer(2.0)` 与 Action engine 的独立 `Timer(2.0)` 都读取同一共享状态。Scene Engine 和 Demo pipeline 因共享 `is_busy` 互斥；Asset/Articraft 面板不写入这一状态，但仍会受 Gradio 队列限制。 + +共享阶段如下；独立 Scene Engine 将其日志映射到相同的进度条： + +```text +idle → received → started → scene_intake → relations +→ asset_generation → gym_export → config → preview → complete + └──────────────→ failed +``` + +## 环境前置条件与验证 + +SimReady 需要 Blender、trimesh、LLM 配置以及可导入的： + +```text +embodichain.gen_sim.simready_pipeline.cli.start +``` + +SimReady 的 OpenAI-compatible 设置来自环境变量，且不应写入 Git： + +```bash +export SIMREADY_OPENAI_API_KEY='' +export SIMREADY_OPENAI_MODEL='' +export SIMREADY_OPENAI_BASE_URL='' +``` + +Demo/Action 需要 action-agent 模块，特别是： + +```text +embodichain.gen_sim.action_agent_pipeline.cli.run_agent_pipeline +embodichain.gen_sim.action_agent_pipeline.cli.generate_action_agent_config +embodichain.gen_sim.action_agent_pipeline.cli.run_agent +``` + +独立 Scene engine 还需要： + +```text +python -m embodichain scene-engine +embodichain/gen_sim/scene_engine/cli/preview.py +embodichain/gen_sim/scene_engine_config.json +``` + +Articulation 还需要 Git（首次 clone）、Conda、`ARTICRAFT_CONDA_ENV` 和 Codex CLI。生成请求会交给本机 Codex CLI 执行，因此只应提交可信请求。 + +每次修改后至少执行： + +```bash +python -m py_compile \ + gradio_app.py app_config.py app_state.py app_commands.py \ + app_processes.py app_media.py app_workflows.py app_ui.py \ + app_asset_engine.py app_articraft.py app_services.py + +env -u HTTP_PROXY -u HTTPS_PROXY -u ALL_PROXY \ + -u http_proxy -u https_proxy -u all_proxy \ + conda run -n embodichain python -c \ + "from app_ui import build_demo; assert build_demo() is not None" +``` + +手动检查： + +1. SimReady 上传简单网格后能显示输入预览；执行后显示 SimReady 输出或明确错误。 +2. Articulation 环境检查能报告 checkout、Conda 和 Codex 状态；成功生成后有 zip、记录目录和 Viser 或明确的预览错误。 +3. Scene engine 从图像生成 `scene_export/scene_config.json`，并在 `8080` 显示 Viser；它不应改写 `gym_project/current`。 +4. Demo 初始生成成功后才替换 `current`；失败时旧场景仍可用。 +5. Action engine 在没有 `current` Gym/action 配置或缺少 CLI 时给出预检错误；任务更新后通过 Demo 的 `Change task only` 重建配置。 +6. Demo 的 Auto/Interact/Parallel Simulation 行为不因 Debug 面板切换而改变；Reset/Stop 能终止其对应的进程组。 diff --git a/embodichain/gen_sim/gradio_ui/random_input.py b/embodichain/gen_sim/gradio_ui/random_input.py new file mode 100644 index 000000000..1ef138cdf --- /dev/null +++ b/embodichain/gen_sim/gradio_ui/random_input.py @@ -0,0 +1,540 @@ +# ---------------------------------------------------------------------------- +# 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 +import os +import uuid +from dataclasses import asdict, dataclass, replace +from pathlib import Path + +import numpy as np + +EMBODICHAIN_ROOT = Path( + os.environ.get("EMBODICHAIN_ROOT", "/home/dex/桌面/EmbodiChain") +).expanduser() +APP_ROOT = Path(__file__).resolve().parent +IMAGE_DIR = Path( + os.environ.get( + "AUTO_IMAGE_DIR", + str(EMBODICHAIN_ROOT / "gym_project/action_agent_pipeline/auto_images"), + ) +).expanduser() +AUTO_IMAGE_DIR_IS_CONFIGURED = "AUTO_IMAGE_DIR" in os.environ +FALLBACK_IMAGE_DIR = ( + EMBODICHAIN_ROOT / "gym_project/action_agent_pipeline/baseline_image_input" +) +PREBUILT_SCENE_DIR = Path( + os.environ.get("AUTO_PREBUILT_SCENE_DIR", str(APP_ROOT / "scenes")) +).expanduser() +GENERATED_IMAGE_DIR = Path( + os.environ.get("AUTO_GENERATED_IMAGE_DIR", "./tmp_img/auto") +).expanduser() +IMAGE_API_KEY = os.environ.get("AUTO_IMAGE_API_KEY", os.environ.get("ARK_API_KEY", "")) +IMAGE_API_URL = os.environ.get( + "AUTO_IMAGE_API_URL", + "https://ark.cn-beijing.volces.com/api/v3", +) +IMAGE_MODEL = os.environ.get("AUTO_IMAGE_MODEL", "doubao-seedream-4-5-251128") +IMAGE_SIZE = os.environ.get("AUTO_IMAGE_SIZE", "2848x1600") +IMAGE_PROMPT = ( + "以原图为基础参考，严格保持原图的相机视角、拍摄距离、透视关系、画面构图、背景环境、桌面材质、桌面纹理、光影方向、阴影位置和整体明暗关系。" + "原图中已有物体的类别、大小、形状、轮廓、空间位置、朝向、部件数量和结构比例必须保持不变。" + "如果提供了 Scene description，只根据该描述在桌面上新增对应背景物体；新增物体必须放在描述指定的位置，尺寸、透视、遮挡和阴影要与原图自然一致。" + "不要删除原有物体，不要移动原有物体，不要改变主任务物体的结构。高清细节，真实自然。" +) + +TASK_DESCRIPTIONS: dict[tuple[int, int], str] = { + (0, 0): "用双臂把两侧的罐头和瓶子放到篮子里", + (0, 1): "用双臂把两侧的方块放到篮子里", + (0, 2): "用双臂把两侧的方块和纸杯放到篮子里", + (0, 3): "用双臂把两侧的方块和苹果放到篮子里", + (1, 0): "用双臂把塑料水盆往前移动", + (1, 1): "用双臂把木棍往前移动", + (1, 2): "用双臂往把苹果和魔方放入盘子，然后用双臂端起盘子", + (1, 3): "用双臂把托盘往前移动", + (2, 0): "用双臂把两侧的香蕉放到盘子里，然后用双臂端起盘子", + (2, 1): "用双臂把两侧的罐头扶正", + (2, 2): "用双臂把两侧的瓶子和罐头扶正", + (2, 3): "用双臂把两侧的罐头扶正", + (3, 0): "把桌面上的物体按照方块按照从右往左的顺序叠起来", + (3, 1): "把桌面上的物体按照右边的方块，左边的方块，纸杯的顺序叠起来", + (3, 2): "把纸杯叠放到爆米花桶上，把蓝色耳机叠放到爆米花桶上", + (3, 3): "把纸杯叠放到爆米花桶上，把固体胶叠放到爆米花桶上", + (4, 0): "把桌面上的方块摆成一排", + (4, 1): "把桌面上的物体按照瓶子，方块排成一排", + (4, 2): "把桌面上的罐头摆成一排", + (4, 3): "把桌面上的物体按照瓶子，罐头，方块的顺序摆成一排", +} + +TASK_DESCRIPTIONS_EN: dict[tuple[int, int], str] = { + ( + 0, + 0, + ): "Use both arms to place the cans and bottles on both sides into the basket.", + (0, 1): "Use both arms to place the blocks on both sides into the basket.", + ( + 0, + 2, + ): "Use both arms to place the blocks and paper cups on both sides into the basket.", + ( + 0, + 3, + ): "Use both arms to place the blocks and apples on both sides into the basket.", + (1, 0): "Use both arms to move the plastic basin forward.", + (1, 1): "Use both arms to move the wooden stick forward.", + ( + 1, + 2, + ): "Use both arms to place the apple and Rubik's Cube onto the plate, then use both arms to lift the plate.", + (1, 3): "Use both arms to move the tray forward.", + ( + 2, + 0, + ): "Use both arms to place the bananas on both sides onto the tray, then use both arms to lift the tray.", + (2, 1): "Use both arms to set the cans on both sides upright.", + (2, 2): "Use both arms to set the bottles and cans on both sides upright.", + (2, 3): "Use both arms to set the cans on both sides upright.", + (3, 0): "Stack the blocks on the table in order from right to left.", + ( + 3, + 1, + ): "Stack the objects on the table in this order: right block, left block, paper cup.", + ( + 3, + 2, + ): "Stack the paper cup on the popcorn bucket, then stack the blue headphones on the popcorn bucket.", + ( + 3, + 3, + ): "Stack the paper cup on the popcorn bucket, then stack the glue stick on the popcorn bucket.", + (4, 0): "Arrange the blocks on the table in a row.", + (4, 1): "Arrange the objects on the table in a row in this order: bottle, block.", + (4, 2): "Arrange the cans on the table in a row.", + ( + 4, + 3, + ): "Arrange the objects on the table in a row in this order: bottle, can, block.", +} + +RELATION_PATTERN = { + (0, 0): ["at the left side of the can", "at the right side of the bottle"], + (0, 1): [ + "at the left side of the left cheese cube", + "at the right side of the right cheese cube", + ], + (0, 2): ["at the left side of the cube", "at the right side of the cup"], + (0, 3): ["at the left side of the cube", "at the right side of the apple"], + (1, 0): [], + (1, 1): [], + (1, 2): [], + (1, 3): [], + (2, 0): [ + "at the left side of the left bottle", + "at the right side of the right bottle", + ], + (2, 1): [ + "at the left side of the left soda can", + "at the right side of the right soda can", + ], + (2, 2): ["at the left side of the bottle", "at the right side of the can"], + (2, 3): ["at the left side of the paper cup", "at the right side of the soda can"], + (3, 0): [], + (3, 1): [], + (3, 2): [], + (3, 3): [], + (4, 0): [], + (4, 1): [], + (4, 2): [], + (4, 3): [], +} + +AREA_PATTERN = [ + "at the left side of the table", + "at the right side of the table", + "at the front of the table", + "at the front right corner of the table", + "at the front left corner of the table", +] + +OBJECT_LIST = [ + "cup", + "potted plant", + "clock", + "book", + "pen", + "bottle", + "soda can", + "photo frame", + "apple", + "peach", + "bread", + "chocolate bar", + "cookie", + "penholder", + "desk lamp", + "stapler", + "headphones", + "desk calendar", + "eyeglasses", + "fan", + "bluetooth speaker", + "table mirror", + "computer mouse", + "keyboard", +] + +CHINESE_OBJECT_NAMES = { + "cup": "杯子", + "potted plant": "盆栽", + "clock": "时钟", + "book": "书", + "bottle": "瓶子", + "soda can": "易拉罐", + "photo frame": "相框", + "apple": "苹果", + "peach": "桃子", + "bread": "小面包", + "chocolate bar": "巧克力棒", + "cookie": "饼干", + "penholder": "笔筒", + "desk lamp": "小台灯", + "stapler": "订书机", + "headphones": "耳机", + "small desk calendar": "小台历", + "eyeglasses": "眼镜", + "fan": "小风扇", + "bluetooth speaker": "蓝牙音箱", + "computer mouse": "鼠标", +} + + +CHINESE_SPATIAL_RELATIONS = { + "at the left side of the can": "罐头左侧", + "at the right side of the bottle": "瓶子右侧", + "at the left side of the left cheese cube": "左侧奶酪方块左侧", + "at the right side of the right cheese cube": "右侧奶酪方块右侧", + "at the left side of the cube": "方块左侧", + "at the right side of the cup": "杯子右侧", + "at the right side of the apple": "苹果右侧", + "at the left side of the left bottle": "左侧瓶子左侧", + "at the right side of the right bottle": "右侧瓶子右侧", + "at the left side of the left soda can": "左侧易拉罐左侧", + "at the right side of the right soda can": "右侧易拉罐右侧", + "at the left side of the bottle": "瓶子左侧", + "at the right side of the can": "罐头右侧", + "at the left side of the paper cup": "纸杯左侧", + "at the right side of the soda can": "易拉罐右侧", + "at the left side of the table": "桌子左侧", + "at the right side of the table": "桌子右侧", + "at the front of the table": "桌子前侧", + "at the front right corner of the table": "桌子右前角", + "at the front left corner of the table": "桌子左前角", + "on the table": "桌面上", +} + + +@dataclass(frozen=True) +class AutoInput: + task_index: tuple[int, int] + base_image_path: Path + prebuilt_scene_dir: Path | None + image_path: Path | None + task_description: str + scene_description: str + + def to_json_dict(self) -> dict[str, object]: + value = asdict(self) + value["task_index"] = list(self.task_index) + value["base_image_path"] = self.base_image_path.as_posix() + value["prebuilt_scene_dir"] = ( + self.prebuilt_scene_dir.as_posix() if self.prebuilt_scene_dir else None + ) + value["image_path"] = self.image_path.as_posix() if self.image_path else None + return value + + +def task_id(task_index: tuple[int, int]) -> str: + return f"task{task_index[0]}_{task_index[1]}" + + +def parse_task_id(value: str) -> tuple[int, int] | None: + stem = Path(value).stem + if not stem.startswith("task"): + return None + parts = stem[4:].split("_", maxsplit=1) + if len(parts) != 2: + return None + try: + return int(parts[0]), int(parts[1]) + except ValueError: + return None + + +def auto_image_directories() -> tuple[Path, ...]: + """Return image sources in precedence order for the Auto loop. + + A user-supplied ``AUTO_IMAGE_DIR`` is authoritative. With the default + directory, retain compatibility with deployments that have the checked-in + ``baseline_image_input`` set but have not created ``auto_images`` yet. + """ + directories = [IMAGE_DIR] + if not AUTO_IMAGE_DIR_IS_CONFIGURED and FALLBACK_IMAGE_DIR != IMAGE_DIR: + directories.append(FALLBACK_IMAGE_DIR) + return tuple(directories) + + +def available_auto_task_indices() -> tuple[tuple[int, int], ...]: + """Return only task variants whose input image and clean scene can be resolved.""" + return tuple( + task_index + for task_index in TASK_DESCRIPTIONS + if any( + (image_dir / f"{task_id(task_index)}.png").is_file() + for image_dir in auto_image_directories() + ) + and get_prebuilt_scene_dir(task_index).is_dir() + ) + + +def random_task(rng: np.random.Generator) -> tuple[int, int]: + available_tasks = available_auto_task_indices() + if not available_tasks: + expected = ", ".join(str(path) for path in auto_image_directories()) + raise FileNotFoundError( + "No Auto input images were found. Add task_.png " + f"files to: {expected}" + ) + return available_tasks[int(rng.integers(0, len(available_tasks)))] + + +def get_base_image_path(task_index: tuple[int, int]) -> Path: + filename = f"{task_id(task_index)}.png" + for image_dir in auto_image_directories(): + candidate = image_dir / filename + if candidate.is_file(): + return candidate + return IMAGE_DIR / filename + + +def get_prebuilt_scene_dir(task_index: tuple[int, int]) -> Path: + return PREBUILT_SCENE_DIR / task_id(task_index) + + +def get_task_description(task_index: tuple[int, int], *, language: str = "zh") -> str: + descriptions = TASK_DESCRIPTIONS_EN if language == "en" else TASK_DESCRIPTIONS + try: + return descriptions[task_index] + except KeyError as exc: + raise KeyError(f"No task description configured for task{task_index}") from exc + + +def image_to_base64(path: Path) -> str: + ext = path.suffix.lower() + if ext in (".jpg", ".jpeg"): + mime = "image/jpeg" + elif ext == ".png": + mime = "image/png" + else: + raise ValueError(f"Not supported: {ext}, only jpg/jpeg/png are supported") + with path.open("rb") as file: + b64_str = base64.b64encode(file.read()).decode("utf-8") + return f"data:{mime};base64,{b64_str}" + + +def build_image_prompt(scene_description: str = "") -> str: + scene_description = (scene_description or "").strip() + if not scene_description: + return IMAGE_PROMPT + return ( + f"{IMAGE_PROMPT}\n\n" + "Scene description:\n" + f"{scene_description}\n\n" + "严格执行 Scene description 中的新增物体和空间位置要求。" + ) + + +def create_image_input( + task_index: tuple[int, int], + *, + scene_description: str = "", + output_dir: Path = GENERATED_IMAGE_DIR, +) -> Path: + base_image_path = get_base_image_path(task_index) + if not base_image_path.is_file(): + raise FileNotFoundError(f"Base auto image not found: {base_image_path}") + + from volcenginesdkarkruntime import Ark + import requests + + image_base64 = image_to_base64(base_image_path) + client = Ark(api_key=IMAGE_API_KEY, base_url=IMAGE_API_URL) + response = client.images.generate( + model=IMAGE_MODEL, + prompt=build_image_prompt(scene_description), + image=image_base64, + size=IMAGE_SIZE, + response_format="url", + watermark=False, + ) + image_url = response.data[0].url + resp = requests.get(image_url, timeout=60) + resp.raise_for_status() + output_dir.mkdir(parents=True, exist_ok=True) + output_path = ( + output_dir + / f"auto_task{task_index[0]}_{task_index[1]}_{uuid.uuid4().hex[:12]}.png" + ) + output_path.write_bytes(resp.content) + return output_path + + +def create_text_input( + task_index: tuple[int, int], + rng: np.random.Generator, + *, + language: str = "en", + min_background_objects: int = 0, +) -> str: + text_parts: list[str] = [] + if task_index[0] == 5 and min_background_objects == 0: + return "" + + mu, sigma = 1.0, 1.0 + raw = rng.normal(loc=mu, scale=sigma) + num_background_objects = int(np.clip(np.round(raw), 0, 3)) + if min_background_objects > 0: + num_background_objects = max(num_background_objects, min_background_objects) + + if num_background_objects == 0: + return "" + + selected_objects = rng.choice( + OBJECT_LIST, + size=num_background_objects, + replace=False, + ).tolist() + + spatial_candidates = [] + spatial_candidates.extend(RELATION_PATTERN.get(task_index, [])) + spatial_candidates.extend(AREA_PATTERN) + spatial_candidates.append("on the table") + + for obj in selected_objects: + selected_spatial = rng.choice(spatial_candidates) + if language == "zh": + chinese_object = CHINESE_OBJECT_NAMES.get(obj, obj) + chinese_relation = CHINESE_SPATIAL_RELATIONS.get( + selected_spatial, + selected_spatial, + ) + text_parts.append(f"将一个{chinese_object}放在{chinese_relation}。") + else: + article = "an" if obj[0].lower() in {"a", "e", "i", "o", "u"} else "a" + text_parts.append(f"Place {article} {obj} {selected_spatial}.") + + return " ".join(text_parts) + + +def generate_auto_scene_description( + *, + rng: np.random.Generator | None = None, + task_index: tuple[int, int] | None = None, + language: str = "en", + ensure_scene: bool = False, +) -> str: + rng = rng or np.random.default_rng() + task_index = task_index or random_task(rng) + return create_text_input( + task_index, + rng, + language=language, + min_background_objects=1 if ensure_scene else 0, + ) + + +def generate_auto_text_input( + *, + rng: np.random.Generator | None = None, + task_index: tuple[int, int] | None = None, + language: str = "en", + ensure_scene: bool = False, + include_scene: bool = True, +) -> AutoInput: + rng = rng or np.random.default_rng() + task_index = task_index or random_task(rng) + base_image_path = get_base_image_path(task_index) + prebuilt_scene_dir = get_prebuilt_scene_dir(task_index) + if not base_image_path.is_file(): + raise FileNotFoundError(f"Base auto image not found: {base_image_path}") + if not prebuilt_scene_dir.is_dir(): + raise FileNotFoundError(f"Prebuilt scene not found: {prebuilt_scene_dir}") + return AutoInput( + task_index=task_index, + base_image_path=base_image_path, + prebuilt_scene_dir=prebuilt_scene_dir, + image_path=None, + task_description=get_task_description(task_index, language=language), + scene_description=( + generate_auto_scene_description( + rng=rng, + task_index=task_index, + language=language, + ensure_scene=ensure_scene, + ) + if include_scene + else "" + ), + ) + + +def generate_auto_image( + auto_input: AutoInput, + *, + output_dir: Path = GENERATED_IMAGE_DIR, +) -> AutoInput: + image_path = create_image_input( + auto_input.task_index, + scene_description=auto_input.scene_description, + output_dir=output_dir, + ) + return replace(auto_input, image_path=image_path) + + +def generate_auto_input( + *, + rng: np.random.Generator | None = None, + task_index: tuple[int, int] | None = None, + output_dir: Path = GENERATED_IMAGE_DIR, + language: str = "en", +) -> AutoInput: + auto_input = generate_auto_text_input( + rng=rng, + task_index=task_index, + language=language, + ) + return generate_auto_image(auto_input, output_dir=output_dir) + + +def main() -> None: + auto_input = generate_auto_input() + print(json.dumps(auto_input.to_json_dict(), ensure_ascii=False, indent=2)) + + +if __name__ == "__main__": + main() From 693e1ebf96b77dabeb6b04be71d4e58536f79b9f Mon Sep 17 00:00:00 2001 From: PengXuanchao Date: Mon, 3 Aug 2026 17:29:02 +0800 Subject: [PATCH 25/46] add conda support --- .../gen_sim/gradio_ui/app_articraft.py | 95 ++++++++++++++++++- 1 file changed, 94 insertions(+), 1 deletion(-) diff --git a/embodichain/gen_sim/gradio_ui/app_articraft.py b/embodichain/gen_sim/gradio_ui/app_articraft.py index 85d8544a0..4fa9bba69 100644 --- a/embodichain/gen_sim/gradio_ui/app_articraft.py +++ b/embodichain/gen_sim/gradio_ui/app_articraft.py @@ -60,6 +60,9 @@ _VISER_START_TIMEOUT_SECONDS = 15.0 _VISER_STOP_TIMEOUT_SECONDS = 5.0 +_ARTICRAFT_PYTHON_VERSION = "3.12" +_CONDA_ENVIRONMENT_SETUP_TIMEOUT_SECONDS = 1_200 +_articraft_environment_lock = threading.Lock() def _command_path(name: str) -> str | None: @@ -109,6 +112,92 @@ def _articraft_conda_environment_exists() -> bool: return False +def _ensure_articraft_conda_environment() -> tuple[bool, str]: + """Create and populate the Articraft Conda environment when it is absent. + + Articraft currently supports Python 3.11 and 3.12, while the Gradio process + can use a different interpreter. The setup therefore creates an isolated + Python 3.12 environment and installs the checked-out project's runtime + dependencies into it. + + Returns: + Whether the environment is ready and a status message suitable for the + Gradio configuration panel. + """ + conda = _conda_path() + if not conda: + return False, "Conda is not on PATH. Set CONDA_EXE to the conda executable." + + with _articraft_environment_lock: + if _articraft_conda_environment_exists(): + return True, f"Conda environment: {ARTICRAFT_CONDA_ENV} (already exists)" + + create_command = [ + conda, + "create", + "--yes", + "--name", + ARTICRAFT_CONDA_ENV, + f"python={_ARTICRAFT_PYTHON_VERSION}", + "pip", + ] + try: + created = subprocess.run( + create_command, + cwd=ARTICRAFT_ROOT, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + timeout=_CONDA_ENVIRONMENT_SETUP_TIMEOUT_SECONDS, + check=False, + ) + except (OSError, subprocess.TimeoutExpired) as exc: + return False, f"Unable to create Conda environment: {exc}" + if created.returncode and not _articraft_conda_environment_exists(): + return False, ( + "Conda environment creation failed: " + f"{_short_output(created, limit=3000)}" + ) + + for install_args, description in ( + ( + ["python", "-m", "pip", "install", "--upgrade", "pip"], + "upgrade pip", + ), + (["python", "-m", "pip", "install", "."], "install Articraft dependencies"), + ): + install_command = [ + conda, + "run", + "--no-capture-output", + "--name", + ARTICRAFT_CONDA_ENV, + *install_args, + ] + try: + installed = subprocess.run( + install_command, + cwd=ARTICRAFT_ROOT, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + timeout=_CONDA_ENVIRONMENT_SETUP_TIMEOUT_SECONDS, + check=False, + ) + except (OSError, subprocess.TimeoutExpired) as exc: + return False, f"Unable to {description}: {exc}" + if installed.returncode: + return ( + False, + f"Unable to {description}: {_short_output(installed, limit=3000)}", + ) + + return True, ( + f"Created Conda environment: {ARTICRAFT_CONDA_ENV} " + f"(Python {_ARTICRAFT_PYTHON_VERSION})" + ) + + def _run_check( command: list[str], *, timeout: int = 45 ) -> subprocess.CompletedProcess[str]: @@ -199,10 +288,13 @@ def _prepare_articraft_checkout() -> tuple[bool, str]: def configure_articraft_environment() -> str: - """Clone the checkout, then verify the Conda environment and Codex.""" + """Clone the checkout, prepare its Conda environment, and verify Codex.""" checkout_ready, checkout_message = _prepare_articraft_checkout() if not checkout_ready: return "**Articulation is not ready.**\n\n- " + checkout_message + environment_ready, environment_message = _ensure_articraft_conda_environment() + if not environment_ready: + return "**Articulation is not ready.**\n\n- " + environment_message try: for directory in ( ARTICRAFT_OUTPUT_ROOT, @@ -218,6 +310,7 @@ def configure_articraft_environment() -> str: f"- {error}" for error in errors ) details.insert(0, checkout_message) + details.insert(1, environment_message) details.extend( ( f"Shared output: `{ARTICRAFT_OUTPUT_ROOT}`", From ab2ad309216eeee0adf2ce68ca02477fde6f3398 Mon Sep 17 00:00:00 2001 From: PengXuanchao Date: Mon, 3 Aug 2026 17:36:03 +0800 Subject: [PATCH 26/46] add scene engine --- embodichain/gen_sim/scene_engine/__init__.py | 19 + .../gen_sim/scene_engine/cli/__init__.py | 19 + .../gen_sim/scene_engine/cli/preview.py | 235 +++ embodichain/gen_sim/scene_engine/cli/start.py | 93 + .../gen_sim/scene_engine/clients/__init__.py | 19 + .../clients/geometry_generation.py | 458 +++++ .../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 | 121 ++ .../scene_engine/pipeline/scene_export.py | 220 +++ .../scene_engine/pipeline/scene_generation.py | 574 ++++++ .../pipeline/scene_segmentation.py | 479 +++++ .../pipeline/scene_understanding.py | 254 +++ .../scene_engine/pipeline/utils/__init__.py | 19 + .../pipeline/utils/scene_generation_utils.py | 1542 +++++++++++++++++ .../utils/scene_segmentation_utils.py | 340 ++++ .../gen_sim/scene_engine/utils/__init__.py | 19 + .../gen_sim/scene_engine/utils/logger.py | 38 + 27 files changed, 5152 insertions(+) create mode 100644 embodichain/gen_sim/scene_engine/__init__.py 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/scene_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/__init__.py b/embodichain/gen_sim/scene_engine/__init__.py new file mode 100644 index 000000000..015c41510 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/__init__.py @@ -0,0 +1,19 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +__all__: list[str] = [] diff --git a/embodichain/gen_sim/scene_engine/cli/__init__.py b/embodichain/gen_sim/scene_engine/cli/__init__.py new file mode 100644 index 000000000..015c41510 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/cli/__init__.py @@ -0,0 +1,19 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +__all__: list[str] = [] diff --git a/embodichain/gen_sim/scene_engine/cli/preview.py b/embodichain/gen_sim/scene_engine/cli/preview.py new file mode 100644 index 000000000..e0dcf884a --- /dev/null +++ b/embodichain/gen_sim/scene_engine/cli/preview.py @@ -0,0 +1,235 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path +import time +from collections.abc import Sequence +from typing import Any + +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.cfg import LightCfg, MeshCfg, RigidObjectCfg +from embodichain.lab.visualization import ( + VisualizationCfg, + add_viser_args_to_parser, + visualization_cfg_from_args, +) + + +def preview_scene_export( + *, + output_root: str | Path, + device: str = "cpu", + headless: bool = False, + visualization: VisualizationCfg | None = None, +) -> None: + """Load ``scene_export/scene_config.json`` and preview its table and assets. + + Args: + output_root: Scene Engine output root containing ``scene_export/``. + device: Simulation device, for example ``"cpu"`` or ``"cuda"``. + headless: Load and validate the scene without an interactive preview. + visualization: Optional live-visualization configuration. + """ + resolved_output_root = Path(output_root).expanduser().resolve() + config_path = resolved_output_root / "scene_export" / "scene_config.json" + if not config_path.is_file(): + raise FileNotFoundError(f"Scene config not found: {config_path}") + + try: + scene_config = json.loads(config_path.read_text(encoding="utf-8")) + except json.JSONDecodeError as exc: + raise ValueError(f"Scene config is not valid JSON: {config_path}") from exc + if not isinstance(scene_config, dict): + raise ValueError("Scene config must be a JSON object.") + if scene_config.get("format") != "embodichain.scene-export/v1": + raise ValueError( + "Expected an EmbodiChain scene export " + "(format='embodichain.scene-export/v1')." + ) + + sim = SimulationManager( + SimulationManagerCfg( + width=1920, + height=1080, + headless=headless, + physics_dt=1.0 / 100.0, + sim_device=device, + visualization=( + VisualizationCfg() if visualization is None else visualization + ), + ) + ) + try: + if sim.is_use_gpu_physics: + sim.init_gpu_physics() + _add_lights(sim) + _add_objects( + sim=sim, + entries=_config_entries(scene_config, "background"), + config_dir=config_path.parent, + label="table", + ) + _add_objects( + sim=sim, + entries=_config_entries(scene_config, "rigid_object"), + config_dir=config_path.parent, + label="asset", + ) + + is_viser = sim.sim_config.visualization.backend == "viser" + if headless and not is_viser: + sim.update(step=1) + print(f"Loaded scene export headlessly: {config_path}") + return + + if is_viser: + sim.update(step=1) + print(f"Previewing in Viser: {config_path}") + else: + print(f"Previewing: {config_path}") + sim.open_window() + print("Close with Ctrl-C.") + while True: + time.sleep(0.1) + except KeyboardInterrupt: + print("Stopping preview.") + finally: + sim.destroy(exit_process=False) + SimulationManager.flush_cleanup_queue() + + +def _config_entries( + scene_config: dict[str, Any], + field_name: str, +) -> list[dict[str, Any]]: + entries = scene_config.get(field_name, []) + if not isinstance(entries, list) or not all( + isinstance(entry, dict) for entry in entries + ): + raise ValueError( + f"Scene config field {field_name!r} must be a list of objects." + ) + return entries + + +def _add_lights(sim: SimulationManager) -> None: + for index in range(8): + angle = 2.0 * math.pi * index / 8 + sim.add_light( + LightCfg( + uid=f"light_{index + 1}", + intensity=80.0, + radius=600, + init_pos=[5.0 * math.cos(angle), 5.0 * math.sin(angle), 8.0], + ) + ) + + +def _add_objects( + *, + sim: SimulationManager, + entries: list[dict[str, Any]], + config_dir: Path, + label: str, +) -> None: + """Add exported meshes as static bodies so previewing does not re-simulate them.""" + for entry in entries: + uid = entry.get("uid") + shape = entry.get("shape") + if not isinstance(uid, str) or not uid: + raise ValueError(f"Scene {label} has no valid uid.") + if not isinstance(shape, dict) or not isinstance(shape.get("fpath"), str): + raise ValueError(f"Scene {label} {uid!r} has no shape.fpath.") + if shape.get("shape_type") != "Mesh": + raise ValueError( + f"Scene {label} {uid!r} must use shape_type='Mesh' for preview." + ) + + 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, + acd_method="vhacd", # Use vhacd by default. + ) + ) + print(f"[{label}] {uid}: pos={init_pos} rot={init_rot} scale={body_scale}") + + +def _vector3(value: object, *, field_name: str) -> list[float]: + if not isinstance(value, list) or len(value) != 3: + raise ValueError(f"Scene config field {field_name!r} must be a length-3 list.") + try: + return [float(item) for item in value] + except (TypeError, ValueError) as exc: + raise ValueError(f"Scene config field {field_name!r} must be numeric.") from exc + + +def main(argv: Sequence[str] | None = None) -> None: + parser = argparse.ArgumentParser( + prog="embodichain preview-scene", + description="Preview a Scene Engine scene export in EmbodiChain simulation.", + ) + parser.add_argument( + "output_root", + type=Path, + help="Scene Engine output root containing scene_export/.", + ) + parser.add_argument( + "--device", + default="cpu", + help="Simulation device, for example cpu or cuda.", + ) + parser.add_argument( + "--headless", + action="store_true", + help="Load and validate the exported scene without opening a window.", + ) + add_viser_args_to_parser(parser) + args = parser.parse_args(argv) + preview_scene_export( + output_root=args.output_root, + device=args.device, + headless=args.headless, + visualization=visualization_cfg_from_args(args), + ) + + +if __name__ == "__main__": + main() diff --git a/embodichain/gen_sim/scene_engine/cli/start.py b/embodichain/gen_sim/scene_engine/cli/start.py new file mode 100644 index 000000000..427e1a2f8 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/cli/start.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 + +import argparse +from collections.abc import Sequence +from pathlib import Path + +from embodichain.gen_sim.scene_engine.pipeline.generate import generate_scene_from_image + +_SUPPORTED_IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png"} + + +def cli_scene_engine( + image: str | Path, + output_root: str | Path, + *, + 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}") + 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, + # 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!") + + +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( + "--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", + ) + 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(argv) + + cli_scene_engine(args.image, args.output_root, config_path=args.config) + + +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..6044d37e8 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/clients/geometry_generation.py @@ -0,0 +1,458 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from contextlib import ExitStack +import json +from pathlib import Path +import time +from typing import Any + +import requests + +_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_objects_path: str, + session: requests.Session | None = None, + ) -> None: + self._base_url = base_url.rstrip("/") + self._timeout_s = timeout_s + self._max_attempts = max_attempts + self._health_path = health_path + self._generate_objects_path = generate_objects_path + self._session = session or requests.Session() + + @classmethod + def from_config( + cls, + config_path: str | Path | None = None, + ) -> "GeometryGenerationClient": + return cls(**_load_config(config_path)) + + def check_health(self) -> None: + last_error: Exception | None = None + for _ in range(self._max_attempts): + try: + response = self._session.get( + self._url(self._health_path), + timeout=10, # Use a shorter timeout for avoiding long waits. + ) + response.raise_for_status() + response_data = response.json() + if ( + not isinstance(response_data, dict) + or response_data.get("ok") is not True + ): + raise RuntimeError( + "Geometry Generation Server health response does not contain ok=true." + ) + return + except (requests.RequestException, ValueError, RuntimeError) as exc: + last_error = exc + + assert last_error is not None + raise RuntimeError( + "Geometry Generation Server health check failed after " + f"{self._max_attempts} attempts." + ) from last_error + + def close(self) -> None: + self._session.close() + + def generate_objects( + self, + *, + image_path: str | Path, + object_masks: list[tuple[str, Path]], + output_root: str | Path, + ) -> tuple[dict[str, Any], list[dict[str, Any]]]: + """Generate objects through the geometry server's mask-list endpoint. + + The 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. + resolved_image_path = Path(image_path).expanduser().resolve() + if not resolved_image_path.is_file(): + raise FileNotFoundError( + f"Geometry generation input not found: {resolved_image_path}" + ) + if not object_masks: + raise ValueError("Geometry generation object_masks must not be empty.") + object_ids = [object_id for object_id, _ in object_masks] + if len(set(object_ids)) != len(object_ids): + raise ValueError("Geometry generation object_ids must be unique.") + + resolved_object_masks: list[tuple[str, Path]] = [] + for object_id, mask_path in object_masks: + resolved_mask_path = Path(mask_path).expanduser().resolve() + if not resolved_mask_path.is_file(): + raise FileNotFoundError( + f"Geometry generation mask not found: {resolved_mask_path}" + ) + resolved_object_masks.append((object_id, resolved_mask_path)) + + # Send one multipart image + masks request, 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. + 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_objects( + self, + *, + image_path: Path, + object_masks: list[tuple[str, Path]], + ) -> tuple[dict[str, Any], list[dict[str, Any]]]: + last_error: Exception | None = None + for _ in range(self._max_attempts): + try: + # This stack manages the context of multiple open files, ensuring they are closed after the request. + with ExitStack() as stack: + image_file = stack.enter_context(image_path.open("rb")) + mask_files = [ + stack.enter_context(mask_path.open("rb")) + for _, mask_path in object_masks + ] + response = self._session.post( + self._url(self._generate_objects_path), + files=[ + ( + "image", + ( + image_path.name, + image_file, + _image_content_type(image_path), + ), + ), + *[ + ( + "masks", + (f"{object_id}.png", mask_file, "image/png"), + ) + for (object_id, _), mask_file in zip( + object_masks, + mask_files, + ) + ], + ], + timeout=self._timeout_s, + ) + response.raise_for_status() + try: + response_data = response.json() + except ValueError as exc: + raise RuntimeError( + "Geometry Generation Server response is not valid JSON." + ) from exc + response_data = self._wait_for_task_if_needed(response_data) + response_objects = _parse_objects_response( + response_data, + object_ids=[object_id for object_id, _ in object_masks], + ) + return response_data, response_objects + except (requests.RequestException, RuntimeError) as exc: + last_error = exc + + assert last_error is not None + raise RuntimeError( + "Geometry Generation Server request failed after " + f"{self._max_attempts} attempts." + ) from last_error + + def _wait_for_task_if_needed(self, response_data: object) -> dict[str, Any]: + """Poll a queued 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): + try: + response = self._session.get( + self._mesh_url(mesh_path), + timeout=self._timeout_s, + ) + response.raise_for_status() + glb_bytes = response.content + if not glb_bytes.startswith(b"glTF"): + raise RuntimeError( + "Geometry Generation Server returned invalid GLB content." + ) + output_path.write_bytes(glb_bytes) + return + except (requests.RequestException, RuntimeError) as exc: + last_error = exc + + assert last_error is not None + raise RuntimeError( + "Geometry Generation Server GLB download failed after " + f"{self._max_attempts} attempts: {mesh_path}" + ) from last_error + + def _mesh_url(self, mesh_path: str) -> str: + if mesh_path.startswith(("http://", "https://")): + return mesh_path + return self._url(mesh_path) + + def _url(self, path: str) -> str: + return f"{self._base_url}/{path.lstrip('/')}" + + +def _parse_objects_response( + response_data: object, + *, + object_ids: list[str], +) -> list[dict[str, Any]]: + if not isinstance(response_data, dict): + raise RuntimeError("Geometry Generation Server response must be a JSON object.") + if response_data.get("ok") is not True: + raise RuntimeError( + "Geometry Generation Server request failed: " + f"{response_data.get('error', 'ok is not true')}" + ) + result = response_data.get("result") + if not isinstance(result, dict): + raise RuntimeError( + "Geometry Generation Server response must contain a result object." + ) + response_objects = result.get("objects") + if not isinstance(response_objects, list) or len(response_objects) != len( + object_ids + ): + raise RuntimeError( + "Geometry Generation Server response object count does not match masks." + ) + + parsed_objects: list[dict[str, Any]] = [] + for index, (object_id, response_object) in enumerate( + zip(object_ids, response_objects) + ): + if not isinstance(response_object, dict): + raise RuntimeError( + f"Geometry Generation Server object {index} must be a JSON object." + ) + if response_object.get("name") != object_id: + raise RuntimeError( + "Geometry Generation Server object name does not match its " + f"requested id: {object_id!r}." + ) + mesh_path = response_object.get("mesh") + if not isinstance(mesh_path, str) or not mesh_path: + raise RuntimeError( + f"Geometry Generation Server object {index} has no mesh path." + ) + parsed_objects.append( + { + "mesh": mesh_path, + "rotation_quaternion_wxyz": _parse_numeric_list( + response_object.get("rotation_quaternion_wxyz"), + expected_length=4, + field_name=f"objects[{index}].rotation_quaternion_wxyz", + ), + "translation": _parse_numeric_list( + response_object.get("translation"), + expected_length=3, + field_name=f"objects[{index}].translation", + ), + "scale": _parse_numeric_list( + response_object.get("scale"), + expected_length=3, + field_name=f"objects[{index}].scale", + ), + } + ) + return parsed_objects + + +def _parse_numeric_list( + value: object, + *, + expected_length: int, + field_name: str, +) -> list[float]: + if not isinstance(value, list) or len(value) != expected_length: + raise RuntimeError( + f"Geometry Generation Server response field {field_name} is invalid." + ) + try: + return [float(item) for item in value] + except (TypeError, ValueError) as exc: + raise RuntimeError( + f"Geometry Generation Server response field {field_name} must be numeric." + ) from exc + + +def _image_content_type(image_path: Path) -> str: + if image_path.suffix.lower() in {".jpg", ".jpeg"}: + return "image/jpeg" + return "image/png" + + +def _load_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_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_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_objects_path": config["generate_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..642901ab3 --- /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_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..5819ce68e --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/generate.py @@ -0,0 +1,121 @@ +# ---------------------------------------------------------------------------- +# 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.scene_export import export_scene + + +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 + ) + 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 + 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 + ) + 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 + log_stage_start("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") + + return scene diff --git a/embodichain/gen_sim/scene_engine/pipeline/scene_export.py b/embodichain/gen_sim/scene_engine/pipeline/scene_export.py new file mode 100644 index 000000000..7593a3c95 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/scene_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( + *, + 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/scene_generation.py b/embodichain/gen_sim/scene_engine/pipeline/scene_generation.py new file mode 100644 index 000000000..6db70010f --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/scene_generation.py @@ -0,0 +1,574 @@ +# ---------------------------------------------------------------------------- +# 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_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..42237711a --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_generation_utils.py @@ -0,0 +1,1542 @@ +# ---------------------------------------------------------------------------- +# 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 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"}) + + +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 = 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 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], + *, + 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..3685ec8ae --- /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 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( + 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 93096a5e62c751dce30406417ad5a47ddcaf1087 Mon Sep 17 00:00:00 2001 From: PengXuanchao Date: Mon, 3 Aug 2026 17:54:16 +0800 Subject: [PATCH 27/46] add .env --- embodichain/__main__.py | 5 + embodichain/gen_sim/.env.example | 39 ++++++++ embodichain/gen_sim/env.py | 95 ++++++++++++++++++ embodichain/gen_sim/gradio_ui/app_config.py | 45 +++++---- .../gen_sim/gradio_ui/app_workflows.py | 6 -- .../gradio_visualization_architecture.md | 9 +- embodichain/gen_sim/gradio_ui/random_input.py | 6 +- embodichain/gen_sim/scene_engine/cli/start.py | 30 ++---- .../clients/geometry_generation.py | 98 ++++++------------- .../clients/image_segmentation.py | 98 ++++++------------- .../configs/scene_engine_config.json | 25 ----- .../gen_sim/scene_engine/llms/load_config.py | 38 +++---- .../llms/openai_compatible_client.py | 8 +- .../gen_sim/scene_engine/pipeline/generate.py | 20 ++-- 14 files changed, 264 insertions(+), 258 deletions(-) create mode 100644 embodichain/gen_sim/.env.example create mode 100644 embodichain/gen_sim/env.py delete mode 100644 embodichain/gen_sim/scene_engine/configs/scene_engine_config.json diff --git a/embodichain/__main__.py b/embodichain/__main__.py index 0af9abae6..70050b26a 100644 --- a/embodichain/__main__.py +++ b/embodichain/__main__.py @@ -51,6 +51,11 @@ 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 from an input image using configured services.", + ), Command( name="preview-asset", target="embodichain.lab.scripts.preview_asset:cli", diff --git a/embodichain/gen_sim/.env.example b/embodichain/gen_sim/.env.example new file mode 100644 index 000000000..9a7352a6c --- /dev/null +++ b/embodichain/gen_sim/.env.example @@ -0,0 +1,39 @@ +# Shared GenSim configuration +# Copy this file to .env and set deployment-specific values. Values exported +# by the shell, container, or CI environment take precedence over this file. + +# Common OpenAI-compatible LLM endpoint used by Scene Engine. +OPENAI_API_KEY="" +OPENAI_MODEL="" +OPENAI_BASE_URL="" +SCENE_ENGINE_OPENAI_DEFAULT_QUERY="{}" +OPENAI_MAX_ATTEMPTS=3 + +# Scene Engine services. +SCENE_ENGINE_IMAGE_SEGMENTATION_BASE_URL="" +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="" +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" + +# Gradio application and local workbench settings. +EMBODICHAIN_ROOT="" +GRADIO_SERVER_NAME="0.0.0.0" +GRADIO_SERVER_PORT=7860 +SCENE_ENGINE_VISER_PORT=8080 +ARTICRAFT_VISER_PORT=8081 +ARTICRAFT_ROOT="" +ARTICRAFT_REPOSITORY_URL="https://github.com/mattzh72/articraft.git" +ARTICRAFT_CONDA_ENV="articraft" +ARTICRAFT_OUTPUT_ROOT="" + +# Optional SimReady endpoint. These values are mapped to OPENAI_* only for +# SimReady subprocesses, leaving Scene Engine settings unchanged. +SIMREADY_OPENAI_API_KEY="" +SIMREADY_OPENAI_MODEL="" +SIMREADY_OPENAI_BASE_URL="" diff --git a/embodichain/gen_sim/env.py b/embodichain/gen_sim/env.py new file mode 100644 index 000000000..be5bc5b4d --- /dev/null +++ b/embodichain/gen_sim/env.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. +# ---------------------------------------------------------------------------- + +"""Load the shared GenSim environment configuration.""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import MutableMapping + +__all__ = ["find_gen_sim_env_file", "load_gen_sim_env"] + + +def find_gen_sim_env_file() -> Path: + """Return the configured shared ``.env`` file path. + + ``EMBODICHAIN_ENV_FILE`` is useful for deployments that keep secrets outside + the source tree. Otherwise GenSim uses ``embodichain/gen_sim/.env``. The + repository-root ``.env`` remains a backward-compatible fallback. + """ + configured_path = os.environ.get("EMBODICHAIN_ENV_FILE") + if configured_path: + return Path(configured_path).expanduser().resolve() + default_path = Path(__file__).resolve().parent / ".env" + if default_path.is_file(): + return default_path + + +def load_gen_sim_env(env: MutableMapping[str, str] | None = None) -> Path | None: + """Load missing variables from the shared GenSim ``.env`` file. + + Existing process environment variables are never overwritten so container, + CI, and shell-provided settings retain precedence over the local file. + + Args: + env: Environment mapping to populate. Defaults to :data:`os.environ`. + + Returns: + The loaded path, or ``None`` when no local ``.env`` file exists. + + Raises: + ValueError: If the file contains an invalid ``KEY=VALUE`` entry. + """ + target_env = os.environ if env is None else env + env_path = find_gen_sim_env_file() + if not env_path.is_file(): + return None + + for line_number, raw_line in enumerate( + env_path.read_text(encoding="utf-8").splitlines(), start=1 + ): + parsed = _parse_env_line(raw_line) + if parsed is None: + continue + key, value = parsed + if not key.isidentifier(): + raise ValueError( + f"Invalid environment variable name at {env_path}:{line_number}: {key!r}" + ) + target_env.setdefault(key, value) + return env_path + + +def _parse_env_line(line: str) -> tuple[str, str] | None: + """Parse one conventional dotenv line without requiring a third-party package.""" + stripped = line.strip() + if not stripped or stripped.startswith("#"): + return None + if stripped.startswith("export "): + stripped = stripped.removeprefix("export ").lstrip() + if "=" not in stripped: + raise ValueError(f"Expected KEY=VALUE entry, got: {line!r}") + + key, value = stripped.split("=", maxsplit=1) + key = key.strip() + value = value.strip() + if len(value) >= 2 and value[0] == value[-1] and value[0] in {"'", '"'}: + value = value[1:-1] + elif " #" in value: + value = value.split(" #", maxsplit=1)[0].rstrip() + return key, value diff --git a/embodichain/gen_sim/gradio_ui/app_config.py b/embodichain/gen_sim/gradio_ui/app_config.py index 7de206cf2..2be27b736 100644 --- a/embodichain/gen_sim/gradio_ui/app_config.py +++ b/embodichain/gen_sim/gradio_ui/app_config.py @@ -14,10 +14,11 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Central configuration for the Gradio application. +"""Shared settings and helpers for the Gradio application. -Keep deployment-specific paths, UI copy, and CLI command definitions here so -application modules do not embed environment-specific values. +Deployment-specific values are read from ``embodichain/gen_sim/.env``; this +module keeps UI constants, path derivation, and CLI command definitions close +to the application code. """ from __future__ import annotations @@ -26,6 +27,10 @@ from pathlib import Path from typing import Any +from embodichain.gen_sim.env import load_gen_sim_env + +load_gen_sim_env() + PROXY_ENV_KEYS = ( "HTTP_PROXY", "HTTPS_PROXY", @@ -38,13 +43,19 @@ ) DIRECT_NO_PROXY_VALUE = "*" + +def _getenv(name: str, default: str) -> str: + """Read a non-empty shared ``.env`` value, falling back to ``default``.""" + return os.environ.get(name) or default + + # SimReady uses an OpenAI-compatible multimodal endpoint. Configure these # values here for a local deployment, or provide the matching SIMREADY_* env # vars before launch. Keep the API key out of commits; an empty value leaves # any inherited OPENAI_* variables and SimReady's own JSON configuration intact. -SIMREADY_OPENAI_API_KEY = os.environ.get("SIMREADY_OPENAI_API_KEY", "") -SIMREADY_OPENAI_MODEL = os.environ.get("SIMREADY_OPENAI_MODEL", "") -SIMREADY_OPENAI_BASE_URL = os.environ.get("SIMREADY_OPENAI_BASE_URL", "") +SIMREADY_OPENAI_API_KEY = _getenv("SIMREADY_OPENAI_API_KEY", "") +SIMREADY_OPENAI_MODEL = _getenv("SIMREADY_OPENAI_MODEL", "") +SIMREADY_OPENAI_BASE_URL = _getenv("SIMREADY_OPENAI_BASE_URL", "") def configure_direct_network_env(env: Any = None) -> None: @@ -74,7 +85,7 @@ def configure_simready_llm_env(env: Any = None) -> None: APP_ROOT = Path(__file__).resolve().parent EMBODICHAIN_ROOT = Path( - os.environ.get("EMBODICHAIN_ROOT", "/home/dex/桌面/EmbodiChain") + _getenv("EMBODICHAIN_ROOT", str(Path(__file__).resolve().parents[3])) ).expanduser() ASSETS_DIR = APP_ROOT / "assets" DEXFORCE_LOGO = ASSETS_DIR / "dexforce.png" @@ -82,25 +93,22 @@ def configure_simready_llm_env(env: Any = None) -> None: DEBUG_ENGINE_ROOT = APP_ROOT / ".debug_engine" DEBUG_ASSET_ENGINE_ROOT = DEBUG_ENGINE_ROOT / "assets" ARTICRAFT_ROOT = Path( - os.environ.get("ARTICRAFT_ROOT", str(APP_ROOT / ".articraft")) + _getenv("ARTICRAFT_ROOT", str(APP_ROOT / ".articraft")) ).expanduser() -ARTICRAFT_REPOSITORY_URL = os.environ.get( +ARTICRAFT_REPOSITORY_URL = _getenv( "ARTICRAFT_REPOSITORY_URL", "https://github.com/mattzh72/articraft.git" ) -ARTICRAFT_CONDA_ENV = os.environ.get("ARTICRAFT_CONDA_ENV", "articraft") +ARTICRAFT_CONDA_ENV = _getenv("ARTICRAFT_CONDA_ENV", "articraft") # Keep every Articraft record, copied reference image, log, and downloadable # result bundle under one app-owned directory rather than the source checkout. ARTICRAFT_OUTPUT_ROOT = Path( - os.environ.get("ARTICRAFT_OUTPUT_ROOT", str(DEBUG_ENGINE_ROOT / "articraft")) + _getenv("ARTICRAFT_OUTPUT_ROOT", str(DEBUG_ENGINE_ROOT / "articraft")) ).expanduser() DEBUG_SCENE_ENGINE_ROOT = DEBUG_ENGINE_ROOT / "scenes" -SCENE_ENGINE_CONFIG = ( - EMBODICHAIN_ROOT / "embodichain" / "gen_sim" / "scene_engine_config.json" -) -SCENE_ENGINE_VISER_PORT = int(os.environ.get("SCENE_ENGINE_VISER_PORT", "8080")) +SCENE_ENGINE_VISER_PORT = int(_getenv("SCENE_ENGINE_VISER_PORT", "8080")) # Articulation previews run as a separate Viser process from scene previews, # so they need their own externally configurable port. -ARTICRAFT_VISER_PORT = int(os.environ.get("ARTICRAFT_VISER_PORT", "8081")) +ARTICRAFT_VISER_PORT = int(_getenv("ARTICRAFT_VISER_PORT", "8081")) SCENE_ID = "current" GYM_PROJECT_ROOT = EMBODICHAIN_ROOT / "gym_project" @@ -264,7 +272,6 @@ def configure_simready_llm_env(env: Any = None) -> None: "single_num_envs": "1", }, # Scene Engine is dispatched by EmbodiChain's registered top-level CLI. - # The scene_engine package itself has no __main__.py in this checkout. "scene_engine": { "module": "embodichain", "base_args": ("scene-engine",), @@ -298,6 +305,6 @@ def configure_simready_llm_env(env: Any = None) -> None: "action_graph_execution", ) -SERVER_NAME = os.environ.get("GRADIO_SERVER_NAME", "0.0.0.0") -SERVER_PORT = int(os.environ.get("GRADIO_SERVER_PORT", "7860")) +SERVER_NAME = _getenv("GRADIO_SERVER_NAME", "0.0.0.0") +SERVER_PORT = int(_getenv("GRADIO_SERVER_PORT", "7860")) DEFAULT_CONCURRENCY_LIMIT = 1 diff --git a/embodichain/gen_sim/gradio_ui/app_workflows.py b/embodichain/gen_sim/gradio_ui/app_workflows.py index 412443ae4..efd2960c8 100644 --- a/embodichain/gen_sim/gradio_ui/app_workflows.py +++ b/embodichain/gen_sim/gradio_ui/app_workflows.py @@ -2043,10 +2043,6 @@ def run_scene_engine(image_value: str | np.ndarray | Image.Image): preview_html = "" try: scene_hash, output_root, image_path = _prepare_scene_engine_input(image_value) - if not SCENE_ENGINE_CONFIG.is_file(): - raise FileNotFoundError( - f"Scene Engine config not found: {SCENE_ENGINE_CONFIG}" - ) except Exception as exc: with runtime_lock: set_runtime_phase_locked("failed") @@ -2093,8 +2089,6 @@ def run_scene_engine(image_value: str | np.ndarray | Image.Image): str(image_path), "--output_root", str(output_root), - "--config", - str(SCENE_ENGINE_CONFIG), ] with runtime_lock: runtime.log_lines.append("$ " + " ".join(command)) diff --git a/embodichain/gen_sim/gradio_ui/gradio_visualization_architecture.md b/embodichain/gen_sim/gradio_ui/gradio_visualization_architecture.md index 6fd293aed..712d56954 100644 --- a/embodichain/gen_sim/gradio_ui/gradio_visualization_architecture.md +++ b/embodichain/gen_sim/gradio_ui/gradio_visualization_architecture.md @@ -20,7 +20,8 @@ app_workflows.py ────► EmbodiChain Scene Engine CLI + Viser ├──────────────► app_processes.py 子进程、环境、日志和阶段检测 ├──────────────► app_state.py 共享 RuntimeState、锁和计时 ├──────────────► app_media.py 视频、数据集预览和日志归档 - └──────────────► app_config.py 路径、端口、文案和固定参数 + └──────────────► app_config.py UI 常量、路径推导和命令定义 + └──────────► ../.env 部署路径、端口和服务凭据 ``` | 模块 | 职责 | @@ -34,7 +35,8 @@ app_workflows.py ────► EmbodiChain Scene Engine CLI + Viser | `app_state.py` | `RuntimeState`、互斥锁、进度阶段、运行 token 和耗时统计。 | | `app_commands.py` | prompt2scene、动作配置和 `run_agent` 的参数构造。 | | `app_media.py` | 观众视频、LeRobot 数据预览、组合视频和运行日志归档。 | -| `app_config.py` | 路径、环境变量、端口、UI 文案、引擎模式和 CLI 固定参数。 | +| `app_config.py` | UI 文案、引擎模式、路径推导和 CLI 固定参数；部署值从 `.env` 读取。 | +| `../.env` | Gradio 与 Scene Engine 共用的路径、端口、LLM 和服务端点配置；不提交凭据。 | ## 启动、路径和网络环境 @@ -172,7 +174,6 @@ image → python -m embodichain scene-engine --image --output_root - --config /embodichain/gen_sim/scene_engine_config.json → /scene_export/scene_config.json → preview.py --viser --viser-host 0.0.0.0 --viser-port 8080 → Gradio iframe @@ -247,7 +248,7 @@ embodichain.gen_sim.action_agent_pipeline.cli.run_agent ```text python -m embodichain scene-engine embodichain/gen_sim/scene_engine/cli/preview.py -embodichain/gen_sim/scene_engine_config.json +.env ``` Articulation 还需要 Git（首次 clone）、Conda、`ARTICRAFT_CONDA_ENV` 和 Codex CLI。生成请求会交给本机 Codex CLI 执行，因此只应提交可信请求。 diff --git a/embodichain/gen_sim/gradio_ui/random_input.py b/embodichain/gen_sim/gradio_ui/random_input.py index 1ef138cdf..618cdd170 100644 --- a/embodichain/gen_sim/gradio_ui/random_input.py +++ b/embodichain/gen_sim/gradio_ui/random_input.py @@ -25,8 +25,12 @@ import numpy as np +from embodichain.gen_sim.env import load_gen_sim_env + +load_gen_sim_env() + EMBODICHAIN_ROOT = Path( - os.environ.get("EMBODICHAIN_ROOT", "/home/dex/桌面/EmbodiChain") + os.environ.get("EMBODICHAIN_ROOT") or str(Path(__file__).resolve().parents[3]) ).expanduser() APP_ROOT = Path(__file__).resolve().parent IMAGE_DIR = Path( diff --git a/embodichain/gen_sim/scene_engine/cli/start.py b/embodichain/gen_sim/scene_engine/cli/start.py index 427e1a2f8..f36c3eb1f 100644 --- a/embodichain/gen_sim/scene_engine/cli/start.py +++ b/embodichain/gen_sim/scene_engine/cli/start.py @@ -20,18 +20,14 @@ from collections.abc import Sequence from pathlib import Path -from embodichain.gen_sim.scene_engine.pipeline.generate import generate_scene_from_image - _SUPPORTED_IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png"} def cli_scene_engine( image: str | Path, output_root: str | Path, - *, - config_path: str | Path | None = None, ) -> None: - """Generate one scene using an optional user-owned service configuration.""" + """Generate one scene using the shared GenSim ``.env`` configuration.""" resolved_image_path = Path(image).expanduser().resolve() if not resolved_image_path.exists(): raise FileNotFoundError(f"Image input not found: {resolved_image_path}") @@ -45,15 +41,16 @@ def cli_scene_engine( resolved_output_root = Path(output_root).expanduser().resolve() resolved_output_root.mkdir(parents=True, exist_ok=True) + # Importing the generation pipeline initializes simulation dependencies. + # Keeping it here lets ``embodichain scene-engine --help`` work in a + # lightweight CLI-only environment. + from embodichain.gen_sim.scene_engine.pipeline.generate import ( + generate_scene_from_image, + ) + 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!") @@ -75,18 +72,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 config containing the llm, " - "image_segmentation, and geometry_generation service settings." - ), - ) 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 6044d37e8..8ff3df14b 100644 --- a/embodichain/gen_sim/scene_engine/clients/geometry_generation.py +++ b/embodichain/gen_sim/scene_engine/clients/geometry_generation.py @@ -19,15 +19,14 @@ 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.env import load_gen_sim_env class GeometryGenerationClient: @@ -51,11 +50,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_env(cls) -> "GeometryGenerationClient": + """Create a client from the shared GenSim ``.env`` configuration.""" + return cls(**_load_config()) def check_health(self) -> None: last_error: Exception | None = None @@ -390,69 +387,30 @@ 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.") - - required_keys = ( - "base_url", - "timeout_s", - "max_attempts", - "health_path", - "generate_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}") +def _load_config() -> dict[str, Any]: + load_gen_sim_env() + prefix = "SCENE_ENGINE_GEOMETRY_GENERATION_" + return { + "base_url": _read_required_string(f"{prefix}BASE_URL"), + "timeout_s": _read_positive_int(f"{prefix}TIMEOUT_S"), + "max_attempts": _read_positive_int(f"{prefix}MAX_ATTEMPTS"), + "health_path": _read_required_string(f"{prefix}HEALTH_PATH"), + "generate_objects_path": _read_required_string(f"{prefix}OBJECTS_PATH"), + } - 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." - ) +def _read_required_string(name: str) -> str: + value = os.getenv(name, "").strip() + if not value: + raise ValueError(f"Missing required environment variable: {name}") + return value - string_keys = ( - "base_url", - "health_path", - "generate_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_objects_path": config["generate_objects_path"].strip(), - } +def _read_positive_int(name: str) -> int: + try: + value = int(os.getenv(name, "")) + except ValueError as exc: + raise ValueError(f"{name} must be an integer.") from exc + if value < 1: + raise ValueError(f"{name} must be at least 1.") + return 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..ea838cca9 100644 --- a/embodichain/gen_sim/scene_engine/clients/image_segmentation.py +++ b/embodichain/gen_sim/scene_engine/clients/image_segmentation.py @@ -17,15 +17,13 @@ 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.env import load_gen_sim_env class ImageSegmentationClient: @@ -48,12 +46,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_env(cls) -> "ImageSegmentationClient": + """Create a client from the shared GenSim ``.env`` configuration.""" + return cls(**_load_config()) def check_health(self) -> None: last_error: requests.RequestException | None = None @@ -136,68 +131,35 @@ 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}") +def _load_config() -> dict[str, Any]: + load_gen_sim_env() + prefix = "SCENE_ENGINE_IMAGE_SEGMENTATION_" + return { + "base_url": _read_required_string(f"{prefix}BASE_URL"), + "timeout_s": _read_positive_int(f"{prefix}TIMEOUT_S"), + "max_attempts": _read_positive_int(f"{prefix}MAX_ATTEMPTS"), + "health_path": _read_required_string(f"{prefix}HEALTH_PATH"), + "segment_single_object_path": _read_required_string( + f"{prefix}SINGLE_OBJECT_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." - ) +def _read_required_string(name: str) -> str: + value = os.getenv(name, "").strip() + if not value: + raise ValueError(f"Missing required environment variable: {name}") + return value - 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 _read_positive_int(name: str) -> int: + try: + value = int(os.getenv(name, "")) + except ValueError as exc: + raise ValueError(f"{name} must be an integer.") from exc + if value < 1: + raise ValueError(f"{name} must be at least 1.") + return value def _extract_rle_masks(response_data: dict[str, Any]) -> list[dict[str, Any]]: 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/llms/load_config.py b/embodichain/gen_sim/scene_engine/llms/load_config.py index f2a786399..6d75e306c 100644 --- a/embodichain/gen_sim/scene_engine/llms/load_config.py +++ b/embodichain/gen_sim/scene_engine/llms/load_config.py @@ -20,12 +20,9 @@ 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.env import load_gen_sim_env @dataclass(frozen=True) @@ -39,32 +36,23 @@ 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 LLM settings from the shared ``.env`` and process environment.""" + load_gen_sim_env() try: - raw_config = json.loads(resolved_config_path.read_text(encoding="utf-8")) + default_query = json.loads(os.getenv("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 be valid JSON." ) 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) + api_key = os.getenv("OPENAI_API_KEY", "") + model = os.getenv("OPENAI_MODEL", "") + base_url = os.getenv("OPENAI_BASE_URL", "") + max_attempts = os.getenv("OPENAI_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 { @@ -80,9 +68,9 @@ def load_llm_config(config_path: str | Path | None = None) -> LLMConfig: 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 + 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(), 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..9a9037a65 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_env(cls) -> "OpenAICompatibleVLM": + """Create a client from the shared GenSim ``.env`` configuration.""" + 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 5819ce68e..732dbdd85 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generate.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generate.py @@ -47,17 +47,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.""" + """Generate the initial core scene state from an input image and ``.env``.""" 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_env() scene = Scene() # 1. Scene Understanding @@ -72,10 +68,8 @@ def generate_scene_from_image( # 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 - ) + # Load the environment configuration and verify the segmentation service. + image_segmentation_client = ImageSegmentationClient.from_env() try: image_segmentation_client.check_health() # Error raising will happen internally. scene = segment_scene( @@ -91,10 +85,8 @@ def generate_scene_from_image( # 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 - ) + # Load the environment configuration and verify the geometry service. + geometry_generation_client = GeometryGenerationClient.from_env() try: geometry_generation_client.check_health() # Error raising will happen internally. scene = generate_scene_and_refine( From 93abb7a30411eef9f993f2a7300f856355a788f3 Mon Sep 17 00:00:00 2001 From: PengXuanchao Date: Mon, 3 Aug 2026 18:11:11 +0800 Subject: [PATCH 28/46] add clean process --- .../gen_sim/gradio_ui/app_articraft.py | 27 ++- .../gen_sim/gradio_ui/app_processes.py | 188 +++++++++++++++--- embodichain/gen_sim/gradio_ui/gradio_app.py | 14 +- 3 files changed, 187 insertions(+), 42 deletions(-) diff --git a/embodichain/gen_sim/gradio_ui/app_articraft.py b/embodichain/gen_sim/gradio_ui/app_articraft.py index 4fa9bba69..d7c2c81fc 100644 --- a/embodichain/gen_sim/gradio_ui/app_articraft.py +++ b/embodichain/gen_sim/gradio_ui/app_articraft.py @@ -49,7 +49,12 @@ ARTICRAFT_ROOT, ARTICRAFT_VISER_PORT, ) -from app_processes import read_process_output, start_pipeline, terminate_process_group +from app_processes import ( + read_process_output, + register_managed_process, + start_pipeline, + terminate_process_group, +) __all__ = [ "build_articraft_panel", @@ -727,15 +732,17 @@ def generate_articraft_asset(prompt_value: str, image_value: Any): ), "" try: - process = subprocess.Popen( - codex_command, - cwd=ARTICRAFT_ROOT, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - text=True, - bufsize=1, - start_new_session=True, - env=os.environ.copy(), + process = register_managed_process( + subprocess.Popen( + codex_command, + cwd=ARTICRAFT_ROOT, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + start_new_session=True, + env=os.environ.copy(), + ) ) except Exception as exc: yield None, record_dir.as_posix(), f"**Codex could not start:** {exc}", "\n".join( diff --git a/embodichain/gen_sim/gradio_ui/app_processes.py b/embodichain/gen_sim/gradio_ui/app_processes.py index 993ac076c..58302e114 100644 --- a/embodichain/gen_sim/gradio_ui/app_processes.py +++ b/embodichain/gen_sim/gradio_ui/app_processes.py @@ -23,13 +23,30 @@ import signal import subprocess import sys +import threading import time from pathlib import Path from app_config import * # noqa: F403 - process settings are central configuration. from app_state import PHASES +__all__ = [ + "build_pipeline_env", + "build_run_agent_command", + "detect_phase_from_files", + "force_stop_all_child_processes", + "read_process_output", + "register_managed_process", + "run_agent_cli_supports_robot_profile", + "start_pipeline", + "terminate_process_group", + "update_phase_from_log", +] + _RUN_AGENT_SUPPORTS_ROBOT_PROFILE: bool | None = None +_managed_processes: dict[int, subprocess.Popen[str]] = {} +_managed_processes_lock = threading.Lock() +_shutdown_requested = False def run_agent_cli_supports_robot_profile() -> bool: @@ -74,15 +91,17 @@ def build_run_agent_command( def start_pipeline(command: list[str]) -> subprocess.Popen[str]: env = build_pipeline_env() env["PYTHONUNBUFFERED"] = "1" - return subprocess.Popen( - command, - cwd=EMBODICHAIN_ROOT, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - text=True, - bufsize=1, - start_new_session=True, - env=env, + return register_managed_process( + subprocess.Popen( + command, + cwd=EMBODICHAIN_ROOT, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + start_new_session=True, + env=env, + ) ) @@ -93,28 +112,147 @@ def build_pipeline_env() -> dict[str, str]: return env -def terminate_process_group(process: subprocess.Popen[str]) -> None: - if process.poll() is not None: - return +def register_managed_process( + process: subprocess.Popen[str], +) -> subprocess.Popen[str]: + """Register a UI-owned subprocess for application-shutdown cleanup. + + Processes must be registered immediately after they are created. If Gradio + shutdown has already begun, the new process is stopped before this function + returns so a callback cannot leave an orphan behind. + """ + with _managed_processes_lock: + if not _shutdown_requested: + _managed_processes[process.pid] = process + return process + + terminate_process_group(process) + return process + + +def _unregister_managed_process(process: subprocess.Popen[str]) -> None: + with _managed_processes_lock: + _managed_processes.pop(process.pid, None) + + +def _child_process_ids(parent_pid: int) -> set[int]: + """Return a snapshot of every descendant of ``parent_pid`` on POSIX.""" try: - os.killpg(process.pid, signal.SIGTERM) - except ProcessLookupError: - return - except Exception: - process.terminate() + result = subprocess.run( + ["ps", "-eo", "pid=,ppid="], + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + text=True, + timeout=2, + check=False, + ) + except (OSError, subprocess.TimeoutExpired): + return set() + + children_by_parent: dict[int, set[int]] = {} + for line in result.stdout.splitlines(): + fields = line.split() + if len(fields) != 2 or not all(field.isdecimal() for field in fields): + continue + pid, ppid = (int(field) for field in fields) + children_by_parent.setdefault(ppid, set()).add(pid) + + descendants: set[int] = set() + pending = list(children_by_parent.get(parent_pid, set())) + while pending: + pid = pending.pop() + if pid in descendants: + continue + descendants.add(pid) + pending.extend(children_by_parent.get(pid, set())) + return descendants + + +def _force_stop_process_ids(process_ids: set[int]) -> None: + """Stop unregistered child PIDs, escalating from SIGTERM to SIGKILL.""" + process_ids.discard(os.getpid()) + for pid in process_ids: + try: + os.kill(pid, signal.SIGTERM) + except ProcessLookupError: + continue + except PermissionError: + continue deadline = time.monotonic() + PROCESS_STOP_TIMEOUT_S - while time.monotonic() < deadline: - if process.poll() is not None: - return - time.sleep(0.2) + remaining = set(process_ids) + while remaining and time.monotonic() < deadline: + remaining = {pid for pid in remaining if _process_is_running(pid)} + if remaining: + time.sleep(0.1) + + for pid in remaining: + try: + os.kill(pid, signal.SIGKILL) + except ProcessLookupError: + continue + except PermissionError: + continue + +def _process_is_running(pid: int) -> bool: try: - os.killpg(process.pid, signal.SIGKILL) + os.kill(pid, 0) except ProcessLookupError: - return - except Exception: - process.kill() + return False + except PermissionError: + return True + try: + status = Path(f"/proc/{pid}/stat").read_text().rsplit(")", maxsplit=1)[1] + except (FileNotFoundError, IndexError, PermissionError): + return True + return not status.lstrip().startswith("Z") + + +def force_stop_all_child_processes() -> None: + """Force-stop every subprocess owned by the Gradio application. + + Registered processes are stopped by their isolated process groups, which + also stops their descendants. A second descendant scan catches short-lived + or legacy subprocesses that were not registered explicitly. + """ + global _shutdown_requested + with _managed_processes_lock: + _shutdown_requested = True + managed_processes = tuple(_managed_processes.values()) + child_process_ids = _child_process_ids(os.getpid()) + + for process in managed_processes: + terminate_process_group(process) + + _force_stop_process_ids(child_process_ids) + + +def terminate_process_group(process: subprocess.Popen[str]) -> None: + try: + if process.poll() is not None: + return + try: + os.killpg(process.pid, signal.SIGTERM) + except ProcessLookupError: + return + except Exception: + process.terminate() + + deadline = time.monotonic() + PROCESS_STOP_TIMEOUT_S + while time.monotonic() < deadline: + if process.poll() is not None: + return + time.sleep(0.2) + + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + return + except Exception: + process.kill() + finally: + _unregister_managed_process(process) def detect_phase_from_files(current_key: str, paths: ScenePaths) -> str: diff --git a/embodichain/gen_sim/gradio_ui/gradio_app.py b/embodichain/gen_sim/gradio_ui/gradio_app.py index c33a8de28..a924d1c41 100644 --- a/embodichain/gen_sim/gradio_ui/gradio_app.py +++ b/embodichain/gen_sim/gradio_ui/gradio_app.py @@ -24,7 +24,6 @@ import signal -from app_articraft import stop_articraft_viser_preview from app_config import ( ASSETS_DIR, DEBUG_ENGINE_ROOT, @@ -33,23 +32,24 @@ SERVER_NAME, SERVER_PORT, ) +from app_processes import force_stop_all_child_processes from app_services import build_demo __all__ = ["main"] -def _stop_child_previews() -> None: - """Release UI-owned preview subprocesses without masking app shutdown.""" +def _stop_child_processes() -> None: + """Force-stop UI-owned subprocesses without masking app shutdown.""" try: - stop_articraft_viser_preview() + force_stop_all_child_processes() except Exception: # Shutdown must not be blocked by an already-exited preview process. pass def _handle_shutdown_signal(signum: int, _frame: object) -> None: - """Terminate the Articraft Viser child before leaving the Gradio process.""" - _stop_child_previews() + """Terminate UI subprocesses before leaving the Gradio process.""" + _stop_child_processes() if signum == signal.SIGINT: raise KeyboardInterrupt raise SystemExit(128 + signum) @@ -78,7 +78,7 @@ def main() -> None: ], ) finally: - _stop_child_previews() + _stop_child_processes() if __name__ == "__main__": 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 29/46] 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 30/46] 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 31/46] 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 e4bc9e31dd078c711396c5c6eff6146ff7bfae48 Mon Sep 17 00:00:00 2001 From: PengXuanchao Date: Tue, 4 Aug 2026 09:51:01 +0800 Subject: [PATCH 32/46] add env config --- .../gen_sim/gradio_ui/app_articraft.py | 2 +- .../gen_sim/gradio_ui/app_asset_engine.py | 7 +- embodichain/gen_sim/gradio_ui/app_config.py | 91 ++------------ embodichain/gen_sim/gradio_ui/app_env.py | 112 ++++++++++++++++++ embodichain/gen_sim/gradio_ui/app_media.py | 1 + .../gen_sim/gradio_ui/app_processes.py | 7 +- .../gen_sim/gradio_ui/app_workflows.py | 1 + embodichain/gen_sim/gradio_ui/gradio_app.py | 4 +- .../gradio_visualization_architecture.md | 6 +- 9 files changed, 135 insertions(+), 96 deletions(-) create mode 100644 embodichain/gen_sim/gradio_ui/app_env.py diff --git a/embodichain/gen_sim/gradio_ui/app_articraft.py b/embodichain/gen_sim/gradio_ui/app_articraft.py index d7c2c81fc..4a425c90b 100644 --- a/embodichain/gen_sim/gradio_ui/app_articraft.py +++ b/embodichain/gen_sim/gradio_ui/app_articraft.py @@ -42,7 +42,7 @@ import gradio as gr -from app_config import ( +from app_env import ( ARTICRAFT_CONDA_ENV, ARTICRAFT_OUTPUT_ROOT, ARTICRAFT_REPOSITORY_URL, diff --git a/embodichain/gen_sim/gradio_ui/app_asset_engine.py b/embodichain/gen_sim/gradio_ui/app_asset_engine.py index 3c021ec90..f63212a46 100644 --- a/embodichain/gen_sim/gradio_ui/app_asset_engine.py +++ b/embodichain/gen_sim/gradio_ui/app_asset_engine.py @@ -35,11 +35,8 @@ import trimesh from app_articraft import build_articraft_panel -from app_config import ( - DEBUG_ASSET_ENGINE_ROOT, - EMBODICHAIN_ROOT, - SIMREADY_MESH_SUFFIXES, -) +from app_config import DEBUG_ASSET_ENGINE_ROOT, SIMREADY_MESH_SUFFIXES +from app_env import EMBODICHAIN_ROOT from app_processes import read_process_output, start_pipeline diff --git a/embodichain/gen_sim/gradio_ui/app_config.py b/embodichain/gen_sim/gradio_ui/app_config.py index 2be27b736..21085ea53 100644 --- a/embodichain/gen_sim/gradio_ui/app_config.py +++ b/embodichain/gen_sim/gradio_ui/app_config.py @@ -14,104 +14,29 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Shared settings and helpers for the Gradio application. +"""Static settings and helpers for the Gradio application. -Deployment-specific values are read from ``embodichain/gen_sim/.env``; this -module keeps UI constants, path derivation, and CLI command definitions close -to the application code. +Deployment-specific settings live in :mod:`app_env`, which reads the shared +``embodichain/gen_sim/.env`` file. This module keeps UI constants, path +derivation, and CLI command definitions close to the application code. """ from __future__ import annotations -import os from pathlib import Path -from typing import Any - -from embodichain.gen_sim.env import load_gen_sim_env - -load_gen_sim_env() - -PROXY_ENV_KEYS = ( - "HTTP_PROXY", - "HTTPS_PROXY", - "ALL_PROXY", - "FTP_PROXY", - "http_proxy", - "https_proxy", - "all_proxy", - "ftp_proxy", -) -DIRECT_NO_PROXY_VALUE = "*" - - -def _getenv(name: str, default: str) -> str: - """Read a non-empty shared ``.env`` value, falling back to ``default``.""" - return os.environ.get(name) or default - - -# SimReady uses an OpenAI-compatible multimodal endpoint. Configure these -# values here for a local deployment, or provide the matching SIMREADY_* env -# vars before launch. Keep the API key out of commits; an empty value leaves -# any inherited OPENAI_* variables and SimReady's own JSON configuration intact. -SIMREADY_OPENAI_API_KEY = _getenv("SIMREADY_OPENAI_API_KEY", "") -SIMREADY_OPENAI_MODEL = _getenv("SIMREADY_OPENAI_MODEL", "") -SIMREADY_OPENAI_BASE_URL = _getenv("SIMREADY_OPENAI_BASE_URL", "") - - -def configure_direct_network_env(env: Any = None) -> None: - """Disable proxy inheritance for local pipeline and Gradio processes.""" - if env is None: - env = os.environ - for key in PROXY_ENV_KEYS: - env.pop(key, None) - env["NO_PROXY"] = DIRECT_NO_PROXY_VALUE - env["no_proxy"] = DIRECT_NO_PROXY_VALUE - env.setdefault("GRADIO_ANALYTICS_ENABLED", "False") - - -def configure_simready_llm_env(env: Any = None) -> None: - """Map app-level SimReady settings to the upstream CLI's environment.""" - if env is None: - env = os.environ - configured_values = { - "OPENAI_API_KEY": SIMREADY_OPENAI_API_KEY, - "OPENAI_MODEL": SIMREADY_OPENAI_MODEL, - "OPENAI_BASE_URL": SIMREADY_OPENAI_BASE_URL, - } - for key, value in configured_values.items(): - if value: - env[key] = value +import app_env APP_ROOT = Path(__file__).resolve().parent -EMBODICHAIN_ROOT = Path( - _getenv("EMBODICHAIN_ROOT", str(Path(__file__).resolve().parents[3])) -).expanduser() ASSETS_DIR = APP_ROOT / "assets" DEXFORCE_LOGO = ASSETS_DIR / "dexforce.png" INTERACT_RANDOM_PREVIEW_DIR = APP_ROOT / ".gradio_previews" DEBUG_ENGINE_ROOT = APP_ROOT / ".debug_engine" DEBUG_ASSET_ENGINE_ROOT = DEBUG_ENGINE_ROOT / "assets" -ARTICRAFT_ROOT = Path( - _getenv("ARTICRAFT_ROOT", str(APP_ROOT / ".articraft")) -).expanduser() -ARTICRAFT_REPOSITORY_URL = _getenv( - "ARTICRAFT_REPOSITORY_URL", "https://github.com/mattzh72/articraft.git" -) -ARTICRAFT_CONDA_ENV = _getenv("ARTICRAFT_CONDA_ENV", "articraft") -# Keep every Articraft record, copied reference image, log, and downloadable -# result bundle under one app-owned directory rather than the source checkout. -ARTICRAFT_OUTPUT_ROOT = Path( - _getenv("ARTICRAFT_OUTPUT_ROOT", str(DEBUG_ENGINE_ROOT / "articraft")) -).expanduser() DEBUG_SCENE_ENGINE_ROOT = DEBUG_ENGINE_ROOT / "scenes" -SCENE_ENGINE_VISER_PORT = int(_getenv("SCENE_ENGINE_VISER_PORT", "8080")) -# Articulation previews run as a separate Viser process from scene previews, -# so they need their own externally configurable port. -ARTICRAFT_VISER_PORT = int(_getenv("ARTICRAFT_VISER_PORT", "8081")) SCENE_ID = "current" -GYM_PROJECT_ROOT = EMBODICHAIN_ROOT / "gym_project" +GYM_PROJECT_ROOT = app_env.EMBODICHAIN_ROOT / "gym_project" ACTION_AGENT_ROOT = GYM_PROJECT_ROOT / "action_agent_pipeline" IMAGE_DIR = ACTION_AGENT_ROOT / "images" AUTO_LOG_DIR = ACTION_AGENT_ROOT / "auto_logs" @@ -119,7 +44,7 @@ def configure_simready_llm_env(env: Any = None) -> None: PROMPT2SCENE_ROOT = GYM_PROJECT_ROOT / SCENE_ID CONFIG_DIR = ACTION_AGENT_ROOT / "configs" / SCENE_ID FAST_GYM_CONFIG = CONFIG_DIR / "fast_gym_config.json" -OUTPUTS_DIR = EMBODICHAIN_ROOT / "outputs" +OUTPUTS_DIR = app_env.EMBODICHAIN_ROOT / "outputs" CURRENT_GYM_EXPORT_DIR = PROMPT2SCENE_ROOT / "gym_export" CURRENT_GYM_EXPORT_CONFIG = CURRENT_GYM_EXPORT_DIR / "gym_config.json" GRADIO_SCENE_DIR = CONFIG_DIR / "gradio_scene" @@ -305,6 +230,4 @@ def configure_simready_llm_env(env: Any = None) -> None: "action_graph_execution", ) -SERVER_NAME = _getenv("GRADIO_SERVER_NAME", "0.0.0.0") -SERVER_PORT = int(_getenv("GRADIO_SERVER_PORT", "7860")) DEFAULT_CONCURRENCY_LIMIT = 1 diff --git a/embodichain/gen_sim/gradio_ui/app_env.py b/embodichain/gen_sim/gradio_ui/app_env.py new file mode 100644 index 000000000..3e450c0e7 --- /dev/null +++ b/embodichain/gen_sim/gradio_ui/app_env.py @@ -0,0 +1,112 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Environment-backed deployment settings for the Gradio application.""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any + +from embodichain.gen_sim.env import load_gen_sim_env + +__all__ = [ + "ARTICRAFT_CONDA_ENV", + "ARTICRAFT_OUTPUT_ROOT", + "ARTICRAFT_REPOSITORY_URL", + "ARTICRAFT_ROOT", + "ARTICRAFT_VISER_PORT", + "DIRECT_NO_PROXY_VALUE", + "EMBODICHAIN_ROOT", + "PROXY_ENV_KEYS", + "SCENE_ENGINE_VISER_PORT", + "SERVER_NAME", + "SERVER_PORT", + "SIMREADY_OPENAI_API_KEY", + "SIMREADY_OPENAI_BASE_URL", + "SIMREADY_OPENAI_MODEL", + "configure_direct_network_env", + "configure_simready_llm_env", +] + +load_gen_sim_env() + +APP_ROOT = Path(__file__).resolve().parent +DEBUG_ENGINE_ROOT = APP_ROOT / ".debug_engine" +PROXY_ENV_KEYS = ( + "HTTP_PROXY", + "HTTPS_PROXY", + "ALL_PROXY", + "FTP_PROXY", + "http_proxy", + "https_proxy", + "all_proxy", + "ftp_proxy", +) +DIRECT_NO_PROXY_VALUE = "*" + + +def _getenv(name: str, default: str) -> str: + """Read a non-empty shared ``.env`` value, falling back to ``default``.""" + return os.environ.get(name) or default + + +EMBODICHAIN_ROOT = Path( + _getenv("EMBODICHAIN_ROOT", str(Path(__file__).resolve().parents[3])) +).expanduser() +ARTICRAFT_ROOT = Path( + _getenv("ARTICRAFT_ROOT", str(APP_ROOT / ".articraft")) +).expanduser() +ARTICRAFT_REPOSITORY_URL = _getenv( + "ARTICRAFT_REPOSITORY_URL", "https://github.com/mattzh72/articraft.git" +) +ARTICRAFT_CONDA_ENV = _getenv("ARTICRAFT_CONDA_ENV", "articraft") +ARTICRAFT_OUTPUT_ROOT = Path( + _getenv("ARTICRAFT_OUTPUT_ROOT", str(DEBUG_ENGINE_ROOT / "articraft")) +).expanduser() +SCENE_ENGINE_VISER_PORT = int(_getenv("SCENE_ENGINE_VISER_PORT", "8080")) +ARTICRAFT_VISER_PORT = int(_getenv("ARTICRAFT_VISER_PORT", "8081")) +SERVER_NAME = _getenv("GRADIO_SERVER_NAME", "0.0.0.0") +SERVER_PORT = int(_getenv("GRADIO_SERVER_PORT", "7860")) +SIMREADY_OPENAI_API_KEY = _getenv("SIMREADY_OPENAI_API_KEY", "") +SIMREADY_OPENAI_MODEL = _getenv("SIMREADY_OPENAI_MODEL", "") +SIMREADY_OPENAI_BASE_URL = _getenv("SIMREADY_OPENAI_BASE_URL", "") + + +def configure_direct_network_env(env: Any = None) -> None: + """Disable proxy inheritance for local pipeline and Gradio processes.""" + if env is None: + env = os.environ + for key in PROXY_ENV_KEYS: + env.pop(key, None) + env["NO_PROXY"] = DIRECT_NO_PROXY_VALUE + env["no_proxy"] = DIRECT_NO_PROXY_VALUE + env.setdefault("GRADIO_ANALYTICS_ENABLED", "False") + + +def configure_simready_llm_env(env: Any = None) -> None: + """Map app-level SimReady settings to the upstream CLI's environment.""" + if env is None: + env = os.environ + configured_values = { + "OPENAI_API_KEY": SIMREADY_OPENAI_API_KEY, + "OPENAI_MODEL": SIMREADY_OPENAI_MODEL, + "OPENAI_BASE_URL": SIMREADY_OPENAI_BASE_URL, + } + for key, value in configured_values.items(): + if value: + env[key] = value diff --git a/embodichain/gen_sim/gradio_ui/app_media.py b/embodichain/gen_sim/gradio_ui/app_media.py index 41271c0f5..5c4813c49 100644 --- a/embodichain/gen_sim/gradio_ui/app_media.py +++ b/embodichain/gen_sim/gradio_ui/app_media.py @@ -32,6 +32,7 @@ from PIL import Image, ImageDraw from app_config import * # noqa: F403 - media paths and limits are configuration. +from app_env import EMBODICHAIN_ROOT from app_state import format_timing_lines, runtime, runtime_lock, snapshot_timing_locked diff --git a/embodichain/gen_sim/gradio_ui/app_processes.py b/embodichain/gen_sim/gradio_ui/app_processes.py index 58302e114..ff052b7d5 100644 --- a/embodichain/gen_sim/gradio_ui/app_processes.py +++ b/embodichain/gen_sim/gradio_ui/app_processes.py @@ -27,7 +27,12 @@ import time from pathlib import Path -from app_config import * # noqa: F403 - process settings are central configuration. +from app_config import COMMANDS, PROCESS_STOP_TIMEOUT_S +from app_env import ( + EMBODICHAIN_ROOT, + configure_direct_network_env, + configure_simready_llm_env, +) from app_state import PHASES __all__ = [ diff --git a/embodichain/gen_sim/gradio_ui/app_workflows.py b/embodichain/gen_sim/gradio_ui/app_workflows.py index efd2960c8..2b77f5e5c 100644 --- a/embodichain/gen_sim/gradio_ui/app_workflows.py +++ b/embodichain/gen_sim/gradio_ui/app_workflows.py @@ -47,6 +47,7 @@ parse_task_id, ) from app_config import * # noqa: F403 - services intentionally consume central config. +from app_env import SCENE_ENGINE_VISER_PORT, configure_direct_network_env from app_processes import ( build_pipeline_env, build_run_agent_command, diff --git a/embodichain/gen_sim/gradio_ui/gradio_app.py b/embodichain/gen_sim/gradio_ui/gradio_app.py index a924d1c41..7dce2e1a8 100644 --- a/embodichain/gen_sim/gradio_ui/gradio_app.py +++ b/embodichain/gen_sim/gradio_ui/gradio_app.py @@ -28,10 +28,8 @@ ASSETS_DIR, DEBUG_ENGINE_ROOT, DEFAULT_CONCURRENCY_LIMIT, - EMBODICHAIN_ROOT, - SERVER_NAME, - SERVER_PORT, ) +from app_env import EMBODICHAIN_ROOT, SERVER_NAME, SERVER_PORT from app_processes import force_stop_all_child_processes from app_services import build_demo diff --git a/embodichain/gen_sim/gradio_ui/gradio_visualization_architecture.md b/embodichain/gen_sim/gradio_ui/gradio_visualization_architecture.md index 712d56954..694d5485b 100644 --- a/embodichain/gen_sim/gradio_ui/gradio_visualization_architecture.md +++ b/embodichain/gen_sim/gradio_ui/gradio_visualization_architecture.md @@ -21,6 +21,7 @@ app_workflows.py ────► EmbodiChain Scene Engine CLI + Viser ├──────────────► app_state.py 共享 RuntimeState、锁和计时 ├──────────────► app_media.py 视频、数据集预览和日志归档 └──────────────► app_config.py UI 常量、路径推导和命令定义 + └──────────────► app_env.py 部署配置读取 └──────────► ../.env 部署路径、端口和服务凭据 ``` @@ -35,7 +36,8 @@ app_workflows.py ────► EmbodiChain Scene Engine CLI + Viser | `app_state.py` | `RuntimeState`、互斥锁、进度阶段、运行 token 和耗时统计。 | | `app_commands.py` | prompt2scene、动作配置和 `run_agent` 的参数构造。 | | `app_media.py` | 观众视频、LeRobot 数据预览、组合视频和运行日志归档。 | -| `app_config.py` | UI 文案、引擎模式、路径推导和 CLI 固定参数；部署值从 `.env` 读取。 | +| `app_config.py` | UI 文案、引擎模式、路径推导和 CLI 固定参数。 | +| `app_env.py` | 从 `.env` 读取 Gradio、Articraft 和 SimReady 的部署值，并保留未配置时的默认值。 | | `../.env` | Gradio 与 Scene Engine 共用的路径、端口、LLM 和服务端点配置；不提交凭据。 | ## 启动、路径和网络环境 @@ -257,7 +259,7 @@ Articulation 还需要 Git（首次 clone）、Conda、`ARTICRAFT_CONDA_ENV` 和 ```bash python -m py_compile \ - gradio_app.py app_config.py app_state.py app_commands.py \ + gradio_app.py app_config.py app_env.py app_state.py app_commands.py \ app_processes.py app_media.py app_workflows.py app_ui.py \ app_asset_engine.py app_articraft.py app_services.py 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 33/46] 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 34/46] 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 fbe72ee4c08bdb3678160f674ee54cd088533c45 Mon Sep 17 00:00:00 2001 From: PengXuanchao Date: Tue, 4 Aug 2026 12:23:26 +0800 Subject: [PATCH 35/46] add reset --- embodichain/gen_sim/.env.example | 4 +- embodichain/gen_sim/env.py | 12 +- .../gen_sim/gradio_ui/app_articraft.py | 194 ++++++++++++++++-- .../gen_sim/gradio_ui/app_asset_engine.py | 179 ++++++++++++---- embodichain/gen_sim/gradio_ui/app_env.py | 9 +- embodichain/gen_sim/gradio_ui/app_state.py | 2 + embodichain/gen_sim/gradio_ui/app_ui.py | 19 +- .../gen_sim/gradio_ui/app_workflows.py | 126 ++++++++++-- .../gradio_visualization_architecture.md | 8 +- embodichain/gen_sim/gradio_ui/random_input.py | 6 +- 10 files changed, 462 insertions(+), 97 deletions(-) diff --git a/embodichain/gen_sim/.env.example b/embodichain/gen_sim/.env.example index 9a7352a6c..6544d0710 100644 --- a/embodichain/gen_sim/.env.example +++ b/embodichain/gen_sim/.env.example @@ -21,8 +21,8 @@ SCENE_ENGINE_GEOMETRY_GENERATION_MAX_ATTEMPTS=3 SCENE_ENGINE_GEOMETRY_GENERATION_HEALTH_PATH="/health" SCENE_ENGINE_GEOMETRY_GENERATION_OBJECTS_PATH="/generate_multiple_objects" -# Gradio application and local workbench settings. -EMBODICHAIN_ROOT="" +# Gradio application and local workbench settings. The EmbodiChain repository +# root is derived automatically from the installed source tree. GRADIO_SERVER_NAME="0.0.0.0" GRADIO_SERVER_PORT=7860 SCENE_ENGINE_VISER_PORT=8080 diff --git a/embodichain/gen_sim/env.py b/embodichain/gen_sim/env.py index be5bc5b4d..44b772568 100644 --- a/embodichain/gen_sim/env.py +++ b/embodichain/gen_sim/env.py @@ -22,7 +22,17 @@ from pathlib import Path from typing import MutableMapping -__all__ = ["find_gen_sim_env_file", "load_gen_sim_env"] +__all__ = ["find_gen_sim_env_file", "get_embodichain_root", "load_gen_sim_env"] + + +def get_embodichain_root() -> Path: + """Return the repository containing the installed GenSim source tree. + + The Gradio app always launches its local tools from this directory. It is + derived from this module instead of a machine-specific dotenv value, so a + checkout continues to work after it is moved or cloned elsewhere. + """ + return Path(__file__).resolve().parents[2] def find_gen_sim_env_file() -> Path: diff --git a/embodichain/gen_sim/gradio_ui/app_articraft.py b/embodichain/gen_sim/gradio_ui/app_articraft.py index 4a425c90b..a7fd81f9e 100644 --- a/embodichain/gen_sim/gradio_ui/app_articraft.py +++ b/embodichain/gen_sim/gradio_ui/app_articraft.py @@ -60,6 +60,7 @@ "build_articraft_panel", "configure_articraft_environment", "generate_articraft_asset", + "reset_articraft_asset", "stop_articraft_viser_preview", ] @@ -68,6 +69,111 @@ _ARTICRAFT_PYTHON_VERSION = "3.12" _CONDA_ENVIRONMENT_SETUP_TIMEOUT_SECONDS = 1_200 _articraft_environment_lock = threading.Lock() +_articraft_generation_lock = threading.Lock() +_articraft_generation_process: subprocess.Popen[str] | None = None +_articraft_generation_token: str | None = None +_ARTICRAFT_IDLE_PREVIEW = ( + "
" + "The interactive Viser articulation preview will appear here after generation." + "
" +) + + +def _begin_articraft_generation() -> str: + """Invalidate the previous generation and return a new ownership token.""" + global _articraft_generation_process, _articraft_generation_token + with _articraft_generation_lock: + previous_process = _articraft_generation_process + _articraft_generation_process = None + token = uuid.uuid4().hex + _articraft_generation_token = token + if previous_process is not None: + terminate_process_group(previous_process) + return token + + +def _articraft_generation_is_active( + token: str, process: subprocess.Popen[str] | None = None +) -> bool: + with _articraft_generation_lock: + return _articraft_generation_token == token and ( + process is None or _articraft_generation_process is process + ) + + +def _set_articraft_generation_process( + token: str, process: subprocess.Popen[str] +) -> bool: + global _articraft_generation_process + with _articraft_generation_lock: + if _articraft_generation_token != token: + return False + _articraft_generation_process = process + return True + + +def _finish_articraft_generation_process( + token: str, process: subprocess.Popen[str] +) -> None: + global _articraft_generation_process + with _articraft_generation_lock: + if ( + _articraft_generation_token == token + and _articraft_generation_process is process + ): + _articraft_generation_process = None + + +def _run_articraft_generation_check( + command: list[str], *, token: str, timeout: int +) -> subprocess.CompletedProcess[str] | None: + """Run one Articraft CLI gate so Reset can stop its whole process group.""" + process = register_managed_process( + subprocess.Popen( + command, + cwd=ARTICRAFT_ROOT, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + start_new_session=True, + env=os.environ.copy(), + ) + ) + if not _set_articraft_generation_process(token, process): + terminate_process_group(process) + return None + try: + stdout, _ = process.communicate(timeout=timeout) + except subprocess.TimeoutExpired: + terminate_process_group(process) + raise + finally: + _finish_articraft_generation_process(token, process) + if not _articraft_generation_is_active(token): + return None + return subprocess.CompletedProcess(command, process.returncode, stdout) + + +def reset_articraft_asset(): + """Clear Articraft inputs/results and stop its command and Viser processes.""" + global _articraft_generation_process, _articraft_generation_token + with _articraft_generation_lock: + process = _articraft_generation_process + _articraft_generation_process = None + _articraft_generation_token = None + if process is not None: + terminate_process_group(process) + stop_articraft_viser_preview() + return ( + "**Environment:** not checked.", + "", + None, + None, + "", + "**Status:** waiting for a description.", + "", + _ARTICRAFT_IDLE_PREVIEW, + ) def _command_path(name: str) -> str | None: @@ -656,9 +762,11 @@ def _build_codex_prompt( def generate_articraft_asset(prompt_value: str, image_value: Any): """Initialize a record, let Codex author it, and expose one result bundle.""" + token = _begin_articraft_generation() prompt = (prompt_value or "").strip() if not prompt: - yield None, "", "**Input error:** enter a description of the articulated object.", "", "" + if _articraft_generation_is_active(token): + yield None, "", "**Input error:** enter a description of the articulated object.", "", "" return details, errors, codex = _check_requirements() @@ -666,7 +774,8 @@ def generate_articraft_asset(prompt_value: str, image_value: Any): message = ( "\n".join(f"- {error}" for error in errors) or "Codex CLI is unavailable." ) - yield None, "", f"**Articulation is not ready.**\n\n{message}", "", "" + if _articraft_generation_is_active(token): + yield None, "", f"**Articulation is not ready.**\n\n{message}", "", "" return record_id = _record_id() @@ -688,7 +797,11 @@ def generate_articraft_asset(prompt_value: str, image_value: Any): prompt, ) log_lines.append("$ " + " ".join(init_command[:-1]) + " ") - initialized = _run_check(init_command, timeout=90) + initialized = _run_articraft_generation_check( + init_command, token=token, timeout=90 + ) + if initialized is None: + return log_lines.append(_short_output(initialized, limit=4000)) if initialized.returncode: yield None, "", "**Articraft record initialization failed.**", "\n".join( @@ -697,7 +810,11 @@ def generate_articraft_asset(prompt_value: str, image_value: Any): return model_path = _active_model_path(record_dir) except Exception as exc: - yield None, "", f"**Setup failed:** {exc}", "\n".join(log_lines), "" + if _articraft_generation_is_active(token): + yield None, "", f"**Setup failed:** {exc}", "\n".join(log_lines), "" + return + + if not _articraft_generation_is_active(token): return final_message = run_root / "codex_final_message.txt" @@ -744,10 +861,14 @@ def generate_articraft_asset(prompt_value: str, image_value: Any): env=os.environ.copy(), ) ) + if not _set_articraft_generation_process(token, process): + terminate_process_group(process) + return except Exception as exc: - yield None, record_dir.as_posix(), f"**Codex could not start:** {exc}", "\n".join( - log_lines - ), "" + if _articraft_generation_is_active(token): + yield None, record_dir.as_posix(), f"**Codex could not start:** {exc}", "\n".join( + log_lines + ), "" return output_queue: queue.Queue[str] = queue.Queue() @@ -756,6 +877,8 @@ def generate_articraft_asset(prompt_value: str, image_value: Any): ) reader.start() while process.poll() is None: + if not _articraft_generation_is_active(token, process): + return try: while True: log_lines.append(output_queue.get_nowait()) @@ -765,12 +888,18 @@ def generate_articraft_asset(prompt_value: str, image_value: Any): log_lines[-240:] ), "" time.sleep(0.75) - reader.join(timeout=2) try: - while True: - log_lines.append(output_queue.get_nowait()) - except queue.Empty: - pass + reader.join(timeout=2) + try: + while True: + log_lines.append(output_queue.get_nowait()) + except queue.Empty: + pass + finally: + _finish_articraft_generation_process(token, process) + + if not _articraft_generation_is_active(token): + return if final_message.is_file(): final_text = final_message.read_text(encoding="utf-8", errors="replace").strip() @@ -800,7 +929,11 @@ def generate_articraft_asset(prompt_value: str, image_value: Any): "", ) try: - checked = _run_check(check_command, timeout=300) + checked = _run_articraft_generation_check( + check_command, token=token, timeout=300 + ) + if checked is None: + return log_lines.append(_short_output(checked, limit=5000)) except Exception as exc: yield ( @@ -845,7 +978,11 @@ def generate_articraft_asset(prompt_value: str, image_value: Any): "", ) try: - compiled = _run_check(compile_command, timeout=300) + compiled = _run_articraft_generation_check( + compile_command, token=token, timeout=300 + ) + if compiled is None: + return log_lines.append(_short_output(compiled, limit=5000)) except Exception as exc: yield ( @@ -886,7 +1023,11 @@ def generate_articraft_asset(prompt_value: str, image_value: Any): ) log_lines.append("$ " + " ".join(finalize_command)) try: - finalized = _run_check(finalize_command, timeout=300) + finalized = _run_articraft_generation_check( + finalize_command, token=token, timeout=300 + ) + if finalized is None: + return log_lines.append(_short_output(finalized, limit=5000)) except Exception as exc: yield ( @@ -907,6 +1048,8 @@ def generate_articraft_asset(prompt_value: str, image_value: Any): ) return + if not _articraft_generation_is_active(token): + return try: materialized, archive = _make_result_bundle(record_id) status = ( @@ -938,6 +1081,7 @@ def build_articraft_panel() -> None: with gr.Row(): configure_button = gr.Button("Configure Articulation & check Codex") generate_button = gr.Button("Generate articulation", variant="primary") + reset_button = gr.Button("Reset Articulation", variant="stop") environment_status = gr.Markdown("**Environment:** not checked.") with gr.Row(): prompt = gr.Textbox( @@ -958,11 +1102,7 @@ def build_articraft_panel() -> None: record_folder = gr.Textbox( label="Articulation record folder", interactive=False ) - articulation_preview = gr.HTML( - "
" - "The interactive Viser articulation preview will appear here after generation." - "
" - ) + articulation_preview = gr.HTML(_ARTICRAFT_IDLE_PREVIEW) generation_status = gr.Markdown("**Status:** waiting for a description.") generation_log = gr.Textbox( label="Codex / Articraft log", lines=14, interactive=False @@ -982,3 +1122,17 @@ def build_articraft_panel() -> None: articulation_preview, ], ) + reset_button.click( + reset_articraft_asset, + outputs=[ + environment_status, + prompt, + image, + output_file, + record_folder, + generation_status, + generation_log, + articulation_preview, + ], + queue=False, + ) diff --git a/embodichain/gen_sim/gradio_ui/app_asset_engine.py b/embodichain/gen_sim/gradio_ui/app_asset_engine.py index f63212a46..5acfbfb76 100644 --- a/embodichain/gen_sim/gradio_ui/app_asset_engine.py +++ b/embodichain/gen_sim/gradio_ui/app_asset_engine.py @@ -26,7 +26,10 @@ import queue import shutil +import subprocess +import sys import threading +import time import uuid from pathlib import Path from typing import Any, Iterable @@ -37,7 +40,57 @@ from app_articraft import build_articraft_panel from app_config import DEBUG_ASSET_ENGINE_ROOT, SIMREADY_MESH_SUFFIXES from app_env import EMBODICHAIN_ROOT -from app_processes import read_process_output, start_pipeline +from app_processes import read_process_output, start_pipeline, terminate_process_group + +_simready_run_lock = threading.Lock() +_simready_process: subprocess.Popen[str] | None = None +_simready_run_token: str | None = None +_SIMREADY_IDLE_STATUS = "**Status:** waiting for an asset." + + +def _begin_simready_run() -> str: + """Invalidate any previous SimReady run and return a new ownership token.""" + global _simready_process, _simready_run_token + with _simready_run_lock: + previous_process = _simready_process + _simready_process = None + token = uuid.uuid4().hex + _simready_run_token = token + if previous_process is not None: + terminate_process_group(previous_process) + return token + + +def _simready_run_is_active( + token: str, process: subprocess.Popen[str] | None = None +) -> bool: + with _simready_run_lock: + return _simready_run_token == token and ( + process is None or _simready_process is process + ) + + +def _finish_simready_run( + token: str, process: subprocess.Popen[str] | None = None +) -> None: + global _simready_process + with _simready_run_lock: + if _simready_run_token == token and ( + process is None or _simready_process is process + ): + _simready_process = None + + +def reset_simready_asset(): + """Clear SimReady widgets and terminate the process group for its active run.""" + global _simready_process, _simready_run_token + with _simready_run_lock: + process = _simready_process + _simready_process = None + _simready_run_token = None + if process is not None: + terminate_process_group(process) + return None, "rigid_object", None, None, None, _SIMREADY_IDLE_STATUS, "" def _as_paths(value: Any) -> list[Path]: @@ -129,9 +182,12 @@ def _find_simready_output(output_root: Path) -> Path: def run_simready_asset(upload_value: Any, category: str): """Run one upstream SimReady job and stream concise subprocess progress.""" + global _simready_process + token = _begin_simready_run() category = (category or "").strip() if not category: - yield None, None, None, "**Input error:** enter an asset category.", "" + if _simready_run_is_active(token): + yield None, None, None, "**Input error:** enter an asset category.", "" return try: uploads = _as_paths(upload_value) @@ -142,11 +198,12 @@ def run_simready_asset(upload_value: Any, category: str): source_mesh = _safe_copy_uploads(uploads, input_dir) input_preview = _export_preview(source_mesh, run_root / "input_preview.glb") except Exception as exc: - yield None, None, None, f"**Input error:** {exc}", "" + if _simready_run_is_active(token): + yield None, None, None, f"**Input error:** {exc}", "" return command = [ - __import__("sys").executable, + sys.executable, "-m", "embodichain.gen_sim.simready_pipeline.cli.start", "--input_dir", @@ -157,6 +214,8 @@ def run_simready_asset(upload_value: Any, category: str): category, ] log_lines = ["$ " + " ".join(command)] + if not _simready_run_is_active(token): + return yield input_preview.as_posix(), None, None, "**SimReady is running…**", "\n".join( log_lines ) @@ -164,53 +223,69 @@ def run_simready_asset(upload_value: Any, category: str): try: process = start_pipeline(command) except Exception as exc: - yield input_preview.as_posix(), None, None, f"**Pipeline start failed:** {exc}", "\n".join( - log_lines - ) + if _simready_run_is_active(token): + yield input_preview.as_posix(), None, None, f"**Pipeline start failed:** {exc}", "\n".join( + log_lines + ) return - output_queue: queue.Queue[str] = queue.Queue() - reader = threading.Thread( - target=read_process_output, args=(process, output_queue), daemon=True - ) - reader.start() - while process.poll() is None: + with _simready_run_lock: + owns_process = _simready_run_token == token + if owns_process: + _simready_process = process + if not owns_process: + terminate_process_group(process) + return + + try: + output_queue: queue.Queue[str] = queue.Queue() + reader = threading.Thread( + target=read_process_output, args=(process, output_queue), daemon=True + ) + reader.start() + while process.poll() is None: + if not _simready_run_is_active(token, process): + return + try: + while True: + log_lines.append(output_queue.get_nowait()) + except queue.Empty: + pass + # Keep the browser responsive while the Blender/LLM stages run. + yield input_preview.as_posix(), None, None, "**SimReady is running…**", "\n".join( + log_lines[-160:] + ) + time.sleep(0.5) + reader.join(timeout=1) try: while True: log_lines.append(output_queue.get_nowait()) except queue.Empty: pass - # Keep the browser responsive while the Blender/LLM stages run. - yield input_preview.as_posix(), None, None, "**SimReady is running…**", "\n".join( - log_lines[-160:] - ) - __import__("time").sleep(0.5) - reader.join(timeout=1) - try: - while True: - log_lines.append(output_queue.get_nowait()) - except queue.Empty: - pass - - if process.returncode != 0: - yield input_preview.as_posix(), None, None, f"**SimReady failed** (exit code {process.returncode}).", "\n".join( - log_lines[-220:] - ) - return - try: - result = _find_simready_output(output_root) - preview = ( - result - if result.suffix.lower() == ".glb" - else _export_preview(result, run_root / "output_preview.glb") - ) - yield input_preview.as_posix(), preview.as_posix(), result.as_posix(), "**SimReady completed.**", "\n".join( - log_lines[-220:] - ) - except Exception as exc: - yield input_preview.as_posix(), None, None, f"**Output error:** {exc}", "\n".join( - log_lines[-220:] - ) + if not _simready_run_is_active(token, process): + return + + if process.returncode != 0: + yield input_preview.as_posix(), None, None, f"**SimReady failed** (exit code {process.returncode}).", "\n".join( + log_lines[-220:] + ) + return + try: + result = _find_simready_output(output_root) + preview = ( + result + if result.suffix.lower() == ".glb" + else _export_preview(result, run_root / "output_preview.glb") + ) + yield input_preview.as_posix(), preview.as_posix(), result.as_posix(), "**SimReady completed.**", "\n".join( + log_lines[-220:] + ) + except Exception as exc: + yield input_preview.as_posix(), None, None, f"**Output error:** {exc}", "\n".join( + log_lines[-220:] + ) + finally: + _finish_simready_run(token, process) def build_asset_engine_panel() -> dict[str, Any]: @@ -258,10 +333,11 @@ def build_asset_engine_panel() -> dict[str, Any]: ) with gr.Row(): run_button = gr.Button("Run SimReady", variant="primary") + reset_button = gr.Button("Reset SimReady", variant="stop") output_file = gr.File( label="SimReady asset output", interactive=False ) - status = gr.Markdown("**Status:** waiting for an asset.") + status = gr.Markdown(_SIMREADY_IDLE_STATUS) log = gr.Textbox(label="Pipeline log", lines=10, interactive=False) with gr.Tab("Articulation"): build_articraft_panel() @@ -277,4 +353,17 @@ def build_asset_engine_panel() -> dict[str, Any]: inputs=[uploads, category], outputs=[input_model, output_model, output_file, status, log], ) + reset_button.click( + reset_simready_asset, + outputs=[ + uploads, + category, + input_model, + output_model, + output_file, + status, + log, + ], + queue=False, + ) return {"panel": panel} diff --git a/embodichain/gen_sim/gradio_ui/app_env.py b/embodichain/gen_sim/gradio_ui/app_env.py index 3e450c0e7..f9227e25e 100644 --- a/embodichain/gen_sim/gradio_ui/app_env.py +++ b/embodichain/gen_sim/gradio_ui/app_env.py @@ -22,7 +22,7 @@ from pathlib import Path from typing import Any -from embodichain.gen_sim.env import load_gen_sim_env +from embodichain.gen_sim.env import get_embodichain_root, load_gen_sim_env __all__ = [ "ARTICRAFT_CONDA_ENV", @@ -65,9 +65,10 @@ def _getenv(name: str, default: str) -> str: return os.environ.get(name) or default -EMBODICHAIN_ROOT = Path( - _getenv("EMBODICHAIN_ROOT", str(Path(__file__).resolve().parents[3])) -).expanduser() +# The repository root must follow this checkout, not a machine-specific .env +# value. Its path is shared with child processes through their working +# directory, so deriving it once here keeps every Debug workflow relocatable. +EMBODICHAIN_ROOT = get_embodichain_root() ARTICRAFT_ROOT = Path( _getenv("ARTICRAFT_ROOT", str(APP_ROOT / ".articraft")) ).expanduser() diff --git a/embodichain/gen_sim/gradio_ui/app_state.py b/embodichain/gen_sim/gradio_ui/app_state.py index 260469af1..9ad6b9dac 100644 --- a/embodichain/gen_sim/gradio_ui/app_state.py +++ b/embodichain/gen_sim/gradio_ui/app_state.py @@ -58,7 +58,9 @@ class RuntimeState: language: str = LANGUAGE_EN process: subprocess.Popen[str] | None = None sim_process: subprocess.Popen[str] | None = None + scene_engine_process: subprocess.Popen[str] | None = None scene_preview_process: subprocess.Popen[str] | None = None + scene_engine_is_running: bool = False sim_started: bool = False sim_finished: bool = False sim_returncode: int | None = None diff --git a/embodichain/gen_sim/gradio_ui/app_ui.py b/embodichain/gen_sim/gradio_ui/app_ui.py index 8e1a10b4b..92e30eeed 100644 --- a/embodichain/gen_sim/gradio_ui/app_ui.py +++ b/embodichain/gen_sim/gradio_ui/app_ui.py @@ -100,7 +100,13 @@ def build_demo() -> gr.Blocks: format="png", height=300, ) - debug_scene_run = gr.Button("Generate scene", variant="primary") + with gr.Row(): + debug_scene_run = gr.Button( + "Generate scene", variant="primary" + ) + debug_scene_reset = gr.Button( + "Reset Scene Engine", variant="stop" + ) with gr.Column(scale=2): debug_scene_progress = gr.Slider( 0, @@ -343,6 +349,17 @@ def build_demo() -> gr.Blocks: debug_scene_preview, ], ) + debug_scene_reset.click( + reset_scene_engine, + outputs=[ + debug_scene_image, + debug_scene_progress, + debug_scene_status, + debug_scene_output, + debug_scene_preview, + ], + queue=False, + ) debug_action_load.click( action_engine_snapshot, outputs=[ diff --git a/embodichain/gen_sim/gradio_ui/app_workflows.py b/embodichain/gen_sim/gradio_ui/app_workflows.py index 2b77f5e5c..3544d7bd1 100644 --- a/embodichain/gen_sim/gradio_ui/app_workflows.py +++ b/embodichain/gen_sim/gradio_ui/app_workflows.py @@ -1258,6 +1258,8 @@ def run_generate( process = start_pipeline(command) except Exception as exc: with runtime_lock: + if runtime.run_token != token: + return runtime.is_busy = False runtime.process = None runtime.phase_key = "failed" @@ -2038,48 +2040,103 @@ def _viser_iframe(port: int, scene_hash: str) -> str: ) +def reset_scene_engine(): + """Clear Scene Engine widgets and stop its generator and Viser process groups.""" + with runtime_lock: + generator_process = runtime.scene_engine_process + preview_process = runtime.scene_preview_process + owns_runtime = runtime.scene_engine_is_running + other_workflow_running = not owns_runtime and runtime.is_busy + if owns_runtime or not runtime.is_busy: + if owns_runtime: + runtime.run_token = uuid.uuid4().hex + if runtime.process is generator_process: + runtime.process = None + runtime.is_busy = False + set_runtime_phase_locked("idle") + runtime.status = "Scene Engine reset." + runtime.image_path = None + runtime.last_error = None + runtime.log_lines.clear() + clear_run_timing_locked() + runtime.scene_engine_process = None + runtime.scene_preview_process = None + runtime.scene_engine_is_running = False + + for process in {generator_process, preview_process}: + if process is not None: + terminate_process_group(process) + + return ( + None, + PHASES["idle"].progress, + format_status( + "Scene Engine reset." + if not other_workflow_running + else "Scene Engine preview reset; another workflow is still running." + ), + "", + "
" + "The Viser preview will appear here after generation." + "
", + ) + + def run_scene_engine(image_value: str | np.ndarray | Image.Image): """Generate one image-conditioned scene and expose its Viser preview.""" output_root: Path | None = None preview_html = "" - try: - scene_hash, output_root, image_path = _prepare_scene_engine_input(image_value) - except Exception as exc: - with runtime_lock: - set_runtime_phase_locked("failed") - runtime.status = f"Input error: {exc}" - runtime.last_error = str(exc) - yield _scene_engine_updates(output_root, preview_html) - return - - old_preview: subprocess.Popen[str] | None = None - busy_message: str | None = None + token = uuid.uuid4().hex with runtime_lock: if runtime.is_busy: runtime.status = "Another pipeline is already running." runtime.last_error = runtime.status busy_message = runtime.status else: - old_preview = runtime.scene_preview_process - runtime.scene_preview_process = None - token = uuid.uuid4().hex runtime.run_token = token runtime.is_busy = True + runtime.scene_engine_is_running = True set_runtime_phase_locked("received") - runtime.status = ( - f"Image saved. Generating Scene Engine output {scene_hash}." - ) + runtime.status = "Preparing Scene Engine input." runtime.last_error = None - runtime.image_path = image_path runtime.log_lines.clear() clear_run_timing_locked() + busy_message = None if busy_message is not None: yield _scene_engine_updates(output_root, preview_html) return + try: + scene_hash, output_root, image_path = _prepare_scene_engine_input(image_value) + except Exception as exc: + with runtime_lock: + if runtime.run_token != token: + return + runtime.is_busy = False + runtime.scene_engine_is_running = False + set_runtime_phase_locked("failed") + runtime.status = f"Input error: {exc}" + runtime.last_error = str(exc) + yield _scene_engine_updates(output_root, preview_html) + return + + old_preview: subprocess.Popen[str] | None = None + old_generator: subprocess.Popen[str] | None = None + with runtime_lock: + if runtime.run_token != token: + return + old_preview = runtime.scene_preview_process + old_generator = runtime.scene_engine_process + runtime.scene_engine_process = None + runtime.scene_preview_process = None + runtime.status = f"Image saved. Generating Scene Engine output {scene_hash}." + runtime.image_path = image_path + if old_preview is not None: terminate_process_group(old_preview) + if old_generator is not None: + terminate_process_group(old_generator) command = [ sys.executable, @@ -2100,6 +2157,7 @@ def run_scene_engine(image_value: str | np.ndarray | Image.Image): except Exception as exc: with runtime_lock: runtime.is_busy = False + runtime.scene_engine_is_running = False set_runtime_phase_locked("failed") runtime.status = f"Scene Engine start failed: {exc}" runtime.last_error = str(exc) @@ -2115,6 +2173,7 @@ def run_scene_engine(image_value: str | np.ndarray | Image.Image): terminate_process_group(process) return runtime.process = process + runtime.scene_engine_process = process start_run_timing_locked("started") set_runtime_phase_locked("started") runtime.status = "Scene Engine generation started." @@ -2123,6 +2182,11 @@ def run_scene_engine(image_value: str | np.ndarray | Image.Image): while process.poll() is None: drained = drain_output_queue(output_queue) with runtime_lock: + if ( + runtime.run_token != token + or runtime.scene_engine_process is not process + ): + return for line in drained: runtime.log_lines.append(line) set_runtime_phase_locked( @@ -2136,12 +2200,15 @@ def run_scene_engine(image_value: str | np.ndarray | Image.Image): reader.join(timeout=1.0) with runtime_lock: + if runtime.run_token != token or runtime.scene_engine_process is not process: + return for line in drain_output_queue(output_queue): runtime.log_lines.append(line) set_runtime_phase_locked( _scene_engine_phase_from_log(line, runtime.phase_key) ) runtime.process = None + runtime.scene_engine_process = None scene_export = output_root / "scene_export" / "scene_config.json" if process.returncode != 0 or not scene_export.is_file(): @@ -2151,7 +2218,10 @@ def run_scene_engine(image_value: str | np.ndarray | Image.Image): else f"Scene Engine did not create {scene_export}." ) with runtime_lock: + if runtime.run_token != token: + return runtime.is_busy = False + runtime.scene_engine_is_running = False set_runtime_phase_locked("failed") runtime.status = detail runtime.last_error = detail @@ -2173,7 +2243,10 @@ def run_scene_engine(image_value: str | np.ndarray | Image.Image): preview_process = start_pipeline(preview_command) except Exception as exc: with runtime_lock: + if runtime.run_token != token: + return runtime.is_busy = False + runtime.scene_engine_is_running = False set_runtime_phase_locked("failed") runtime.status = f"Viser preview start failed: {exc}" runtime.last_error = str(exc) @@ -2181,6 +2254,10 @@ def run_scene_engine(image_value: str | np.ndarray | Image.Image): return with runtime_lock: + if runtime.run_token != token: + terminate_process_group(preview_process) + return + runtime.scene_preview_process = preview_process runtime.log_lines.append("$ " + " ".join(preview_command)) set_runtime_phase_locked("preview") runtime.status = "Starting Viser preview..." @@ -2189,7 +2266,11 @@ def run_scene_engine(image_value: str | np.ndarray | Image.Image): if not _wait_for_viser(port, preview_process): terminate_process_group(preview_process) with runtime_lock: + if runtime.run_token != token: + return + runtime.scene_preview_process = None runtime.is_busy = False + runtime.scene_engine_is_running = False set_runtime_phase_locked("failed") runtime.status = "Viser preview did not start." runtime.last_error = runtime.status @@ -2198,8 +2279,15 @@ def run_scene_engine(image_value: str | np.ndarray | Image.Image): preview_html = _viser_iframe(port, scene_hash) with runtime_lock: + if ( + runtime.run_token != token + or runtime.scene_preview_process is not preview_process + ): + terminate_process_group(preview_process) + return runtime.scene_preview_process = preview_process runtime.is_busy = False + runtime.scene_engine_is_running = False set_runtime_phase_locked("complete") runtime.status = "Scene generated successfully. Viser preview is ready." runtime.last_error = None diff --git a/embodichain/gen_sim/gradio_ui/gradio_visualization_architecture.md b/embodichain/gen_sim/gradio_ui/gradio_visualization_architecture.md index 694d5485b..f8cc2fe80 100644 --- a/embodichain/gen_sim/gradio_ui/gradio_visualization_architecture.md +++ b/embodichain/gen_sim/gradio_ui/gradio_visualization_architecture.md @@ -50,7 +50,7 @@ conda run -n embodichain python gradio_app.py | 变量 | 默认值 | 用途 | | --- | --- | --- | -| `EMBODICHAIN_ROOT` | `/home/dex/桌面/EmbodiChain` | EmbodiChain 根目录。 | +| EmbodiChain root | 自动从 `embodichain/gen_sim/env.py` 的源码位置推导 | EmbodiChain 根目录；不再从 `.env` 配置。 | | `GRADIO_SERVER_NAME` | `0.0.0.0` | Gradio 监听地址。 | | `GRADIO_SERVER_PORT` | `7860` | Gradio 监听端口。 | | `SCENE_ENGINE_VISER_PORT` | `8080` | 独立 Scene Engine 的 Viser 端口。 | @@ -147,6 +147,8 @@ python -m embodichain.gen_sim.simready_pipeline.cli.start \ 处理函数以 generator 持续返回最近的 stdout；完成时优先预览 `asset_simready.glb`，只有 OBJ 时再转为 GLB。此路径不依赖 DexSim。 +`Reset SimReady` 会清空上传、类别、预览、下载项和日志，并按进程组终止正在运行的 SimReady CLI 及其子进程。 + ### Articulation：Articraft + Codex Articulation 标签页根据文本和可选参考图生成一个可下载的 articulated asset。先点击环境检查：若 `ARTICRAFT_ROOT` 不存在，应用会 clone `ARTICRAFT_REPOSITORY_URL`；随后检查 Conda、指定的 Articraft 环境和 Codex CLI。该操作会创建 checkout 和 `.debug_engine/articraft/` 中的输出目录，现有的非 Articraft 目录不会被覆盖。 @@ -166,6 +168,8 @@ description + optional image 产物、记录和参考图均在 `ARTICRAFT_OUTPUT_ROOT` 下，不能直接当作 Demo 的 Gym 场景或 SimReady 资产；若要进入后续仿真，需要另行定义并实现转换/导入流程。Articraft Viser 每次成功预览会终止旧的 Articraft 预览进程，再以 `0.0.0.0:` 启动新进程。 +`Reset Articulation` 会清空描述、参考图、记录与下载结果，终止当前 Articraft/Codex 命令进程组，并关闭该面板启动的 Viser。 + ## 独立 Scene engine 和 Viser Scene engine 只接收图像。上传图像会先进行 EXIF 归正并转为 RGB PNG，以 PNG 字节的 SHA-256 前 16 位作为目录名；相同图像会复用同一目录： @@ -183,6 +187,8 @@ image 当 `scene_export/scene_config.json` 存在且生成进程返回成功时，应用才启动 Viser。iframe 使用 Gradio 页面当前的协议和主机名转向 Viser 端口，因此从其他设备访问时，浏览器必须能访问该端口。每次新 Scene Engine 任务开始前会终止旧的 Scene Viser 进程。输出目录会显示在 UI 中，便于检查 hash 命名的场景导出。 +`Reset Scene Engine` 会清空图像、进度、输出目录和 iframe，并终止当前生成命令与 Scene Viser 的进程组；运行 token 会使已经失效的生成器停止回写界面。 + ## Action engine：Gym 场景契约 Action engine 不接收裸 GLB。普通 GLB 只有渲染数据，而 DexSim 还需要碰撞、物理参数、初始位姿、资源相对路径和 action 配置。当前实现的前置条件是： diff --git a/embodichain/gen_sim/gradio_ui/random_input.py b/embodichain/gen_sim/gradio_ui/random_input.py index 618cdd170..9ce553992 100644 --- a/embodichain/gen_sim/gradio_ui/random_input.py +++ b/embodichain/gen_sim/gradio_ui/random_input.py @@ -25,13 +25,11 @@ import numpy as np -from embodichain.gen_sim.env import load_gen_sim_env +from embodichain.gen_sim.env import get_embodichain_root, load_gen_sim_env load_gen_sim_env() -EMBODICHAIN_ROOT = Path( - os.environ.get("EMBODICHAIN_ROOT") or str(Path(__file__).resolve().parents[3]) -).expanduser() +EMBODICHAIN_ROOT = get_embodichain_root() APP_ROOT = Path(__file__).resolve().parent IMAGE_DIR = Path( os.environ.get( 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 36/46] 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 37/46] 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 38/46] 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 39/46] 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 40/46] 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 41/46] 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 42/46] 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 43/46] 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 44539f3779b778519402438630805efc1f34d96c Mon Sep 17 00:00:00 2001 From: PengXuanchao Date: Wed, 5 Aug 2026 14:18:57 +0800 Subject: [PATCH 44/46] fix bug --- embodichain/gen_sim/gradio_ui/app_workflows.py | 1 + 1 file changed, 1 insertion(+) diff --git a/embodichain/gen_sim/gradio_ui/app_workflows.py b/embodichain/gen_sim/gradio_ui/app_workflows.py index 3544d7bd1..9154ab84e 100644 --- a/embodichain/gen_sim/gradio_ui/app_workflows.py +++ b/embodichain/gen_sim/gradio_ui/app_workflows.py @@ -2232,6 +2232,7 @@ def run_scene_engine(image_value: str | np.ndarray | Image.Image): preview_command = [ sys.executable, COMMANDS["scene_engine"]["preview_script"], + "--output_root", str(output_root), "--viser", "--viser-host", From 9b0c5ebe555c6916112919a49313c92db8f6ab9f Mon Sep 17 00:00:00 2001 From: PengXuanchao Date: Wed, 5 Aug 2026 14:55:03 +0800 Subject: [PATCH 45/46] run black --- 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 3e3c50bab8541434d0e0a1194f6ea0e12eef385e Mon Sep 17 00:00:00 2001 From: PengXuanchao Date: Thu, 6 Aug 2026 19:16:10 +0800 Subject: [PATCH 46/46] add log --- embodichain/gen_sim/gradio_ui/app_processes.py | 16 ++++++++++++++-- embodichain/gen_sim/gradio_ui/app_workflows.py | 9 ++++++++- 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/embodichain/gen_sim/gradio_ui/app_processes.py b/embodichain/gen_sim/gradio_ui/app_processes.py index ff052b7d5..a0f414f3a 100644 --- a/embodichain/gen_sim/gradio_ui/app_processes.py +++ b/embodichain/gen_sim/gradio_ui/app_processes.py @@ -329,8 +329,20 @@ def update_phase_from_log(line: str, current_key: str) -> str: def read_process_output( process: subprocess.Popen[str], output_queue: queue.Queue[str], + log_path: Path | None = None, ) -> None: + """Forward merged subprocess output to the UI queue and an optional log.""" if process.stdout is None: return - for line in process.stdout: - output_queue.put(line.rstrip()) + log_file = log_path.open("a", encoding="utf-8") if log_path is not None else None + try: + for line in process.stdout: + output_queue.put(line.rstrip()) + if log_file is not None: + log_file.write(line) + if not line.endswith("\n"): + log_file.write("\n") + log_file.flush() + finally: + if log_file is not None: + log_file.close() diff --git a/embodichain/gen_sim/gradio_ui/app_workflows.py b/embodichain/gen_sim/gradio_ui/app_workflows.py index 9154ab84e..4895454eb 100644 --- a/embodichain/gen_sim/gradio_ui/app_workflows.py +++ b/embodichain/gen_sim/gradio_ui/app_workflows.py @@ -2148,6 +2148,11 @@ def run_scene_engine(image_value: str | np.ndarray | Image.Image): "--output_root", str(output_root), ] + scene_engine_log = output_root / "scene_engine.log" + scene_engine_log.write_text( + "$ " + " ".join(command) + "\n", + encoding="utf-8", + ) with runtime_lock: runtime.log_lines.append("$ " + " ".join(command)) yield _scene_engine_updates(output_root, preview_html) @@ -2166,7 +2171,9 @@ def run_scene_engine(image_value: str | np.ndarray | Image.Image): output_queue: queue.Queue[str] = queue.Queue() reader = threading.Thread( - target=read_process_output, args=(process, output_queue), daemon=True + target=read_process_output, + args=(process, output_queue, scene_engine_log), + daemon=True, ) with runtime_lock: if runtime.run_token != token: