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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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/53] 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: From e5d9d20bc15eead4d600f5506a4204e48502bc12 Mon Sep 17 00:00:00 2001 From: PengXuanchao Date: Mon, 10 Aug 2026 10:46:19 +0800 Subject: [PATCH 47/53] delete demo --- .gitignore | 4 +- embodichain/__main__.py | 5 + .../gen_sim/gradio_ui/app_asset_engine.py | 10 +- embodichain/gen_sim/gradio_ui/app_commands.py | 126 +- embodichain/gen_sim/gradio_ui/app_config.py | 167 +- embodichain/gen_sim/gradio_ui/app_env.py | 4 +- embodichain/gen_sim/gradio_ui/app_media.py | 655 +--- .../gen_sim/gradio_ui/app_processes.py | 78 +- embodichain/gen_sim/gradio_ui/app_services.py | 4 +- embodichain/gen_sim/gradio_ui/app_state.py | 133 +- embodichain/gen_sim/gradio_ui/app_ui.py | 669 +--- .../gen_sim/gradio_ui/app_workflows.py | 3180 ++--------------- embodichain/gen_sim/gradio_ui/gradio_app.py | 12 +- .../gradio_visualization_architecture.md | 109 +- embodichain/gen_sim/gradio_ui/random_input.py | 542 --- 15 files changed, 516 insertions(+), 5182 deletions(-) delete mode 100644 embodichain/gen_sim/gradio_ui/random_input.py diff --git a/.gitignore b/.gitignore index e61027a6b..17b7d5c41 100644 --- a/.gitignore +++ b/.gitignore @@ -208,9 +208,9 @@ scripts/benchmark/rl/reports/* .worktrees/ # Local gym project workspace /gym_project/ -.debug_engine/ +.gen_sim/ # 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/.gen_sim/ /embodichain/gen_sim/gradio_ui/__pycache__/ diff --git a/embodichain/__main__.py b/embodichain/__main__.py index 0e73bd076..975e7649d 100644 --- a/embodichain/__main__.py +++ b/embodichain/__main__.py @@ -96,6 +96,11 @@ class Command: target="embodichain.workspace_cache_cli:main", help="Inspect and clean workspace analyzer caches.", ), + Command( + name="analyze-workspace", + target="embodichain.lab.scripts.analyze_workspace:cli", + help="Analyze a robot's reachable workspace from a URDF/USD asset.", + ), ) diff --git a/embodichain/gen_sim/gradio_ui/app_asset_engine.py b/embodichain/gen_sim/gradio_ui/app_asset_engine.py index 5acfbfb76..4aa57bcc9 100644 --- a/embodichain/gen_sim/gradio_ui/app_asset_engine.py +++ b/embodichain/gen_sim/gradio_ui/app_asset_engine.py @@ -14,7 +14,7 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Standalone SimReady asset-engine workflow used by Debug mode. +"""Standalone SimReady asset-engine workflow used by the engine workspace. The upstream SimReady CLI works on a directory, while Gradio uploads files. This adapter creates an isolated directory for every run, keeps material @@ -38,7 +38,7 @@ import trimesh from app_articraft import build_articraft_panel -from app_config import DEBUG_ASSET_ENGINE_ROOT, SIMREADY_MESH_SUFFIXES +from app_config import GEN_SIM_ASSET_ROOT, SIMREADY_MESH_SUFFIXES from app_env import EMBODICHAIN_ROOT from app_processes import read_process_output, start_pipeline, terminate_process_group @@ -151,7 +151,7 @@ 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" + preview = GEN_SIM_ASSET_ROOT / "previews" / f"{uuid.uuid4().hex}.glb" _export_preview(source, preview) return ( preview.as_posix(), @@ -192,7 +192,7 @@ def run_simready_asset(upload_value: Any, category: str): try: uploads = _as_paths(upload_value) _mesh_path(uploads) - run_root = DEBUG_ASSET_ENGINE_ROOT / "runs" / uuid.uuid4().hex + run_root = GEN_SIM_ASSET_ROOT / "runs" / uuid.uuid4().hex input_dir = run_root / "input" output_root = run_root / "output" source_mesh = _safe_copy_uploads(uploads, input_dir) @@ -289,7 +289,7 @@ def run_simready_asset(upload_value: Any, category: str): def build_asset_engine_panel() -> dict[str, Any]: - """Create the Debug Asset-engine panel and return its event endpoints.""" + """Create the 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." diff --git a/embodichain/gen_sim/gradio_ui/app_commands.py b/embodichain/gen_sim/gradio_ui/app_commands.py index 7dc1f4a1c..1c14861b2 100644 --- a/embodichain/gen_sim/gradio_ui/app_commands.py +++ b/embodichain/gen_sim/gradio_ui/app_commands.py @@ -14,30 +14,26 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""CLI command builders for EmbodiChain pipelines.""" +"""CLI command builder for the Action engine.""" from __future__ import annotations import sys -from typing import Protocol from app_config import ( + AGENT_CONFIG, COMMANDS, + FAST_GYM_CONFIG, ROBOT_PROFILE_FRANKA, ROBOT_PROFILE_UR5, ROBOT_PROFILE_UR10, SCENE_ID, ) +__all__ = ["build_run_agent_command"] -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: +def _robot_profile_cli_value(robot_profile: str | None) -> str | None: return { ROBOT_PROFILE_FRANKA: "franka", ROBOT_PROFILE_UR5: "dual_ur5", @@ -45,112 +41,12 @@ def robot_profile_cli_value(robot_profile: str | None) -> str | None: }.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]: + """Build the DexSim command for the existing ``current`` Gym scene.""" agent = COMMANDS["agent"] command = [ sys.executable, @@ -159,15 +55,13 @@ def build_run_agent_command( "--task_name", SCENE_ID, "--gym_config", - str(paths.fast_gym_config), + str(FAST_GYM_CONFIG), "--agent_config", - str(paths.agent_config), + str(AGENT_CONFIG), *agent["base_args"], "--num_envs", - agent["parallel_num_envs"] if parallel_env else agent["single_num_envs"], + agent["single_num_envs"], ] - if parallel_env: - command.extend(agent["parallel_args"]) - if supports_robot_profile and (profile := robot_profile_cli_value(robot_profile)): + 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 index 21085ea53..06f696e01 100644 --- a/embodichain/gen_sim/gradio_ui/app_config.py +++ b/embodichain/gen_sim/gradio_ui/app_config.py @@ -14,12 +14,7 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Static settings and helpers for the Gradio application. - -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. -""" +"""Static settings for the engine-only Gradio application.""" from __future__ import annotations @@ -30,45 +25,23 @@ APP_ROOT = Path(__file__).resolve().parent 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" -DEBUG_SCENE_ENGINE_ROOT = DEBUG_ENGINE_ROOT / "scenes" -SCENE_ID = "current" +GEN_SIM_ROOT = APP_ROOT / ".gen_sim" +GEN_SIM_ASSET_ROOT = GEN_SIM_ROOT / "assets" +GEN_SIM_SCENE_ROOT = GEN_SIM_ROOT / "scenes" +SCENE_ID = "current" 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" -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" +AGENT_CONFIG = CONFIG_DIR / "agent_config.json" + 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" -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" +VIDEO_SUFFIXES = {".mp4", ".avi", ".mov", ".mkv", ".webm"} 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 +DEFAULT_CONCURRENCY_LIMIT = 1 -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" @@ -78,125 +51,32 @@ (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", + "single_video_preview": "DexSim Video 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. "scene_engine": { "module": "embodichain", "base_args": ("scene-engine",), @@ -207,27 +87,12 @@ PHASE_DEFINITIONS = { "idle": (0, "Idle"), "received": (5, "Input received"), - "started": (10, "Local pipeline started"), + "started": (10, "Scene generation 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"), + "relations": (35, "Scene segmentation"), + "asset_generation": (55, "Geometry generation"), + "gym_export": (75, "Scene export"), + "preview": (90, "Preview generation"), "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", -) - -DEFAULT_CONCURRENCY_LIMIT = 1 diff --git a/embodichain/gen_sim/gradio_ui/app_env.py b/embodichain/gen_sim/gradio_ui/app_env.py index f9227e25e..ac1464a5d 100644 --- a/embodichain/gen_sim/gradio_ui/app_env.py +++ b/embodichain/gen_sim/gradio_ui/app_env.py @@ -46,7 +46,7 @@ load_gen_sim_env() APP_ROOT = Path(__file__).resolve().parent -DEBUG_ENGINE_ROOT = APP_ROOT / ".debug_engine" +GEN_SIM_ROOT = APP_ROOT / ".gen_sim" PROXY_ENV_KEYS = ( "HTTP_PROXY", "HTTPS_PROXY", @@ -77,7 +77,7 @@ def _getenv(name: str, default: str) -> str: ) ARTICRAFT_CONDA_ENV = _getenv("ARTICRAFT_CONDA_ENV", "articraft") ARTICRAFT_OUTPUT_ROOT = Path( - _getenv("ARTICRAFT_OUTPUT_ROOT", str(DEBUG_ENGINE_ROOT / "articraft")) + _getenv("ARTICRAFT_OUTPUT_ROOT", str(GEN_SIM_ROOT / "articraft")) ).expanduser() SCENE_ENGINE_VISER_PORT = int(_getenv("SCENE_ENGINE_VISER_PORT", "8080")) ARTICRAFT_VISER_PORT = int(_getenv("ARTICRAFT_VISER_PORT", "8081")) diff --git a/embodichain/gen_sim/gradio_ui/app_media.py b/embodichain/gen_sim/gradio_ui/app_media.py index 5c4813c49..417498de7 100644 --- a/embodichain/gen_sim/gradio_ui/app_media.py +++ b/embodichain/gen_sim/gradio_ui/app_media.py @@ -14,177 +14,44 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Run-log archival and video/dataset preview generation.""" +"""Media helpers used by the Action and Articulation engines.""" 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 OUTPUTS_DIR, VIDEO_SUFFIXES -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 +__all__ = [ + "articraft_viser_preview_cli", + "latest_audience_output_video", + "run_articraft_viser_preview", +] -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]: +def _collect_audience_output_videos() -> list[Path]: if not OUTPUTS_DIR.is_dir(): return [] - return sorted( + videos = [ 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 [ + audience_videos = [ path for path in videos if "audience" in path.relative_to(OUTPUTS_DIR).as_posix().lower() ] + return audience_videos or videos def latest_audience_output_video(min_mtime_ns: int | None = None) -> Path | None: + """Return the newest DexSim video created after the requested time.""" latest_path: Path | None = None latest_mtime = -1 - for path in collect_audience_output_videos(): + for path in _collect_audience_output_videos(): try: mtime = path.stat().st_mtime_ns except OSError: @@ -197,496 +64,8 @@ def latest_audience_output_video(min_mtime_ns: int | None = None) -> Path | None 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. - """ + """Load an Articraft URDF and publish its initial topology to Viser.""" from embodichain.lab.scripts import preview_asset from embodichain.lab.sim.sim_manager import SimulationManager from embodichain.utils.logger import log_info @@ -710,11 +89,7 @@ def run_articraft_viser_preview(args: argparse.Namespace) -> None: 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``. - """ + """Run the Articraft-aware variant of the generic preview-asset CLI.""" from embodichain.lab.scripts import preview_asset parser = preview_asset._create_parser() diff --git a/embodichain/gen_sim/gradio_ui/app_processes.py b/embodichain/gen_sim/gradio_ui/app_processes.py index a0f414f3a..11bad0a5a 100644 --- a/embodichain/gen_sim/gradio_ui/app_processes.py +++ b/embodichain/gen_sim/gradio_ui/app_processes.py @@ -14,7 +14,7 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Pipeline subprocess execution and progress detection.""" +"""Subprocess execution and lifecycle management for the engine workspace.""" from __future__ import annotations @@ -33,19 +33,16 @@ configure_direct_network_env, configure_simready_llm_env, ) -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 @@ -80,14 +77,11 @@ def run_agent_cli_supports_robot_profile() -> bool: return _RUN_AGENT_SUPPORTS_ROBOT_PROFILE -def build_run_agent_command( - paths: ScenePaths, *, parallel_env: bool = False, robot_profile: str | None = None -) -> list[str]: +def build_run_agent_command(*, robot_profile: str | None = None) -> list[str]: + """Build the Action-engine command for the existing current scene.""" 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(), ) @@ -260,72 +254,6 @@ def terminate_process_group(process: subprocess.Popen[str]) -> None: _unregister_managed_process(process) -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], diff --git a/embodichain/gen_sim/gradio_ui/app_services.py b/embodichain/gen_sim/gradio_ui/app_services.py index 36b7fccdf..a3956eb4a 100644 --- a/embodichain/gen_sim/gradio_ui/app_services.py +++ b/embodichain/gen_sim/gradio_ui/app_services.py @@ -23,6 +23,6 @@ from __future__ import annotations -from app_ui import build_demo +from app_ui import build_app -__all__ = ["build_demo"] +__all__ = ["build_app"] diff --git a/embodichain/gen_sim/gradio_ui/app_state.py b/embodichain/gen_sim/gradio_ui/app_state.py index 9ad6b9dac..cf9df4067 100644 --- a/embodichain/gen_sim/gradio_ui/app_state.py +++ b/embodichain/gen_sim/gradio_ui/app_state.py @@ -14,26 +14,20 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Shared, thread-safe runtime state and timing helpers.""" +"""Shared runtime state for the Scene and Action engines.""" 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, -) +from app_config import PHASE_DEFINITIONS + +__all__ = ["PHASES", "Phase", "runtime", "runtime_lock", "set_runtime_phase_locked"] @dataclass(frozen=True) @@ -49,141 +43,24 @@ class Phase: 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_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 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) + """Set the current UI phase while the caller holds ``runtime_lock``.""" 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 index 92e30eeed..a2352d28c 100644 --- a/embodichain/gen_sim/gradio_ui/app_ui.py +++ b/embodichain/gen_sim/gradio_ui/app_ui.py @@ -14,7 +14,7 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Gradio layout and event bindings. +"""Gradio layout and event bindings for the engine workspace. The workflow layer supplies all callbacks; this module only owns presentation and wires components to those callbacks. @@ -22,28 +22,35 @@ 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 +import gradio as gr +from app_asset_engine import build_asset_engine_panel +from app_config import ( + DEBUG_ENGINE_ACTION, + DEBUG_ENGINE_ASSET, + DEBUG_ENGINE_SCENE, + DEBUG_ENGINES, + DEFAULT_ROBOT_PROFILE, + DEXFORCE_LOGO, + LANGUAGE_EN, + ROBOT_PROFILES, + UI_TEXT, +) +from app_workflows import ( + format_status, + preview_saved_scene, + refresh_saved_scenes, + reset_scene_engine, + run_action_engine_from_current, + run_scene_engine, + ui_snapshot, +) -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), - ) +__all__ = ["build_app"] -def select_debug_engine(selected_engine: str): - """Expose an explicit active state without starting any pipeline.""" +def select_engine(selected_engine: str): + """Show the selected engine panel without starting a pipeline.""" button_updates = tuple( gr.update(variant="primary" if engine == selected_engine else "secondary") for engine, _ in DEBUG_ENGINES @@ -57,122 +64,22 @@ def select_debug_engine(selected_engine: str): 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 + """Adapt the shared runtime snapshot to the Action-engine status widgets.""" + video, task, progress, status, _initial, _edited, _objects = ui_snapshot() + return video, task, progress, status def run_action_engine_panel(task_text: str, robot_profile: str | None): + """Run the Action engine and return its latest UI snapshot.""" 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, - ) - 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, - 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(): +def build_app() -> gr.Blocks: + """Build the engine-only Gradio application.""" + with gr.Blocks(title="EmbodiChain Gradio") as app: + if DEXFORCE_LOGO.is_file(): + with gr.Row(equal_height=True): gr.Image( value=str(DEXFORCE_LOGO), show_label=False, @@ -180,154 +87,114 @@ def build_demo() -> gr.Blocks: 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, - ) + with gr.Row(): + 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") - 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), + 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." ) - edited_model = gr.Model3D( - label=UI_TEXT[LANGUAGE_EN]["edited_preview"], - height=520, - clear_color=(0.94, 0.94, 0.94, 1.0), + with gr.Row(): + with gr.Column(scale=1): + scene_image = gr.Image( + label=UI_TEXT[LANGUAGE_EN]["input_image"], + sources=["upload", "webcam"], + type="filepath", + format="png", + height=300, + ) + with gr.Row(): + scene_run = gr.Button("Generate scene", variant="primary") + scene_reset = gr.Button("Reset Scene Engine", variant="stop") + with gr.Column(scale=2): + scene_progress = gr.Slider( + 0, + 100, + value=0, + step=1, + label=UI_TEXT[LANGUAGE_EN]["progress"], + interactive=False, + ) + scene_status = gr.Markdown(format_status("Idle.")) + scene_output = gr.Textbox( + label="Scene output directory (hash-named)", + interactive=False, + ) + 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\n" + "Select a generated Scene Engine export to inspect it in Viser. " + "Scene selection is currently independent from DexSim execution." ) - 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", - ) + with gr.Row(): + with gr.Column(scale=1): + action_scene_list = gr.Dropdown( + choices=[], + value=None, + label="Generated scenes", + info="Complete scenes stored under .gen_sim/scenes.", + ) + action_scene_refresh = gr.Button("Refresh scenes") + action_scene_status = gr.Markdown( + "**Scene list:** open Action engine or refresh to load scenes." + ) + action_task = gr.Textbox( + label="Task description", + placeholder="e.g. Put the bottle on the table", + ) + action_robot = gr.Radio( + choices=ROBOT_PROFILES, + value=DEFAULT_ROBOT_PROFILE, + label=UI_TEXT[LANGUAGE_EN]["robot"], + ) + action_run = gr.Button("Run DexSim", variant="primary") + with gr.Column(scale=2): + action_scene = gr.HTML( + "
" + "Select a generated scene to preview it." + "
" + ) + action_video = gr.Video( + label=UI_TEXT[LANGUAGE_EN]["single_video_preview"], + height=320, + autoplay=True, + loop=True, + ) + action_current_task = gr.Textbox( + label=UI_TEXT[LANGUAGE_EN]["current_task"], + interactive=False, + ) + action_progress = gr.Slider( + 0, + 100, + value=0, + step=1, + label=UI_TEXT[LANGUAGE_EN]["progress"], + interactive=False, + ) + action_status = gr.Markdown( + format_status("Load or generate a scene first.") + ) + action_refresh_timer = gr.Timer(2.0) - 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, - ), + (engine for engine, _label in DEBUG_ENGINES), + (asset_engine_button, scene_engine_button, action_engine_button), ): button.click( - select_debug_engine, + select_engine, inputs=[gr.State(engine)], outputs=[ asset_engine_button, @@ -339,248 +206,60 @@ def build_demo() -> gr.Blocks: ], queue=False, ) - debug_scene_run.click( + + action_engine_button.click( + refresh_saved_scenes, + inputs=[action_scene_list], + outputs=[action_scene_list, action_scene_status], + queue=False, + ) + + scene_run.click( run_scene_engine, - inputs=[debug_scene_image], - outputs=[ - debug_scene_progress, - debug_scene_status, - debug_scene_output, - debug_scene_preview, - ], + inputs=[scene_image], + outputs=[scene_progress, scene_status, scene_output, scene_preview], ) - debug_scene_reset.click( + scene_reset.click( reset_scene_engine, outputs=[ - debug_scene_image, - debug_scene_progress, - debug_scene_status, - debug_scene_output, - debug_scene_preview, + scene_image, + scene_progress, + scene_status, + scene_output, + scene_preview, ], queue=False, ) - debug_action_load.click( - action_engine_snapshot, - outputs=[ - debug_action_video, - debug_action_current_task, - debug_action_progress, - debug_action_status, - debug_action_scene, - ], + action_scene_refresh.click( + refresh_saved_scenes, + inputs=[action_scene_list], + outputs=[action_scene_list, action_scene_status], queue=False, ) - debug_action_run.click( + action_scene_list.change( + preview_saved_scene, + inputs=[action_scene_list], + outputs=[action_scene, action_scene_status], + ) + action_run.click( run_action_engine_panel, - inputs=[debug_action_task, debug_action_robot], + inputs=[action_task, action_robot], outputs=[ - debug_action_video, - debug_action_current_task, - debug_action_progress, - debug_action_status, - debug_action_scene, + action_video, + action_current_task, + action_progress, + action_status, ], ) - debug_action_refresh_timer.tick( + action_refresh_timer.tick( action_engine_snapshot, outputs=[ - debug_action_video, - debug_action_current_task, - debug_action_progress, - debug_action_status, - debug_action_scene, + action_video, + action_current_task, + action_progress, + action_status, ], 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 + + return app diff --git a/embodichain/gen_sim/gradio_ui/app_workflows.py b/embodichain/gen_sim/gradio_ui/app_workflows.py index 4895454eb..e6122ed28 100644 --- a/embodichain/gen_sim/gradio_ui/app_workflows.py +++ b/embodichain/gen_sim/gradio_ui/app_workflows.py @@ -14,1943 +14,68 @@ # limitations under the License. # ---------------------------------------------------------------------------- +"""Scene- and Action-engine workflows for the Gradio workspace.""" + from __future__ import annotations import hashlib import html +import importlib.util 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, + +import gradio as gr +import numpy as np +from PIL import Image, ImageOps + +from app_config import ( + AGENT_CONFIG, + COMMANDS, + GEN_SIM_SCENE_ROOT, + FAST_GYM_CONFIG, ) -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_media import latest_audience_output_video 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. +from app_state import PHASES, Phase, runtime, runtime_lock, set_runtime_phase_locked + +__all__ = [ + "format_status", + "preview_saved_scene", + "refresh_saved_scenes", + "reset_scene_engine", + "run_action_engine_from_current", + "run_scene_engine", + "ui_snapshot", +] 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: - if runtime.run_token != token: - return - 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() +def _drain_output_queue(output_queue: queue.Queue[str]) -> list[str]: + lines: list[str] = [] 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) + lines.append(output_queue.get_nowait()) + except queue.Empty: + return lines 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.""" + """Map Scene Engine stage names to the shared progress UI.""" text = line.lower() mapping = ( ("scene understanding", "scene_intake"), @@ -2004,7 +129,7 @@ def _prepare_scene_engine_input( 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 = GEN_SIM_SCENE_ROOT / scene_hash output_root.mkdir(parents=True, exist_ok=True) image_path = output_root / "input.png" image_path.write_bytes(image_bytes.getvalue()) @@ -2012,7 +137,6 @@ def _prepare_scene_engine_input( 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: @@ -2025,40 +149,159 @@ def _wait_for_viser(port: int, process: subprocess.Popen[str]) -> bool: return False -def _viser_iframe(port: int, scene_hash: str) -> str: - """Embed the Viser service using the same hostname as the Gradio page.""" - srcdoc = ( - "" - ) +def _viser_iframe(port: int, scene_hash: str) -> str: + srcdoc = ( + "" + ) + return ( + "
Viser preview: " + f"{html.escape(scene_hash)}" + f"
" + ) + + +def _saved_scene_root(scene_name: str) -> Path: + """Resolve a scene-list value without allowing paths outside the store.""" + if not scene_name or Path(scene_name).name != scene_name: + raise ValueError("Select a valid generated scene.") + scene_store = GEN_SIM_SCENE_ROOT.resolve() + scene_root = (scene_store / scene_name).resolve() + if scene_root.parent != scene_store: + raise ValueError("Selected scene must stay within the generated scene store.") + config_path = scene_root / "scene_export" / "scene_config.json" + if not config_path.is_file(): + raise FileNotFoundError(f"Scene export is incomplete: {config_path}") + try: + scene_config = json.loads(config_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise ValueError(f"Scene config is invalid: {config_path}") from exc + if not isinstance(scene_config, dict) or scene_config.get("format") != ( + "embodichain.scene-export/v1" + ): + raise ValueError(f"Unsupported scene export: {config_path}") + return scene_root + + +def saved_scene_choices() -> list[tuple[str, str]]: + """List complete Scene Engine exports, newest first.""" + if not GEN_SIM_SCENE_ROOT.is_dir(): + return [] + choices: list[tuple[int, str, str]] = [] + for scene_root in GEN_SIM_SCENE_ROOT.iterdir(): + if not scene_root.is_dir(): + continue + config_path = scene_root / "scene_export" / "scene_config.json" + if not config_path.is_file(): + continue + try: + scene_config = json.loads(config_path.read_text(encoding="utf-8")) + if not isinstance(scene_config, dict) or scene_config.get("format") != ( + "embodichain.scene-export/v1" + ): + continue + scene_id = scene_config.get("scene_id") + label = ( + f"{scene_root.name} · {scene_id}" + if isinstance(scene_id, str) and scene_id + else scene_root.name + ) + modified_ns = config_path.stat().st_mtime_ns + except (OSError, json.JSONDecodeError): + continue + choices.append((modified_ns, label, scene_root.name)) + choices.sort(reverse=True) + return [(label, value) for _modified_ns, label, value in choices] + + +def refresh_saved_scenes(selected_scene: str | None = None): + """Refresh the Action-engine scene list without selecting a scene implicitly.""" + choices = saved_scene_choices() + values = {value for _label, value in choices} + value = selected_scene if selected_scene in values else None + status = ( + f"**Scene list:** {len(choices)} generated scene(s) available." + if choices + else "**Scene list:** no complete generated scenes found." + ) + return gr.update(choices=choices, value=value), status + + +def preview_saved_scene(scene_name: str | None): + """Start a Viser preview for the explicitly selected generated scene.""" + idle_preview = ( + "
" + "Select a generated scene to preview it." + "
" + ) + if not scene_name: + return idle_preview, "**Scene preview:** no scene selected." + + try: + scene_root = _saved_scene_root(scene_name) + except (ValueError, FileNotFoundError) as exc: + return idle_preview, f"**Scene preview error:** {exc}" + + with runtime_lock: + if runtime.scene_engine_is_running: + return idle_preview, "**Scene preview:** Scene Engine is still running." + old_preview = runtime.scene_preview_process + runtime.scene_preview_process = None + + if old_preview is not None: + terminate_process_group(old_preview) + + port = SCENE_ENGINE_VISER_PORT + preview_command = [ + sys.executable, + COMMANDS["scene_engine"]["preview_script"], + "--output_root", + str(scene_root), + "--viser", + "--viser-host", + "0.0.0.0", + "--viser-port", + str(port), + ] + try: + preview_process = start_pipeline(preview_command) + except Exception as exc: + return idle_preview, f"**Scene preview error:** {exc}" + + with runtime_lock: + runtime.scene_preview_process = preview_process + if not _wait_for_viser(port, preview_process): + terminate_process_group(preview_process) + with runtime_lock: + if runtime.scene_preview_process is preview_process: + runtime.scene_preview_process = None + return idle_preview, "**Scene preview error:** Viser did not start." + return ( - f"
Viser preview: {html.escape(scene_hash)}" - f"" - "
" + _viser_iframe(port, scene_name), + f"**Scene preview:** `{scene_name}` is ready.", ) def reset_scene_engine(): - """Clear Scene Engine widgets and stop its generator and Viser process groups.""" + """Clear Scene Engine widgets and stop its owned 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 + action_running = runtime.sim_process is not None + if owns_runtime: + runtime.run_token = uuid.uuid4().hex runtime.is_busy = False + if not action_running: 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.image_path = None runtime.scene_engine_process = None runtime.scene_preview_process = None runtime.scene_engine_is_running = False @@ -2067,14 +310,15 @@ def reset_scene_engine(): if process is not None: terminate_process_group(process) + message = ( + "Scene Engine reset." + if not action_running + else "Scene Engine preview reset; Action Engine is still running." + ) 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." - ), + format_status(message), "", "
" "The Viser preview will appear here after generation." @@ -2089,7 +333,7 @@ def run_scene_engine(image_value: str | np.ndarray | Image.Image): token = uuid.uuid4().hex with runtime_lock: if runtime.is_busy: - runtime.status = "Another pipeline is already running." + runtime.status = "Another engine is already running." runtime.last_error = runtime.status busy_message = runtime.status else: @@ -2100,7 +344,6 @@ def run_scene_engine(image_value: str | np.ndarray | Image.Image): runtime.status = "Preparing Scene Engine input." runtime.last_error = None runtime.log_lines.clear() - clear_run_timing_locked() busy_message = None if busy_message is not None: @@ -2121,22 +364,16 @@ def run_scene_engine(image_value: str | np.ndarray | Image.Image): 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, @@ -2149,10 +386,7 @@ def run_scene_engine(image_value: str | np.ndarray | Image.Image): str(output_root), ] scene_engine_log = output_root / "scene_engine.log" - scene_engine_log.write_text( - "$ " + " ".join(command) + "\n", - encoding="utf-8", - ) + 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) @@ -2179,15 +413,13 @@ def run_scene_engine(image_value: str | np.ndarray | Image.Image): if runtime.run_token != token: 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." reader.start() while process.poll() is None: - drained = drain_output_queue(output_queue) + drained = _drain_output_queue(output_queue) with runtime_lock: if ( runtime.run_token != token @@ -2209,12 +441,11 @@ def run_scene_engine(image_value: str | np.ndarray | Image.Image): 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): + 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" @@ -2293,7 +524,6 @@ def run_scene_engine(image_value: str | np.ndarray | Image.Image): ): 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") @@ -2302,475 +532,66 @@ def run_scene_engine(image_value: str | np.ndarray | Image.Image): yield _scene_engine_updates(output_root, preview_html) +def _action_scene_is_available() -> bool: + return FAST_GYM_CONFIG.is_file() and AGENT_CONFIG.is_file() + + +def _action_agent_cli_is_available() -> bool: + try: + return importlib.util.find_spec(COMMANDS["agent"]["module"]) is not None + except (ImportError, ModuleNotFoundError, ValueError): + return False + + 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.""" + """Launch DexSim for the existing ``current`` Gym scene.""" 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 + failure = "Enter a task description first." + elif not _action_scene_is_available(): + failure = "Current Gym scene/config is unavailable." + elif runtime.is_busy or runtime.sim_process is not None: + failure = "Another engine is already running." + elif not _action_agent_cli_is_available(): + failure = "Action-agent CLI is unavailable in this environment." else: + failure = None token = uuid.uuid4().hex runtime.run_token = token + runtime.is_busy = True 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 + runtime.log_lines.clear() + set_runtime_phase_locked("started") - 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) - ) + if failure is not None: + runtime.status = failure + runtime.last_error = failure + if failure is None: + error = _launch_current_simulation(token, robot_profile=robot_profile) + if error: 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", - ) + runtime.status = error + runtime.last_error = error + return ui_snapshot() -def launch_current_simulation( +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, - ) + command = build_run_agent_command(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}" + return f"DexSim launch failed: {exc}" output_queue: queue.Queue[str] = queue.Queue() reader = threading.Thread( @@ -2779,18 +600,8 @@ def launch_current_simulation( 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, - ), + target=_monitor_simulation, + args=(token, process, output_queue, reader, started_at_ns), daemon=True, ) @@ -2800,10 +611,6 @@ def launch_current_simulation( 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: @@ -2815,90 +622,40 @@ def launch_current_simulation( return None -def monitor_simulation( +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)) + _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 + _append_simulation_logs(token, process, _drain_output_queue(output_queue)) + display_video = latest_audience_output_video(min_mtime_ns=started_at_ns) - 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.is_busy = False 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}" - ) + set_runtime_phase_locked("complete") + runtime.status = "DexSim simulation finished successfully." + runtime.last_error = None + if display_video is None: + runtime.log_lines.append("No simulation preview video was found.") 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, - ) + set_runtime_phase_locked("failed") + runtime.status = f"DexSim exited with return code {process.returncode}." + runtime.last_error = runtime.status -def append_simulation_logs( +def _append_simulation_logs( token: str, process: subprocess.Popen[str], lines: list[str], @@ -2906,408 +663,12 @@ def append_simulation_logs( 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, - ) + if runtime.run_token == token and runtime.sim_process is process: + runtime.log_lines.extend(lines) def ui_snapshot(extra_status: str | None = None): + """Return the current Action-engine widget values.""" with runtime_lock: phase = PHASES.get(runtime.phase_key, PHASES["idle"]) video_value = None @@ -3315,35 +676,18 @@ def ui_snapshot(extra_status: str | None = 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: + if 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, @@ -3354,52 +698,9 @@ def ui_snapshot(extra_status: str | None = None): 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), + None, + None, + None, ) @@ -3410,6 +711,7 @@ def format_status( busy: bool = False, last_error: str | None = None, ) -> str: + """Format an engine status for display in Gradio.""" if phase is None: phase = PHASES["idle"] state = "running" if busy else "ready" diff --git a/embodichain/gen_sim/gradio_ui/gradio_app.py b/embodichain/gen_sim/gradio_ui/gradio_app.py index 7dce2e1a8..e0713836c 100644 --- a/embodichain/gen_sim/gradio_ui/gradio_app.py +++ b/embodichain/gen_sim/gradio_ui/gradio_app.py @@ -26,12 +26,12 @@ from app_config import ( ASSETS_DIR, - DEBUG_ENGINE_ROOT, + GEN_SIM_ROOT, DEFAULT_CONCURRENCY_LIMIT, ) 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 +from app_services import build_app __all__ = ["main"] @@ -62,17 +62,17 @@ def _install_shutdown_handlers() -> None: 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) + app = build_app() + app.queue(default_concurrency_limit=DEFAULT_CONCURRENCY_LIMIT) _install_shutdown_handlers() try: - demo.launch( + app.launch( server_name=SERVER_NAME, server_port=SERVER_PORT, allowed_paths=[ str(EMBODICHAIN_ROOT), str(ASSETS_DIR), - str(DEBUG_ENGINE_ROOT), + str(GEN_SIM_ROOT), ], ) finally: diff --git a/embodichain/gen_sim/gradio_ui/gradio_visualization_architecture.md b/embodichain/gen_sim/gradio_ui/gradio_visualization_architecture.md index f8cc2fe80..59d8fe666 100644 --- a/embodichain/gen_sim/gradio_ui/gradio_visualization_architecture.md +++ b/embodichain/gen_sim/gradio_ui/gradio_visualization_architecture.md @@ -1,6 +1,6 @@ # Gradio 可视化系统架构 -本文档以当前代码为准，描述 Gradio Demo、Debug 下的三个引擎，以及它们与 EmbodiChain、SimReady、Articraft 和 DexSim 的边界。`gradio_app.py` 只负责启动；界面、资产工作流、场景工作流和进程管理分散在专用模块中。 +本文档以当前代码为准，描述 Gradio 中的三个引擎，以及它们与 EmbodiChain、SimReady、Articraft 和 DexSim 的边界。`gradio_app.py` 只负责启动；界面、资产工作流、场景工作流和进程管理分散在专用模块中。 ## 架构总览 @@ -11,11 +11,11 @@ gradio_app.py 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 + │ └► action-agent `run_agent` + DexSim ├──────────────► app_commands.py 命令构造 ├──────────────► app_processes.py 子进程、环境、日志和阶段检测 ├──────────────► app_state.py 共享 RuntimeState、锁和计时 @@ -28,14 +28,14 @@ app_workflows.py ────► EmbodiChain Scene Engine CLI + Viser | 模块 | 职责 | | --- | --- | | `gradio_app.py` | 唯一启动入口；校验 `EMBODICHAIN_ROOT`，创建 Blocks，设置队列和本地文件访问路径。 | -| `app_ui.py` | Demo/Debug 布局、引擎面板切换和回调绑定；不实现 pipeline。 | +| `app_ui.py` | 顶部图标、引擎面板切换和回调绑定；不实现 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_workflows.py` | Scene Engine 工作流、Action Engine 所需的共享状态、GLB 预览和 DexSim。 | +| `app_processes.py` | 子进程环境、进程组终止、stdout 读取和 pipeline 阶段检测。 | | `app_state.py` | `RuntimeState`、互斥锁、进度阶段、运行 token 和耗时统计。 | -| `app_commands.py` | prompt2scene、动作配置和 `run_agent` 的参数构造。 | -| `app_media.py` | 观众视频、LeRobot 数据预览、组合视频和运行日志归档。 | +| `app_commands.py` | Action engine 的 `run_agent` 参数构造。 | +| `app_media.py` | DexSim 观众视频发现和 Articraft Viser CLI 适配。 | | `app_config.py` | UI 文案、引擎模式、路径推导和 CLI 固定参数。 | | `app_env.py` | 从 `.env` 读取 Gradio、Articraft 和 SimReady 的部署值，并保留未配置时的默认值。 | | `../.env` | Gradio 与 Scene Engine 共用的路径、端口、LLM 和服务端点配置；不提交凭据。 | @@ -57,69 +57,22 @@ conda run -n embodichain python gradio_app.py | `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。 | +| `ARTICRAFT_OUTPUT_ROOT` | `<项目>/.gen_sim/articraft` | Articraft 记录、运行日志和导出 bundle。 | -`demo.launch()` 仅开放 EmbodiChain 根目录、`assets/` 和 `.debug_engine/` 给浏览器读取。pipeline 子进程由 `build_pipeline_env()` 创建环境：它清除代理变量、设置 `NO_PROXY=no_proxy=*`、关闭 Gradio analytics，并把非空的 SimReady 配置映射为 `OPENAI_*`。这不会改写启动 Gradio 的父进程环境。 +`app.launch()` 仅开放 EmbodiChain 根目录、`assets/` 和 `.gen_sim/` 给浏览器读取。pipeline 子进程由 `build_pipeline_env()` 创建环境：它清除代理变量、设置 `NO_PROXY=no_proxy=*`、关闭 Gradio analytics，并把非空的 SimReady 配置映射为 `OPENAI_*`。这不会改写启动 Gradio 的父进程环境。 ## 页面与引擎 -顶部的 `Demo` / `Debug` 只切换可见面板，不会启动任务；切换后共享运行状态保留。Debug 有三个按钮：`Asset_engine`、`Scene_engine`、`Action_engine`。它们的实际输入和产物并不完全相同： +页面顶部保留 DexForce 图标，并直接显示 `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/` | 是 | +| Asset engine / SimReady | 一个网格、可选材质附件、类别 | 输入 GLB、SimReady GLB、原始输出下载 | `.gen_sim/assets/runs//` | 否 | +| Asset engine / Articulation | 文字、可选参考图 | URDF articulation 的 Viser、zip 下载 | `.gen_sim/articraft/` | 否 | +| Scene engine | 一张图片 | Scene Engine 的 Viser | `.gen_sim/scenes//` | 否 | +| Action engine | 已生成场景列表、任务、机器人 | 选中场景的 Viser 和 DexSim 视频 | 场景预览来自 `.gen_sim/scenes/`；DexSim 暂沿用现有命令 | 是 | -因此，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。 +因此，Scene engine 是独立的图像条件场景生成器；它不会提升、复制或转换输出到 `gym_project/current`。Action engine 只消费已有的 `current` Gym 场景。界面中的 “Scene engine” 文案表达的是所需场景类型，并不意味着独立 Scene Engine 输出已自动连到 Action engine。 ## Asset engine @@ -129,7 +82,7 @@ SimReady CLI 接收目录，而 Gradio 接收上传文件。上传文件会复 ```text mesh + sidecar files - → .debug_engine/assets/runs//input/ + → .gen_sim/assets/runs//input/ → trimesh 导出 input_preview.glb → SimReady CLI → output/**/asset_simready.glb（优先）或 asset_simready.obj @@ -151,7 +104,7 @@ python -m embodichain.gen_sim.simready_pipeline.cli.start \ ### Articulation：Articraft + Codex -Articulation 标签页根据文本和可选参考图生成一个可下载的 articulated asset。先点击环境检查：若 `ARTICRAFT_ROOT` 不存在，应用会 clone `ARTICRAFT_REPOSITORY_URL`；随后检查 Conda、指定的 Articraft 环境和 Codex CLI。该操作会创建 checkout 和 `.debug_engine/articraft/` 中的输出目录，现有的非 Articraft 目录不会被覆盖。 +Articulation 标签页根据文本和可选参考图生成一个可下载的 articulated asset。先点击环境检查：若 `ARTICRAFT_ROOT` 不存在，应用会 clone `ARTICRAFT_REPOSITORY_URL`；随后检查 Conda、指定的 Articraft 环境和 Codex CLI。该操作会创建 checkout 和 `.gen_sim/articraft/` 中的输出目录，现有的非 Articraft 目录不会被覆盖。 生成流程： @@ -166,7 +119,7 @@ description + optional image → exports/.zip + Viser articulation preview ``` -产物、记录和参考图均在 `ARTICRAFT_OUTPUT_ROOT` 下，不能直接当作 Demo 的 Gym 场景或 SimReady 资产；若要进入后续仿真，需要另行定义并实现转换/导入流程。Articraft Viser 每次成功预览会终止旧的 Articraft 预览进程，再以 `0.0.0.0:` 启动新进程。 +产物、记录和参考图均在 `ARTICRAFT_OUTPUT_ROOT` 下，不能直接当作 Action engine 的 Gym 场景或 SimReady 资产；若要进入后续仿真，需要另行定义并实现转换/导入流程。Articraft Viser 每次成功预览会终止旧的 Articraft 预览进程，再以 `0.0.0.0:` 启动新进程。 `Reset Articulation` 会清空描述、参考图、记录与下载结果，终止当前 Articraft/Codex 命令进程组，并关闭该面板启动的 Viser。 @@ -176,7 +129,7 @@ Scene engine 只接收图像。上传图像会先进行 EXIF 归正并转为 RGB ```text image - → .debug_engine/scenes//input.png + → .gen_sim/scenes//input.png → python -m embodichain scene-engine --image --output_root @@ -199,7 +152,9 @@ 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`，或者实现显式的配置再生成步骤。 +进入 Action engine 或点击 `Refresh scenes` 会扫描 `.gen_sim/scenes/`，只列出包含 `scene_export/scene_config.json` 的完整场景。列表不会自动选中场景；用户显式选择后，右侧通过 Viser 展示该场景。当前场景选择只负责可视化，尚未传递给 DexSim 命令。 + +点击 `Run DexSim` 仍会检查任务、现有 `current` Gym/action 配置、运行占用和可导入的 `embodichain.gen_sim.action_agent_pipeline.cli.run_agent`，再以当前配置调用 `run_agent`。 运行命令的核心参数为： @@ -208,16 +163,16 @@ 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> + --regenerate --renderer fast-rt --num_envs 1 ``` -并行模式额外传入 arena 和数据保存过滤参数。`--robot-profile` 仅在通过 `run_agent --help` 探测到该参数时加入。DexSim 完成后会寻找 audience 视频和 LeRobot 数据集；单环境可组合两种预览视频。 +`--robot-profile` 仅在通过 `run_agent --help` 探测到该参数时加入。DexSim 完成后会寻找本次运行产生的 audience 视频并显示在 Action engine 中。 ## 共享状态、并发和进度 -Demo、独立 Scene engine 和 Action engine 共享 `RuntimeState` 与 `runtime_lock`，其中包含运行 token、pipeline/DexSim/Scene Viser 进程、输入、预览、日志、阶段和计时。运行 token 用于丢弃过期线程的更新。Articraft Viser 使用单独的锁和进程引用；SimReady 使用自己的同步 generator。 +Scene engine 和 Action engine 共享 `RuntimeState` 与 `runtime_lock`，其中包含运行 token、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 队列限制。 +`app.queue(default_concurrency_limit=1)` 将队列中的高成本回调串行化。Action engine 的 `Timer(2.0)` 读取共享状态。Asset/Articraft 面板不写入这一状态，但仍会受 Gradio 队列限制。 共享阶段如下；独立 Scene Engine 将其日志映射到相同的进度条： @@ -243,11 +198,9 @@ export SIMREADY_OPENAI_MODEL='' export SIMREADY_OPENAI_BASE_URL='' ``` -Demo/Action 需要 action-agent 模块，特别是： +Action engine 只需要 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 ``` @@ -272,7 +225,7 @@ python -m py_compile \ 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" + "from app_ui import build_app; assert build_app() is not None" ``` 手动检查： @@ -280,6 +233,4 @@ env -u HTTP_PROXY -u HTTPS_PROXY -u ALL_PROXY \ 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 能终止其对应的进程组。 +4. Action engine 在没有 `current` Gym/action 配置或缺少 CLI 时给出预检错误。 diff --git a/embodichain/gen_sim/gradio_ui/random_input.py b/embodichain/gen_sim/gradio_ui/random_input.py deleted file mode 100644 index 9ce553992..000000000 --- a/embodichain/gen_sim/gradio_ui/random_input.py +++ /dev/null @@ -1,542 +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 base64 -import json -import os -import uuid -from dataclasses import asdict, dataclass, replace -from pathlib import Path - -import numpy as np - -from embodichain.gen_sim.env import get_embodichain_root, load_gen_sim_env - -load_gen_sim_env() - -EMBODICHAIN_ROOT = get_embodichain_root() -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 f508f1892141fd8bb97106cce44bb0fed89106e1 Mon Sep 17 00:00:00 2001 From: PengXuanchao Date: Mon, 10 Aug 2026 11:02:13 +0800 Subject: [PATCH 48/53] fix simready config --- .../gen_sim/gradio_ui/app_asset_engine.py | 2 +- .../gen_sim/gradio_ui/app_processes.py | 29 ++++++++++++++++--- .../gradio_visualization_architecture.md | 2 +- 3 files changed, 27 insertions(+), 6 deletions(-) diff --git a/embodichain/gen_sim/gradio_ui/app_asset_engine.py b/embodichain/gen_sim/gradio_ui/app_asset_engine.py index 4aa57bcc9..e94adca26 100644 --- a/embodichain/gen_sim/gradio_ui/app_asset_engine.py +++ b/embodichain/gen_sim/gradio_ui/app_asset_engine.py @@ -221,7 +221,7 @@ def run_simready_asset(upload_value: Any, category: str): ) try: - process = start_pipeline(command) + process = start_pipeline(command, use_simready_llm=True) except Exception as exc: if _simready_run_is_active(token): yield input_preview.as_posix(), None, None, f"**Pipeline start failed:** {exc}", "\n".join( diff --git a/embodichain/gen_sim/gradio_ui/app_processes.py b/embodichain/gen_sim/gradio_ui/app_processes.py index 11bad0a5a..15358c451 100644 --- a/embodichain/gen_sim/gradio_ui/app_processes.py +++ b/embodichain/gen_sim/gradio_ui/app_processes.py @@ -87,8 +87,20 @@ def build_run_agent_command(*, robot_profile: str | None = None) -> list[str]: ) -def start_pipeline(command: list[str]) -> subprocess.Popen[str]: - env = build_pipeline_env() +def start_pipeline( + command: list[str], *, use_simready_llm: bool = False +) -> subprocess.Popen[str]: + """Start a managed pipeline subprocess with its scoped dotenv settings. + + Args: + command: Command and arguments to execute. + use_simready_llm: Whether to map the dotenv ``SIMREADY_OPENAI_*`` values + to the upstream SimReady CLI's ``OPENAI_*`` variable names. + + Returns: + The registered subprocess. + """ + env = build_pipeline_env(use_simready_llm=use_simready_llm) env["PYTHONUNBUFFERED"] = "1" return register_managed_process( subprocess.Popen( @@ -104,10 +116,19 @@ def start_pipeline(command: list[str]) -> subprocess.Popen[str]: ) -def build_pipeline_env() -> dict[str, str]: +def build_pipeline_env(*, use_simready_llm: bool = False) -> dict[str, str]: + """Build a child environment from the shared GenSim dotenv configuration. + + Args: + use_simready_llm: Whether to apply the SimReady-specific LLM mapping. + + Returns: + A copy of the loaded process environment configured for the child. + """ env = os.environ.copy() configure_direct_network_env(env) - configure_simready_llm_env(env) + if use_simready_llm: + configure_simready_llm_env(env) return env diff --git a/embodichain/gen_sim/gradio_ui/gradio_visualization_architecture.md b/embodichain/gen_sim/gradio_ui/gradio_visualization_architecture.md index 59d8fe666..6c98118f4 100644 --- a/embodichain/gen_sim/gradio_ui/gradio_visualization_architecture.md +++ b/embodichain/gen_sim/gradio_ui/gradio_visualization_architecture.md @@ -59,7 +59,7 @@ conda run -n embodichain python gradio_app.py | `ARTICRAFT_CONDA_ENV` | `articraft` | 运行 Articraft CLI 的 Conda 环境。 | | `ARTICRAFT_OUTPUT_ROOT` | `<项目>/.gen_sim/articraft` | Articraft 记录、运行日志和导出 bundle。 | -`app.launch()` 仅开放 EmbodiChain 根目录、`assets/` 和 `.gen_sim/` 给浏览器读取。pipeline 子进程由 `build_pipeline_env()` 创建环境：它清除代理变量、设置 `NO_PROXY=no_proxy=*`、关闭 Gradio analytics，并把非空的 SimReady 配置映射为 `OPENAI_*`。这不会改写启动 Gradio 的父进程环境。 +`app.launch()` 仅开放 EmbodiChain 根目录、`assets/` 和 `.gen_sim/` 给浏览器读取。pipeline 子进程由 `build_pipeline_env()` 从共享 `.env` 创建环境：它清除代理变量、设置 `NO_PROXY=no_proxy=*` 并关闭 Gradio analytics。只有 SimReady 子进程会额外把非空的 `SIMREADY_OPENAI_*` 映射为其上游 CLI 需要的 `OPENAI_*`；Scene Engine、DexSim、Viser 和 Articraft 直接继承 `.env` 中的原始配置。这不会改写启动 Gradio 的父进程环境。 ## 页面与引擎 From 14bbf5284c6b21a829afea6687c567fab46b5187 Mon Sep 17 00:00:00 2001 From: PengXuanchao Date: Mon, 10 Aug 2026 11:37:49 +0800 Subject: [PATCH 49/53] fix config and session isolation --- embodichain/gen_sim/.env.example | 5 +- embodichain/gen_sim/env.py | 11 +- .../gen_sim/gradio_ui/app_articraft.py | 383 +++++++++--------- .../gen_sim/gradio_ui/app_asset_engine.py | 130 +++--- embodichain/gen_sim/gradio_ui/app_env.py | 103 ++++- .../gen_sim/gradio_ui/app_processes.py | 230 ++++++++++- embodichain/gen_sim/gradio_ui/app_ui.py | 4 +- .../gen_sim/gradio_ui/app_workflows.py | 17 +- embodichain/gen_sim/gradio_ui/gradio_app.py | 36 +- .../gradio_visualization_architecture.md | 18 +- pyproject.toml | 1 + tests/gen_sim/__init__.py | 17 + tests/gen_sim/gradio_ui/__init__.py | 17 + tests/gen_sim/gradio_ui/test_app_articraft.py | 59 +++ tests/gen_sim/gradio_ui/test_app_env.py | 72 ++++ tests/gen_sim/gradio_ui/test_app_processes.py | 102 +++++ tests/gen_sim/scene_engine/__init__.py | 17 + .../scene_engine/test_scene_engine_config.py | 100 ----- tests/gen_sim/simready_pipeline/__init__.py | 17 + tests/gen_sim/test_gen_sim_env.py | 65 +++ 20 files changed, 1015 insertions(+), 389 deletions(-) create mode 100644 tests/gen_sim/__init__.py create mode 100644 tests/gen_sim/gradio_ui/__init__.py create mode 100644 tests/gen_sim/gradio_ui/test_app_articraft.py create mode 100644 tests/gen_sim/gradio_ui/test_app_env.py create mode 100644 tests/gen_sim/gradio_ui/test_app_processes.py create mode 100644 tests/gen_sim/scene_engine/__init__.py delete mode 100644 tests/gen_sim/scene_engine/test_scene_engine_config.py create mode 100644 tests/gen_sim/simready_pipeline/__init__.py create mode 100644 tests/gen_sim/test_gen_sim_env.py diff --git a/embodichain/gen_sim/.env.example b/embodichain/gen_sim/.env.example index 6544d0710..ee293c134 100644 --- a/embodichain/gen_sim/.env.example +++ b/embodichain/gen_sim/.env.example @@ -23,8 +23,11 @@ SCENE_ENGINE_GEOMETRY_GENERATION_OBJECTS_PATH="/generate_multiple_objects" # 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_NAME="127.0.0.1" GRADIO_SERVER_PORT=7860 +# Both values are required when GRADIO_SERVER_NAME is not a loopback address. +GRADIO_AUTH_USERNAME="" +GRADIO_AUTH_PASSWORD="" SCENE_ENGINE_VISER_PORT=8080 ARTICRAFT_VISER_PORT=8081 ARTICRAFT_ROOT="" diff --git a/embodichain/gen_sim/env.py b/embodichain/gen_sim/env.py index 44b772568..aee913a25 100644 --- a/embodichain/gen_sim/env.py +++ b/embodichain/gen_sim/env.py @@ -35,12 +35,14 @@ def get_embodichain_root() -> Path: return Path(__file__).resolve().parents[2] -def find_gen_sim_env_file() -> Path: +def find_gen_sim_env_file() -> Path | None: """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. + the source tree. Otherwise GenSim uses ``embodichain/gen_sim/.env``. + + Returns: + The configured or default dotenv path, or ``None`` when neither exists. """ configured_path = os.environ.get("EMBODICHAIN_ENV_FILE") if configured_path: @@ -48,6 +50,7 @@ def find_gen_sim_env_file() -> Path: default_path = Path(__file__).resolve().parent / ".env" if default_path.is_file(): return default_path + return None def load_gen_sim_env(env: MutableMapping[str, str] | None = None) -> Path | None: @@ -67,7 +70,7 @@ def load_gen_sim_env(env: MutableMapping[str, str] | None = None) -> Path | None """ target_env = os.environ if env is None else env env_path = find_gen_sim_env_file() - if not env_path.is_file(): + if env_path is None or not env_path.is_file(): return None for line_number, raw_line in enumerate( diff --git a/embodichain/gen_sim/gradio_ui/app_articraft.py b/embodichain/gen_sim/gradio_ui/app_articraft.py index a7fd81f9e..dc3069246 100644 --- a/embodichain/gen_sim/gradio_ui/app_articraft.py +++ b/embodichain/gen_sim/gradio_ui/app_articraft.py @@ -29,13 +29,13 @@ import json import shutil import html -import signal import socket import subprocess import sys import threading import time import uuid +from collections.abc import Iterator from datetime import datetime, timezone from pathlib import Path from typing import Any @@ -48,16 +48,25 @@ ARTICRAFT_REPOSITORY_URL, ARTICRAFT_ROOT, ARTICRAFT_VISER_PORT, + EMBODICHAIN_ROOT, + validate_gradio_artifact_root, ) from app_processes import ( + SessionProcessRegistry, + build_codex_env, + build_pipeline_env, + get_request_session_id, read_process_output, + redact_sensitive_text, register_managed_process, start_pipeline, terminate_process_group, ) +from embodichain.gen_sim.env import find_gen_sim_env_file __all__ = [ "build_articraft_panel", + "cleanup_articraft_session", "configure_articraft_environment", "generate_articraft_asset", "reset_articraft_asset", @@ -65,13 +74,10 @@ ] _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() -_articraft_generation_lock = threading.Lock() -_articraft_generation_process: subprocess.Popen[str] | None = None -_articraft_generation_token: str | None = None +_articraft_runs = SessionProcessRegistry() _ARTICRAFT_IDLE_PREVIEW = ( "
" "The interactive Viser articulation preview will appear here after generation." @@ -79,53 +85,8 @@ ) -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 + command: list[str], *, session_id: 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( @@ -136,10 +97,10 @@ def _run_articraft_generation_check( stderr=subprocess.STDOUT, text=True, start_new_session=True, - env=os.environ.copy(), + env=build_pipeline_env(), ) ) - if not _set_articraft_generation_process(token, process): + if not _articraft_runs.attach(session_id, token, process): terminate_process_group(process) return None try: @@ -148,22 +109,27 @@ def _run_articraft_generation_check( terminate_process_group(process) raise finally: - _finish_articraft_generation_process(token, process) - if not _articraft_generation_is_active(token): + _articraft_runs.finish(session_id, token, process) + if not _articraft_runs.is_active(session_id, token): return None - return subprocess.CompletedProcess(command, process.returncode, stdout) + return subprocess.CompletedProcess( + command, + process.returncode, + redact_sensitive_text(stdout or ""), + ) -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() +def reset_articraft_asset(request: gr.Request) -> tuple[Any, ...]: + """Clear Articraft state and stop only the requesting session's processes. + + Args: + request: Gradio request for the browser session initiating Reset. + + Returns: + Reset values for all Articraft panel widgets. + """ + session_id = get_request_session_id(request) + cleanup_articraft_session(session_id) return ( "**Environment:** not checked.", "", @@ -176,6 +142,16 @@ def reset_articraft_asset(): ) +def cleanup_articraft_session(session_id: str) -> None: + """Stop Articraft generation and preview processes for one session. + + Args: + session_id: Stable Gradio session identifier. + """ + _articraft_runs.reset(session_id) + stop_articraft_viser_preview(session_id) + + 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") @@ -334,7 +310,10 @@ 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 ( + isolation_error = _articraft_isolation_error() + if isolation_error: + errors.append(isolation_error) + elif not ( ARTICRAFT_ROOT.is_dir() and (ARTICRAFT_ROOT / ".git").exists() and (ARTICRAFT_ROOT / "pyproject.toml").is_file() @@ -367,6 +346,8 @@ def _check_requirements() -> tuple[list[str], list[str], str | None]: def _prepare_articraft_checkout() -> tuple[bool, str]: """Clone the configured checkout when absent, without overwriting a directory.""" + if isolation_error := _articraft_isolation_error(): + return False, isolation_error if ARTICRAFT_ROOT.exists(): if (ARTICRAFT_ROOT / ".git").exists() and ( ARTICRAFT_ROOT / "pyproject.toml" @@ -398,6 +379,27 @@ def _prepare_articraft_checkout() -> tuple[bool, str]: return True, f"Cloned .articraft from {ARTICRAFT_REPOSITORY_URL}" +def _articraft_isolation_error() -> str | None: + """Return an error when Codex roots could contain deployment secrets.""" + checkout = ARTICRAFT_ROOT.expanduser().resolve() + repository = EMBODICHAIN_ROOT.resolve() + if checkout == repository or repository.is_relative_to(checkout): + return ( + "ARTICRAFT_ROOT must be a dedicated nested or external Git checkout, " + "not the EmbodiChain repository or one of its parents." + ) + env_path = find_gen_sim_env_file() + if env_path is not None and env_path.resolve().is_relative_to(checkout): + return "ARTICRAFT_ROOT must not contain the shared GenSim dotenv file." + try: + output_root = validate_gradio_artifact_root(ARTICRAFT_OUTPUT_ROOT) + except ValueError as exc: + return str(exc) + if env_path is not None and env_path.resolve().is_relative_to(output_root): + return "ARTICRAFT_OUTPUT_ROOT must not contain the shared GenSim dotenv file." + return None + + def configure_articraft_environment() -> str: """Clone the checkout, prepare its Conda environment, and verify Codex.""" checkout_ready, checkout_message = _prepare_articraft_checkout() @@ -484,11 +486,11 @@ def _make_result_bundle(record_id: str) -> tuple[Path, Path]: return materialized, archive -def _articraft_viser_iframe(record_id: str) -> str: +def _articraft_viser_iframe(record_id: str, port: int) -> str: """Embed the Articulation Viser service through the Gradio page hostname.""" srcdoc = ( "" + f"window.top.location.hostname + ':{port}');" ) escaped_record_id = html.escape(record_id) return ( @@ -502,34 +504,45 @@ def _articraft_viser_iframe(record_id: str) -> str: class _ArticraftViserPreview: - """Own the single Articraft Viser process and its dedicated TCP port.""" + """Own an isolated Articraft Viser process for each Gradio session.""" - def __init__(self, port: int) -> None: - self._port = port + def __init__(self, preferred_port: int) -> None: + self._preferred_port = preferred_port self._lock = threading.Lock() - self._process: subprocess.Popen[str] | None = None + self._processes: dict[str, tuple[subprocess.Popen[str], int]] = {} - def start(self, urdf_path: Path, record_id: str) -> str: - """Replace the active preview with a verified preview of one URDF.""" + def start(self, session_id: str, urdf_path: Path, record_id: str) -> str: + """Replace one session's 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): + previous = self._processes.pop(session_id, None) + if previous is not None: + terminate_process_group(previous[0]) + port = self._select_available_port() + process = start_pipeline(self._command(urdf_path, port)) + if not self._wait_until_owned(process, port): terminate_process_group(process) raise RuntimeError("New Articraft Viser preview did not bind its port.") - self._process = process - return _articraft_viser_iframe(record_id) + self._processes[session_id] = (process, port) + return _articraft_viser_iframe(record_id, port) - def stop(self) -> None: - """Stop the preview process, if this panel started one.""" + def stop(self, session_id: str | None = None) -> None: + """Stop one session's preview, or every preview during shutdown.""" with self._lock: - self._stop_managed_process() + if session_id is None: + processes = tuple( + process for process, _port in self._processes.values() + ) + self._processes.clear() + else: + current = self._processes.pop(session_id, None) + processes = () if current is None else (current[0],) + for process in processes: + terminate_process_group(process) - def _command(self, urdf_path: Path) -> list[str]: + def _command(self, urdf_path: Path, port: int) -> list[str]: return [ sys.executable, str(Path(__file__).with_name("app_media.py")), @@ -542,133 +555,69 @@ def _command(self, urdf_path: Path) -> list[str]: "--viser-host", "0.0.0.0", "--viser-port", - str(self._port), + str(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: + def _wait_until_owned(self, process: subprocess.Popen[str], port: int) -> 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 + with socket.create_connection(("127.0.0.1", port), timeout=0.2): + return True except OSError: pass time.sleep(0.25) return False - def _port_is_available(self) -> bool: + def _select_available_port(self) -> int: + if self._port_is_available(self._preferred_port): + return self._preferred_port probe = socket.socket(socket.AF_INET, socket.SOCK_STREAM) try: - probe.bind(("0.0.0.0", self._port)) + probe.bind(("0.0.0.0", 0)) + return int(probe.getsockname()[1]) + finally: + probe.close() + + @staticmethod + def _port_is_available(port: int) -> bool: + probe = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + try: + probe.bind(("0.0.0.0", 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: +def stop_articraft_viser_preview(session_id: str | None = None) -> 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() + _articraft_viser_preview.stop(session_id) atexit.register(stop_articraft_viser_preview) -def _start_articraft_viser_preview(materialized: Path, record_id: str) -> str: +def _start_articraft_viser_preview( + session_id: str, 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) + return _articraft_viser_preview.start( + session_id, + materialized / "model.urdf", + record_id, + ) def _external_check_is_unsupported(result: subprocess.CompletedProcess[str]) -> bool: @@ -760,12 +709,26 @@ def _build_codex_prompt( 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.""" - token = _begin_articraft_generation() +def generate_articraft_asset( + prompt_value: str, + image_value: Any, + request: gr.Request, +) -> Iterator[tuple[Any, ...]]: + """Initialize a record, let Codex author it, and expose one result bundle. + + Args: + prompt_value: Requested articulated-object description. + image_value: Optional Gradio reference-image value. + request: Gradio request identifying the owning browser session. + + Yields: + Updated artifact, status, log, and Viser preview values for the panel. + """ + session_id = get_request_session_id(request) + token = _articraft_runs.begin(session_id) prompt = (prompt_value or "").strip() if not prompt: - if _articraft_generation_is_active(token): + if _articraft_runs.is_active(session_id, token): yield None, "", "**Input error:** enter a description of the articulated object.", "", "" return @@ -774,7 +737,7 @@ def generate_articraft_asset(prompt_value: str, image_value: Any): message = ( "\n".join(f"- {error}" for error in errors) or "Codex CLI is unavailable." ) - if _articraft_generation_is_active(token): + if _articraft_runs.is_active(session_id, token): yield None, "", f"**Articulation is not ready.**\n\n{message}", "", "" return @@ -798,7 +761,10 @@ def generate_articraft_asset(prompt_value: str, image_value: Any): ) log_lines.append("$ " + " ".join(init_command[:-1]) + " ") initialized = _run_articraft_generation_check( - init_command, token=token, timeout=90 + init_command, + session_id=session_id, + token=token, + timeout=90, ) if initialized is None: return @@ -810,11 +776,11 @@ def generate_articraft_asset(prompt_value: str, image_value: Any): return model_path = _active_model_path(record_dir) except Exception as exc: - if _articraft_generation_is_active(token): + if _articraft_runs.is_active(session_id, token): yield None, "", f"**Setup failed:** {exc}", "\n".join(log_lines), "" return - if not _articraft_generation_is_active(token): + if not _articraft_runs.is_active(session_id, token): return final_message = run_root / "codex_final_message.txt" @@ -823,6 +789,11 @@ def generate_articraft_asset(prompt_value: str, image_value: Any): "exec", "--sandbox", "workspace-write", + "--ephemeral", + "--ignore-user-config", + "--ignore-rules", + "-c", + 'web_search="disabled"', "--color", "never", "-C", @@ -858,14 +829,14 @@ def generate_articraft_asset(prompt_value: str, image_value: Any): text=True, bufsize=1, start_new_session=True, - env=os.environ.copy(), + env=build_codex_env(), ) ) - if not _set_articraft_generation_process(token, process): + if not _articraft_runs.attach(session_id, token, process): terminate_process_group(process) return except Exception as exc: - if _articraft_generation_is_active(token): + if _articraft_runs.is_active(session_id, token): yield None, record_dir.as_posix(), f"**Codex could not start:** {exc}", "\n".join( log_lines ), "" @@ -873,11 +844,14 @@ def generate_articraft_asset(prompt_value: str, image_value: Any): 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), + kwargs={"redact_sensitive": True}, + daemon=True, ) reader.start() while process.poll() is None: - if not _articraft_generation_is_active(token, process): + if not _articraft_runs.is_active(session_id, token, process): return try: while True: @@ -896,15 +870,17 @@ def generate_articraft_asset(prompt_value: str, image_value: Any): except queue.Empty: pass finally: - _finish_articraft_generation_process(token, process) + _articraft_runs.finish(session_id, token, process) - if not _articraft_generation_is_active(token): + if not _articraft_runs.is_active(session_id, token): return 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) + log_lines.append( + "\nCodex final response:\n" + redact_sensitive_text(final_text) + ) if process.returncode: yield None, record_dir.as_posix(), f"**Codex generation failed** (exit code {process.returncode}).", "\n".join( log_lines[-300:] @@ -930,7 +906,10 @@ def generate_articraft_asset(prompt_value: str, image_value: Any): ) try: checked = _run_articraft_generation_check( - check_command, token=token, timeout=300 + check_command, + session_id=session_id, + token=token, + timeout=300, ) if checked is None: return @@ -979,7 +958,10 @@ def generate_articraft_asset(prompt_value: str, image_value: Any): ) try: compiled = _run_articraft_generation_check( - compile_command, token=token, timeout=300 + compile_command, + session_id=session_id, + token=token, + timeout=300, ) if compiled is None: return @@ -1024,7 +1006,10 @@ def generate_articraft_asset(prompt_value: str, image_value: Any): log_lines.append("$ " + " ".join(finalize_command)) try: finalized = _run_articraft_generation_check( - finalize_command, token=token, timeout=300 + finalize_command, + session_id=session_id, + token=token, + timeout=300, ) if finalized is None: return @@ -1048,7 +1033,7 @@ def generate_articraft_asset(prompt_value: str, image_value: Any): ) return - if not _articraft_generation_is_active(token): + if not _articraft_runs.is_active(session_id, token): return try: materialized, archive = _make_result_bundle(record_id) @@ -1057,7 +1042,11 @@ def generate_articraft_asset(prompt_value: str, image_value: Any): f"- Record: `{record_dir}`\n- Compiled output: `{materialized}`\n- Downloadable bundle: `{archive}`" ) try: - preview_html = _start_articraft_viser_preview(materialized, record_id) + preview_html = _start_articraft_viser_preview( + session_id, + materialized, + record_id, + ) status += "\n- Interactive Viser preview: ready" except Exception as exc: preview_html = "" diff --git a/embodichain/gen_sim/gradio_ui/app_asset_engine.py b/embodichain/gen_sim/gradio_ui/app_asset_engine.py index e94adca26..261e6865a 100644 --- a/embodichain/gen_sim/gradio_ui/app_asset_engine.py +++ b/embodichain/gen_sim/gradio_ui/app_asset_engine.py @@ -26,73 +26,64 @@ import queue import shutil -import subprocess import sys -import threading import time import uuid +from collections.abc import Iterator from pathlib import Path from typing import Any, Iterable import gradio as gr import trimesh -from app_articraft import build_articraft_panel +from app_articraft import build_articraft_panel, cleanup_articraft_session from app_config import GEN_SIM_ASSET_ROOT, SIMREADY_MESH_SUFFIXES -from app_env import EMBODICHAIN_ROOT -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 +from app_processes import ( + SessionProcessRegistry, + get_request_session_id, + read_process_output, + start_pipeline, + terminate_process_group, +) + +__all__ = [ + "build_asset_engine_panel", + "cleanup_asset_engine_session", + "prepare_asset_input_preview", + "reset_simready_asset", + "run_simready_asset", +] + +_simready_runs = SessionProcessRegistry() _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 reset_simready_asset( + request: gr.Request, +) -> tuple[None, str, None, None, None, str, str]: + """Clear SimReady widgets and stop only the requesting session's run. + Args: + request: Gradio request for the browser session initiating Reset. -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) + Returns: + Reset values for the SimReady panel widgets. + """ + _simready_runs.reset(get_request_session_id(request)) return None, "rigid_object", None, None, None, _SIMREADY_IDLE_STATUS, "" +def cleanup_asset_engine_session(request: gr.Request) -> None: + """Stop Asset-engine subprocesses owned by a disconnected session. + + Args: + request: Gradio request for the disconnecting browser session. + """ + session_id = get_request_session_id(request) + _simready_runs.reset(session_id) + cleanup_articraft_session(session_id) + + def _as_paths(value: Any) -> list[Path]: if value is None: return [] @@ -180,13 +171,26 @@ def _find_simready_output(output_root: Path) -> Path: return candidates[0] -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() +def run_simready_asset( + upload_value: Any, + category: str, + request: gr.Request, +) -> Iterator[tuple[Any, ...]]: + """Run one upstream SimReady job and stream concise subprocess progress. + + Args: + upload_value: Gradio upload value containing the asset and sidecars. + category: SimReady asset category. + request: Gradio request identifying the owning browser session. + + Yields: + Updated preview, output, status, and log values for the panel. + """ + session_id = get_request_session_id(request) + token = _simready_runs.begin(session_id) category = (category or "").strip() if not category: - if _simready_run_is_active(token): + if _simready_runs.is_active(session_id, token): yield None, None, None, "**Input error:** enter an asset category.", "" return try: @@ -198,7 +202,7 @@ 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: - if _simready_run_is_active(token): + if _simready_runs.is_active(session_id, token): yield None, None, None, f"**Input error:** {exc}", "" return @@ -214,7 +218,7 @@ def run_simready_asset(upload_value: Any, category: str): category, ] log_lines = ["$ " + " ".join(command)] - if not _simready_run_is_active(token): + if not _simready_runs.is_active(session_id, token): return yield input_preview.as_posix(), None, None, "**SimReady is running…**", "\n".join( log_lines @@ -223,17 +227,13 @@ def run_simready_asset(upload_value: Any, category: str): try: process = start_pipeline(command, use_simready_llm=True) except Exception as exc: - if _simready_run_is_active(token): + if _simready_runs.is_active(session_id, token): yield input_preview.as_posix(), None, None, f"**Pipeline start failed:** {exc}", "\n".join( log_lines ) return - with _simready_run_lock: - owns_process = _simready_run_token == token - if owns_process: - _simready_process = process - if not owns_process: + if not _simready_runs.attach(session_id, token, process): terminate_process_group(process) return @@ -244,7 +244,7 @@ def run_simready_asset(upload_value: Any, category: str): ) reader.start() while process.poll() is None: - if not _simready_run_is_active(token, process): + if not _simready_runs.is_active(session_id, token, process): return try: while True: @@ -262,7 +262,7 @@ def run_simready_asset(upload_value: Any, category: str): log_lines.append(output_queue.get_nowait()) except queue.Empty: pass - if not _simready_run_is_active(token, process): + if not _simready_runs.is_active(session_id, token, process): return if process.returncode != 0: @@ -285,7 +285,7 @@ def run_simready_asset(upload_value: Any, category: str): log_lines[-220:] ) finally: - _finish_simready_run(token, process) + _simready_runs.finish(session_id, token, process) def build_asset_engine_panel() -> dict[str, Any]: diff --git a/embodichain/gen_sim/gradio_ui/app_env.py b/embodichain/gen_sim/gradio_ui/app_env.py index ac1464a5d..4a1bc54ad 100644 --- a/embodichain/gen_sim/gradio_ui/app_env.py +++ b/embodichain/gen_sim/gradio_ui/app_env.py @@ -32,6 +32,8 @@ "ARTICRAFT_VISER_PORT", "DIRECT_NO_PROXY_VALUE", "EMBODICHAIN_ROOT", + "GRADIO_AUTH_PASSWORD", + "GRADIO_AUTH_USERNAME", "PROXY_ENV_KEYS", "SCENE_ENGINE_VISER_PORT", "SERVER_NAME", @@ -39,8 +41,12 @@ "SIMREADY_OPENAI_API_KEY", "SIMREADY_OPENAI_BASE_URL", "SIMREADY_OPENAI_MODEL", + "build_gradio_allowed_paths", + "build_gradio_blocked_paths", "configure_direct_network_env", "configure_simready_llm_env", + "get_gradio_auth", + "validate_gradio_artifact_root", ] load_gen_sim_env() @@ -81,8 +87,10 @@ def _getenv(name: str, default: str) -> str: ).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_NAME = _getenv("GRADIO_SERVER_NAME", "127.0.0.1") SERVER_PORT = int(_getenv("GRADIO_SERVER_PORT", "7860")) +GRADIO_AUTH_USERNAME = _getenv("GRADIO_AUTH_USERNAME", "") +GRADIO_AUTH_PASSWORD = _getenv("GRADIO_AUTH_PASSWORD", "") 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", "") @@ -111,3 +119,96 @@ def configure_simready_llm_env(env: Any = None) -> None: for key, value in configured_values.items(): if value: env[key] = value + + +def get_gradio_auth( + server_name: str = SERVER_NAME, + username: str = GRADIO_AUTH_USERNAME, + password: str = GRADIO_AUTH_PASSWORD, +) -> tuple[str, str] | None: + """Validate deployment exposure and return Gradio credentials. + + Args: + server_name: Interface address used by the Gradio server. + username: Optional HTTP basic-auth username. + password: Optional HTTP basic-auth password. + + Returns: + A ``(username, password)`` tuple, or ``None`` for a local-only server. + + Raises: + ValueError: If credentials are incomplete or a non-loopback server has + no authentication configured. + """ + has_username = bool(username) + has_password = bool(password) + if has_username != has_password: + raise ValueError( + "Set both GRADIO_AUTH_USERNAME and GRADIO_AUTH_PASSWORD, or neither." + ) + if has_username and has_password: + return username, password + if server_name.strip().lower() not in {"127.0.0.1", "localhost", "::1"}: + raise ValueError( + "A non-loopback GRADIO_SERVER_NAME requires Gradio authentication." + ) + return None + + +def build_gradio_allowed_paths(*roots: Path) -> list[str]: + """Resolve the explicit static and generated roots Gradio may serve. + + Args: + *roots: Static-resource or generated-artifact directories. + + Returns: + Sorted, de-duplicated absolute path strings. + """ + return sorted({str(path.expanduser().resolve()) for path in roots}) + + +def build_gradio_blocked_paths(env_path: Path | None) -> list[str]: + """Resolve repository metadata and dotenv paths Gradio must never serve. + + Args: + env_path: Active shared dotenv path, if one exists. + + Returns: + Sorted, de-duplicated absolute path strings. + """ + blocked = { + EMBODICHAIN_ROOT / ".env", + EMBODICHAIN_ROOT / ".git", + EMBODICHAIN_ROOT / "embodichain" / "gen_sim" / ".env", + } + if env_path is not None: + blocked.add(env_path) + return build_gradio_allowed_paths(*blocked) + + +def validate_gradio_artifact_root(root: Path) -> Path: + """Reject an artifact setting broad enough to expose the repository. + + Args: + root: Configured directory containing generated artifacts. + + Returns: + The normalized artifact directory. + + Raises: + ValueError: If the directory is the repository or one of its ancestors. + """ + resolved_root = root.expanduser().resolve() + repository = EMBODICHAIN_ROOT.resolve() + if resolved_root == repository or repository.is_relative_to(resolved_root): + raise ValueError( + "ARTICRAFT_OUTPUT_ROOT must be a dedicated artifact directory, not " + "the EmbodiChain repository or one of its parents." + ) + return resolved_root + + +# Gradio imports its HTTP client during application module loading. Remove the +# same unsupported or credential-bearing proxies that child pipelines exclude +# before importing any view module. +configure_direct_network_env() diff --git a/embodichain/gen_sim/gradio_ui/app_processes.py b/embodichain/gen_sim/gradio_ui/app_processes.py index 15358c451..ac92c8ad4 100644 --- a/embodichain/gen_sim/gradio_ui/app_processes.py +++ b/embodichain/gen_sim/gradio_ui/app_processes.py @@ -25,6 +25,7 @@ import sys import threading import time +import uuid from pathlib import Path from app_config import COMMANDS, PROCESS_STOP_TIMEOUT_S @@ -35,20 +36,186 @@ ) __all__ = [ + "SessionProcessRegistry", + "build_codex_env", "build_pipeline_env", "build_run_agent_command", "force_stop_all_child_processes", + "get_request_session_id", "read_process_output", "register_managed_process", "run_agent_cli_supports_robot_profile", "start_pipeline", "terminate_process_group", + "redact_sensitive_text", ] _RUN_AGENT_SUPPORTS_ROBOT_PROFILE: bool | None = None _managed_processes: dict[int, subprocess.Popen[str]] = {} _managed_processes_lock = threading.Lock() _shutdown_requested = False +_CODEX_ENV_ALLOWLIST = { + "CODEX_HOME", + "COLORTERM", + "HOME", + "LANG", + "LC_ALL", + "LC_CTYPE", + "LOGNAME", + "PATH", + "REQUESTS_CA_BUNDLE", + "SHELL", + "SSL_CERT_DIR", + "SSL_CERT_FILE", + "TEMP", + "TERM", + "TMP", + "TMPDIR", + "USER", + "XDG_CACHE_HOME", + "XDG_CONFIG_HOME", + "XDG_DATA_HOME", +} +_SENSITIVE_ENV_MARKERS = ( + "API_KEY", + "CREDENTIAL", + "PASSWORD", + "SECRET", + "TOKEN", +) + + +def get_request_session_id(request: object) -> str: + """Return the stable session identifier supplied by Gradio. + + Args: + request: Gradio request object injected into an event callback. + + Returns: + The non-empty Gradio session hash. + + Raises: + RuntimeError: If the callback was invoked without a session hash. + """ + session_id = getattr(request, "session_hash", None) + if not isinstance(session_id, str) or not session_id: + raise RuntimeError("This operation requires an active Gradio session.") + return session_id + + +class SessionProcessRegistry: + """Track one replaceable subprocess for each Gradio session. + + A registry instance belongs to one workflow, such as SimReady or Articraft. + Resetting one session can therefore never invalidate or terminate another + session's run. + """ + + def __init__(self) -> None: + self._lock = threading.Lock() + self._runs: dict[str, tuple[str, subprocess.Popen[str] | None]] = {} + + def begin(self, session_id: str) -> str: + """Start a new logical run for one session. + + Args: + session_id: Stable Gradio session identifier. + + Returns: + A new ownership token for the run. + """ + token = uuid.uuid4().hex + with self._lock: + previous = self._runs.get(session_id) + self._runs[session_id] = (token, None) + if previous is not None and previous[1] is not None: + terminate_process_group(previous[1]) + return token + + def is_active( + self, + session_id: str, + token: str, + process: subprocess.Popen[str] | None = None, + ) -> bool: + """Return whether a run still owns its session slot. + + Args: + session_id: Stable Gradio session identifier. + token: Token returned by :meth:`begin`. + process: Optional process that must also match the registered child. + + Returns: + ``True`` when the token and optional process still match. + """ + with self._lock: + current = self._runs.get(session_id) + return ( + current is not None + and current[0] == token + and (process is None or current[1] is process) + ) + + def attach( + self, + session_id: str, + token: str, + process: subprocess.Popen[str], + ) -> bool: + """Attach a subprocess to an active session run. + + Args: + session_id: Stable Gradio session identifier. + token: Token returned by :meth:`begin`. + process: Managed subprocess started for the run. + + Returns: + ``True`` if the token still owns the session slot. + """ + with self._lock: + current = self._runs.get(session_id) + if current is None or current[0] != token: + return False + self._runs[session_id] = (token, process) + return True + + def finish( + self, + session_id: str, + token: str, + process: subprocess.Popen[str], + ) -> None: + """Clear a finished subprocess while keeping its logical run active. + + Args: + session_id: Stable Gradio session identifier. + token: Token returned by :meth:`begin`. + process: Subprocess that has finished. + """ + with self._lock: + current = self._runs.get(session_id) + if current == (token, process): + self._runs[session_id] = (token, None) + + def reset(self, session_id: str) -> None: + """Invalidate and terminate only one session's process. + + Args: + session_id: Stable Gradio session identifier. + """ + with self._lock: + current = self._runs.pop(session_id, None) + if current is not None and current[1] is not None: + terminate_process_group(current[1]) + + def reset_all(self) -> None: + """Invalidate and terminate every process tracked by this registry.""" + with self._lock: + runs = tuple(self._runs.values()) + self._runs.clear() + for _token, process in runs: + if process is not None: + terminate_process_group(process) def run_agent_cli_supports_robot_profile() -> bool: @@ -132,6 +299,50 @@ def build_pipeline_env(*, use_simready_llm: bool = False) -> dict[str, str]: return env +def build_codex_env() -> dict[str, str]: + """Build a credential-minimized environment for user-directed Codex runs. + + The Codex CLI may still use its own login state through ``CODEX_HOME`` or + the normal user configuration directory, but GenSim service credentials + and dotenv-specific settings are not inherited by the command sandbox. + + Returns: + An allowlisted child-process environment. + + .. attention:: + Deployments that authenticate Codex exclusively through + ``OPENAI_API_KEY`` must use ``codex login`` or another isolated Codex + credential store instead. Passing the server key to a user-directed + process would recreate the disclosure boundary this function removes. + """ + return { + key: value + for key, value in os.environ.items() + if key in _CODEX_ENV_ALLOWLIST and value + } + + +def redact_sensitive_text(text: str) -> str: + """Replace known environment credential values in UI-bound output. + + Args: + text: Subprocess output or a final message that may contain credentials. + + Returns: + Text with non-trivial sensitive environment values replaced. + """ + redacted = text + for key, value in os.environ.items(): + upper_key = key.upper() + if ( + value + and len(value) >= 4 + and any(marker in upper_key for marker in _SENSITIVE_ENV_MARKERS) + ): + redacted = redacted.replace(value, "[REDACTED]") + return redacted + + def register_managed_process( process: subprocess.Popen[str], ) -> subprocess.Popen[str]: @@ -279,17 +490,28 @@ def read_process_output( process: subprocess.Popen[str], output_queue: queue.Queue[str], log_path: Path | None = None, + *, + redact_sensitive: bool = False, ) -> None: - """Forward merged subprocess output to the UI queue and an optional log.""" + """Forward merged subprocess output to the UI queue and an optional log. + + Args: + process: Child process whose merged stdout should be consumed. + output_queue: Destination for individual output lines. + log_path: Optional file receiving the same output. + redact_sensitive: Whether to redact known environment credentials before + forwarding or persisting each line. + """ if process.stdout is None: return 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()) + output_line = redact_sensitive_text(line) if redact_sensitive else line + output_queue.put(output_line.rstrip()) if log_file is not None: - log_file.write(line) - if not line.endswith("\n"): + log_file.write(output_line) + if not output_line.endswith("\n"): log_file.write("\n") log_file.flush() finally: diff --git a/embodichain/gen_sim/gradio_ui/app_ui.py b/embodichain/gen_sim/gradio_ui/app_ui.py index a2352d28c..62af2145f 100644 --- a/embodichain/gen_sim/gradio_ui/app_ui.py +++ b/embodichain/gen_sim/gradio_ui/app_ui.py @@ -24,7 +24,7 @@ import gradio as gr -from app_asset_engine import build_asset_engine_panel +from app_asset_engine import build_asset_engine_panel, cleanup_asset_engine_session from app_config import ( DEBUG_ENGINE_ACTION, DEBUG_ENGINE_ASSET, @@ -262,4 +262,6 @@ def build_app() -> gr.Blocks: queue=False, ) + app.unload(cleanup_asset_engine_session) + return app diff --git a/embodichain/gen_sim/gradio_ui/app_workflows.py b/embodichain/gen_sim/gradio_ui/app_workflows.py index e6122ed28..ca30fb5ac 100644 --- a/embodichain/gen_sim/gradio_ui/app_workflows.py +++ b/embodichain/gen_sim/gradio_ui/app_workflows.py @@ -24,6 +24,7 @@ import io import json import queue +import shutil import socket import subprocess import sys @@ -39,6 +40,7 @@ from app_config import ( AGENT_CONFIG, COMMANDS, + GEN_SIM_ROOT, GEN_SIM_SCENE_ROOT, FAST_GYM_CONFIG, ) @@ -635,7 +637,20 @@ def _monitor_simulation( reader.join(timeout=1.0) _append_simulation_logs(token, process, _drain_output_queue(output_queue)) - display_video = latest_audience_output_video(min_mtime_ns=started_at_ns) + source_video = latest_audience_output_video(min_mtime_ns=started_at_ns) + display_video: Path | None = None + if source_video is not None: + destination = GEN_SIM_ROOT / "action_videos" / token / source_video.name + try: + destination.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(source_video, destination) + display_video = destination + except OSError as exc: + _append_simulation_logs( + token, + process, + [f"Could not copy the simulation preview into the workspace: {exc}"], + ) with runtime_lock: if runtime.run_token != token or runtime.sim_process is not process: diff --git a/embodichain/gen_sim/gradio_ui/gradio_app.py b/embodichain/gen_sim/gradio_ui/gradio_app.py index e0713836c..17d0b79d1 100644 --- a/embodichain/gen_sim/gradio_ui/gradio_app.py +++ b/embodichain/gen_sim/gradio_ui/gradio_app.py @@ -26,12 +26,22 @@ from app_config import ( ASSETS_DIR, - GEN_SIM_ROOT, DEFAULT_CONCURRENCY_LIMIT, + GEN_SIM_ROOT, +) +from app_env import ( + ARTICRAFT_OUTPUT_ROOT, + EMBODICHAIN_ROOT, + SERVER_NAME, + SERVER_PORT, + build_gradio_allowed_paths, + build_gradio_blocked_paths, + get_gradio_auth, + validate_gradio_artifact_root, ) -from app_env import EMBODICHAIN_ROOT, SERVER_NAME, SERVER_PORT from app_processes import force_stop_all_child_processes from app_services import build_app +from embodichain.gen_sim.env import find_gen_sim_env_file __all__ = ["main"] @@ -59,6 +69,20 @@ def _install_shutdown_handlers() -> None: signal.signal(signal.SIGTERM, _handle_shutdown_signal) +def _allowed_paths() -> list[str]: + """Return only static assets and workspace-generated artifact roots.""" + return build_gradio_allowed_paths( + ASSETS_DIR, + GEN_SIM_ROOT, + validate_gradio_artifact_root(ARTICRAFT_OUTPUT_ROOT), + ) + + +def _blocked_paths() -> list[str]: + """Return source-control and dotenv paths that Gradio must never serve.""" + return build_gradio_blocked_paths(find_gen_sim_env_file()) + + def main() -> None: if not EMBODICHAIN_ROOT.is_dir(): raise FileNotFoundError(f"EmbodiChain root not found: {EMBODICHAIN_ROOT}") @@ -69,11 +93,9 @@ def main() -> None: app.launch( server_name=SERVER_NAME, server_port=SERVER_PORT, - allowed_paths=[ - str(EMBODICHAIN_ROOT), - str(ASSETS_DIR), - str(GEN_SIM_ROOT), - ], + auth=get_gradio_auth(), + allowed_paths=_allowed_paths(), + blocked_paths=_blocked_paths(), ) finally: _stop_child_processes() diff --git a/embodichain/gen_sim/gradio_ui/gradio_visualization_architecture.md b/embodichain/gen_sim/gradio_ui/gradio_visualization_architecture.md index 6c98118f4..186586efe 100644 --- a/embodichain/gen_sim/gradio_ui/gradio_visualization_architecture.md +++ b/embodichain/gen_sim/gradio_ui/gradio_visualization_architecture.md @@ -51,15 +51,17 @@ conda run -n embodichain python gradio_app.py | 变量 | 默认值 | 用途 | | --- | --- | --- | | EmbodiChain root | 自动从 `embodichain/gen_sim/env.py` 的源码位置推导 | EmbodiChain 根目录；不再从 `.env` 配置。 | -| `GRADIO_SERVER_NAME` | `0.0.0.0` | Gradio 监听地址。 | +| `GRADIO_SERVER_NAME` | `127.0.0.1` | Gradio 监听地址；非回环地址必须启用认证。 | | `GRADIO_SERVER_PORT` | `7860` | Gradio 监听端口。 | +| `GRADIO_AUTH_USERNAME` | 空 | 非本机部署的 Gradio 用户名。 | +| `GRADIO_AUTH_PASSWORD` | 空 | 非本机部署的 Gradio 密码。 | | `SCENE_ENGINE_VISER_PORT` | `8080` | 独立 Scene Engine 的 Viser 端口。 | -| `ARTICRAFT_VISER_PORT` | `8081` | Articraft 关节预览的 Viser 端口。 | +| `ARTICRAFT_VISER_PORT` | `8081` | Articraft 关节预览的首选 Viser 端口；占用时为会话分配其他可用端口。 | | `ARTICRAFT_ROOT` | `<项目>/.articraft` | Articraft checkout。 | | `ARTICRAFT_CONDA_ENV` | `articraft` | 运行 Articraft CLI 的 Conda 环境。 | | `ARTICRAFT_OUTPUT_ROOT` | `<项目>/.gen_sim/articraft` | Articraft 记录、运行日志和导出 bundle。 | -`app.launch()` 仅开放 EmbodiChain 根目录、`assets/` 和 `.gen_sim/` 给浏览器读取。pipeline 子进程由 `build_pipeline_env()` 从共享 `.env` 创建环境：它清除代理变量、设置 `NO_PROXY=no_proxy=*` 并关闭 Gradio analytics。只有 SimReady 子进程会额外把非空的 `SIMREADY_OPENAI_*` 映射为其上游 CLI 需要的 `OPENAI_*`；Scene Engine、DexSim、Viser 和 Articraft 直接继承 `.env` 中的原始配置。这不会改写启动 Gradio 的父进程环境。 +`app.launch()` 仅开放 UI 静态资源、`.gen_sim/` 生成物和配置的 Articraft 输出目录，并显式禁止 `.env`、`.git/` 等敏感路径。pipeline 子进程由 `build_pipeline_env()` 从共享 `.env` 创建环境：它清除代理变量、设置 `NO_PROXY=no_proxy=*` 并关闭 Gradio analytics。只有 SimReady 子进程会额外把非空的 `SIMREADY_OPENAI_*` 映射为其上游 CLI 需要的 `OPENAI_*`；Scene Engine、DexSim、Viser 和 Articraft 直接继承 `.env` 中的原始配置。Codex 作为用户指令驱动的子进程，使用独立登录状态和最小化环境，不继承 GenSim 服务凭据。 ## 页面与引擎 @@ -100,7 +102,7 @@ python -m embodichain.gen_sim.simready_pipeline.cli.start \ 处理函数以 generator 持续返回最近的 stdout；完成时优先预览 `asset_simready.glb`，只有 OBJ 时再转为 GLB。此路径不依赖 DexSim。 -`Reset SimReady` 会清空上传、类别、预览、下载项和日志，并按进程组终止正在运行的 SimReady CLI 及其子进程。 +`Reset SimReady` 会清空当前浏览器会话的上传、类别、预览、下载项和日志，并仅按进程组终止该会话正在运行的 SimReady CLI 及其子进程。 ### Articulation：Articraft + Codex @@ -119,9 +121,9 @@ description + optional image → exports/.zip + Viser articulation preview ``` -产物、记录和参考图均在 `ARTICRAFT_OUTPUT_ROOT` 下，不能直接当作 Action engine 的 Gym 场景或 SimReady 资产；若要进入后续仿真，需要另行定义并实现转换/导入流程。Articraft Viser 每次成功预览会终止旧的 Articraft 预览进程，再以 `0.0.0.0:` 启动新进程。 +产物、记录和参考图均在 `ARTICRAFT_OUTPUT_ROOT` 下，不能直接当作 Action engine 的 Gym 场景或 SimReady 资产；若要进入后续仿真，需要另行定义并实现转换/导入流程。Articraft Viser 按 Gradio 会话管理预览进程：首先尝试 `ARTICRAFT_VISER_PORT`，若已被占用则分配其他可用端口，不会查找或发送信号给占用端口的外部进程。 -`Reset Articulation` 会清空描述、参考图、记录与下载结果，终止当前 Articraft/Codex 命令进程组，并关闭该面板启动的 Viser。 +`Reset Articulation` 会清空当前会话的描述、参考图、记录与下载结果，终止该会话的 Articraft/Codex 命令进程组，并关闭该会话启动的 Viser。 ## 独立 Scene engine 和 Viser @@ -170,7 +172,7 @@ python -m embodichain.gen_sim.action_agent_pipeline.cli.run_agent \ ## 共享状态、并发和进度 -Scene engine 和 Action engine 共享 `RuntimeState` 与 `runtime_lock`，其中包含运行 token、DexSim/Scene Viser 进程、输入、预览、日志和阶段。运行 token 用于丢弃过期线程的更新。Articraft Viser 使用单独的锁和进程引用；SimReady 使用自己的同步 generator。 +Scene engine 和 Action engine 共享 `RuntimeState` 与 `runtime_lock`，其中包含运行 token、DexSim/Scene Viser 进程、输入、预览、日志和阶段。运行 token 用于丢弃过期线程的更新。SimReady 和 Articraft 使用按 `request.session_hash` 隔离的进程注册表；Reset 和页面卸载只清理所属会话，应用退出时再统一清理全部已注册子进程。 `app.queue(default_concurrency_limit=1)` 将队列中的高成本回调串行化。Action engine 的 `Timer(2.0)` 读取共享状态。Asset/Articraft 面板不写入这一状态，但仍会受 Gradio 队列限制。 @@ -212,7 +214,7 @@ embodichain/gen_sim/scene_engine/cli/preview.py .env ``` -Articulation 还需要 Git（首次 clone）、Conda、`ARTICRAFT_CONDA_ENV` 和 Codex CLI。生成请求会交给本机 Codex CLI 执行，因此只应提交可信请求。 +Articulation 还需要 Git（首次 clone）、Conda、`ARTICRAFT_CONDA_ENV` 和已通过独立凭据存储完成登录的 Codex CLI。Codex 子进程不继承 `.env` 中的 API key、token 或密码，输出在返回浏览器前还会按已知敏感环境值脱敏。生成请求会交给本机 Codex CLI 执行，因此默认只在本机可信工作台中使用。 每次修改后至少执行： diff --git a/pyproject.toml b/pyproject.toml index af9177192..65c7b2d41 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,6 +54,7 @@ dependencies = [ [project.optional-dependencies] gensim = [ "bpy", + "gradio>=6.19,<7", "pyrender==0.1.45", "requests", "Pillow", diff --git a/tests/gen_sim/__init__.py b/tests/gen_sim/__init__.py new file mode 100644 index 000000000..355d915ff --- /dev/null +++ b/tests/gen_sim/__init__.py @@ -0,0 +1,17 @@ +# ---------------------------------------------------------------------------- +# 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 diff --git a/tests/gen_sim/gradio_ui/__init__.py b/tests/gen_sim/gradio_ui/__init__.py new file mode 100644 index 000000000..355d915ff --- /dev/null +++ b/tests/gen_sim/gradio_ui/__init__.py @@ -0,0 +1,17 @@ +# ---------------------------------------------------------------------------- +# 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 diff --git a/tests/gen_sim/gradio_ui/test_app_articraft.py b/tests/gen_sim/gradio_ui/test_app_articraft.py new file mode 100644 index 000000000..bfc68e1ca --- /dev/null +++ b/tests/gen_sim/gradio_ui/test_app_articraft.py @@ -0,0 +1,59 @@ +# ---------------------------------------------------------------------------- +# 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 socket +import sys +from pathlib import Path + +import pytest + +GRADIO_UI_ROOT = ( + Path(__file__).resolve().parents[3] / "embodichain" / "gen_sim" / "gradio_ui" +) +sys.path.insert(0, str(GRADIO_UI_ROOT)) + +import app_articraft # noqa: E402 + + +def test_preview_selects_another_port_when_preferred_port_is_occupied() -> None: + listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + listener.bind(("127.0.0.1", 0)) + listener.listen() + occupied_port = int(listener.getsockname()[1]) + previews = app_articraft._ArticraftViserPreview(occupied_port) + try: + selected_port = previews._select_available_port() + finally: + listener.close() + + assert selected_port != occupied_port + assert selected_port > 0 + + +def test_codex_checkout_cannot_be_the_embodichain_repository( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + app_articraft, + "ARTICRAFT_ROOT", + app_articraft.EMBODICHAIN_ROOT, + ) + + assert "dedicated nested or external Git checkout" in ( + app_articraft._articraft_isolation_error() or "" + ) diff --git a/tests/gen_sim/gradio_ui/test_app_env.py b/tests/gen_sim/gradio_ui/test_app_env.py new file mode 100644 index 000000000..27231eb51 --- /dev/null +++ b/tests/gen_sim/gradio_ui/test_app_env.py @@ -0,0 +1,72 @@ +# ---------------------------------------------------------------------------- +# 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 sys +from pathlib import Path + +import pytest + +GRADIO_UI_ROOT = ( + Path(__file__).resolve().parents[3] / "embodichain" / "gen_sim" / "gradio_ui" +) +sys.path.insert(0, str(GRADIO_UI_ROOT)) + +import app_env # noqa: E402 + + +def test_local_gradio_server_does_not_require_authentication() -> None: + assert app_env.get_gradio_auth("127.0.0.1", "", "") is None + + +def test_remote_gradio_server_requires_authentication() -> None: + with pytest.raises(ValueError, match="requires Gradio authentication"): + app_env.get_gradio_auth("0.0.0.0", "", "") + + +def test_remote_gradio_server_accepts_complete_credentials() -> None: + assert app_env.get_gradio_auth("0.0.0.0", "workspace", "secret") == ( + "workspace", + "secret", + ) + + +def test_partial_gradio_credentials_are_rejected() -> None: + with pytest.raises(ValueError, match="Set both"): + app_env.get_gradio_auth("127.0.0.1", "workspace", "") + + +def test_gradio_file_access_excludes_repository_and_blocks_dotenv( + tmp_path: Path, +) -> None: + generated_root = tmp_path / "generated" + static_root = tmp_path / "static" + external_env = tmp_path / "deployment.env" + + allowed = app_env.build_gradio_allowed_paths(generated_root, static_root) + blocked = app_env.build_gradio_blocked_paths(external_env) + + assert str(app_env.EMBODICHAIN_ROOT.resolve()) not in allowed + assert str(generated_root.resolve()) in allowed + assert str(static_root.resolve()) in allowed + assert str(external_env.resolve()) in blocked + assert str((app_env.EMBODICHAIN_ROOT / ".git").resolve()) in blocked + + +def test_repository_cannot_be_configured_as_artifact_root() -> None: + with pytest.raises(ValueError, match="dedicated artifact directory"): + app_env.validate_gradio_artifact_root(app_env.EMBODICHAIN_ROOT) diff --git a/tests/gen_sim/gradio_ui/test_app_processes.py b/tests/gen_sim/gradio_ui/test_app_processes.py new file mode 100644 index 000000000..a222a3273 --- /dev/null +++ b/tests/gen_sim/gradio_ui/test_app_processes.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 sys +from pathlib import Path +from typing import Any + +import pytest + +GRADIO_UI_ROOT = ( + Path(__file__).resolve().parents[3] / "embodichain" / "gen_sim" / "gradio_ui" +) +sys.path.insert(0, str(GRADIO_UI_ROOT)) + +import app_processes # noqa: E402 + + +class FakeProcess: + """Minimal subprocess stand-in for ownership tests.""" + + +def test_reset_stops_only_the_requesting_session( + monkeypatch: pytest.MonkeyPatch, +) -> None: + registry = app_processes.SessionProcessRegistry() + first_process = FakeProcess() + second_process = FakeProcess() + stopped: list[Any] = [] + monkeypatch.setattr(app_processes, "terminate_process_group", stopped.append) + + first_token = registry.begin("first-session") + second_token = registry.begin("second-session") + assert registry.attach("first-session", first_token, first_process) + assert registry.attach("second-session", second_token, second_process) + + registry.reset("second-session") + + assert stopped == [second_process] + assert registry.is_active("first-session", first_token, first_process) + assert not registry.is_active("second-session", second_token, second_process) + + +def test_new_run_replaces_only_the_same_sessions_process( + monkeypatch: pytest.MonkeyPatch, +) -> None: + registry = app_processes.SessionProcessRegistry() + previous_process = FakeProcess() + stopped: list[Any] = [] + monkeypatch.setattr(app_processes, "terminate_process_group", stopped.append) + + previous_token = registry.begin("session") + assert registry.attach("session", previous_token, previous_process) + + replacement_token = registry.begin("session") + + assert stopped == [previous_process] + assert not registry.is_active("session", previous_token) + assert registry.is_active("session", replacement_token) + + +def test_codex_environment_excludes_service_credentials( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("PATH", "/usr/bin") + monkeypatch.setenv("HOME", "/tmp/codex-user") + monkeypatch.setenv("OPENAI_API_KEY", "server-api-key") + monkeypatch.setenv("SIMREADY_OPENAI_API_KEY", "simready-api-key") + monkeypatch.setenv("SCENE_ENGINE_IMAGE_SEGMENTATION_BASE_URL", "https://service") + + child_env = app_processes.build_codex_env() + + assert child_env["PATH"] == "/usr/bin" + assert child_env["HOME"] == "/tmp/codex-user" + assert "OPENAI_API_KEY" not in child_env + assert "SIMREADY_OPENAI_API_KEY" not in child_env + assert "SCENE_ENGINE_IMAGE_SEGMENTATION_BASE_URL" not in child_env + + +def test_sensitive_output_values_are_redacted( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("SERVICE_TOKEN", "sensitive-token-value") + + assert ( + app_processes.redact_sensitive_text("credential=sensitive-token-value") + == "credential=[REDACTED]" + ) diff --git a/tests/gen_sim/scene_engine/__init__.py b/tests/gen_sim/scene_engine/__init__.py new file mode 100644 index 000000000..355d915ff --- /dev/null +++ b/tests/gen_sim/scene_engine/__init__.py @@ -0,0 +1,17 @@ +# ---------------------------------------------------------------------------- +# 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 diff --git a/tests/gen_sim/scene_engine/test_scene_engine_config.py b/tests/gen_sim/scene_engine/test_scene_engine_config.py deleted file mode 100644 index 01914e144..000000000 --- a/tests/gen_sim/scene_engine/test_scene_engine_config.py +++ /dev/null @@ -1,100 +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 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/simready_pipeline/__init__.py b/tests/gen_sim/simready_pipeline/__init__.py new file mode 100644 index 000000000..355d915ff --- /dev/null +++ b/tests/gen_sim/simready_pipeline/__init__.py @@ -0,0 +1,17 @@ +# ---------------------------------------------------------------------------- +# 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 diff --git a/tests/gen_sim/test_gen_sim_env.py b/tests/gen_sim/test_gen_sim_env.py new file mode 100644 index 000000000..215a12813 --- /dev/null +++ b/tests/gen_sim/test_gen_sim_env.py @@ -0,0 +1,65 @@ +# ---------------------------------------------------------------------------- +# 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 import env as gen_sim_env + + +def test_missing_default_env_file_is_optional( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + module_path = tmp_path / "embodichain" / "gen_sim" / "env.py" + monkeypatch.delenv("EMBODICHAIN_ENV_FILE", raising=False) + monkeypatch.setattr(gen_sim_env, "__file__", str(module_path)) + + assert gen_sim_env.find_gen_sim_env_file() is None + assert gen_sim_env.load_gen_sim_env({}) is None + + +def test_missing_configured_env_file_is_optional( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + env_path = tmp_path / "missing.env" + monkeypatch.setenv("EMBODICHAIN_ENV_FILE", str(env_path)) + + assert gen_sim_env.find_gen_sim_env_file() == env_path.resolve() + assert gen_sim_env.load_gen_sim_env({}) is None + + +def test_shell_values_take_precedence_over_dotenv( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + env_path = tmp_path / ".env" + env_path.write_text( + "OPENAI_MODEL=dotenv-model\nOPENAI_API_KEY=dotenv-key\n", + encoding="utf-8", + ) + monkeypatch.setenv("EMBODICHAIN_ENV_FILE", str(env_path)) + target_env = {"OPENAI_MODEL": "shell-model"} + + assert gen_sim_env.load_gen_sim_env(target_env) == env_path.resolve() + assert target_env == { + "OPENAI_MODEL": "shell-model", + "OPENAI_API_KEY": "dotenv-key", + } From 0d78765d87d6dce3721f6fd9bad4a13b1c4000c1 Mon Sep 17 00:00:00 2001 From: PengXuanchao Date: Mon, 10 Aug 2026 14:17:18 +0800 Subject: [PATCH 50/53] change the logic of engine isolation and add the stop function in action engine --- embodichain/gen_sim/.env.example | 1 + embodichain/gen_sim/gradio_ui/app_env.py | 2 + embodichain/gen_sim/gradio_ui/app_state.py | 53 +- embodichain/gen_sim/gradio_ui/app_ui.py | 49 +- .../gen_sim/gradio_ui/app_workflows.py | 497 ++++++++++++------ .../gradio_visualization_architecture.md | 23 +- tests/gen_sim/gradio_ui/test_app_workflows.py | 254 +++++++++ 7 files changed, 684 insertions(+), 195 deletions(-) create mode 100644 tests/gen_sim/gradio_ui/test_app_workflows.py diff --git a/embodichain/gen_sim/.env.example b/embodichain/gen_sim/.env.example index ee293c134..ba46bef48 100644 --- a/embodichain/gen_sim/.env.example +++ b/embodichain/gen_sim/.env.example @@ -30,6 +30,7 @@ GRADIO_AUTH_USERNAME="" GRADIO_AUTH_PASSWORD="" SCENE_ENGINE_VISER_PORT=8080 ARTICRAFT_VISER_PORT=8081 +ACTION_ENGINE_VISER_PORT=8082 ARTICRAFT_ROOT="" ARTICRAFT_REPOSITORY_URL="https://github.com/mattzh72/articraft.git" ARTICRAFT_CONDA_ENV="articraft" diff --git a/embodichain/gen_sim/gradio_ui/app_env.py b/embodichain/gen_sim/gradio_ui/app_env.py index 4a1bc54ad..cbba7de1f 100644 --- a/embodichain/gen_sim/gradio_ui/app_env.py +++ b/embodichain/gen_sim/gradio_ui/app_env.py @@ -25,6 +25,7 @@ from embodichain.gen_sim.env import get_embodichain_root, load_gen_sim_env __all__ = [ + "ACTION_ENGINE_VISER_PORT", "ARTICRAFT_CONDA_ENV", "ARTICRAFT_OUTPUT_ROOT", "ARTICRAFT_REPOSITORY_URL", @@ -87,6 +88,7 @@ def _getenv(name: str, default: str) -> str: ).expanduser() SCENE_ENGINE_VISER_PORT = int(_getenv("SCENE_ENGINE_VISER_PORT", "8080")) ARTICRAFT_VISER_PORT = int(_getenv("ARTICRAFT_VISER_PORT", "8081")) +ACTION_ENGINE_VISER_PORT = int(_getenv("ACTION_ENGINE_VISER_PORT", "8082")) SERVER_NAME = _getenv("GRADIO_SERVER_NAME", "127.0.0.1") SERVER_PORT = int(_getenv("GRADIO_SERVER_PORT", "7860")) GRADIO_AUTH_USERNAME = _getenv("GRADIO_AUTH_USERNAME", "") diff --git a/embodichain/gen_sim/gradio_ui/app_state.py b/embodichain/gen_sim/gradio_ui/app_state.py index cf9df4067..d16c68b05 100644 --- a/embodichain/gen_sim/gradio_ui/app_state.py +++ b/embodichain/gen_sim/gradio_ui/app_state.py @@ -14,20 +14,26 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Shared runtime state for the Scene and Action engines.""" +"""Per-session runtime state for the Scene and Action engines.""" from __future__ import annotations -import subprocess import threading -import uuid from collections import deque from dataclasses import dataclass, field from pathlib import Path from app_config import PHASE_DEFINITIONS -__all__ = ["PHASES", "Phase", "runtime", "runtime_lock", "set_runtime_phase_locked"] +__all__ = [ + "PHASES", + "Phase", + "RuntimeState", + "SessionRuntimeRegistry", + "runtime_lock", + "runtime_registry", + "set_runtime_phase_locked", +] @dataclass(frozen=True) @@ -41,11 +47,9 @@ class Phase: @dataclass class RuntimeState: + """Mutable Scene/Action UI state owned by one Gradio session.""" + is_busy: bool = False - run_token: str = field(default_factory=lambda: uuid.uuid4().hex) - 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 phase_key: str = "idle" status: str = "Idle." @@ -57,10 +61,39 @@ class RuntimeState: log_lines: deque[str] = field(default_factory=deque) -runtime = RuntimeState() +class SessionRuntimeRegistry: + """Own one Scene/Action UI runtime for each Gradio session hash.""" + + def __init__(self) -> None: + self._lock = threading.Lock() + self._states: dict[str, RuntimeState] = {} + + def get(self, session_id: str) -> RuntimeState: + """Return the existing state for a session, creating it when absent. + + Args: + session_id: Stable Gradio session identifier. + + Returns: + Runtime state owned exclusively by ``session_id``. + """ + with self._lock: + return self._states.setdefault(session_id, RuntimeState()) + + def reset(self, session_id: str) -> None: + """Discard only one session's UI runtime state. + + Args: + session_id: Stable Gradio session identifier. + """ + with self._lock: + self._states.pop(session_id, None) + + +runtime_registry = SessionRuntimeRegistry() runtime_lock = threading.Lock() -def set_runtime_phase_locked(new_phase_key: str) -> None: +def set_runtime_phase_locked(runtime: RuntimeState, new_phase_key: str) -> None: """Set the current UI phase while the caller holds ``runtime_lock``.""" runtime.phase_key = new_phase_key diff --git a/embodichain/gen_sim/gradio_ui/app_ui.py b/embodichain/gen_sim/gradio_ui/app_ui.py index 62af2145f..2869bc46a 100644 --- a/embodichain/gen_sim/gradio_ui/app_ui.py +++ b/embodichain/gen_sim/gradio_ui/app_ui.py @@ -36,13 +36,16 @@ ROBOT_PROFILES, UI_TEXT, ) +from app_processes import get_request_session_id from app_workflows import ( + cleanup_workflow_session, format_status, preview_saved_scene, refresh_saved_scenes, reset_scene_engine, run_action_engine_from_current, run_scene_engine, + stop_action_engine, ui_snapshot, ) @@ -63,16 +66,29 @@ def select_engine(selected_engine: str): ) -def action_engine_snapshot(): - """Adapt the shared runtime snapshot to the Action-engine status widgets.""" - video, task, progress, status, _initial, _edited, _objects = ui_snapshot() +def action_engine_snapshot(request: gr.Request) -> tuple[object, ...]: + """Adapt this session's runtime snapshot to the Action status widgets.""" + session_id = get_request_session_id(request) + video, task, progress, status, _initial, _edited, _objects = ui_snapshot(session_id) return video, task, progress, status -def run_action_engine_panel(task_text: str, robot_profile: str | None): +def run_action_engine_panel( + task_text: str, + robot_profile: str | None, + request: gr.Request, +) -> tuple[object, ...]: """Run the Action engine and return its latest UI snapshot.""" - run_action_engine_from_current(task_text, robot_profile) - return action_engine_snapshot() + video, task, progress, status, _initial, _edited, _objects = ( + run_action_engine_from_current(task_text, robot_profile, request) + ) + return video, task, progress, status + + +def cleanup_app_session(request: gr.Request) -> None: + """Stop every engine process owned by a disconnected Gradio session.""" + cleanup_workflow_session(request) + cleanup_asset_engine_session(request) def build_app() -> gr.Blocks: @@ -159,7 +175,12 @@ def build_app() -> gr.Blocks: value=DEFAULT_ROBOT_PROFILE, label=UI_TEXT[LANGUAGE_EN]["robot"], ) - action_run = gr.Button("Run DexSim", variant="primary") + with gr.Row(): + action_run = gr.Button("Run DexSim", variant="primary") + action_stop = gr.Button( + "Stop Action Engine", + variant="stop", + ) with gr.Column(scale=2): action_scene = gr.HTML( "
" @@ -251,6 +272,18 @@ def build_app() -> gr.Blocks: action_status, ], ) + action_stop.click( + stop_action_engine, + outputs=[ + action_scene, + action_scene_status, + action_video, + action_current_task, + action_progress, + action_status, + ], + queue=False, + ) action_refresh_timer.tick( action_engine_snapshot, outputs=[ @@ -262,6 +295,6 @@ def build_app() -> gr.Blocks: queue=False, ) - app.unload(cleanup_asset_engine_session) + app.unload(cleanup_app_session) return app diff --git a/embodichain/gen_sim/gradio_ui/app_workflows.py b/embodichain/gen_sim/gradio_ui/app_workflows.py index ca30fb5ac..ce81400ae 100644 --- a/embodichain/gen_sim/gradio_ui/app_workflows.py +++ b/embodichain/gen_sim/gradio_ui/app_workflows.py @@ -30,7 +30,7 @@ import sys import threading import time -import uuid +from collections.abc import Iterator from pathlib import Path import gradio as gr @@ -44,28 +44,54 @@ GEN_SIM_SCENE_ROOT, FAST_GYM_CONFIG, ) -from app_env import SCENE_ENGINE_VISER_PORT, configure_direct_network_env +from app_env import ( + ACTION_ENGINE_VISER_PORT, + SCENE_ENGINE_VISER_PORT, + configure_direct_network_env, +) from app_media import latest_audience_output_video from app_processes import ( + SessionProcessRegistry, build_run_agent_command, + get_request_session_id, read_process_output, start_pipeline, terminate_process_group, ) -from app_state import PHASES, Phase, runtime, runtime_lock, set_runtime_phase_locked +from app_state import ( + PHASES, + Phase, + RuntimeState, + runtime_lock, + runtime_registry, + set_runtime_phase_locked, +) __all__ = [ + "cleanup_workflow_session", "format_status", "preview_saved_scene", "refresh_saved_scenes", "reset_scene_engine", "run_action_engine_from_current", "run_scene_engine", + "stop_action_engine", "ui_snapshot", ] configure_direct_network_env() +_scene_runs = SessionProcessRegistry() +_action_runs = SessionProcessRegistry() +_action_preview_runs = SessionProcessRegistry() +_preview_start_lock = threading.Lock() + +_ACTION_IDLE_PREVIEW = ( + "
" + "Select a generated scene to preview it." + "
" +) + def _drain_output_queue(output_queue: queue.Queue[str]) -> list[str]: lines: list[str] = [] @@ -93,6 +119,7 @@ def _scene_engine_phase_from_log(line: str, current_key: str) -> str: def _scene_engine_updates( + runtime: RuntimeState, output_root: Path | None = None, preview_html: str | None = None, ) -> tuple[int, str, str | None, str]: @@ -151,6 +178,18 @@ def _wait_for_viser(port: int, process: subprocess.Popen[str]) -> bool: return False +def _select_available_port(preferred_port: int) -> int: + """Return the preferred Viser port, or an ephemeral port when occupied.""" + for port in (preferred_port, 0): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener: + try: + listener.bind(("127.0.0.1", port)) + except OSError: + continue + return int(listener.getsockname()[1]) + raise RuntimeError("Could not allocate a local Viser port.") + + def _viser_iframe(port: int, scene_hash: str) -> str: srcdoc = ( "