From 4038b31047ee3071518b290aaac1d859be64393b Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Mon, 10 Aug 2026 22:51:02 +0800 Subject: [PATCH 01/29] Align the doc, test, and client with the newest .env setup (image segmentation client) --- .../source/features/generative_sim/scene_engine.md | 2 +- .../scene_engine/clients/image_segmentation.py | 14 +++++++------- tests/gen_sim/scene_engine/test_clients.py | 10 +++++----- 3 files changed, 13 insertions(+), 13 deletions(-) diff --git a/docs/source/features/generative_sim/scene_engine.md b/docs/source/features/generative_sim/scene_engine.md index 03e4b3fa4..a8792faab 100644 --- a/docs/source/features/generative_sim/scene_engine.md +++ b/docs/source/features/generative_sim/scene_engine.md @@ -43,7 +43,7 @@ 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_IMAGE_SEGMENTATION_BY_PROMPT_PATH="/segment_by_prompt" SCENE_ENGINE_GEOMETRY_GENERATION_BASE_URL="http://host:port" SCENE_ENGINE_GEOMETRY_GENERATION_TIMEOUT_S=600 diff --git a/embodichain/gen_sim/scene_engine/clients/image_segmentation.py b/embodichain/gen_sim/scene_engine/clients/image_segmentation.py index 2c56af91a..8414b9627 100644 --- a/embodichain/gen_sim/scene_engine/clients/image_segmentation.py +++ b/embodichain/gen_sim/scene_engine/clients/image_segmentation.py @@ -36,14 +36,14 @@ def __init__( timeout_s: int, max_attempts: int, health_path: str, - segment_single_object_path: str, + segment_by_prompt_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._segment_by_prompt_path = segment_by_prompt_path self._session = session or requests.Session() @classmethod @@ -96,7 +96,7 @@ def segment_single_object( try: with resolved_image_path.open("rb") as image_file: response = self._session.post( - self._url(self._segment_single_object_path), + self._url(self._segment_by_prompt_path), data={"prompt": prompt}, files={"image": (resolved_image_path.name, image_file)}, timeout=self._timeout_s, @@ -138,7 +138,7 @@ def _load_dotenv_config() -> dict[str, Any]: "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", + "SCENE_ENGINE_IMAGE_SEGMENTATION_BY_PROMPT_PATH", ) try: timeout_s = int(values["SCENE_ENGINE_IMAGE_SEGMENTATION_TIMEOUT_S"]) @@ -165,7 +165,7 @@ def _load_dotenv_config() -> dict[str, Any]: string_keys = ( "SCENE_ENGINE_IMAGE_SEGMENTATION_BASE_URL", "SCENE_ENGINE_IMAGE_SEGMENTATION_HEALTH_PATH", - "SCENE_ENGINE_IMAGE_SEGMENTATION_SINGLE_OBJECT_PATH", + "SCENE_ENGINE_IMAGE_SEGMENTATION_BY_PROMPT_PATH", ) for key in string_keys: if not values[key].strip(): @@ -176,8 +176,8 @@ def _load_dotenv_config() -> dict[str, Any]: "timeout_s": timeout_s, "max_attempts": max_attempts, "health_path": values["SCENE_ENGINE_IMAGE_SEGMENTATION_HEALTH_PATH"].strip(), - "segment_single_object_path": values[ - "SCENE_ENGINE_IMAGE_SEGMENTATION_SINGLE_OBJECT_PATH" + "segment_by_prompt_path": values[ + "SCENE_ENGINE_IMAGE_SEGMENTATION_BY_PROMPT_PATH" ].strip(), } diff --git a/tests/gen_sim/scene_engine/test_clients.py b/tests/gen_sim/scene_engine/test_clients.py index 513f0c5c0..a2ccacb21 100644 --- a/tests/gen_sim/scene_engine/test_clients.py +++ b/tests/gen_sim/scene_engine/test_clients.py @@ -79,7 +79,7 @@ def test_clients_load_their_required_dotenv_values( "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", + "SCENE_ENGINE_IMAGE_SEGMENTATION_BY_PROMPT_PATH": "/segment_by_prompt", } llm_values = { "OPENAI_API_KEY": "test-key", @@ -107,7 +107,7 @@ def test_clients_load_their_required_dotenv_values( 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 segmentation_client._segment_by_prompt_path == "/segment_by_prompt" assert llm_client_config.default_query == {"api-version": "1"} assert llm_client_config.base_url == "http://llm/v1" @@ -146,7 +146,7 @@ def test_service_health_checks_use_the_configured_health_path() -> None: timeout_s=30, max_attempts=1, health_path="/health", - segment_single_object_path="/predict", + segment_by_prompt_path="/segment_by_prompt", session=segmentation_session, ) @@ -170,7 +170,7 @@ def test_segmentation_client_posts_prompt_and_returns_rle_masks(tmp_path: Path) timeout_s=30, max_attempts=1, health_path="/health", - segment_single_object_path="/predict", + segment_by_prompt_path="/segment_by_prompt", session=session, ) @@ -178,7 +178,7 @@ def test_segmentation_client_posts_prompt_and_returns_rle_masks(tmp_path: Path) rle_mask ] assert session.post_call is not None - assert session.post_call["url"] == "http://segment/predict" + assert session.post_call["url"] == "http://segment/segment_by_prompt" assert session.post_call["data"] == {"prompt": "table"} From e435426c5a669f0052f2200606a2af870ad542eb Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Tue, 11 Aug 2026 15:09:49 +0800 Subject: [PATCH 02/29] Add image generation client + update .env + update doc + test files --- .../features/generative_sim/scene_engine.md | 10 +- .../scene_engine/clients/image_generation.py | 172 ++++++++++++++++++ tests/gen_sim/scene_engine/test_clients.py | 106 ++++++++++- 3 files changed, 285 insertions(+), 3 deletions(-) create mode 100644 embodichain/gen_sim/scene_engine/clients/image_generation.py diff --git a/docs/source/features/generative_sim/scene_engine.md b/docs/source/features/generative_sim/scene_engine.md index a8792faab..c9f569ee2 100644 --- a/docs/source/features/generative_sim/scene_engine.md +++ b/docs/source/features/generative_sim/scene_engine.md @@ -29,8 +29,8 @@ python -m embodichain scene-engine \ ## Configuration -Scene Engine reads the LLM, segmentation, and geometry-generation settings -from `embodichain/gen_sim/.env`: +Scene Engine reads the LLM, segmentation, image-generation, and +geometry-generation settings from `embodichain/gen_sim/.env`: ```bash OPENAI_API_KEY="your-api-key" @@ -45,6 +45,12 @@ SCENE_ENGINE_IMAGE_SEGMENTATION_MAX_ATTEMPTS=3 SCENE_ENGINE_IMAGE_SEGMENTATION_HEALTH_PATH="/health" SCENE_ENGINE_IMAGE_SEGMENTATION_BY_PROMPT_PATH="/segment_by_prompt" +SCENE_ENGINE_IMAGE_GENERATION_BASE_URL="http://host:port" +SCENE_ENGINE_IMAGE_GENERATION_TIMEOUT_S=120 +SCENE_ENGINE_IMAGE_GENERATION_MAX_ATTEMPTS=3 +SCENE_ENGINE_IMAGE_GENERATION_HEALTH_PATH="/health" +SCENE_ENGINE_IMAGE_GENERATION_BY_PROMPT_PATH="/generate_image_by_prompt" + SCENE_ENGINE_GEOMETRY_GENERATION_BASE_URL="http://host:port" SCENE_ENGINE_GEOMETRY_GENERATION_TIMEOUT_S=600 SCENE_ENGINE_GEOMETRY_GENERATION_MAX_ATTEMPTS=3 diff --git a/embodichain/gen_sim/scene_engine/clients/image_generation.py b/embodichain/gen_sim/scene_engine/clients/image_generation.py new file mode 100644 index 000000000..4286e26f8 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/clients/image_generation.py @@ -0,0 +1,172 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import requests + +from embodichain.gen_sim.scene_engine.configs.environment import ( + read_scene_engine_env_values, +) + + +class ImageGenerationClient: + """Manage the Image Generation Server connection.""" + + def __init__( + self, + *, + base_url: str, + timeout_s: int, + max_attempts: int, + health_path: str, + generate_image_by_prompt_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_image_by_prompt_path = generate_image_by_prompt_path + self._session = session or requests.Session() + + @classmethod + def from_dotenv(cls) -> "ImageGenerationClient": + """Create a client from its required ``gen_sim/.env`` settings.""" + return cls(**_load_dotenv_config()) + + def check_health(self) -> None: + last_error: requests.RequestException | RuntimeError | 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( + "Image 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( + "Image Generation Server health check failed after " + f"{self._max_attempts} attempts." + ) from last_error + + def close(self) -> None: + self._session.close() + + def generate_image_by_prompt( + self, + *, + prompt: str, + output_path: str | Path, + ) -> Path: + """Generate one PNG image from ``prompt`` and save it to ``output_path``.""" + prompt = prompt.strip() + if not prompt: + raise ValueError("Image generation prompt must not be empty.") + + resolved_output_path = Path(output_path).expanduser().resolve() + resolved_output_path.parent.mkdir(parents=True, exist_ok=True) + + last_error: Exception | None = None + for _ in range(self._max_attempts): + try: + response = self._session.post( + self._url(self._generate_image_by_prompt_path), + json={"prompt": prompt}, + timeout=self._timeout_s, + ) + response.raise_for_status() + content_type = response.headers.get("content-type", "").split(";")[0] + if content_type != "image/png": + raise RuntimeError( + "Image Generation Server response is not a PNG image." + ) + resolved_output_path.write_bytes(response.content) + return resolved_output_path + except (requests.RequestException, RuntimeError) as exc: + last_error = exc + + assert last_error is not None + raise RuntimeError( + "Image Generation Server request failed after " + f"{self._max_attempts} attempts." + ) from last_error + + def _url(self, path: str) -> str: + return f"{self._base_url}/{path.lstrip('/')}" + + +def _load_dotenv_config() -> dict[str, Any]: + values = read_scene_engine_env_values( + "SCENE_ENGINE_IMAGE_GENERATION_BASE_URL", + "SCENE_ENGINE_IMAGE_GENERATION_TIMEOUT_S", + "SCENE_ENGINE_IMAGE_GENERATION_MAX_ATTEMPTS", + "SCENE_ENGINE_IMAGE_GENERATION_HEALTH_PATH", + "SCENE_ENGINE_IMAGE_GENERATION_BY_PROMPT_PATH", + ) + try: + timeout_s = int(values["SCENE_ENGINE_IMAGE_GENERATION_TIMEOUT_S"]) + except (TypeError, ValueError) as exc: + raise ValueError( + "SCENE_ENGINE_IMAGE_GENERATION_TIMEOUT_S must be an integer." + ) from exc + if timeout_s < 1: + raise ValueError("SCENE_ENGINE_IMAGE_GENERATION_TIMEOUT_S must be at least 1.") + + try: + max_attempts = int(values["SCENE_ENGINE_IMAGE_GENERATION_MAX_ATTEMPTS"]) + except (TypeError, ValueError) as exc: + raise ValueError( + "SCENE_ENGINE_IMAGE_GENERATION_MAX_ATTEMPTS must be an integer." + ) from exc + if max_attempts < 1: + raise ValueError( + "SCENE_ENGINE_IMAGE_GENERATION_MAX_ATTEMPTS must be at least 1." + ) + + string_keys = ( + "SCENE_ENGINE_IMAGE_GENERATION_BASE_URL", + "SCENE_ENGINE_IMAGE_GENERATION_HEALTH_PATH", + "SCENE_ENGINE_IMAGE_GENERATION_BY_PROMPT_PATH", + ) + for key in string_keys: + if not values[key].strip(): + raise ValueError(f"Scene Engine .env key {key} must be a non-empty string.") + + return { + "base_url": values["SCENE_ENGINE_IMAGE_GENERATION_BASE_URL"].strip(), + "timeout_s": timeout_s, + "max_attempts": max_attempts, + "health_path": values["SCENE_ENGINE_IMAGE_GENERATION_HEALTH_PATH"].strip(), + "generate_image_by_prompt_path": values[ + "SCENE_ENGINE_IMAGE_GENERATION_BY_PROMPT_PATH" + ].strip(), + } diff --git a/tests/gen_sim/scene_engine/test_clients.py b/tests/gen_sim/scene_engine/test_clients.py index a2ccacb21..948db9c73 100644 --- a/tests/gen_sim/scene_engine/test_clients.py +++ b/tests/gen_sim/scene_engine/test_clients.py @@ -23,6 +23,7 @@ import pytest from embodichain.gen_sim.scene_engine.clients import geometry_generation +from embodichain.gen_sim.scene_engine.clients import image_generation from embodichain.gen_sim.scene_engine.clients import image_segmentation from embodichain.gen_sim.scene_engine.llms import load_config @@ -30,9 +31,16 @@ class _Response: """Minimal successful HTTP response used by client unit tests.""" - def __init__(self, payload: object, *, content: bytes = b"") -> None: + def __init__( + self, + payload: object, + *, + content: bytes = b"", + headers: dict[str, str] | None = None, + ) -> None: self._payload = payload self.content = content + self.headers = headers or {} def raise_for_status(self) -> None: return None @@ -81,6 +89,13 @@ def test_clients_load_their_required_dotenv_values( "SCENE_ENGINE_IMAGE_SEGMENTATION_HEALTH_PATH": "/health", "SCENE_ENGINE_IMAGE_SEGMENTATION_BY_PROMPT_PATH": "/segment_by_prompt", } + image_generation_values = { + "SCENE_ENGINE_IMAGE_GENERATION_BASE_URL": "http://image-generation/", + "SCENE_ENGINE_IMAGE_GENERATION_TIMEOUT_S": "120", + "SCENE_ENGINE_IMAGE_GENERATION_MAX_ATTEMPTS": "2", + "SCENE_ENGINE_IMAGE_GENERATION_HEALTH_PATH": "/health", + "SCENE_ENGINE_IMAGE_GENERATION_BY_PROMPT_PATH": "/generate_image_by_prompt", + } llm_values = { "OPENAI_API_KEY": "test-key", "OPENAI_MODEL": "test-model", @@ -96,18 +111,29 @@ def test_clients_load_their_required_dotenv_values( "read_scene_engine_env_values", lambda *_: segmentation_values, ) + monkeypatch.setattr( + image_generation, + "read_scene_engine_env_values", + lambda *_: image_generation_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() + image_generation_client = image_generation.ImageGenerationClient.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_by_prompt_path == "/segment_by_prompt" + assert image_generation_client._base_url == "http://image-generation" + assert ( + image_generation_client._generate_image_by_prompt_path + == "/generate_image_by_prompt" + ) assert llm_client_config.default_query == {"api-version": "1"} assert llm_client_config.base_url == "http://llm/v1" @@ -149,12 +175,90 @@ def test_service_health_checks_use_the_configured_health_path() -> None: segment_by_prompt_path="/segment_by_prompt", session=segmentation_session, ) + image_generation_session = _Session(get_payload={"ok": True}) + image_generation_client = image_generation.ImageGenerationClient( + base_url="http://image-generation", + timeout_s=120, + max_attempts=1, + health_path="/health", + generate_image_by_prompt_path="/generate_image_by_prompt", + session=image_generation_session, + ) geometry_client.check_health() segmentation_client.check_health() + image_generation_client.check_health() assert geometry_session.get_calls == [("http://geometry/health", 10)] assert segmentation_session.get_calls == [("http://segment/health", 30)] + assert image_generation_session.get_calls == [ + ("http://image-generation/health", 10) + ] + + +def test_image_generation_client_posts_prompt_and_writes_png( + tmp_path: Path, +) -> None: + png_bytes = b"\x89PNG\r\n\x1a\nimage" + + class ImageGenerationSession(_Session): + def post(self, url: str, **kwargs: object) -> _Response: + self.post_call = {"url": url, **kwargs} + return _Response( + {}, + content=png_bytes, + headers={"content-type": "image/png"}, + ) + + session = ImageGenerationSession(get_payload={"ok": True}) + client = image_generation.ImageGenerationClient( + base_url="http://image-generation", + timeout_s=120, + max_attempts=1, + health_path="/health", + generate_image_by_prompt_path="/generate_image_by_prompt", + session=session, + ) + + output_path = client.generate_image_by_prompt( + prompt="a red mug on a wooden table", + output_path=tmp_path / "generated.png", + ) + + assert output_path.read_bytes() == png_bytes + assert session.post_call is not None + assert session.post_call["url"] == ( + "http://image-generation/generate_image_by_prompt" + ) + assert session.post_call["json"] == {"prompt": "a red mug on a wooden table"} + + +def test_image_generation_client_rejects_non_png_response(tmp_path: Path) -> None: + class ImageGenerationSession(_Session): + def post(self, url: str, **kwargs: object) -> _Response: + self.post_call = {"url": url, **kwargs} + return _Response( + {"ok": False, "error": "failed"}, + content=b'{"ok": false}', + headers={"content-type": "application/json"}, + ) + + session = ImageGenerationSession(get_payload={"ok": True}) + client = image_generation.ImageGenerationClient( + base_url="http://image-generation", + timeout_s=120, + max_attempts=1, + health_path="/health", + generate_image_by_prompt_path="/generate_image_by_prompt", + session=session, + ) + + with pytest.raises(RuntimeError, match="request failed after 1 attempts") as exc: + client.generate_image_by_prompt( + prompt="a red mug on a wooden table", + output_path=tmp_path / "generated.png", + ) + assert "response is not a PNG image" in str(exc.value.__cause__) def test_segmentation_client_posts_prompt_and_returns_rle_masks(tmp_path: Path) -> None: From 7af223fc4e6f9d411567e025002c7646d3987a1e Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Tue, 11 Aug 2026 20:21:29 +0800 Subject: [PATCH 03/29] Reformat the code, add scene edit basic framework, notice that the docs and tests may not be updated --- embodichain/gen_sim/scene_engine/cli/start.py | 26 +- .../gen_sim/scene_engine/pipeline/edit.py | 89 +++++++ .../scene_engine/pipeline/editing/__init__.py | 19 ++ .../editing/scene_edit_understanding.py | 43 ++++ .../gen_sim/scene_engine/pipeline/generate.py | 4 +- .../pipeline/generation/__init__.py | 19 ++ .../{ => generation}/scene_generation.py | 0 .../{ => generation}/scene_understanding.py | 0 .../pipeline/utils/scene_importer.py | 233 ++++++++++++++++++ tests/gen_sim/scene_engine/test_scene_edit.py | 133 ++++++++++ .../scene_engine/test_scene_understanding.py | 2 +- 11 files changed, 564 insertions(+), 4 deletions(-) create mode 100644 embodichain/gen_sim/scene_engine/pipeline/edit.py create mode 100644 embodichain/gen_sim/scene_engine/pipeline/editing/__init__.py create mode 100644 embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py create mode 100644 embodichain/gen_sim/scene_engine/pipeline/generation/__init__.py rename embodichain/gen_sim/scene_engine/pipeline/{ => generation}/scene_generation.py (100%) rename embodichain/gen_sim/scene_engine/pipeline/{ => generation}/scene_understanding.py (100%) create mode 100644 embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py create mode 100644 tests/gen_sim/scene_engine/test_scene_edit.py diff --git a/embodichain/gen_sim/scene_engine/cli/start.py b/embodichain/gen_sim/scene_engine/cli/start.py index 59454b09f..edced21c4 100644 --- a/embodichain/gen_sim/scene_engine/cli/start.py +++ b/embodichain/gen_sim/scene_engine/cli/start.py @@ -21,6 +21,7 @@ from pathlib import Path from embodichain.gen_sim.scene_engine.pipeline.generate import generate_scene_from_image +from embodichain.gen_sim.scene_engine.pipeline.edit import edit_scene _SUPPORTED_IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png"} @@ -28,6 +29,8 @@ def cli_scene_engine( image: str | Path, output_root: str | Path, + *, + edit_prompt: str | None = None, ) -> None: """Generate one scene using the required ``gen_sim/.env`` settings.""" resolved_image_path = Path(image).expanduser().resolve() @@ -41,12 +44,27 @@ def cli_scene_engine( ) resolved_output_root = Path(output_root).expanduser().resolve() + # If this scene needs editing. + if edit_prompt is not None: + edit_prompt = edit_prompt.strip() + if not edit_prompt: + raise ValueError("Edit prompt must not be empty.") + if not resolved_output_root.is_dir() or not any(resolved_output_root.iterdir()): + raise ValueError( + "Output root must exist and contain files when edit_prompt is provided." + ) + resolved_output_root.mkdir(parents=True, exist_ok=True) generate_scene_from_image( image_path=resolved_image_path, output_root=resolved_output_root, ) + if edit_prompt is not None: + edit_scene( + output_root=resolved_output_root, + edit_prompt=edit_prompt, + ) print("Successfully completed!") @@ -68,9 +86,15 @@ def main(argv: Sequence[str] | None = None) -> None: required=True, help="Path to the output directory", ) + parser.add_argument( + "--edit_prompt", + type=str, + default=None, + help="Optional text instruction for editing an existing output root", + ) args = parser.parse_args(argv) - cli_scene_engine(args.image, args.output_root) + cli_scene_engine(args.image, args.output_root, edit_prompt=args.edit_prompt) if __name__ == "__main__": diff --git a/embodichain/gen_sim/scene_engine/pipeline/edit.py b/embodichain/gen_sim/scene_engine/pipeline/edit.py new file mode 100644 index 000000000..c055decf2 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/edit.py @@ -0,0 +1,89 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF 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.llms.openai_compatible_client import ( + OpenAICompatibleVLM, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_importer import ( + SceneExportImporter, +) +from embodichain.gen_sim.scene_engine.pipeline.editing.scene_edit_understanding import ( + understand_scene_edit, +) +from embodichain.utils.logger import log_info + + +def edit_scene( + *, + output_root: str | Path, + edit_prompt: str, +) -> None: + """Apply one text edit instruction to an existing Scene Engine output.""" + + # Initialize the VLM client that will interpret the edit instruction. + vlm_client = OpenAICompatibleVLM.from_dotenv() + scene_importer = SceneExportImporter(output_root=output_root) + # Validate scene_export, write scene.json, and return Scene; failures raise before editing. + scene = scene_importer.import_scene() + + # 1. Edit Understanding + log_info("Starting Edit Understanding") + edit_plan = understand_scene_edit( + scene=scene, + edit_prompt=edit_prompt, + output_root=output_root, # Has already been resolved. + vlm_client=vlm_client, + ) + log_info("Completed Edit Understanding") + + # 2. Scene Graph Update(or Initialization) + log_info("Starting Scene Graph Update") + # updated_scene_graph = update_scene_graph( + # scene=scene, + # edit_plan=edit_plan, + # output_root=output_root, + # ) + log_info("Completed Scene Graph Update") + + # 3. Prepare Objects. + log_info("Preparing Objects if necessary") + # scene = prepare_objects( + # scene=scene, + # output_root=output_root, + # ) + log_info("Completed Preparing Objects") + + # 4. Layout Editing + log_info("Starting Layout Editing") + # scene = edit_layout( + # scene=scene, + # edit_plan=edit_plan, + # scene_graph=updated_scene_graph, + # output_root=output_root, + # ) + log_info("Completed Layout Editing") + + # 5. Scene Export + # Re export the scene to the same output format, + # and delete some temporary files or folders. + log_info("Starting Scene Export") + log_info("Completed Scene Export") + + raise NotImplementedError("Scene editing is not implemented yet.") diff --git a/embodichain/gen_sim/scene_engine/pipeline/editing/__init__.py b/embodichain/gen_sim/scene_engine/pipeline/editing/__init__.py new file mode 100644 index 000000000..015c41510 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/editing/__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/editing/scene_edit_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py new file mode 100644 index 000000000..c5db9531b --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py @@ -0,0 +1,43 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF 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, +) + + +def understand_scene_edit( + *, + scene: Scene, + edit_prompt: str, + output_root: str | Path, + vlm_client: OpenAICompatibleVLM, +) -> dict[str, object]: + """Understand one text edit instruction for an existing scene.""" + edit_prompt = edit_prompt.strip() + if not edit_prompt: + raise ValueError("Edit prompt must not be empty.") + + return { + "edit_prompt": edit_prompt, + "operations": [], + } diff --git a/embodichain/gen_sim/scene_engine/pipeline/generate.py b/embodichain/gen_sim/scene_engine/pipeline/generate.py index 773a897bf..17e0a69dc 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generate.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generate.py @@ -26,12 +26,12 @@ GeometryGenerationClient, ) -from embodichain.gen_sim.scene_engine.pipeline.scene_understanding import ( +from embodichain.gen_sim.scene_engine.pipeline.generation.scene_understanding import ( understand_scene, ) from embodichain.utils.logger import log_info -from embodichain.gen_sim.scene_engine.pipeline.scene_generation import ( +from embodichain.gen_sim.scene_engine.pipeline.generation.scene_generation import ( generate_scene_and_refine, ) from embodichain.gen_sim.scene_engine.pipeline.utils.scene_exporter import SceneExporter diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/__init__.py b/embodichain/gen_sim/scene_engine/pipeline/generation/__init__.py new file mode 100644 index 000000000..015c41510 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/__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/scene_generation.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py similarity index 100% rename from embodichain/gen_sim/scene_engine/pipeline/scene_generation.py rename to embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py diff --git a/embodichain/gen_sim/scene_engine/pipeline/scene_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py similarity index 100% rename from embodichain/gen_sim/scene_engine/pipeline/scene_understanding.py rename to embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py new file mode 100644 index 000000000..716ce036d --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py @@ -0,0 +1,233 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF 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 numpy as np +from scipy.spatial.transform import Rotation + +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_object import ( + ObjectPhysics, + SceneObject, +) +from embodichain.utils.logger import log_info + +_Y_UP_TO_Z_UP_ROTATION = np.array( + [ + [1.0, 0.0, 0.0], + [0.0, 0.0, -1.0], + [0.0, 1.0, 0.0], + ], + dtype=float, +) +_Z_UP_TO_Y_UP_ROTATION = _Y_UP_TO_Z_UP_ROTATION.T + + +class SceneExportImporter: + """Import an editable ``Scene`` from an exported Scene Engine directory.""" + + def __init__( + self, + *, + output_root: str | Path, + ) -> None: + self.output_root = Path(output_root).expanduser().resolve() + self.scene_export_root = self.output_root / "scene_export" + self.mesh_assets_root = self.scene_export_root / "mesh_assets" + self.scene_config_path = self.scene_export_root / "scene_config.json" + self.scene_json_path = self.scene_export_root / "scene.json" + + def import_scene(self) -> Scene: + """Validate the scene export, write ``scene.json``, and return a ``Scene``.""" + # Editing only runs on an existing Scene Engine output directory. + if not self.output_root.is_dir() or not any(self.output_root.iterdir()): + raise ValueError( + "Output root must exist and contain files when edit_prompt is provided." + ) + + # The editor consumes the portable scene export and its copied GLB assets. + if not self.scene_export_root.is_dir(): + raise FileNotFoundError( + f"Scene export directory not found: {self.scene_export_root}" + ) + if not self.mesh_assets_root.is_dir(): + raise FileNotFoundError( + f"Scene mesh assets directory not found: {self.mesh_assets_root}" + ) + if not self.scene_config_path.is_file(): + raise FileNotFoundError(f"Scene config not found: {self.scene_config_path}") + + try: + scene_config = json.loads( + self.scene_config_path.read_text(encoding="utf-8") + ) + except json.JSONDecodeError as exc: + raise ValueError( + f"Scene config is not valid JSON: {self.scene_config_path}" + ) from exc + if not isinstance(scene_config, dict): + raise ValueError("Scene config must be a JSON object.") + + scene = self._scene_from_config(scene_config) + if self.scene_json_path.exists(): + self.scene_json_path.unlink() + self.scene_json_path.write_text( + json.dumps(scene.to_dict(), indent=2, ensure_ascii=False) + "\n", + encoding="utf-8", + ) + log_info(f"Imported scene JSON: {self.scene_json_path}") + return scene + + def _scene_from_config(self, scene_config: dict[str, Any]) -> Scene: + """Build a y-up ``Scene`` from the z-up scene-export config.""" + # The table is the required support object for scene-edit operations. + background = scene_config.get("background", []) + if not isinstance(background, list): + raise ValueError("Scene config background must be a list.") + table_entry = next( + ( + scene_object + for scene_object in background + if isinstance(scene_object, dict) and scene_object.get("uid") == "table" + ), + None, + ) + if table_entry is None: + raise ValueError("Scene config background must contain a table entry.") + + rigid_object_entries = scene_config.get("rigid_object", []) + if not isinstance(rigid_object_entries, list): + raise ValueError("Scene config rigid_object must be a list.") + + return Scene( + objects=[ + self._scene_object_from_export_entry(table_entry, kind="table"), + *[ + self._scene_object_from_export_entry(entry, kind="asset") + for entry in rigid_object_entries + ], + ] + ) + + def _scene_object_from_export_entry( + self, + entry: object, + *, + kind: str, + ) -> SceneObject: + """Convert one z-up scene-export entry back to a y-up ``SceneObject``.""" + if not isinstance(entry, dict): + raise ValueError("Scene config entries must be objects.") + uid = entry.get("uid") + if not isinstance(uid, str) or not uid: + raise ValueError("Scene config entries must contain a valid uid.") + + glb_path = self._resolve_export_glb_path(entry, uid=uid) + pos_z_up = self._vector3( + entry.get("init_pos", [0.0, 0.0, 0.0]), + field_name=f"{uid}.init_pos", + ) + rot_z_up = self._vector3( + entry.get("init_rot", [0.0, 0.0, 0.0]), + field_name=f"{uid}.init_rot", + ) + scale = self._vector3( + entry.get("body_scale", [1.0, 1.0, 1.0]), + field_name=f"{uid}.body_scale", + ) + + pos_y_up = _Z_UP_TO_Y_UP_ROTATION @ np.asarray(pos_z_up, dtype=float) + rotation_z_up = Rotation.from_euler("XYZ", rot_z_up, degrees=True).as_matrix() + rotation_y_up = ( + _Z_UP_TO_Y_UP_ROTATION @ rotation_z_up @ _Z_UP_TO_Y_UP_ROTATION.T + ) + rot_y_up = Rotation.from_matrix(rotation_y_up).as_euler("xyz", degrees=True) + + return SceneObject( + id=uid, + kind=kind, # type: ignore[arg-type] + category=uid, + name=uid, + description=str(entry.get("description") or uid), + simready_glb_path=str(glb_path), + rot=rot_y_up.tolist(), + pos=pos_y_up.tolist(), + scale=scale, + physics=ObjectPhysics( + body_type=str(entry.get("body_type", "dynamic")), # type: ignore[arg-type] + attrs=self._physics_attrs(entry.get("attrs", {"mass": 1.0})), + max_convex_hull_num=max(1, int(entry.get("max_convex_hull_num", 32))), + ), + ) + + def _resolve_export_glb_path( + self, + entry: dict[str, Any], + *, + uid: str, + ) -> Path: + """Validate one exported mesh reference and return its absolute GLB path.""" + shape = entry.get("shape") + if not isinstance(shape, dict) or not isinstance(shape.get("fpath"), str): + raise ValueError(f"Scene object {uid!r} must contain shape.fpath.") + fpath = Path(shape["fpath"]) + if fpath.is_absolute(): + raise ValueError(f"Scene object {uid!r} shape.fpath must be relative.") + if fpath.suffix.lower() != ".glb": + raise ValueError(f"Scene object {uid!r} shape.fpath must point to a GLB.") + glb_path = (self.scene_export_root / fpath).resolve() + if self.scene_export_root.resolve() not in glb_path.parents: + raise ValueError( + f"Scene object {uid!r} shape.fpath must stay within " + f"{self.scene_export_root.resolve()}." + ) + if not glb_path.is_file(): + raise FileNotFoundError(f"Scene object {uid!r} GLB not found: {glb_path}") + return glb_path + + @staticmethod + def _vector3(value: object, *, field_name: str) -> list[float]: + """Validate one length-3 numeric vector.""" + if not isinstance(value, list) or len(value) != 3: + raise ValueError( + f"Scene config field {field_name!r} must be a length-3 list." + ) + vector = [float(item) for item in value] + if not np.all(np.isfinite(vector)): + raise ValueError(f"Scene config field {field_name!r} must be finite.") + return vector + + @staticmethod + def _physics_attrs(value: object) -> dict[str, float | int]: + """Validate exported physics attributes.""" + if not isinstance(value, dict) or not value: + raise ValueError("Scene object attrs must be a non-empty object.") + attrs: dict[str, float | int] = {} + for key, item in value.items(): + if not isinstance(key, str) or not isinstance(item, (float, int)): + raise ValueError("Scene object attrs must map strings to numbers.") + attrs[key] = item + return attrs + + +def import_scene_from_output_root(output_root: str | Path) -> Scene: + """Import an editable ``Scene`` from ``scene_export/scene_config.json``.""" + return SceneExportImporter(output_root=output_root).import_scene() diff --git a/tests/gen_sim/scene_engine/test_scene_edit.py b/tests/gen_sim/scene_engine/test_scene_edit.py new file mode 100644 index 000000000..aef5a1462 --- /dev/null +++ b/tests/gen_sim/scene_engine/test_scene_edit.py @@ -0,0 +1,133 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF 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.pipeline.utils.scene_importer import ( + import_scene_from_output_root, +) + + +def _write_scene_export( + output_root: Path, + *, + include_table: bool = True, + include_asset_mesh: bool = True, +) -> None: + scene_export_root = output_root / "scene_export" + table_mesh_path = scene_export_root / "mesh_assets" / "table" / "table.glb" + asset_mesh_path = scene_export_root / "mesh_assets" / "cup" / "cup.glb" + table_mesh_path.parent.mkdir(parents=True) + asset_mesh_path.parent.mkdir(parents=True) + table_mesh_path.write_bytes(b"glTF-table") + if include_asset_mesh: + asset_mesh_path.write_bytes(b"glTF-cup") + + background = [] + if include_table: + background.append( + { + "uid": "table", + "shape": { + "shape_type": "Mesh", + "fpath": "mesh_assets/table/table.glb", + }, + "attrs": {"mass": 1.0}, + "body_type": "kinematic", + "init_pos": [0.0, 0.0, 0.0], + "init_rot": [0.0, 0.0, 0.0], + "body_scale": [1.0, 1.0, 1.0], + "max_convex_hull_num": 16, + } + ) + scene_config = { + "background": background, + "rigid_object": [ + { + "uid": "cup", + "shape": { + "shape_type": "Mesh", + "fpath": "mesh_assets/cup/cup.glb", + }, + "attrs": {"mass": 1.0}, + "body_type": "dynamic", + "init_pos": [1.0, -3.0, 2.0], + "init_rot": [0.0, 0.0, 0.0], + "body_scale": [1.0, 2.0, 3.0], + "max_convex_hull_num": 32, + } + ], + } + (scene_export_root / "scene_config.json").write_text( + json.dumps(scene_config), + encoding="utf-8", + ) + + +def test_import_scene_from_output_root_writes_y_up_scene_json(tmp_path: Path) -> None: + _write_scene_export(tmp_path) + (tmp_path / "scene_export" / "scene.json").write_text( + '{"old": true}', + encoding="utf-8", + ) + + scene = import_scene_from_output_root(tmp_path) + scene_json = json.loads( + (tmp_path / "scene_export" / "scene.json").read_text(encoding="utf-8") + ) + + assert scene.table is not None + assert scene.table.id == "table" + assert scene.assets[0].id == "cup" + assert scene.assets[0].pos == [1.0, 2.0, 3.0] + assert scene.assets[0].scale == [1.0, 2.0, 3.0] + assert scene.assets[0].simready_glb_path == str( + (tmp_path / "scene_export" / "mesh_assets" / "cup" / "cup.glb").resolve() + ) + assert scene_json["objects"][1]["id"] == "cup" + assert scene_json["objects"][1]["pos"] == [1.0, 2.0, 3.0] + + +def test_check_scene_export_for_edit_requires_export_directories( + tmp_path: Path, +) -> None: + with pytest.raises(ValueError, match="Output root"): + import_scene_from_output_root(tmp_path) + + (tmp_path / "scene_export").mkdir() + with pytest.raises(FileNotFoundError, match="mesh assets"): + import_scene_from_output_root(tmp_path) + + +def test_check_scene_export_for_edit_requires_table(tmp_path: Path) -> None: + _write_scene_export(tmp_path, include_table=False) + + with pytest.raises(ValueError, match="table"): + import_scene_from_output_root(tmp_path) + + +def test_check_scene_export_for_edit_requires_rigid_object_glb( + tmp_path: Path, +) -> None: + _write_scene_export(tmp_path, include_asset_mesh=False) + + with pytest.raises(FileNotFoundError, match="cup"): + import_scene_from_output_root(tmp_path) diff --git a/tests/gen_sim/scene_engine/test_scene_understanding.py b/tests/gen_sim/scene_engine/test_scene_understanding.py index 1fdb51f7b..5d9a285c3 100644 --- a/tests/gen_sim/scene_engine/test_scene_understanding.py +++ b/tests/gen_sim/scene_engine/test_scene_understanding.py @@ -23,7 +23,7 @@ import pytest from embodichain.gen_sim.scene_engine.core.scene import Scene -from embodichain.gen_sim.scene_engine.pipeline import scene_understanding +from embodichain.gen_sim.scene_engine.pipeline.generation import scene_understanding def _response(*, asset_name: str = "cup") -> str: From 23a6557934fbde051377db9a156d19677da2105c Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Wed, 12 Aug 2026 12:25:19 +0800 Subject: [PATCH 04/29] Basic design of scene graph --- .../gen_sim/scene_engine/core/scene_graph.py | 352 ++++++++++++++++++ .../gen_sim/scene_engine/test_scene_graph.py | 309 +++++++++++++++ 2 files changed, 661 insertions(+) create mode 100644 embodichain/gen_sim/scene_engine/core/scene_graph.py create mode 100644 tests/gen_sim/scene_engine/test_scene_graph.py diff --git a/embodichain/gen_sim/scene_engine/core/scene_graph.py b/embodichain/gen_sim/scene_engine/core/scene_graph.py new file mode 100644 index 000000000..806c5a0ee --- /dev/null +++ b/embodichain/gen_sim/scene_engine/core/scene_graph.py @@ -0,0 +1,352 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF 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 typing import Literal + +TABLE_OBJECT_ID = "table" + +# 9-grid table regions, treat the table as a 3x3 grid. +TableRegion = Literal[ + "left_back", + "back_center", + "right_back", + "left_center", + "center", + "right_center", + "left_front", + "front_center", + "right_front", +] + +# A on B, then B is the parent node of A. +SupportRelationType = Literal["on"] + +# A PlanarRelation with B, then A and B must have the same parent node. +PlanarRelationType = Literal["left_of", "right_of", "in_front_of", "behind"] +SceneConstraintType = SupportRelationType | PlanarRelationType + + +@dataclass +class SceneGraphNode: + """One object node in the edit-time scene hierarchy.""" + + object_id: str + parent_id: str | None + parent_relation: SupportRelationType | None = None + table_region: TableRegion | None = None + + def __post_init__(self) -> None: + """Validate local node fields before graph-level checks.""" + if not self.object_id: + raise ValueError("object_id must be non-empty.") + # If the node is the table. + if self.object_id == TABLE_OBJECT_ID: + if self.parent_id is not None: + raise ValueError("table must not have a parent.") + if self.parent_relation is not None: + raise ValueError("table must not have a parent relation.") + # If the node is not the table. + elif self.parent_id is None: + raise ValueError("non-table nodes must have a parent.") + elif self.parent_relation not in {None, "on"}: + raise ValueError("non-table nodes must be on their parent.") + + def to_dict(self) -> dict[str, object]: + """Serialize this node for scene graph debugging artifacts.""" + return { + "object_id": self.object_id, + "parent_id": self.parent_id, + "parent_relation": self.parent_relation, + "table_region": self.table_region, + } + + +@dataclass +class SceneGraphRelation: + """One edit-time spatial relation between two non-table objects.""" + + source_id: str + relation: PlanarRelationType + target_id: str + + def __post_init__(self) -> None: + """Validate local relation fields before graph-level checks.""" + if not self.source_id or not self.target_id: + raise ValueError("relation endpoints must be non-empty.") + if self.source_id == self.target_id: + raise ValueError("relation endpoints must be different.") + + def to_dict(self) -> dict[str, object]: + """Serialize this planar relation for scene graph debugging artifacts.""" + return { + "source_id": self.source_id, + "relation": self.relation, + "target_id": self.target_id, + } + + +@dataclass +class SceneGraph: + """Layered support graph plus planar relations for scene editing.""" + + nodes: list[SceneGraphNode] = field(default_factory=list) + relations: list[SceneGraphRelation] = field(default_factory=list) + validate_on_refresh: bool = True # Validate after each automatic refresh. + + def __post_init__(self) -> None: + """Normalize new graphs so downstream stages see canonical constraints.""" + self.refresh() + + def refresh(self) -> None: + """Normalize the graph and optionally validate semantic constraints.""" + # First normalize then validate (if applicable). + self.normalize() + if self.validate_on_refresh: + self.validate() + + def node_by_id(self) -> dict[str, SceneGraphNode]: + """Return nodes keyed by object id, raising on duplicate ids.""" + nodes_by_id: dict[str, SceneGraphNode] = {} + for node in self.nodes: + if node.object_id in nodes_by_id: + raise ValueError(f"Duplicate scene graph node: {node.object_id}") + nodes_by_id[node.object_id] = node + return nodes_by_id + + def normalize(self) -> None: + """Materialize inverse planar relations and remove duplicates.""" + self._materialize_inverse_planar_relations() + self._deduplicate_relations() + + def layer_by_id(self) -> dict[str, int]: + """Return the layer depth of each node inferred from parent links.""" + # Build fast lookup tables before walking the table-rooted tree. + nodes_by_id = self.node_by_id() + children_by_parent = self._children_by_parent() + layers: dict[str, int] = {} + visiting: set[str] = set() + visited: set[str] = set() + + def visit(node_id: str, layer: int) -> None: + # A node already on the recursion path means the parent chain loops. + if node_id in visiting: + raise ValueError(f"Parent cycle detected at node: {node_id}") + if node_id in visited: + return + # Missing parent nodes cannot contribute a valid table-rooted layer. + if node_id not in nodes_by_id: + raise ValueError(f"Parent node does not exist: {node_id}") + + visiting.add(node_id) + layers[node_id] = layer + # Children are exactly one support level above their parent. + for child in children_by_parent.get(node_id, []): + visit(child.object_id, layer + 1) + visiting.remove(node_id) + visited.add(node_id) + + if TABLE_OBJECT_ID not in nodes_by_id: + raise ValueError("Scene graph must contain a table node.") + visit(TABLE_OBJECT_ID, 0) + return layers + + def derive_constraints(self) -> list[dict[str, str]]: + """Return support constraints plus materialized planar relations.""" + self.refresh() + constraints: list[dict[str, str]] = [] + for node in self.nodes: + if node.parent_id is None: + continue + # Parent links become direct support constraints. + constraints.append( + self._constraint_dict( + source_id=node.object_id, + relation=node.parent_relation, + target_id=node.parent_id, + ), + ) + for relation in self.relations: + # Inverse planar relations are already stored during normalization. + constraints.append( + self._constraint_dict( + source_id=relation.source_id, + relation=relation.relation, + target_id=relation.target_id, + ), + ) + return self._deduplicate_constraints(constraints) + + def validate(self) -> None: + """Validate hierarchy, table regions, and planar relation constraints.""" + # id -> Node mapping. + nodes_by_id = self.node_by_id() + # Table node must exist. + if TABLE_OBJECT_ID not in nodes_by_id: + raise ValueError("Scene graph must contain a table node.") + + # Validate the table-rooted support tree before checking sibling relations. + for node in self.nodes: + if node.object_id == TABLE_OBJECT_ID: + if node.table_region is not None: + raise ValueError("table must not have a table_region.") + continue + parent = nodes_by_id.get(node.parent_id) + # Parent must exist, except for the table. (root node) + if parent is None: + raise ValueError(f"Parent node does not exist: {node.parent_id}") + if node.parent_relation is None: + raise ValueError( + f"Node {node.object_id} must define its parent relation." + ) + # Table regions are only valid for objects directly on the table. + if node.table_region is not None and node.parent_id != TABLE_OBJECT_ID: + raise ValueError("table_region is only valid for objects on the table.") + # Get id -> layer mapping. + layers = self.layer_by_id() + if len(layers) != len(nodes_by_id): + raise ValueError("All scene graph nodes must be reachable from the table.") + # Validate planar relations between nodes with the same parent. + for relation in self.relations: + source = nodes_by_id.get(relation.source_id) + target = nodes_by_id.get(relation.target_id) + if source is None or target is None: + raise ValueError("Planar relation endpoint does not exist.") + if source.parent_id != target.parent_id: + raise ValueError("Planar relation endpoints must share one parent.") + if source.parent_relation != "on" or target.parent_relation != "on": + raise ValueError("Planar relation endpoints must be on their parent.") + # Validate and imply planar relation. + self._validate_planar_relation_conflicts() + + def to_dict(self) -> dict[str, object]: + """Serialize the normalized graph state.""" + self.refresh() + return { + "nodes": [node.to_dict() for node in self.nodes], + "relations": [relation.to_dict() for relation in self.relations], + } + + def _children_by_parent(self) -> dict[str, list[SceneGraphNode]]: + children_by_parent: dict[str, list[SceneGraphNode]] = {} + for node in self.nodes: + if node.parent_id is not None: + children_by_parent.setdefault(node.parent_id, []).append(node) + return children_by_parent + + def _deduplicate_relations(self) -> None: + """Remove duplicate planar relations while preserving the first occurrence.""" + deduplicated: list[SceneGraphRelation] = [] + seen: set[tuple[str, PlanarRelationType, str]] = set() + for relation in self.relations: + key = (relation.source_id, relation.relation, relation.target_id) + # Only identical triples are duplicates; inverse relations are both retained. + if key in seen: + continue + seen.add(key) + deduplicated.append(relation) + self.relations = deduplicated + + def _materialize_inverse_planar_relations(self) -> None: + """Add the inverse of every planar relation to the graph.""" + inverse_relations = [ + SceneGraphRelation( + source_id=relation.target_id, + relation=self._inverse_planar_relation(relation.relation), + target_id=relation.source_id, + ) + for relation in self.relations + ] + self.relations.extend(inverse_relations) + + def _deduplicate_constraints( + self, + constraints: list[dict[str, str]], + ) -> list[dict[str, str]]: + deduplicated: list[dict[str, str]] = [] + seen: set[tuple[str, SceneConstraintType, str]] = set() + for constraint in constraints: + key = ( + constraint["source_id"], + constraint["relation"], + constraint["target_id"], + ) + if key in seen: + continue + seen.add(key) + deduplicated.append(constraint) + return deduplicated + + def _constraint_dict( + self, + *, + source_id: str, + relation: SceneConstraintType, + target_id: str, + ) -> dict[str, str]: + return { + "source_id": source_id, + "relation": relation, + "target_id": target_id, + } + + def _validate_planar_relation_conflicts(self) -> None: + implied_relations: dict[tuple[str, str], PlanarRelationType] = {} + for relation in self.relations: + self._add_implied_planar_relation( + implied_relations, + source_id=relation.source_id, + relation=relation.relation, + target_id=relation.target_id, + ) + self._add_implied_planar_relation( + implied_relations, + source_id=relation.target_id, + relation=self._inverse_planar_relation(relation.relation), + target_id=relation.source_id, + ) + + def _add_implied_planar_relation( + self, + implied_relations: dict[tuple[str, str], PlanarRelationType], + *, + source_id: str, + relation: PlanarRelationType, + target_id: str, + ) -> None: + key = (source_id, target_id) + existing_relation = implied_relations.get(key) + if existing_relation is not None and existing_relation != relation: + raise ValueError( + f"Conflicting planar relations: {source_id} " + f"{existing_relation} and {relation} {target_id}" + ) + implied_relations[key] = relation + + @classmethod + def _inverse_planar_relation( + cls, + relation: PlanarRelationType, + ) -> PlanarRelationType: + if relation == "left_of": + return "right_of" + if relation == "right_of": + return "left_of" + if relation == "in_front_of": + return "behind" + return "in_front_of" diff --git a/tests/gen_sim/scene_engine/test_scene_graph.py b/tests/gen_sim/scene_engine/test_scene_graph.py new file mode 100644 index 000000000..970a3fc68 --- /dev/null +++ b/tests/gen_sim/scene_engine/test_scene_graph.py @@ -0,0 +1,309 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF 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 pytest + +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + SceneGraph, + SceneGraphNode, + SceneGraphRelation, +) + + +def test_scene_graph_accepts_layered_on_relations() -> None: + graph = SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="plate", + parent_id="table", + parent_relation="on", + table_region="center", + ), + SceneGraphNode( + object_id="cup", + parent_id="table", + parent_relation="on", + table_region="right_center", + ), + SceneGraphNode( + object_id="spoon", + parent_id="plate", + parent_relation="on", + ), + ], + relations=[ + SceneGraphRelation( + source_id="plate", + relation="left_of", + target_id="cup", + ), + ], + ) + + graph.validate() + assert graph.layer_by_id()["spoon"] == 2 + + +def test_scene_graph_rejects_planar_relations_without_common_parent() -> None: + with pytest.raises(ValueError, match="share one parent"): + SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="plate", + parent_id="table", + parent_relation="on", + ), + SceneGraphNode( + object_id="spoon", + parent_id="plate", + parent_relation="on", + ), + ], + relations=[ + SceneGraphRelation( + source_id="plate", + relation="right_of", + target_id="spoon", + ), + ], + ) + + +def test_scene_graph_rejects_conflicting_planar_relations() -> None: + with pytest.raises(ValueError, match="Conflicting planar relations"): + SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="plate", + parent_id="table", + parent_relation="on", + ), + SceneGraphNode( + object_id="cup", + parent_id="table", + parent_relation="on", + ), + ], + relations=[ + SceneGraphRelation( + source_id="plate", + relation="left_of", + target_id="cup", + ), + SceneGraphRelation( + source_id="cup", + relation="left_of", + target_id="plate", + ), + ], + ) + + +def test_scene_graph_requires_explicit_parent_relation() -> None: + with pytest.raises(ValueError, match="parent relation"): + SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="plate", + parent_id="table", + ), + ], + ) + + +def test_scene_graph_can_skip_validation_during_refresh() -> None: + graph = SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="plate", + parent_id="table", + ), + ], + validate_on_refresh=False, + ) + + with pytest.raises(ValueError, match="parent relation"): + graph.validate() + + +def test_scene_graph_rejects_unsupported_parent_relation() -> None: + with pytest.raises(ValueError, match="must be on their parent"): + SceneGraphNode( + object_id="orange", + parent_id="box", + parent_relation="inside", + ) + + +def test_scene_graph_derives_layers_from_parent_links() -> None: + graph = SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="plate", + parent_id="table", + parent_relation="on", + ), + SceneGraphNode( + object_id="spoon", + parent_id="plate", + parent_relation="on", + ), + ], + ) + + assert graph.layer_by_id() == { + "table": 0, + "plate": 1, + "spoon": 2, + } + + +def test_scene_graph_layer_by_id_requires_table_root() -> None: + graph = SceneGraph(nodes=[], validate_on_refresh=False) + + with pytest.raises(ValueError, match="table node"): + graph.layer_by_id() + + +def test_scene_graph_rejects_table_region_for_non_table_parent() -> None: + with pytest.raises(ValueError, match="only valid for objects on the table"): + SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="plate", + parent_id="table", + parent_relation="on", + ), + SceneGraphNode( + object_id="spoon", + parent_id="plate", + parent_relation="on", + table_region="center", + ), + ], + ) + + +def test_scene_graph_derives_support_and_inverse_planar_constraints() -> None: + graph = SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="plate", + parent_id="table", + parent_relation="on", + ), + SceneGraphNode( + object_id="cup", + parent_id="table", + parent_relation="on", + ), + ], + relations=[ + SceneGraphRelation( + source_id="plate", + relation="left_of", + target_id="cup", + ), + SceneGraphRelation( + source_id="plate", + relation="left_of", + target_id="cup", + ), + ], + ) + + constraints = graph.derive_constraints() + + assert constraints == [ + {"source_id": "plate", "relation": "on", "target_id": "table"}, + {"source_id": "cup", "relation": "on", "target_id": "table"}, + {"source_id": "plate", "relation": "left_of", "target_id": "cup"}, + {"source_id": "cup", "relation": "right_of", "target_id": "plate"}, + ] + + +def test_scene_graph_materializes_inverse_planar_relations() -> None: + graph = SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="plate", + parent_id="table", + parent_relation="on", + ), + SceneGraphNode( + object_id="cup", + parent_id="table", + parent_relation="on", + ), + ], + relations=[ + SceneGraphRelation( + source_id="plate", + relation="left_of", + target_id="cup", + ), + ], + ) + + assert [relation.to_dict() for relation in graph.relations] == [ + {"source_id": "plate", "relation": "left_of", "target_id": "cup"}, + {"source_id": "cup", "relation": "right_of", "target_id": "plate"}, + ] + + +def test_scene_graph_to_dict_serializes_graph_state() -> None: + graph = SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="plate", + parent_id="table", + parent_relation="on", + table_region="center", + ), + ], + ) + + graph_dict = graph.to_dict() + + assert graph_dict == { + "nodes": [ + { + "object_id": "table", + "parent_id": None, + "parent_relation": None, + "table_region": None, + }, + { + "object_id": "plate", + "parent_id": "table", + "parent_relation": "on", + "table_region": "center", + }, + ], + "relations": [], + } From c21e39900a8a516fe8096ee90b43d7ac9438253d Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Wed, 12 Aug 2026 15:41:40 +0800 Subject: [PATCH 05/29] i 1. modify cli/start logic 2. add xy position info in scene object data sturcture 3. add some test files 4. modify scene export and scene import, thus the scene graph will be export and import automatically --- embodichain/gen_sim/scene_engine/cli/start.py | 37 ++--- .../gen_sim/scene_engine/core/scene_object.py | 2 + .../gen_sim/scene_engine/pipeline/edit.py | 28 ++-- .../editing/scene_edit_understanding.py | 5 + .../gen_sim/scene_engine/pipeline/generate.py | 4 +- .../pipeline/generation/scene_generation.py | 26 ++- .../generation/scene_understanding.py | 37 ++++- .../pipeline/utils/scene_exporter.py | 14 ++ .../pipeline/utils/scene_importer.py | 148 +++++++++++++++++- .../test_scene_core_and_export.py | 66 +++++++- tests/gen_sim/scene_engine/test_scene_edit.py | 2 + .../scene_engine/test_scene_engine_config.py | 52 ++++++ .../scene_engine/test_scene_understanding.py | 44 ++++++ 13 files changed, 419 insertions(+), 46 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/cli/start.py b/embodichain/gen_sim/scene_engine/cli/start.py index edced21c4..6d46f32dc 100644 --- a/embodichain/gen_sim/scene_engine/cli/start.py +++ b/embodichain/gen_sim/scene_engine/cli/start.py @@ -27,12 +27,25 @@ def cli_scene_engine( - image: str | Path, + image: str | Path | None, output_root: str | Path, *, edit_prompt: str | None = None, ) -> None: - """Generate one scene using the required ``gen_sim/.env`` settings.""" + """Generate a scene from an image, edit an export, or do both in sequence.""" + resolved_output_root = Path(output_root).expanduser().resolve() + if edit_prompt is not None: + edit_prompt = edit_prompt.strip() + if not edit_prompt: + raise ValueError("Edit prompt must not be empty.") + + if image is None: + if edit_prompt is None: + raise ValueError("Provide --image, --edit_prompt, or both.") + edit_scene(output_root=resolved_output_root, edit_prompt=edit_prompt) + print("Successfully completed!") + return + resolved_image_path = Path(image).expanduser().resolve() if not resolved_image_path.exists(): raise FileNotFoundError(f"Image input not found: {resolved_image_path}") @@ -43,19 +56,7 @@ def cli_scene_engine( "Image input must have one of these extensions: .jpg, .jpeg, .png" ) - resolved_output_root = Path(output_root).expanduser().resolve() - # If this scene needs editing. - if edit_prompt is not None: - edit_prompt = edit_prompt.strip() - if not edit_prompt: - raise ValueError("Edit prompt must not be empty.") - if not resolved_output_root.is_dir() or not any(resolved_output_root.iterdir()): - raise ValueError( - "Output root must exist and contain files when edit_prompt is provided." - ) - resolved_output_root.mkdir(parents=True, exist_ok=True) - generate_scene_from_image( image_path=resolved_image_path, output_root=resolved_output_root, @@ -71,14 +72,14 @@ def cli_scene_engine( def main(argv: Sequence[str] | None = None) -> None: parser = argparse.ArgumentParser( prog="embodichain scene-engine", - description="Generate a Scene Engine export from one input image.", + description="Generate a Scene Engine export, edit one, or do both.", epilog="Service settings are read from embodichain/gen_sim/.env.", ) parser.add_argument( "--image", type=str, - required=True, - help="Path to the required input image file (.jpg, .jpeg, or .png)", + required=False, + help="Optional input image file (.jpg, .jpeg, or .png)", ) parser.add_argument( "--output_root", @@ -90,7 +91,7 @@ def main(argv: Sequence[str] | None = None) -> None: "--edit_prompt", type=str, default=None, - help="Optional text instruction for editing an existing output root", + help="Text instruction for editing an existing or newly generated output root", ) args = parser.parse_args(argv) diff --git a/embodichain/gen_sim/scene_engine/core/scene_object.py b/embodichain/gen_sim/scene_engine/core/scene_object.py index d9a837e87..0a8ed3aea 100644 --- a/embodichain/gen_sim/scene_engine/core/scene_object.py +++ b/embodichain/gen_sim/scene_engine/core/scene_object.py @@ -66,6 +66,7 @@ class SceneObject: 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. + center_xy: list[float] | None = None # Z-up table-frame XY AABB center. physics: ObjectPhysics | None = None # Assigned when SimReady processing succeeds. def to_dict(self) -> dict[str, object]: @@ -81,5 +82,6 @@ def to_dict(self) -> dict[str, object]: "rot": self.rot, "pos": self.pos, "scale": self.scale, + "center_xy": self.center_xy, "physics": self.physics.to_dict() if self.physics is not None else None, } diff --git a/embodichain/gen_sim/scene_engine/pipeline/edit.py b/embodichain/gen_sim/scene_engine/pipeline/edit.py index c055decf2..3b6c95e02 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/edit.py +++ b/embodichain/gen_sim/scene_engine/pipeline/edit.py @@ -16,6 +16,7 @@ from __future__ import annotations +import json from pathlib import Path from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( @@ -41,11 +42,19 @@ def edit_scene( vlm_client = OpenAICompatibleVLM.from_dotenv() scene_importer = SceneExportImporter(output_root=output_root) # Validate scene_export, write scene.json, and return Scene; failures raise before editing. - scene = scene_importer.import_scene() + scene, scene_graph = scene_importer.import_scene_and_graph() + print( + json.dumps( + {"scene": scene.to_dict(), "scene_graph": scene_graph.to_dict()}, + indent=2, + ensure_ascii=False, + ) + ) # 1. Edit Understanding + # And update or initialize the scene graph log_info("Starting Edit Understanding") - edit_plan = understand_scene_edit( + updated_scene_graph = understand_scene_edit( scene=scene, edit_prompt=edit_prompt, output_root=output_root, # Has already been resolved. @@ -53,16 +62,7 @@ def edit_scene( ) log_info("Completed Edit Understanding") - # 2. Scene Graph Update(or Initialization) - log_info("Starting Scene Graph Update") - # updated_scene_graph = update_scene_graph( - # scene=scene, - # edit_plan=edit_plan, - # output_root=output_root, - # ) - log_info("Completed Scene Graph Update") - - # 3. Prepare Objects. + # 2. Prepare Objects. log_info("Preparing Objects if necessary") # scene = prepare_objects( # scene=scene, @@ -70,7 +70,7 @@ def edit_scene( # ) log_info("Completed Preparing Objects") - # 4. Layout Editing + # 3. Layout Editing log_info("Starting Layout Editing") # scene = edit_layout( # scene=scene, @@ -80,7 +80,7 @@ def edit_scene( # ) log_info("Completed Layout Editing") - # 5. Scene Export + # 4. Scene Export # Re export the scene to the same output format, # and delete some temporary files or folders. log_info("Starting Scene Export") diff --git a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py index c5db9531b..48dda64bf 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py @@ -20,6 +20,11 @@ from pathlib import Path from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + SceneGraph, + SceneGraphNode, + TABLE_OBJECT_ID, +) from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( OpenAICompatibleVLM, ) diff --git a/embodichain/gen_sim/scene_engine/pipeline/generate.py b/embodichain/gen_sim/scene_engine/pipeline/generate.py index 17e0a69dc..32aa6550f 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generate.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generate.py @@ -51,7 +51,7 @@ def generate_scene_from_image( # 1. Scene Understanding log_info("Starting Scene Understanding") - scene = understand_scene( + scene, scene_graph = understand_scene( scene=scene, image_path=image_path, output_root=resolved_output_root, @@ -69,6 +69,7 @@ def generate_scene_from_image( image_path=image_path, output_root=resolved_output_root, scene=scene, + scene_graph=scene_graph, geometry_generation_client=geometry_generation_client, ) finally: @@ -79,6 +80,7 @@ def generate_scene_from_image( log_info("Starting Scene Export") scene_exporter = SceneExporter( scene=scene, + scene_graph=scene_graph, output_root=resolved_output_root, ) scene_exporter.export() diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py index d14ce364f..1cfe3f592 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py @@ -28,6 +28,7 @@ GeometryGenerationClient, ) from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_graph import SceneGraph 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, @@ -62,11 +63,14 @@ def generate_scene_and_refine( image_path: str | Path, output_root: str | Path, scene: Scene, + scene_graph: SceneGraph, *, geometry_generation_client: GeometryGenerationClient, ) -> Scene: resolved_image_path = _validate_image_path(image_path) + # Validate the scene graph before layout refinement consumes it. + scene_graph.validate() # Create stage output directory. stage_output_root = Path(output_root).expanduser().resolve() / "scene_generation" if stage_output_root.exists(): @@ -195,16 +199,18 @@ def _generate_coarse_results_from_masks( return None -def _update_scene_final_y_up_layout( +def _update_scene_final_y_up_layout_and_z_up_centers( *, scene: Scene, table_layout: dict[str, object], assets_layout: list[dict[str, object]], + geometry_root: str | Path, ) -> None: - """Copy final y-up layout values into the matching table and asset objects.""" + """Write final y-up layouts and z-up XY centers into the scene.""" if scene.table is None: raise ValueError("Cannot update a final layout without a table.") + # Keep final poses in the y-up layout convention used by exported GLBs. _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() @@ -223,6 +229,17 @@ def _update_scene_final_y_up_layout( f"Final layout is missing scene assets: {sorted(missing_assets)}." ) + # Measure final geometry in z-up so scene edits can compare tabletop XY positions. + table_mesh, assets_aabb_corners_by_id = _measure_table_and_assets_in_z_up_world( + table_layout=table_layout, + assets_layout=assets_layout, + geometry_root=geometry_root, + ) + # Persist AABB centers for future scene-edit object disambiguation. + scene.table.center_xy = table_mesh.bounds[:, :2].mean(axis=0).tolist() + for asset in scene.assets: + asset.center_xy = assets_aabb_corners_by_id[asset.id].mean(axis=0).tolist() + def _copy_y_up_layout_to_scene_object( scene_object: SceneObject, @@ -407,11 +424,12 @@ def _layout_refinement( ) refined_assets_layout = gravity_settler.settle() - # Update the scene data structure with the final y-up layout values. - _update_scene_final_y_up_layout( + # Update the scene data structure with the final layout and spatial metadata. + _update_scene_final_y_up_layout_and_z_up_centers( scene=scene, table_layout=refined_table_layout, assets_layout=refined_assets_layout, + geometry_root=simready_geometry_output_root, ) return refined_table_layout, refined_assets_layout diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py index f91c665e6..00a9ea250 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py @@ -29,6 +29,11 @@ ImageSegmentationClient, ) from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + SceneGraph, + SceneGraphNode, + TABLE_OBJECT_ID, +) from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( OpenAICompatibleVLM, @@ -150,7 +155,7 @@ def understand_scene( *, vlm_client: OpenAICompatibleVLM, json_max_attempts: int = 3, -) -> Scene: +) -> tuple[Scene, SceneGraph]: resolved_image_path = _validate_image_path(image_path) # The output in this stage will keep a JSON which contains @@ -181,12 +186,40 @@ def understand_scene( finally: image_segmentation_client.close() # Kill the session to avoid resource leaks. + # Use the segmented image to initialize the scene graph + # with the help of the VLM client. + # But at here, we do with the simplest way (hard code). + scene_graph = _initialize_scene_graph_from_segmented_scene(scene) + # 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 + (stage_output_root / "scene_graph.json").write_text( + json.dumps(scene_graph.to_dict(), indent=2, ensure_ascii=False) + "\n", + encoding="utf-8", + ) + return scene, scene_graph + + +def _initialize_scene_graph_from_segmented_scene(scene: Scene) -> SceneGraph: + """Build the initial graph assuming every segmented asset rests on the table.""" + if scene.table is None: + raise ValueError("Cannot initialize a scene graph without a table.") + return SceneGraph( + nodes=[ + SceneGraphNode(object_id=TABLE_OBJECT_ID, parent_id=None), + *[ + SceneGraphNode( + object_id=asset.id, + parent_id=TABLE_OBJECT_ID, + parent_relation="on", + ) + for asset in scene.assets + ], + ], + ) def _analyze_image_objects( 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 fb66c30c4..391ccd156 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py @@ -26,6 +26,7 @@ from scipy.spatial.transform import Rotation from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_graph import SceneGraph from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject from embodichain.utils.logger import log_info @@ -46,12 +47,15 @@ def __init__( self, *, scene: Scene, + scene_graph: SceneGraph, output_root: str | Path, ) -> None: self.scene = scene + self.scene_graph = scene_graph self.output_root = Path(output_root).expanduser().resolve() self.export_root = self.output_root / "scene_export" self.scene_config_path: Path | None = None + self.scene_graph_path: Path | None = None def export(self) -> Path: """Write a scene-only config and copy SimReady GLBs into ``mesh_assets``. @@ -72,6 +76,9 @@ def export(self) -> Path: 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.") + self.scene_graph.validate() + if set(self.scene_graph.node_by_id()) != set(object_ids): + raise ValueError("Scene graph nodes must match exported scene object ids.") exported_entries = { scene_object.id: self._copy_scene_object_to_assets( @@ -106,6 +113,12 @@ def export(self) -> Path: encoding="utf-8", ) log_info(f"Exported scene config: {self.scene_config_path}") + self.scene_graph_path = self.export_root / "scene_graph.json" + self.scene_graph_path.write_text( + json.dumps(self.scene_graph.to_dict(), indent=2, ensure_ascii=False) + "\n", + encoding="utf-8", + ) + log_info(f"Exported scene graph: {self.scene_graph_path}") return self.scene_config_path @staticmethod @@ -179,6 +192,7 @@ def _scene_object_config( # 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, + "center_xy": scene_object.center_xy, "max_convex_hull_num": scene_object.physics.max_convex_hull_num, } diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py index 716ce036d..1ae6cdb1b 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py @@ -24,6 +24,11 @@ from scipy.spatial.transform import Rotation from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + SceneGraph, + SceneGraphNode, + SceneGraphRelation, +) from embodichain.gen_sim.scene_engine.core.scene_object import ( ObjectPhysics, SceneObject, @@ -53,10 +58,28 @@ def __init__( self.scene_export_root = self.output_root / "scene_export" self.mesh_assets_root = self.scene_export_root / "mesh_assets" self.scene_config_path = self.scene_export_root / "scene_config.json" + self.scene_graph_path = self.scene_export_root / "scene_graph.json" self.scene_json_path = self.scene_export_root / "scene.json" def import_scene(self) -> Scene: """Validate the scene export, write ``scene.json``, and return a ``Scene``.""" + scene = self._load_scene() + self._write_scene_json(scene) + return scene + + def import_scene_and_graph(self) -> tuple[Scene, SceneGraph]: + """Import a scene and graph after validating the complete edit input.""" + scene = self._load_scene() + scene_graph = self._load_scene_graph() + if set(scene_graph.node_by_id()) != { + scene_object.id for scene_object in scene.objects + }: + raise ValueError("Scene graph nodes must match imported scene object ids.") + self._write_scene_json(scene) + return scene, scene_graph + + def _load_scene(self) -> Scene: + """Validate the exported scene files and restore the ``Scene`` data.""" # Editing only runs on an existing Scene Engine output directory. if not self.output_root.is_dir() or not any(self.output_root.iterdir()): raise ValueError( @@ -86,15 +109,29 @@ def import_scene(self) -> Scene: if not isinstance(scene_config, dict): raise ValueError("Scene config must be a JSON object.") - scene = self._scene_from_config(scene_config) - if self.scene_json_path.exists(): - self.scene_json_path.unlink() + return self._scene_from_config(scene_config) + + def _load_scene_graph(self) -> SceneGraph: + """Read and validate the exported scene graph.""" + if not self.scene_graph_path.is_file(): + raise FileNotFoundError(f"Scene graph not found: {self.scene_graph_path}") + try: + scene_graph_data = json.loads( + self.scene_graph_path.read_text(encoding="utf-8") + ) + except json.JSONDecodeError as exc: + raise ValueError( + f"Scene graph is not valid JSON: {self.scene_graph_path}" + ) from exc + return self._scene_graph_from_data(scene_graph_data) + + def _write_scene_json(self, scene: Scene) -> None: + """Write the restored scene debugging artifact after validation succeeds.""" self.scene_json_path.write_text( json.dumps(scene.to_dict(), indent=2, ensure_ascii=False) + "\n", encoding="utf-8", ) log_info(f"Imported scene JSON: {self.scene_json_path}") - return scene def _scene_from_config(self, scene_config: dict[str, Any]) -> Scene: """Build a y-up ``Scene`` from the z-up scene-export config.""" @@ -127,6 +164,88 @@ def _scene_from_config(self, scene_config: dict[str, Any]) -> Scene: ] ) + @staticmethod + def _scene_graph_from_data(value: object) -> SceneGraph: + """Build a validated ``SceneGraph`` from exported graph JSON.""" + if not isinstance(value, dict) or set(value) != {"nodes", "relations"}: + raise ValueError("Scene graph must contain exactly nodes and relations.") + nodes_value = value["nodes"] + relations_value = value["relations"] + if not isinstance(nodes_value, list) or not isinstance(relations_value, list): + raise ValueError("Scene graph nodes and relations must be lists.") + + nodes = [ + SceneExportImporter._scene_graph_node_from_data(node) + for node in nodes_value + ] + relations = [ + SceneExportImporter._scene_graph_relation_from_data(relation) + for relation in relations_value + ] + return SceneGraph(nodes=nodes, relations=relations) + + @staticmethod + def _scene_graph_node_from_data(value: object) -> SceneGraphNode: + if not isinstance(value, dict) or set(value) != { + "object_id", + "parent_id", + "parent_relation", + "table_region", + }: + raise ValueError("Scene graph nodes must use the serialized node schema.") + object_id = value["object_id"] + parent_id = value["parent_id"] + parent_relation = value["parent_relation"] + table_region = value["table_region"] + if not isinstance(object_id, str) or not isinstance( + parent_id, (str, type(None)) + ): + raise ValueError("Scene graph node ids must be strings or null.") + if parent_relation not in {None, "on"}: + raise ValueError("Scene graph parent_relation must be 'on' or null.") + if table_region not in { + None, + "left_back", + "back_center", + "right_back", + "left_center", + "center", + "right_center", + "left_front", + "front_center", + "right_front", + }: + raise ValueError("Scene graph table_region is invalid.") + return SceneGraphNode( + object_id=object_id, + parent_id=parent_id, + parent_relation=parent_relation, + table_region=table_region, + ) + + @staticmethod + def _scene_graph_relation_from_data(value: object) -> SceneGraphRelation: + if not isinstance(value, dict) or set(value) != { + "source_id", + "relation", + "target_id", + }: + raise ValueError( + "Scene graph relations must use the serialized relation schema." + ) + source_id = value["source_id"] + relation = value["relation"] + target_id = value["target_id"] + if not isinstance(source_id, str) or not isinstance(target_id, str): + raise ValueError("Scene graph relation ids must be strings.") + if relation not in {"left_of", "right_of", "in_front_of", "behind"}: + raise ValueError("Scene graph relation is invalid.") + return SceneGraphRelation( + source_id=source_id, + relation=relation, + target_id=target_id, + ) + def _scene_object_from_export_entry( self, entry: object, @@ -153,6 +272,9 @@ def _scene_object_from_export_entry( entry.get("body_scale", [1.0, 1.0, 1.0]), field_name=f"{uid}.body_scale", ) + center_xy = entry.get("center_xy") + if center_xy is not None: + center_xy = self._vector2(center_xy, field_name=f"{uid}.center_xy") pos_y_up = _Z_UP_TO_Y_UP_ROTATION @ np.asarray(pos_z_up, dtype=float) rotation_z_up = Rotation.from_euler("XYZ", rot_z_up, degrees=True).as_matrix() @@ -171,6 +293,7 @@ def _scene_object_from_export_entry( rot=rot_y_up.tolist(), pos=pos_y_up.tolist(), scale=scale, + center_xy=center_xy, physics=ObjectPhysics( body_type=str(entry.get("body_type", "dynamic")), # type: ignore[arg-type] attrs=self._physics_attrs(entry.get("attrs", {"mass": 1.0})), @@ -193,6 +316,11 @@ def _resolve_export_glb_path( raise ValueError(f"Scene object {uid!r} shape.fpath must be relative.") if fpath.suffix.lower() != ".glb": raise ValueError(f"Scene object {uid!r} shape.fpath must point to a GLB.") + expected_fpath = Path("mesh_assets") / uid / f"{uid}.glb" + if fpath != expected_fpath: + raise ValueError( + f"Scene object {uid!r} shape.fpath must be {expected_fpath.as_posix()!r}." + ) glb_path = (self.scene_export_root / fpath).resolve() if self.scene_export_root.resolve() not in glb_path.parents: raise ValueError( @@ -215,6 +343,18 @@ def _vector3(value: object, *, field_name: str) -> list[float]: raise ValueError(f"Scene config field {field_name!r} must be finite.") return vector + @staticmethod + def _vector2(value: object, *, field_name: str) -> list[float]: + """Validate one length-2 numeric vector.""" + if not isinstance(value, list) or len(value) != 2: + raise ValueError( + f"Scene config field {field_name!r} must be a length-2 list." + ) + vector = [float(item) for item in value] + if not np.all(np.isfinite(vector)): + raise ValueError(f"Scene config field {field_name!r} must be finite.") + return vector + @staticmethod def _physics_attrs(value: object) -> dict[str, float | int]: """Validate exported physics attributes.""" 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 56c12af1c..b7b614d49 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,11 +24,18 @@ import pytest from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + SceneGraph, + SceneGraphNode, +) from embodichain.gen_sim.scene_engine.core.scene_object import ( ObjectPhysics, SceneObject, ) from embodichain.gen_sim.scene_engine.pipeline.utils.scene_exporter import SceneExporter +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_importer import ( + SceneExportImporter, +) def _scene_object( @@ -60,6 +67,24 @@ def _physics(body_type: str) -> ObjectPhysics: ) +def _scene_graph(scene: Scene) -> SceneGraph: + if scene.table is None: + raise ValueError("Test scene must contain a table.") + return SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + *[ + SceneGraphNode( + object_id=asset.id, + parent_id="table", + parent_relation="on", + ) + for asset in scene.assets + ], + ] + ) + + 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") @@ -120,9 +145,12 @@ def test_scene_export_copies_meshes_and_converts_y_up_pose(tmp_path: Path) -> No glb_path=asset_glb, physics=_physics("dynamic"), ) + asset.center_xy = [0.25, -0.5] + scene = Scene(objects=[table, asset]) export_path = SceneExporter( - scene=Scene(objects=[table, asset]), + scene=scene, + scene_graph=_scene_graph(scene), output_root=tmp_path / "output", ).export() exported = json.loads(export_path.read_text(encoding="utf-8")) @@ -136,7 +164,31 @@ def test_scene_export_copies_meshes_and_converts_y_up_pose(tmp_path: Path) -> No 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 entry["center_xy"] == [0.25, -0.5] assert np.allclose(entry["init_rot"], [0.0, 0.0, 0.0]) + assert json.loads((export_path.parent / "scene_graph.json").read_text()) == { + "nodes": [ + { + "object_id": "table", + "parent_id": None, + "parent_relation": None, + "table_region": None, + }, + { + "object_id": "cup", + "parent_id": "table", + "parent_relation": "on", + "table_region": None, + }, + ], + "relations": [], + } + + imported_scene, imported_graph = SceneExportImporter( + output_root=tmp_path / "output" + ).import_scene_and_graph() + assert [asset.id for asset in imported_scene.assets] == ["cup"] + assert imported_graph.to_dict() == _scene_graph(scene).to_dict() def test_scene_export_requires_final_physics(tmp_path: Path) -> None: @@ -145,7 +197,13 @@ def test_scene_export_requires_final_physics(tmp_path: Path) -> None: 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() + SceneExporter( + scene=Scene(objects=[table]), + scene_graph=SceneGraph( + nodes=[SceneGraphNode(object_id="table", parent_id=None)] + ), + output_root=tmp_path, + ).export() def test_scene_export_rejects_backslash_in_object_id(tmp_path: Path) -> None: @@ -165,7 +223,9 @@ def test_scene_export_rejects_backslash_in_object_id(tmp_path: Path) -> None: ) with pytest.raises(ValueError, match="not safe for a GLB filename"): + scene = Scene(objects=[table, unsafe_asset]) SceneExporter( - scene=Scene(objects=[table, unsafe_asset]), + scene=scene, + scene_graph=_scene_graph(scene), output_root=tmp_path / "output", ).export() diff --git a/tests/gen_sim/scene_engine/test_scene_edit.py b/tests/gen_sim/scene_engine/test_scene_edit.py index aef5a1462..fd1a9e933 100644 --- a/tests/gen_sim/scene_engine/test_scene_edit.py +++ b/tests/gen_sim/scene_engine/test_scene_edit.py @@ -72,6 +72,7 @@ def _write_scene_export( "init_pos": [1.0, -3.0, 2.0], "init_rot": [0.0, 0.0, 0.0], "body_scale": [1.0, 2.0, 3.0], + "center_xy": [1.0, -3.0], "max_convex_hull_num": 32, } ], @@ -99,6 +100,7 @@ def test_import_scene_from_output_root_writes_y_up_scene_json(tmp_path: Path) -> assert scene.assets[0].id == "cup" assert scene.assets[0].pos == [1.0, 2.0, 3.0] assert scene.assets[0].scale == [1.0, 2.0, 3.0] + assert scene.assets[0].center_xy == [1.0, -3.0] assert scene.assets[0].simready_glb_path == str( (tmp_path / "scene_export" / "mesh_assets" / "cup" / "cup.glb").resolve() ) diff --git a/tests/gen_sim/scene_engine/test_scene_engine_config.py b/tests/gen_sim/scene_engine/test_scene_engine_config.py index 01914e144..3754210d9 100644 --- a/tests/gen_sim/scene_engine/test_scene_engine_config.py +++ b/tests/gen_sim/scene_engine/test_scene_engine_config.py @@ -87,6 +87,58 @@ def generate_scene(*, image_path: Path, output_root: Path) -> None: } +def test_scene_engine_cli_edits_existing_output_without_an_image( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + captured: dict[str, object] = {} + + def edit_scene(*, output_root: Path, edit_prompt: str) -> None: + captured["output_root"] = output_root + captured["edit_prompt"] = edit_prompt + + monkeypatch.setattr(start, "edit_scene", edit_scene) + output_root = tmp_path / "existing_output" + + start.cli_scene_engine(None, output_root, edit_prompt="move the cup right") + + assert captured == { + "output_root": output_root.resolve(), + "edit_prompt": "move the cup right", + } + + +def test_scene_engine_cli_generates_then_edits_when_both_inputs_exist( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + image_path = tmp_path / "scene.png" + image_path.write_bytes(b"png") + call_order: list[str] = [] + + def generate_scene(*, image_path: Path, output_root: Path) -> None: + call_order.append("generate") + + def edit_scene(*, output_root: Path, edit_prompt: str) -> None: + call_order.append("edit") + + monkeypatch.setattr(start, "generate_scene_from_image", generate_scene) + monkeypatch.setattr(start, "edit_scene", edit_scene) + + start.cli_scene_engine( + image_path, + tmp_path / "output", + edit_prompt="move the cup right", + ) + + assert call_order == ["generate", "edit"] + + +def test_scene_engine_cli_requires_an_image_or_edit_prompt(tmp_path: Path) -> None: + with pytest.raises(ValueError, match="--image, --edit_prompt, or both"): + start.cli_scene_engine(None, tmp_path / "output") + + @pytest.mark.parametrize("image_name", ["missing.png", "scene.gif"]) def test_scene_engine_cli_rejects_invalid_image_inputs( tmp_path: Path, diff --git a/tests/gen_sim/scene_engine/test_scene_understanding.py b/tests/gen_sim/scene_engine/test_scene_understanding.py index 5d9a285c3..9f5ae3884 100644 --- a/tests/gen_sim/scene_engine/test_scene_understanding.py +++ b/tests/gen_sim/scene_engine/test_scene_understanding.py @@ -23,6 +23,7 @@ import pytest from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject from embodichain.gen_sim.scene_engine.pipeline.generation import scene_understanding @@ -83,3 +84,46 @@ def complete(self, **_: object) -> str: assert scene.table is not None assert [asset.id for asset in scene.assets] == ["cup_001"] + + +def test_initial_scene_graph_places_every_asset_on_table() -> None: + scene = Scene( + objects=[ + SceneObject( + id="table", + kind="table", + category="table", + name="wooden table", + description="A wooden table.", + ), + SceneObject( + id="cup_001", + kind="asset", + category="cup", + name="blue cup", + description="A blue cup.", + ), + ], + ) + + scene_graph = scene_understanding._initialize_scene_graph_from_segmented_scene( + scene + ) + + assert scene_graph.to_dict() == { + "nodes": [ + { + "object_id": "table", + "parent_id": None, + "parent_relation": None, + "table_region": None, + }, + { + "object_id": "cup_001", + "parent_id": "table", + "parent_relation": "on", + "table_region": None, + }, + ], + "relations": [], + } From dd7cbaf873caccce98b9fe30c27bb53d44240402 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Wed, 12 Aug 2026 16:25:40 +0800 Subject: [PATCH 06/29] Fixed the prompt of scne understanding: delete the location in description, and add check code --- .../generation/scene_understanding.py | 29 +++++++++++-------- 1 file changed, 17 insertions(+), 12 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py index 00a9ea250..c69ec3719 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py @@ -67,8 +67,8 @@ 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. + keep the same category and name. Do not encode location or spatial context + in any semantic field. 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. @@ -78,8 +78,9 @@ 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. +8. For assets, description contains only visible category, material, color, + texture, shape, and structural details. Do not mention location, the table, + or any relationship to another object. Return JSON only: no Markdown, comments, or prose outside this exact schema: { @@ -92,14 +93,14 @@ { "category": "mug", "name": "blue ceramic mug", - "description": "small blue ceramic mug on the left side of the table" + "description": "small blue ceramic mug with a curved handle" } ] } -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.""" +For two identical blue mugs, output two asset entries with the same category, +name, and description. 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." @@ -363,13 +364,17 @@ def _parse_scene_object_fields( 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. + # Check whether the name and description contain location or relationship words. + if _LOCATION_WORD_PATTERN.search(fields["name"]): raise ValueError( f"VLM JSON key {field_name}.name must not contain location or " "relationship words." ) + if _LOCATION_WORD_PATTERN.search(fields["description"]): + raise ValueError( + f"VLM JSON key {field_name}.description must not contain location or " + "relationship words." + ) return fields From 2dffcfeafe70bf702c663712cf52122fd90ef25c Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Wed, 12 Aug 2026 16:53:06 +0800 Subject: [PATCH 07/29] modify the prompt of scene understanding, to avoid a very long description in objects' names --- .../scene_engine/pipeline/generation/scene_understanding.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py index c69ec3719..cae1c612f 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py @@ -71,9 +71,9 @@ in any semantic field. 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. +6. name is a concise human-readable phrase containing only color, material, + texture, shape, and object details. It may contain spaces, but 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 From c3d2203f20db984ca8f0d411096f2cb94cfbb157 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Wed, 12 Aug 2026 17:31:14 +0800 Subject: [PATCH 08/29] finish the llm understand scene edit --- .../scene_engine/core/scene_edit_plan.py | 239 +++++++++++++++ .../gen_sim/scene_engine/pipeline/edit.py | 7 +- .../editing/scene_edit_understanding.py | 280 +++++++++++++++++- .../scene_engine/test_scene_edit_plan.py | 204 +++++++++++++ .../scene_engine/test_scene_understanding.py | 8 + 5 files changed, 725 insertions(+), 13 deletions(-) create mode 100644 embodichain/gen_sim/scene_engine/core/scene_edit_plan.py create mode 100644 tests/gen_sim/scene_engine/test_scene_edit_plan.py diff --git a/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py b/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py new file mode 100644 index 000000000..26fe96bc9 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py @@ -0,0 +1,239 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF 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 typing import Literal + +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + SceneConstraintType, + SceneGraph, + TABLE_OBJECT_ID, +) + +__all__ = ["SceneEditOperation", "SceneEditPlan"] + +SceneEditOperationType = Literal["add", "move", "delete"] + + +@dataclass(frozen=True) +class SceneEditOperation: + """One normalized edit operation produced from an LLM edit draft.""" + + op: SceneEditOperationType + object_id: str | None = None + target_id: str | None = None + relation: SceneConstraintType | None = None + category: str | None = None + name: str | None = None + description: str | None = None + + def to_dict(self) -> dict[str, object]: + """Serialize one normalized edit operation.""" + return { + "op": self.op, + "object_id": self.object_id, + "target_id": self.target_id, + "relation": self.relation, + "category": self.category, + "name": self.name, + "description": self.description, + } + + +@dataclass +class SceneEditPlan: + """Validated operations against one immutable pre-edit scene state.""" + + scene: Scene + scene_graph: SceneGraph + operations: list[SceneEditOperation] = field(default_factory=list) + + def __post_init__(self) -> None: + """Validate the plan before later stages prepare assets or edit layouts.""" + self.validate() + + def to_dict(self) -> dict[str, object]: + """Serialize the input scene state and normalized edit operations.""" + return { + "scene": self.scene.to_dict(), + "scene_graph": self.scene_graph.to_dict(), + "operations": [operation.to_dict() for operation in self.operations], + } + + def validate(self) -> None: + """Validate object references and edit conflicts against the input scene.""" + # Scene object IDs must remain a one-to-one lookup key for edit operations. + scene_object_ids = {scene_object.id for scene_object in self.scene.objects} + if len(scene_object_ids) != len(self.scene.objects): + raise ValueError("Scene edit input must contain unique object ids.") + # Parent and child checks require the graph to describe this exact scene. + if set(self.scene_graph.node_by_id()) != scene_object_ids: + raise ValueError("Scene edit plan graph nodes must match scene object ids.") + + existing_object_ids = set(scene_object_ids) + added_object_ids: set[str] = set() + # Collect deletions first so other operations cannot target removed objects. + deleted_object_ids = { + operation.object_id + for operation in self.operations + if operation.op == "delete" + } + if None in deleted_object_ids: + raise ValueError("Delete operations must identify an existing object.") + + edited_object_ids: set[str] = set() + # Validate each operation against the unchanged input scene and graph. + for operation in self.operations: + self._validate_operation( + operation=operation, + existing_object_ids=existing_object_ids, + deleted_object_ids=deleted_object_ids, + edited_object_ids=edited_object_ids, + added_object_ids=added_object_ids, + ) + # A removed support object must not leave any child objects orphaned. + self._validate_deleted_subtrees(deleted_object_ids) + + def _validate_operation( + self, + *, + operation: SceneEditOperation, + existing_object_ids: set[str], + deleted_object_ids: set[str], + edited_object_ids: set[str], + added_object_ids: set[str], + ) -> None: + if operation.op == "add": + self._validate_add_operation( + operation, + existing_object_ids, + deleted_object_ids, + added_object_ids, + ) + return + if operation.op not in {"move", "delete"}: + raise ValueError(f"Unsupported scene edit operation: {operation.op!r}") + if operation.object_id not in existing_object_ids: + raise ValueError( + "Move and delete operations must reference existing objects." + ) + if operation.object_id == TABLE_OBJECT_ID: + raise ValueError("The table cannot be moved or deleted.") + # Existing objects accept only one move or delete instruction per plan. + if operation.object_id in edited_object_ids: + raise ValueError("An existing object may have only one edit operation.") + edited_object_ids.add(operation.object_id) + + if operation.op == "delete": + # Delete carries no new metadata or spatial placement. + if any( + value is not None + for value in ( + operation.target_id, + operation.relation, + operation.category, + operation.name, + operation.description, + ) + ): + raise ValueError("Delete operations may only specify object_id.") + return + + self._validate_position_reference( + operation=operation, + existing_object_ids=existing_object_ids, + deleted_object_ids=deleted_object_ids, + ) + if any( + value is not None + for value in (operation.category, operation.name, operation.description) + ): + raise ValueError("Move operations must not declare a new object.") + + @staticmethod + def _validate_add_operation( + operation: SceneEditOperation, + existing_object_ids: set[str], + deleted_object_ids: set[str], + added_object_ids: set[str], + ) -> None: + if not operation.object_id: + raise ValueError("Add operations must have a generated object_id.") + # Generated IDs must not collide with the input scene or this add batch. + if ( + operation.object_id in existing_object_ids + or operation.object_id in added_object_ids + ): + raise ValueError("Add operations must use unique new object ids.") + added_object_ids.add(operation.object_id) + if not all( + isinstance(value, str) and value.strip() + for value in (operation.category, operation.name, operation.description) + ): + raise ValueError("Add operations require category, name, and description.") + SceneEditPlan._validate_position_reference( + operation=operation, + existing_object_ids=existing_object_ids, + deleted_object_ids=deleted_object_ids, + ) + + @staticmethod + def _validate_position_reference( + *, + operation: SceneEditOperation, + existing_object_ids: set[str], + deleted_object_ids: set[str], + ) -> None: + if (operation.target_id is None) != (operation.relation is None): + raise ValueError("target_id and relation must be specified together.") + if operation.target_id is None: + return + if operation.target_id not in existing_object_ids: + raise ValueError("Edit targets must reference existing scene objects.") + # One edit may not position an object relative to a deleted target. + if operation.target_id in deleted_object_ids: + raise ValueError("Edit targets must not reference deleted objects.") + + def _validate_deleted_subtrees(self, deleted_object_ids: set[str]) -> None: + # Index the support graph once before checking every deleted parent. + children_by_parent: dict[str, list[str]] = {} + for node in self.scene_graph.nodes: + if node.parent_id is not None: + children_by_parent.setdefault(node.parent_id, []).append(node.object_id) + + for object_id in deleted_object_ids: + descendants = self._descendant_ids(object_id, children_by_parent) + if not descendants.issubset(deleted_object_ids): + raise ValueError( + "Deleting a parent requires deleting all of its children." + ) + + @staticmethod + def _descendant_ids( + object_id: str, + children_by_parent: dict[str, list[str]], + ) -> set[str]: + descendants: set[str] = set() + # Traverse every support descendant, not only direct children. + pending = list(children_by_parent.get(object_id, [])) + while pending: + child_id = pending.pop() + descendants.add(child_id) + pending.extend(children_by_parent.get(child_id, [])) + return descendants diff --git a/embodichain/gen_sim/scene_engine/pipeline/edit.py b/embodichain/gen_sim/scene_engine/pipeline/edit.py index 3b6c95e02..93001afb6 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/edit.py +++ b/embodichain/gen_sim/scene_engine/pipeline/edit.py @@ -43,6 +43,7 @@ def edit_scene( scene_importer = SceneExportImporter(output_root=output_root) # Validate scene_export, write scene.json, and return Scene; failures raise before editing. scene, scene_graph = scene_importer.import_scene_and_graph() + # Only for debug. print( json.dumps( {"scene": scene.to_dict(), "scene_graph": scene_graph.to_dict()}, @@ -52,12 +53,12 @@ def edit_scene( ) # 1. Edit Understanding - # And update or initialize the scene graph + # Will return an already checked scene edit plan. log_info("Starting Edit Understanding") - updated_scene_graph = understand_scene_edit( + scene_edit_plan = understand_scene_edit( scene=scene, + scene_graph=scene_graph, edit_prompt=edit_prompt, - output_root=output_root, # Has already been resolved. vlm_client=vlm_client, ) log_info("Completed Edit Understanding") diff --git a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py index 48dda64bf..8099eda4b 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py @@ -17,32 +17,292 @@ from __future__ import annotations -from pathlib import Path +import json -from embodichain.gen_sim.scene_engine.core.scene import Scene -from embodichain.gen_sim.scene_engine.core.scene_graph import ( - SceneGraph, - SceneGraphNode, - TABLE_OBJECT_ID, +from embodichain.gen_sim.scene_engine.core.scene_edit_plan import ( + SceneEditOperation, + SceneEditPlan, ) +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_graph import SceneGraph from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( OpenAICompatibleVLM, ) +_EDIT_SYSTEM_PROMPT = """You convert one user instruction into edits for an existing tabletop scene. + +Use an existing object ID only when it appears in the supplied Existing object +IDs list. IDs identify existing objects exactly; never invent, correct, or +renumber them. The table ID is "table" and cannot be moved or deleted. + +Each operation is one of: +1. move: move one existing object. object_id identifies it. target_id and + relation are either both provided or both null. +2. delete: delete one existing object. Only object_id is provided. +3. add: create one new object. object_id must be null. Provide a lower-case + singular snake_case category, name, and description. Multiple add operations + may have the same category and name; their final IDs are assigned by the + program in operation order. target_id and relation are either both provided + or both null. + +For a positioned move or add, target_id must be an Existing object ID and +relation must be one of on, left_of, right_of, in_front_of, or behind. Do not +position a new object relative to another newly added object. + +Each existing object's center_xy is its center position [x, y] in the +table-frame Z-up world coordinate system. Smaller x is left, larger x is right, +larger y is in front, and smaller y is behind. Use center_xy only to disambiguate +references such as "the bottle on the left"; do not output coordinates. Express +the requested position using target_id and one allowed relation instead. + +For every newly added object, category is its lower-case singular snake_case +class. name contains only color, material, texture, shape, and object details. +description contains only visible category, material, color, texture, shape, +and structural details. name and description must not mention position, the +table, or relations to any object. + +Return JSON only: no Markdown, comments, or prose. Every operation must contain +exactly these fields: op, object_id, target_id, relation, category, name, and +description. Use null for every field that does not apply to an operation: +{ + "operations": [ + { + "op": "move", + "object_id": "bottle_001", + "target_id": "book_001", + "relation": "right_of", + "category": null, + "name": null, + "description": null + }, + { + "op": "delete", + "object_id": "cup_001", + "target_id": null, + "relation": null, + "category": null, + "name": null, + "description": null + }, + { + "op": "add", + "object_id": null, + "target_id": null, + "relation": null, + "category": "orange", + "name": "small orange", + "description": "small round orange with a textured peel" + }, + { + "op": "add", + "object_id": null, + "target_id": "book_001", + "relation": "right_of", + "category": "orange", + "name": "small orange", + "description": "small round orange with a textured peel" + } + ] +} +The two orange additions intentionally share category and name. Do not add +fields beyond the required schema.""" + def understand_scene_edit( *, scene: Scene, + scene_graph: SceneGraph, edit_prompt: str, - output_root: str | Path, vlm_client: OpenAICompatibleVLM, -) -> dict[str, object]: + json_max_attempts: int = 3, +) -> SceneEditPlan: """Understand one text edit instruction for an existing scene.""" edit_prompt = edit_prompt.strip() if not edit_prompt: raise ValueError("Edit prompt must not be empty.") + if json_max_attempts < 1: + raise ValueError("json_max_attempts must be at least 1.") + # Give the VLM only the scene metadata needed to identify existing objects. + simplified_scene_info = _simplify_scene_info(scene=scene) + operations = _vlm_understand_scene_edit( + scene=scene, + edit_prompt=edit_prompt, + simplified_scene_info=simplified_scene_info, + vlm_client=vlm_client, + json_max_attempts=json_max_attempts, + ) + + # SceneEditPlan validates all references against the immutable input scene graph. + return SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=operations, + ) + +def _simplify_scene_info(scene: Scene) -> dict[str, object]: + """Return the object metadata needed for edit instruction resolution.""" return { - "edit_prompt": edit_prompt, - "operations": [], + "existing_object_ids": [scene_object.id for scene_object in scene.objects], + "objects": [ + { + "id": scene_object.id, + "category": scene_object.category, + "name": scene_object.name, + "description": scene_object.description, + "center_xy": scene_object.center_xy, + } + for scene_object in scene.objects + ], } + + +def _vlm_understand_scene_edit( + *, + scene: Scene, + edit_prompt: str, + simplified_scene_info: dict[str, object], + vlm_client: OpenAICompatibleVLM, + json_max_attempts: int, +) -> list[SceneEditOperation]: + """Return parsed edit operations from the VLM with assigned add IDs.""" + # Construct user prompt. + user_prompt = ( + f"User edit instruction:\n{edit_prompt}\n\n" + "Existing scene metadata:\n" + f"{json.dumps(simplified_scene_info, indent=2, ensure_ascii=False)}" + ) + last_error: ValueError | None = None + for _ in range(json_max_attempts): + response_text = vlm_client.complete( + system_prompt=_EDIT_SYSTEM_PROMPT, + user_prompt=user_prompt, + ) + try: + value = json.loads(_strip_json_code_fence(response_text)) + return _parse_scene_edit_operations(value, scene=scene) + except (json.JSONDecodeError, ValueError) as exc: + last_error = ValueError(f"VLM returned invalid scene edit JSON: {exc}") + continue + + assert last_error is not None + raise ValueError( + "VLM returned invalid scene edit JSON after " + f"{json_max_attempts} attempts: {last_error}" + ) from last_error + + +def _strip_json_code_fence(response_text: str) -> str: + """Remove one optional Markdown JSON fence from a VLM response.""" + 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 _parse_scene_edit_operations( + value: object, + *, + scene: Scene, +) -> list[SceneEditOperation]: + """Parse the strict VLM edit-draft schema into typed operations with add IDs.""" + if not isinstance(value, dict) or set(value) != {"operations"}: + raise ValueError("Scene edit draft must contain exactly operations.") + # Get and validate a list operation value. + operations_value = value["operations"] + if not isinstance(operations_value, list): + raise ValueError("Scene edit draft operations must be a list.") + + expected_keys = { + "op", + "object_id", + "target_id", + "relation", + "category", + "name", + "description", + } + # Get ids and counts of existing objects to assign new add IDs. + assigned_object_ids = {scene_object.id for scene_object in scene.objects} + category_counts = { + category: sum( + scene_object.category == category for scene_object in scene.objects + ) + for category in {scene_object.category for scene_object in scene.objects} + } + operations: list[SceneEditOperation] = [] + for value in operations_value: + if not isinstance(value, dict) or not isinstance(value.get("op"), str): + raise ValueError("Scene edit operations must contain a string op.") + op = value["op"] + if op not in {"add", "move", "delete"}: + raise ValueError("Scene edit operation op is invalid.") + if set(value) != expected_keys: + raise ValueError("Scene edit operations must use the required schema.") + object_id = _optional_string(value.get("object_id"), field_name="object_id") + category = _optional_string(value.get("category"), field_name="category") + if op == "add": + if object_id is not None: + raise ValueError("VLM add operations must set object_id to null.") + if category is None: + raise ValueError("VLM add operations must provide a category.") + # Add operation should generate new id here. + # Never believe the LLM could always generate a valid id. + object_id = _next_add_object_id( + category=category, + category_counts=category_counts, + assigned_object_ids=assigned_object_ids, + ) + operations.append( + SceneEditOperation( + op=op, + object_id=object_id, + target_id=_optional_string( + value.get("target_id"), field_name="target_id" + ), + relation=_optional_relation(value.get("relation")), + category=category, + name=_optional_string(value.get("name"), field_name="name"), + description=_optional_string( + value.get("description"), field_name="description" + ), + ) + ) + return operations + + +def _next_add_object_id( + *, + category: str, + category_counts: dict[str, int], + assigned_object_ids: set[str], +) -> str: + """Assign the next available ID for one new object category.""" + index = category_counts.get(category, 0) + 1 + object_id = f"{category}_{index:03d}" + # In case the scene have orange_001 and orange_003. + while object_id in assigned_object_ids: + index += 1 + object_id = f"{category}_{index:03d}" + category_counts[category] = index + assigned_object_ids.add(object_id) + return object_id + + +def _optional_string(value: object, *, field_name: str) -> str | None: + if value is None: + return None + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"Scene edit operation {field_name} must be a string or null.") + return value.strip() + + +def _optional_relation(value: object) -> str | None: + if value is None: + return None + if value not in {"on", "left_of", "right_of", "in_front_of", "behind"}: + raise ValueError("Scene edit operation relation is invalid.") + return value diff --git a/tests/gen_sim/scene_engine/test_scene_edit_plan.py b/tests/gen_sim/scene_engine/test_scene_edit_plan.py new file mode 100644 index 000000000..c01049eb6 --- /dev/null +++ b/tests/gen_sim/scene_engine/test_scene_edit_plan.py @@ -0,0 +1,204 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF 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 + +import pytest + +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_edit_plan import ( + SceneEditOperation, + SceneEditPlan, +) +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + SceneGraph, + SceneGraphNode, +) +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.pipeline.editing.scene_edit_understanding import ( + _parse_scene_edit_operations, +) + + +def _scene_and_graph() -> tuple[Scene, SceneGraph]: + scene = Scene( + objects=[ + SceneObject( + id="table", + kind="table", + category="table", + name="wooden table", + description="A wooden table.", + ), + SceneObject( + id="book_001", + kind="asset", + category="book", + name="blue book", + description="A blue book.", + ), + SceneObject( + id="orange_001", + kind="asset", + category="orange", + name="orange", + description="An orange.", + ), + ] + ) + scene_graph = SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="book_001", + parent_id="table", + parent_relation="on", + ), + SceneGraphNode( + object_id="orange_001", + parent_id="book_001", + parent_relation="on", + ), + ] + ) + return scene, scene_graph + + +def test_scene_edit_plan_accepts_add_without_a_position() -> None: + scene, scene_graph = _scene_and_graph() + + plan = SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[ + SceneEditOperation( + op="add", + object_id="cup_001", + category="cup", + name="green cup", + description="A small green ceramic cup.", + ) + ], + ) + + assert len(plan.operations) == 1 + assert plan.to_dict()["operations"] == [ + { + "op": "add", + "object_id": "cup_001", + "target_id": None, + "relation": None, + "category": "cup", + "name": "green cup", + "description": "A small green ceramic cup.", + } + ] + + +def test_scene_edit_plan_accepts_multiple_new_objects_with_the_same_category() -> None: + scene, scene_graph = _scene_and_graph() + + plan = SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[ + SceneEditOperation( + op="add", + object_id="orange_002", + category="orange", + name="small orange", + description="A small round orange with a textured peel.", + ), + SceneEditOperation( + op="add", + object_id="orange_003", + category="orange", + name="large orange", + description="A large round orange with a textured peel.", + ), + ], + ) + + assert [operation.category for operation in plan.operations] == ["orange", "orange"] + + +def test_scene_edit_parser_assigns_ids_to_same_category_adds_in_order() -> None: + scene, _ = _scene_and_graph() + draft = { + "operations": [ + { + "op": "add", + "object_id": None, + "target_id": None, + "relation": None, + "category": "orange", + "name": "small_orange", + "description": "A small round orange with a textured peel.", + }, + { + "op": "add", + "object_id": None, + "target_id": None, + "relation": None, + "category": "orange", + "name": "small_orange", + "description": "A small round orange with a textured peel.", + }, + ] + } + + operations = _parse_scene_edit_operations( + json.loads(json.dumps(draft)), scene=scene + ) + + assert [operation.object_id for operation in operations] == [ + "orange_002", + "orange_003", + ] + + +def test_scene_edit_plan_rejects_targets_outside_the_input_scene() -> None: + scene, scene_graph = _scene_and_graph() + + with pytest.raises(ValueError, match="existing scene objects"): + SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[ + SceneEditOperation( + op="add", + object_id="spoon_001", + target_id="new_orange_001", + relation="left_of", + category="spoon", + name="metal spoon", + description="A metal spoon.", + ) + ], + ) + + +def test_scene_edit_plan_requires_deleting_all_children_of_a_deleted_parent() -> None: + scene, scene_graph = _scene_and_graph() + + with pytest.raises(ValueError, match="all of its children"): + SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[SceneEditOperation(op="delete", object_id="book_001")], + ) diff --git a/tests/gen_sim/scene_engine/test_scene_understanding.py b/tests/gen_sim/scene_engine/test_scene_understanding.py index 9f5ae3884..8b9357f33 100644 --- a/tests/gen_sim/scene_engine/test_scene_understanding.py +++ b/tests/gen_sim/scene_engine/test_scene_understanding.py @@ -63,6 +63,14 @@ def test_image_object_analysis_rejects_location_words_in_object_names() -> None: ) +def test_image_object_analysis_rejects_location_words_in_object_descriptions() -> None: + response = json.loads(_response()) + response["assets"][0]["description"] = "A small ceramic cup on the table." + + with pytest.raises(ValueError, match="description must not contain location"): + scene_understanding._parse_image_object_analysis_response(json.dumps(response)) + + def test_image_object_analysis_retries_then_updates_scene(tmp_path: Path) -> None: class VLM: def __init__(self) -> None: From e2177c0aa748e7d6ef268183bc5377e89ffb9c4e Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Wed, 12 Aug 2026 19:54:11 +0800 Subject: [PATCH 09/29] finish the scene edit plan -> updated scene graph --- .../scene_engine/core/scene_edit_plan.py | 14 +- .../gen_sim/scene_engine/core/scene_graph.py | 152 +++++++++++++++ .../gen_sim/scene_engine/pipeline/edit.py | 35 ++-- .../editing/scene_edit_understanding.py | 97 +++++++++- .../scene_engine/test_scene_edit_plan.py | 183 ++++++++++++++++++ 5 files changed, 459 insertions(+), 22 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py b/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py index 26fe96bc9..a7ac152d7 100644 --- a/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py +++ b/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py @@ -78,6 +78,15 @@ def to_dict(self) -> dict[str, object]: def validate(self) -> None: """Validate object references and edit conflicts against the input scene.""" + # Edit-plan rules: + # - move and delete identify one existing non-table object with object_id. + # - add carries generated object_id plus non-empty category, name, and description. + # - move always supplies target_id and relation; add may omit both. + # - target_id and relation are otherwise supplied together or both absent. + # - every target is from the pre-edit scene; new and deleted objects are invalid targets. + # - an existing object has at most one move or delete operation in one plan. + # - delete carries no placement or new-object metadata and must delete every descendant. + # - these checks validate intent only; they do not mutate Scene or SceneGraph. # Scene object IDs must remain a one-to-one lookup key for edit operations. scene_object_ids = {scene_object.id for scene_object in self.scene.objects} if len(scene_object_ids) != len(self.scene.objects): @@ -88,7 +97,7 @@ def validate(self) -> None: existing_object_ids = set(scene_object_ids) added_object_ids: set[str] = set() - # Collect deletions first so other operations cannot target removed objects. + # Collect deletion intents first so move and add cannot target them regardless of order. deleted_object_ids = { operation.object_id for operation in self.operations @@ -155,6 +164,8 @@ def _validate_operation( raise ValueError("Delete operations may only specify object_id.") return + if operation.target_id is None or operation.relation is None: + raise ValueError("Move operations must specify target_id and relation.") self._validate_position_reference( operation=operation, existing_object_ids=existing_object_ids, @@ -204,6 +215,7 @@ def _validate_position_reference( raise ValueError("target_id and relation must be specified together.") if operation.target_id is None: return + # Targets come only from the pre-edit scene, so new objects cannot be targets. if operation.target_id not in existing_object_ids: raise ValueError("Edit targets must reference existing scene objects.") # One edit may not position an object relative to a deleted target. diff --git a/embodichain/gen_sim/scene_engine/core/scene_graph.py b/embodichain/gen_sim/scene_engine/core/scene_graph.py index 806c5a0ee..514176211 100644 --- a/embodichain/gen_sim/scene_engine/core/scene_graph.py +++ b/embodichain/gen_sim/scene_engine/core/scene_graph.py @@ -129,6 +129,158 @@ def node_by_id(self) -> dict[str, SceneGraphNode]: nodes_by_id[node.object_id] = node return nodes_by_id + def remove_nodes(self, object_ids: set[str]) -> None: + """Remove nodes and their incident planar relations, then validate.""" + # If no node to be removed, return directly. + if not object_ids: + return + if TABLE_OBJECT_ID in object_ids: + raise ValueError("The table cannot be removed from a scene graph.") + unknown_object_ids = object_ids - set(self.node_by_id()) + if unknown_object_ids: + raise ValueError( + f"Cannot remove unknown scene graph nodes: {sorted(unknown_object_ids)}" + ) + + # Removing every incident relation prevents dangling planar endpoints. + self.nodes = [node for node in self.nodes if node.object_id not in object_ids] + self.relations = [ + relation + for relation in self.relations + if relation.source_id not in object_ids + and relation.target_id not in object_ids + ] + # Refresh. + self.refresh() + + def add_node(self, node: SceneGraphNode) -> None: + """Add one node and validate the resulting graph.""" + if node.object_id in self.node_by_id(): + raise ValueError(f"Duplicate scene graph node: {node.object_id}") + self.nodes.append(node) + self.refresh() + + def apply_updates( + self, + *, + deleted_object_ids: set[str], + added_object_ids: list[str], + on_parent_updates: list[tuple[str, str]], + planar_relation_updates: list[tuple[str, PlanarRelationType, str]], + ) -> None: + """Apply one atomic batch of node and relationship updates.""" + if TABLE_OBJECT_ID in deleted_object_ids: + raise ValueError("The table cannot be removed from a scene graph.") + + existing_object_ids = set(self.node_by_id()) + unknown_object_ids = deleted_object_ids - existing_object_ids + if unknown_object_ids: + raise ValueError( + f"Cannot remove unknown scene graph nodes: {sorted(unknown_object_ids)}" + ) + + # Delete all requested nodes before resolving new parents and relations. + self.nodes = [ + node for node in self.nodes if node.object_id not in deleted_object_ids + ] + self.relations = [ + relation + for relation in self.relations + if relation.source_id not in deleted_object_ids + and relation.target_id not in deleted_object_ids + ] + + remaining_object_ids = set(self.node_by_id()) + if len(added_object_ids) != len(set(added_object_ids)): + raise ValueError("Added scene graph node ids must be unique.") + duplicate_object_ids = set(added_object_ids) & remaining_object_ids + if duplicate_object_ids: + raise ValueError( + f"Duplicate scene graph nodes: {sorted(duplicate_object_ids)}" + ) + + # New nodes default to the table; later updates replace that parent when needed. + self.nodes.extend( + SceneGraphNode( + object_id=object_id, + parent_id=TABLE_OBJECT_ID, + parent_relation="on", + ) + for object_id in added_object_ids + ) + + # Apply support-parent changes before planar updates need the final parent. + for object_id, parent_id in on_parent_updates: + self._set_on_parent(object_id=object_id, parent_id=parent_id) + + # Resolve chained planar parent inheritance before adding final relations. + self._resolve_planar_parent_updates(planar_relation_updates) + for source_id, relation, target_id in planar_relation_updates: + self._clear_incident_planar_relations(source_id) + self.relations.append( + SceneGraphRelation( + source_id=source_id, + relation=relation, + target_id=target_id, + ) + ) + + # Normalize inverse relations and reject invalid final graph constraints. + self.refresh() + + def _set_on_parent(self, *, object_id: str, parent_id: str) -> None: + """Replace one node's support parent and stale planar constraints.""" + nodes_by_id = self.node_by_id() + if object_id == TABLE_OBJECT_ID: + raise ValueError("The table cannot be moved onto another object.") + if object_id not in nodes_by_id or parent_id not in nodes_by_id: + raise ValueError( + "Parent updates must reference existing scene graph nodes." + ) + if object_id == parent_id: + raise ValueError("A scene graph node cannot be its own parent.") + + node = nodes_by_id[object_id] + node.parent_id = parent_id + node.parent_relation = "on" + node.table_region = None + self._clear_incident_planar_relations(object_id) + + def _resolve_planar_parent_updates( + self, + planar_relation_updates: list[tuple[str, PlanarRelationType, str]], + ) -> None: + """Make every planar source share its target's final support parent.""" + for _ in range(len(planar_relation_updates)): + changed = False + for source_id, _, target_id in planar_relation_updates: + nodes_by_id = self.node_by_id() + if source_id == TABLE_OBJECT_ID: + raise ValueError("The table cannot have a planar relation.") + if source_id not in nodes_by_id or target_id not in nodes_by_id: + raise ValueError( + "Planar updates must reference existing scene graph nodes." + ) + target_parent_id = nodes_by_id[target_id].parent_id + if target_parent_id is None: + raise ValueError("Planar relation targets must have a parent.") + source = nodes_by_id[source_id] + if source.parent_id != target_parent_id: + source.parent_id = target_parent_id + source.parent_relation = "on" + source.table_region = None + changed = True + if not changed: + return + + def _clear_incident_planar_relations(self, object_id: str) -> None: + """Remove planar constraints invalidated when one node changes parent.""" + self.relations = [ + relation + for relation in self.relations + if relation.source_id != object_id and relation.target_id != object_id + ] + def normalize(self) -> None: """Materialize inverse planar relations and remove duplicates.""" self._materialize_inverse_planar_relations() diff --git a/embodichain/gen_sim/scene_engine/pipeline/edit.py b/embodichain/gen_sim/scene_engine/pipeline/edit.py index 93001afb6..a13ef17fd 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/edit.py +++ b/embodichain/gen_sim/scene_engine/pipeline/edit.py @@ -44,31 +44,40 @@ def edit_scene( # Validate scene_export, write scene.json, and return Scene; failures raise before editing. scene, scene_graph = scene_importer.import_scene_and_graph() # Only for debug. - print( - json.dumps( - {"scene": scene.to_dict(), "scene_graph": scene_graph.to_dict()}, - indent=2, - ensure_ascii=False, - ) - ) + # print( + # json.dumps( + # {"scene": scene.to_dict(), "scene_graph": scene_graph.to_dict()}, + # indent=2, + # ensure_ascii=False, + # ) + # ) # 1. Edit Understanding - # Will return an already checked scene edit plan. + # Will return an already checked scene edit plan + # and a validated updated scene graph. log_info("Starting Edit Understanding") - scene_edit_plan = understand_scene_edit( + scene_edit_plan, updated_scene_graph = understand_scene_edit( scene=scene, scene_graph=scene_graph, edit_prompt=edit_prompt, vlm_client=vlm_client, ) log_info("Completed Edit Understanding") + # Only for debug. + # print( + # json.dumps( + # {"updated scene graph": updated_scene_graph.to_dict()}, + # indent=2, + # ensure_ascii=False, + # ) + # ) # 2. Prepare Objects. log_info("Preparing Objects if necessary") - # scene = prepare_objects( - # scene=scene, - # output_root=output_root, - # ) + scene = prepare_scene_edit_assets( + scene=scene, + scene_edit_plan=scene_edit_plan, + ) log_info("Completed Preparing Objects") # 3. Layout Editing diff --git a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py index 8099eda4b..1eb3a857f 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py @@ -24,7 +24,12 @@ SceneEditPlan, ) from embodichain.gen_sim.scene_engine.core.scene import Scene -from embodichain.gen_sim.scene_engine.core.scene_graph import SceneGraph +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + PlanarRelationType, + SceneGraph, + SceneGraphNode, + SceneGraphRelation, +) from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( OpenAICompatibleVLM, ) @@ -36,8 +41,8 @@ renumber them. The table ID is "table" and cannot be moved or deleted. Each operation is one of: -1. move: move one existing object. object_id identifies it. target_id and - relation are either both provided or both null. +1. move: move one existing object. object_id, target_id, and relation must all + be provided. 2. delete: delete one existing object. Only object_id is provided. 3. add: create one new object. object_id must be null. Provide a lower-case singular snake_case category, name, and description. Multiple add operations @@ -45,9 +50,9 @@ program in operation order. target_id and relation are either both provided or both null. -For a positioned move or add, target_id must be an Existing object ID and -relation must be one of on, left_of, right_of, in_front_of, or behind. Do not -position a new object relative to another newly added object. +For every move and every positioned add, target_id must be an Existing object +ID and relation must be one of on, left_of, right_of, in_front_of, or behind. +Do not position a new object relative to another newly added object. Each existing object's center_xy is its center position [x, y] in the table-frame Z-up world coordinate system. Smaller x is left, larger x is right, @@ -115,7 +120,7 @@ def understand_scene_edit( edit_prompt: str, vlm_client: OpenAICompatibleVLM, json_max_attempts: int = 3, -) -> SceneEditPlan: +) -> tuple[SceneEditPlan, SceneGraph]: """Understand one text edit instruction for an existing scene.""" edit_prompt = edit_prompt.strip() if not edit_prompt: @@ -133,11 +138,87 @@ def understand_scene_edit( ) # SceneEditPlan validates all references against the immutable input scene graph. - return SceneEditPlan( + scene_edit_plan = SceneEditPlan( scene=scene, scene_graph=scene_graph, operations=operations, ) + updated_scene_graph = _build_updated_scene_graph( + scene_graph=scene_graph, + scene_edit_plan=scene_edit_plan, + ) + return scene_edit_plan, updated_scene_graph + + +def _build_updated_scene_graph( + *, + scene_graph: SceneGraph, + scene_edit_plan: SceneEditPlan, +) -> SceneGraph: + """Build and validate the target graph implied by one edit plan.""" + # Copy every mutable graph value so the pre-edit graph remains unchanged. + updated_scene_graph = SceneGraph( + nodes=[ + SceneGraphNode( + object_id=node.object_id, + parent_id=node.parent_id, + parent_relation=node.parent_relation, + table_region=node.table_region, + ) + for node in scene_graph.nodes + ], + relations=[ + SceneGraphRelation( + source_id=relation.source_id, + relation=relation.relation, + target_id=relation.target_id, + ) + for relation in scene_graph.relations + ], + validate_on_refresh=scene_graph.validate_on_refresh, + ) + _apply_scene_edit_plan_to_scene_graph( + scene_graph=updated_scene_graph, + scene_edit_plan=scene_edit_plan, + ) + return updated_scene_graph + + +def _apply_scene_edit_plan_to_scene_graph( + *, + scene_graph: SceneGraph, + scene_edit_plan: SceneEditPlan, +) -> None: + """Apply the target graph updates implied by add and move operations.""" + deleted_object_ids: set[str] = set() + added_object_ids: list[str] = [] + on_parent_updates: list[tuple[str, str]] = [] + planar_relation_updates: list[tuple[str, PlanarRelationType, str]] = [] + for operation in scene_edit_plan.operations: + if operation.op == "delete": + if operation.object_id is not None: + deleted_object_ids.add(operation.object_id) + continue + if operation.object_id is None: + raise ValueError("Add and move operations must have an object_id.") + if operation.op == "add": + added_object_ids.append(operation.object_id) + if operation.target_id is None or operation.relation is None: + continue + if operation.relation == "on": + on_parent_updates.append((operation.object_id, operation.target_id)) + continue + planar_relation_updates.append( + (operation.object_id, operation.relation, operation.target_id) + ) + + # Apply all graph changes atomically so intermediate edit states need not be valid. + scene_graph.apply_updates( + deleted_object_ids=deleted_object_ids, + added_object_ids=added_object_ids, + on_parent_updates=on_parent_updates, + planar_relation_updates=planar_relation_updates, + ) def _simplify_scene_info(scene: Scene) -> dict[str, object]: diff --git a/tests/gen_sim/scene_engine/test_scene_edit_plan.py b/tests/gen_sim/scene_engine/test_scene_edit_plan.py index c01049eb6..30c11c995 100644 --- a/tests/gen_sim/scene_engine/test_scene_edit_plan.py +++ b/tests/gen_sim/scene_engine/test_scene_edit_plan.py @@ -31,8 +31,13 @@ ) from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject from embodichain.gen_sim.scene_engine.pipeline.editing.scene_edit_understanding import ( + _apply_scene_edit_plan_to_scene_graph, + _build_updated_scene_graph, _parse_scene_edit_operations, ) +from embodichain.gen_sim.scene_engine.pipeline.editing.scene_edit_asset_preparation import ( + prepare_scene_edit_assets, +) def _scene_and_graph() -> tuple[Scene, SceneGraph]: @@ -202,3 +207,181 @@ def test_scene_edit_plan_requires_deleting_all_children_of_a_deleted_parent() -> scene_graph=scene_graph, operations=[SceneEditOperation(op="delete", object_id="book_001")], ) + + +def test_scene_edit_plan_requires_a_position_for_move_operations() -> None: + scene, scene_graph = _scene_and_graph() + + with pytest.raises(ValueError, match="must specify target_id and relation"): + SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[SceneEditOperation(op="move", object_id="book_001")], + ) + + +def test_scene_edit_asset_preparation_skips_plans_without_adds() -> None: + scene, scene_graph = _scene_and_graph() + plan = SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[ + SceneEditOperation( + op="move", + object_id="book_001", + target_id="table", + relation="on", + ) + ], + ) + + prepared_scene = prepare_scene_edit_assets( + scene=scene, + scene_edit_plan=plan, + ) + + assert prepared_scene is scene + + +def test_scene_edit_graph_builder_copies_the_pre_edit_graph() -> None: + scene, scene_graph = _scene_and_graph() + plan = SceneEditPlan(scene=scene, scene_graph=scene_graph) + + updated_scene_graph = _build_updated_scene_graph( + scene_graph=scene_graph, + scene_edit_plan=plan, + ) + + assert updated_scene_graph is not scene_graph + assert updated_scene_graph.nodes is not scene_graph.nodes + assert updated_scene_graph.relations is not scene_graph.relations + assert updated_scene_graph.to_dict() == scene_graph.to_dict() + + +def test_scene_edit_graph_builder_removes_deleted_nodes() -> None: + scene, scene_graph = _scene_and_graph() + plan = SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[ + SceneEditOperation(op="delete", object_id="orange_001"), + SceneEditOperation(op="delete", object_id="book_001"), + ], + ) + + updated_scene_graph = _build_updated_scene_graph( + scene_graph=scene_graph, + scene_edit_plan=plan, + ) + + assert set(updated_scene_graph.node_by_id()) == {"table"} + assert set(scene_graph.node_by_id()) == {"table", "book_001", "orange_001"} + + +def test_scene_edit_graph_builder_adds_unpositioned_objects_on_the_table() -> None: + scene, scene_graph = _scene_and_graph() + plan = SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[ + SceneEditOperation( + op="add", + object_id="cup_001", + category="cup", + name="green cup", + description="A small green ceramic cup.", + ) + ], + ) + + updated_scene_graph = _build_updated_scene_graph( + scene_graph=scene_graph, + scene_edit_plan=plan, + ) + + added_node = updated_scene_graph.node_by_id()["cup_001"] + assert added_node.parent_id == "table" + assert added_node.parent_relation == "on" + + +def test_scene_edit_graph_builder_updates_move_on_parent() -> None: + scene, scene_graph = _scene_and_graph() + plan = SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[ + SceneEditOperation( + op="move", + object_id="orange_001", + target_id="table", + relation="on", + ) + ], + ) + + updated_scene_graph = _build_updated_scene_graph( + scene_graph=scene_graph, + scene_edit_plan=plan, + ) + + assert updated_scene_graph.node_by_id()["orange_001"].parent_id == "table" + + +def test_scene_edit_graph_builder_adds_planar_relation_with_target_parent() -> None: + scene, scene_graph = _scene_and_graph() + plan = SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[ + SceneEditOperation( + op="add", + object_id="cup_001", + target_id="book_001", + relation="right_of", + category="cup", + name="green cup", + description="A small green ceramic cup.", + ) + ], + ) + + updated_scene_graph = _build_updated_scene_graph( + scene_graph=scene_graph, + scene_edit_plan=plan, + ) + + assert updated_scene_graph.node_by_id()["cup_001"].parent_id == "table" + assert any( + relation.source_id == "cup_001" + and relation.relation == "right_of" + and relation.target_id == "book_001" + for relation in updated_scene_graph.relations + ) + + +def test_scene_edit_plan_application_adds_new_nodes_before_relationship_updates() -> ( + None +): + scene, scene_graph = _scene_and_graph() + plan = SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[ + SceneEditOperation( + op="add", + object_id="cup_001", + target_id="book_001", + relation="right_of", + category="cup", + name="green cup", + description="A small green ceramic cup.", + ) + ], + ) + + _apply_scene_edit_plan_to_scene_graph( + scene_graph=scene_graph, + scene_edit_plan=plan, + ) + + assert scene_graph.node_by_id()["cup_001"].parent_id == "table" From 809203a7fd8d15c6c8831fd20c87d219a427298e Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Wed, 12 Aug 2026 20:50:46 +0800 Subject: [PATCH 10/29] move client try catch outside the understand_scene --- .../gen_sim/scene_engine/pipeline/generate.py | 22 ++++++++++++++----- .../generation/scene_understanding.py | 21 +++++++----------- 2 files changed, 24 insertions(+), 19 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/pipeline/generate.py b/embodichain/gen_sim/scene_engine/pipeline/generate.py index 32aa6550f..1df9027f9 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generate.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generate.py @@ -25,6 +25,9 @@ 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.pipeline.generation.scene_understanding import ( understand_scene, @@ -51,12 +54,19 @@ def generate_scene_from_image( # 1. Scene Understanding log_info("Starting Scene Understanding") - scene, scene_graph = understand_scene( - scene=scene, - image_path=image_path, - output_root=resolved_output_root, - vlm_client=vlm_client, - ) + # Load .env settings and fail if the Image Segmentation Server is unavailable. + image_segmentation_client = ImageSegmentationClient.from_dotenv() + try: + image_segmentation_client.check_health() + scene, scene_graph = understand_scene( + scene=scene, + image_path=image_path, + output_root=resolved_output_root, + vlm_client=vlm_client, + image_segmentation_client=image_segmentation_client, + ) + finally: + image_segmentation_client.close() # Close the session after scene understanding. log_info("Completed Scene Understanding") # 2. Objects + Coarse Layout Generation diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py index cae1c612f..f2103a945 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py @@ -155,6 +155,7 @@ def understand_scene( output_root: str | Path, *, vlm_client: OpenAICompatibleVLM, + image_segmentation_client: ImageSegmentationClient, json_max_attempts: int = 3, ) -> tuple[Scene, SceneGraph]: @@ -173,19 +174,13 @@ def understand_scene( json_max_attempts=json_max_attempts, ) - # Load .env settings and fail if the Image Segmentation Server is unavailable. - image_segmentation_client = ImageSegmentationClient.from_dotenv() - try: - image_segmentation_client.check_health() # Error raising will happen internally. - _segment_scene( - image_path=resolved_image_path, - stage_output_root=stage_output_root, - scene=scene, - vlm_client=vlm_client, - image_segmentation_client=image_segmentation_client, - ) - finally: - image_segmentation_client.close() # Kill the session to avoid resource leaks. + _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, + ) # Use the segmented image to initialize the scene graph # with the help of the VLM client. From 15ec06dfb659cbb50817558aaaea898c13d97437 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Thu, 13 Aug 2026 17:07:06 +0800 Subject: [PATCH 11/29] finish part of the scene edit: from user prompt to simreadyed asset, but no real-world scale... --- .../scene_engine/core/scene_edit_plan.py | 11 + .../gen_sim/scene_engine/core/scene_graph.py | 38 ++- .../gen_sim/scene_engine/pipeline/edit.py | 64 ++-- .../editing/scene_edit_asset_preparation.py | 309 ++++++++++++++++++ .../editing/scene_edit_understanding.py | 70 +++- .../pipeline/generation/scene_generation.py | 6 +- .../utils/image_segmentation_utils.py | 42 +++ .../pipeline/utils/scene_importer.py | 14 +- ...ene_processor.py => simready_processor.py} | 14 +- .../scene_engine/test_scene_edit_plan.py | 142 +++++++- 10 files changed, 646 insertions(+), 64 deletions(-) create mode 100644 embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py rename embodichain/gen_sim/scene_engine/pipeline/utils/{simready_scene_processor.py => simready_processor.py} (97%) diff --git a/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py b/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py index a7ac152d7..83ad5bd79 100644 --- a/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py +++ b/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py @@ -23,6 +23,7 @@ from embodichain.gen_sim.scene_engine.core.scene_graph import ( SceneConstraintType, SceneGraph, + TableRegion, TABLE_OBJECT_ID, ) @@ -39,6 +40,7 @@ class SceneEditOperation: object_id: str | None = None target_id: str | None = None relation: SceneConstraintType | None = None + table_region: TableRegion | None = None category: str | None = None name: str | None = None description: str | None = None @@ -50,6 +52,7 @@ def to_dict(self) -> dict[str, object]: "object_id": self.object_id, "target_id": self.target_id, "relation": self.relation, + "table_region": self.table_region, "category": self.category, "name": self.name, "description": self.description, @@ -82,6 +85,7 @@ def validate(self) -> None: # - move and delete identify one existing non-table object with object_id. # - add carries generated object_id plus non-empty category, name, and description. # - move always supplies target_id and relation; add may omit both. + # - table_region is only valid with target_id=table and relation=on. # - target_id and relation are otherwise supplied together or both absent. # - every target is from the pre-edit scene; new and deleted objects are invalid targets. # - an existing object has at most one move or delete operation in one plan. @@ -156,6 +160,7 @@ def _validate_operation( for value in ( operation.target_id, operation.relation, + operation.table_region, operation.category, operation.name, operation.description, @@ -213,6 +218,12 @@ def _validate_position_reference( ) -> None: if (operation.target_id is None) != (operation.relation is None): raise ValueError("target_id and relation must be specified together.") + if operation.table_region is not None and ( + operation.target_id != TABLE_OBJECT_ID or operation.relation != "on" + ): + raise ValueError( + "table_region requires target_id='table' and relation='on'." + ) if operation.target_id is None: return # Targets come only from the pre-edit scene, so new objects cannot be targets. diff --git a/embodichain/gen_sim/scene_engine/core/scene_graph.py b/embodichain/gen_sim/scene_engine/core/scene_graph.py index 514176211..c5c49c625 100644 --- a/embodichain/gen_sim/scene_engine/core/scene_graph.py +++ b/embodichain/gen_sim/scene_engine/core/scene_graph.py @@ -21,7 +21,7 @@ TABLE_OBJECT_ID = "table" -# 9-grid table regions, treat the table as a 3x3 grid. +# Static type constraint for the nine regions of the tabletop 3x3 grid. TableRegion = Literal[ "left_back", "back_center", @@ -33,6 +33,20 @@ "front_center", "right_front", ] +# Runtime membership set for validating serialized and user-provided regions. +TABLE_REGIONS = frozenset( + { + "left_back", + "back_center", + "right_back", + "left_center", + "center", + "right_center", + "left_front", + "front_center", + "right_front", + } +) # A on B, then B is the parent node of A. SupportRelationType = Literal["on"] @@ -55,6 +69,8 @@ def __post_init__(self) -> None: """Validate local node fields before graph-level checks.""" if not self.object_id: raise ValueError("object_id must be non-empty.") + if self.table_region not in {None, *TABLE_REGIONS}: + raise ValueError("table_region is invalid.") # If the node is the table. if self.object_id == TABLE_OBJECT_ID: if self.parent_id is not None: @@ -165,7 +181,7 @@ def apply_updates( *, deleted_object_ids: set[str], added_object_ids: list[str], - on_parent_updates: list[tuple[str, str]], + on_parent_updates: list[tuple[str, str, TableRegion | None]], planar_relation_updates: list[tuple[str, PlanarRelationType, str]], ) -> None: """Apply one atomic batch of node and relationship updates.""" @@ -210,8 +226,12 @@ def apply_updates( ) # Apply support-parent changes before planar updates need the final parent. - for object_id, parent_id in on_parent_updates: - self._set_on_parent(object_id=object_id, parent_id=parent_id) + for object_id, parent_id, table_region in on_parent_updates: + self._set_on_parent( + object_id=object_id, + parent_id=parent_id, + table_region=table_region, + ) # Resolve chained planar parent inheritance before adding final relations. self._resolve_planar_parent_updates(planar_relation_updates) @@ -228,7 +248,13 @@ def apply_updates( # Normalize inverse relations and reject invalid final graph constraints. self.refresh() - def _set_on_parent(self, *, object_id: str, parent_id: str) -> None: + def _set_on_parent( + self, + *, + object_id: str, + parent_id: str, + table_region: TableRegion | None = None, + ) -> None: """Replace one node's support parent and stale planar constraints.""" nodes_by_id = self.node_by_id() if object_id == TABLE_OBJECT_ID: @@ -243,7 +269,7 @@ def _set_on_parent(self, *, object_id: str, parent_id: str) -> None: node = nodes_by_id[object_id] node.parent_id = parent_id node.parent_relation = "on" - node.table_region = None + node.table_region = table_region self._clear_incident_planar_relations(object_id) def _resolve_planar_parent_updates( diff --git a/embodichain/gen_sim/scene_engine/pipeline/edit.py b/embodichain/gen_sim/scene_engine/pipeline/edit.py index a13ef17fd..6ba7898d1 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/edit.py +++ b/embodichain/gen_sim/scene_engine/pipeline/edit.py @@ -16,9 +16,17 @@ from __future__ import annotations -import json from pathlib import Path +from embodichain.gen_sim.scene_engine.clients.geometry_generation import ( + GeometryGenerationClient, +) +from embodichain.gen_sim.scene_engine.clients.image_generation import ( + ImageGenerationClient, +) +from embodichain.gen_sim.scene_engine.clients.image_segmentation import ( + ImageSegmentationClient, +) from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( OpenAICompatibleVLM, ) @@ -28,6 +36,9 @@ from embodichain.gen_sim.scene_engine.pipeline.editing.scene_edit_understanding import ( understand_scene_edit, ) +from embodichain.gen_sim.scene_engine.pipeline.editing.scene_edit_asset_preparation import ( + prepare_scene_edit_assets, +) from embodichain.utils.logger import log_info @@ -37,20 +48,14 @@ def edit_scene( edit_prompt: str, ) -> None: """Apply one text edit instruction to an existing Scene Engine output.""" + resolved_output_root = Path(output_root).expanduser().resolve() + resolved_output_root.mkdir(parents=True, exist_ok=True) # Initialize the VLM client that will interpret the edit instruction. vlm_client = OpenAICompatibleVLM.from_dotenv() scene_importer = SceneExportImporter(output_root=output_root) # Validate scene_export, write scene.json, and return Scene; failures raise before editing. scene, scene_graph = scene_importer.import_scene_and_graph() - # Only for debug. - # print( - # json.dumps( - # {"scene": scene.to_dict(), "scene_graph": scene_graph.to_dict()}, - # indent=2, - # ensure_ascii=False, - # ) - # ) # 1. Edit Understanding # Will return an already checked scene edit plan @@ -63,22 +68,31 @@ def edit_scene( vlm_client=vlm_client, ) log_info("Completed Edit Understanding") - # Only for debug. - # print( - # json.dumps( - # {"updated scene graph": updated_scene_graph.to_dict()}, - # indent=2, - # ensure_ascii=False, - # ) - # ) - # 2. Prepare Objects. - log_info("Preparing Objects if necessary") - scene = prepare_scene_edit_assets( - scene=scene, - scene_edit_plan=scene_edit_plan, - ) - log_info("Completed Preparing Objects") + # 2. Prepare Objects + log_info("Starting Objects Preparation") + # Initialize all the clients and then check. + image_generation_client = ImageGenerationClient.from_dotenv() + geometry_generation_client = GeometryGenerationClient.from_dotenv() + image_segmentation_client = ImageSegmentationClient.from_dotenv() + try: + image_generation_client.check_health() + geometry_generation_client.check_health() + image_segmentation_client.check_health() + # Return a list of added SceneObjects assets. + # Now do not support editing the table. + added_assets = prepare_scene_edit_assets( + scene_edit_plan=scene_edit_plan, + output_root=resolved_output_root, + image_generation_client=image_generation_client, + geometry_generation_client=geometry_generation_client, + image_segmentation_client=image_segmentation_client, + ) + finally: + image_generation_client.close() + geometry_generation_client.close() + image_segmentation_client.close() + log_info("Completed Objects Preparation") # 3. Layout Editing log_info("Starting Layout Editing") @@ -96,4 +110,4 @@ def edit_scene( log_info("Starting Scene Export") log_info("Completed Scene Export") - raise NotImplementedError("Scene editing is not implemented yet.") + return None diff --git a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py new file mode 100644 index 000000000..70c02d1a7 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py @@ -0,0 +1,309 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF 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 shutil + +from PIL import Image + +from embodichain.gen_sim.scene_engine.clients.geometry_generation import ( + GeometryGenerationClient, +) +from embodichain.gen_sim.scene_engine.clients.image_generation import ( + ImageGenerationClient, +) +from embodichain.gen_sim.scene_engine.clients.image_segmentation import ( + ImageSegmentationClient, +) +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_edit_plan import SceneEditPlan +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.pipeline.utils.image_segmentation_utils import ( + MaskCandidate, + build_mask_candidates, + invert_mask_if_foreground_is_off_center, + save_binary_mask, + union_overlapping_mask_candidates, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor import ( + SimReadyProcessor, +) + +__all__ = ["prepare_scene_edit_assets"] + + +@dataclass(frozen=True) +class _AddedAssetInfo: + """Semantic information needed while preparing one newly added asset.""" + + object_id: str + category: str + name: str + description: str + + +def prepare_scene_edit_assets( + *, + scene_edit_plan: SceneEditPlan, + output_root: str | Path, + image_generation_client: ImageGenerationClient, + geometry_generation_client: GeometryGenerationClient, + image_segmentation_client: ImageSegmentationClient, +) -> list[SceneObject]: + """Prepare and return SimReady assets required by add operations.""" + # Prepare descriptions for all newly added objects. + added_asset_descriptions = _collect_added_asset_descriptions(scene_edit_plan) + # Skip asset generation when the edit plan only moves or deletes existing objects. + if not added_asset_descriptions: + return [] + + # Recreate this stage only when new assets need image, segmentation, and geometry outputs. + stage_output_root = ( + Path(output_root).expanduser().resolve() / "scene_editing" / "asset_preparation" + ) + if stage_output_root.exists(): + shutil.rmtree(stage_output_root) + stage_output_root.mkdir(parents=True, exist_ok=True) + + generated_asset_images = _generate_added_asset_images( + added_asset_descriptions=added_asset_descriptions, + stage_output_root=stage_output_root, + image_generation_client=image_generation_client, + ) + generated_asset_masks = _segment_generated_added_asset_images( + added_asset_descriptions=added_asset_descriptions, + generated_asset_images=generated_asset_images, + stage_output_root=stage_output_root, + image_segmentation_client=image_segmentation_client, + ) + + generated_asset_glbs = _generate_added_assets_coarse_geometry( + generated_asset_images=generated_asset_images, + generated_asset_masks=generated_asset_masks, + stage_output_root=stage_output_root, + geometry_generation_client=geometry_generation_client, + ) + # Build a list of added SceneObjects. + added_assets = _build_added_scene_objects( + added_asset_descriptions=added_asset_descriptions, + generated_asset_glbs=generated_asset_glbs, + ) + # The temporary scene contains only new assets because the existing table is reused. + tmp_scene = Scene(objects=added_assets) + simready_processor = SimReadyProcessor( + scene=tmp_scene, + coarse_layout_by_id=_coarse_layouts_by_id(generated_asset_glbs), + coarse_geometry_root=stage_output_root / "coarse_geometry", + simready_geometry_root=stage_output_root / "simready_geometry", + ) + # process_assets() validates and processes assets only; it does not require a table. + simready_processor.process_assets() + # Canonical GLBs use identity edit-time poses; layout editing sets them later. + _reset_added_asset_layouts(added_assets) + return added_assets + + +def _build_added_scene_objects( + *, + added_asset_descriptions: list[_AddedAssetInfo], + generated_asset_glbs: list[tuple[str, Path]], +) -> list[SceneObject]: + """Build temporary SceneObjects from generated coarse GLBs.""" + glbs_by_id = dict(generated_asset_glbs) + if len(glbs_by_id) != len(generated_asset_glbs): + raise ValueError("Generated asset GLBs must use unique object ids.") + assets: list[SceneObject] = [] + for asset_info in added_asset_descriptions: + glb_path = glbs_by_id.get(asset_info.object_id) + if glb_path is None: + raise ValueError(f"Generated asset {asset_info.object_id!r} has no GLB.") + assets.append( + SceneObject( + id=asset_info.object_id, + kind="asset", + category=asset_info.category, + name=asset_info.name, + description=asset_info.description, + simready_glb_path=str(glb_path), + ) + ) + return assets + + +def _reset_added_asset_layouts(added_assets: list[SceneObject]) -> None: + """Reset added asset poses after SimReady canonicalization.""" + for asset in added_assets: + asset.rot = [0.0, 0.0, 0.0] + asset.pos = [0.0, 0.0, 0.0] + asset.scale = [1.0, 1.0, 1.0] + + +def _coarse_layouts_by_id( + generated_asset_glbs: list[tuple[str, Path]], +) -> dict[str, dict[str, object]]: + """Build edit-time layouts with fixed identity poses and scale.""" + return { + object_id: { + "rot": [0.0, 0.0, 0.0], + "pos": [0.0, 0.0, 0.0], + "scale": [1.0, 1.0, 1.0], + } + for object_id, _ in generated_asset_glbs + } + + +def _collect_added_asset_descriptions( + scene_edit_plan: SceneEditPlan, +) -> list[_AddedAssetInfo]: + """Return complete semantic information for add operations in plan order.""" + # Existing assets already have SimReady GLBs, so only add operations need assets. + added_asset_descriptions: list[_AddedAssetInfo] = [] + for operation in scene_edit_plan.operations: + if operation.op != "add": + continue + if ( + operation.object_id is None + or operation.category is None + or operation.name is None + or operation.description is None + ): + raise ValueError( + "Add operations must have an object_id, category, name, and description." + ) + added_asset_descriptions.append( + _AddedAssetInfo( + object_id=operation.object_id, + category=operation.category, + name=operation.name, + description=operation.description, + ) + ) + return added_asset_descriptions + + +def _generate_added_asset_images( + *, + added_asset_descriptions: list[_AddedAssetInfo], + stage_output_root: Path, + image_generation_client: ImageGenerationClient, +) -> list[tuple[str, Path]]: + """Generate one stable PNG for each new object description.""" + # Prepare a list. + generated_asset_images: list[tuple[str, Path]] = [] + # Create a subdir. + image_output_root = stage_output_root / "generated_images" + image_output_root.mkdir(parents=True, exist_ok=True) + + for asset_info in added_asset_descriptions: + object_id = asset_info.object_id + # Stable object IDs preserve the image-to-asset mapping across later stages. + image_path = image_generation_client.generate_image_by_prompt( + prompt=asset_info.description, + output_path=image_output_root / f"{object_id}.png", + ) + generated_asset_images.append((object_id, image_path)) + return generated_asset_images + + +def _segment_generated_added_asset_images( + *, + added_asset_descriptions: list[_AddedAssetInfo], + generated_asset_images: list[tuple[str, Path]], + stage_output_root: Path, + image_segmentation_client: ImageSegmentationClient, +) -> list[tuple[str, Path]]: + """Segment each generated image with its description and return binary masks.""" + asset_info_by_id = { + asset_info.object_id: asset_info for asset_info in added_asset_descriptions + } + if len(asset_info_by_id) != len(added_asset_descriptions): + raise ValueError("Added asset descriptions must use unique object ids.") + + masks_output_root = stage_output_root / "generated_masks" + masks_output_root.mkdir(parents=True, exist_ok=True) + generated_asset_masks: list[tuple[str, Path]] = [] + for object_id, image_path in generated_asset_images: + asset_info = asset_info_by_id.get(object_id) + if asset_info is None: + raise ValueError(f"Generated image {object_id!r} has no description.") + + candidates: list[MaskCandidate] = [] + # Retry with simpler semantic prompts when the detailed description is not found. + for prompt in (asset_info.description, asset_info.name, asset_info.category): + candidates = union_overlapping_mask_candidates( + build_mask_candidates( + image_segmentation_client.segment_single_object( + image_path=image_path, + prompt=prompt, + ) + ), + min_iou=0.8, + ) + if candidates: + break + # A single generated object may still yield multiple SAM3 candidates; use the first one. + if not candidates: + raise ValueError( + f"Generated asset {object_id!r} produced no segmentation candidates." + ) + with Image.open(image_path) as image: + image_size = image.size + mask_path = save_binary_mask( + invert_mask_if_foreground_is_off_center(candidates[0]), + image_size=image_size, + output_path=masks_output_root / f"{object_id}_mask.png", + ) + generated_asset_masks.append((object_id, mask_path)) + return generated_asset_masks + + +def _generate_added_assets_coarse_geometry( + *, + generated_asset_images: list[tuple[str, Path]], + generated_asset_masks: list[tuple[str, Path]], + stage_output_root: Path, + geometry_generation_client: GeometryGenerationClient, +) -> list[tuple[str, Path]]: + """Generate one coarse GLB for each generated image and binary mask.""" + masks_by_id = dict(generated_asset_masks) + if len(masks_by_id) != len(generated_asset_masks): + raise ValueError("Generated asset masks must use unique object ids.") + if set(masks_by_id) != {object_id for object_id, _ in generated_asset_images}: + raise ValueError( + "Generated asset images and masks must have matching object ids." + ) + + geometry_output_root = stage_output_root / "coarse_geometry" + geometry_output_root.mkdir(parents=True, exist_ok=True) + generated_asset_glbs: list[tuple[str, Path]] = [] + for object_id, image_path in generated_asset_images: + # Each generated object has its own color image, so it needs an individual request. + geometry_generation_client.generate_objects( + image_path=image_path, + object_masks=[(object_id, masks_by_id[object_id])], + output_root=geometry_output_root, + ) + glb_path = geometry_output_root / f"{object_id}.glb" + if not glb_path.is_file(): + raise FileNotFoundError( + f"Geometry generation did not produce a GLB for {object_id!r}: {glb_path}" + ) + generated_asset_glbs.append((object_id, glb_path)) + return generated_asset_glbs diff --git a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py index 1eb3a857f..e66fd8653 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py @@ -29,6 +29,8 @@ SceneGraph, SceneGraphNode, SceneGraphRelation, + TABLE_REGIONS, + TableRegion, ) from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( OpenAICompatibleVLM, @@ -52,6 +54,17 @@ For every move and every positioned add, target_id must be an Existing object ID and relation must be one of on, left_of, right_of, in_front_of, or behind. +When the target is the tabletop, use target_id "table", relation "on", and set +table_region to one of left_back, back_center, right_back, left_center, center, +right_center, left_front, front_center, or right_front. Do not use a planar +relation with the table. For non-table placement, table_region must be null. +If an add operation has no target_id and relation, it is placed on the table by +default and table_region must be null. +In the tabletop 9-grid, smaller x means left, larger x means right, smaller y +means back, and larger y means front: left_back is the upper-left/back cell, +back_center is the upper-center/back cell, right_back is the upper-right/back +cell, left_center/center/right_center are the middle row, and +left_front/front_center/right_front are the lower/front row. Do not position a new object relative to another newly added object. Each existing object's center_xy is its center position [x, y] in the @@ -67,8 +80,8 @@ table, or relations to any object. Return JSON only: no Markdown, comments, or prose. Every operation must contain -exactly these fields: op, object_id, target_id, relation, category, name, and -description. Use null for every field that does not apply to an operation: +exactly these fields: op, object_id, target_id, relation, table_region, category, +name, and description. Use null for every field that does not apply: { "operations": [ { @@ -76,6 +89,7 @@ "object_id": "bottle_001", "target_id": "book_001", "relation": "right_of", + "table_region": null, "category": null, "name": null, "description": null @@ -85,6 +99,7 @@ "object_id": "cup_001", "target_id": null, "relation": null, + "table_region": null, "category": null, "name": null, "description": null @@ -92,8 +107,9 @@ { "op": "add", "object_id": null, - "target_id": null, - "relation": null, + "target_id": "table", + "relation": "on", + "table_region": "back_center", "category": "orange", "name": "small orange", "description": "small round orange with a textured peel" @@ -103,9 +119,20 @@ "object_id": null, "target_id": "book_001", "relation": "right_of", + "table_region": null, "category": "orange", "name": "small orange", "description": "small round orange with a textured peel" + }, + { + "op": "add", + "object_id": null, + "target_id": null, + "relation": null, + "table_region": null, + "category": "banana", + "name": "yellow banana", + "description": "curved yellow banana with a green stem" } ] } @@ -128,7 +155,10 @@ def understand_scene_edit( if json_max_attempts < 1: raise ValueError("json_max_attempts must be at least 1.") # Give the VLM only the scene metadata needed to identify existing objects. - simplified_scene_info = _simplify_scene_info(scene=scene) + simplified_scene_info = _simplify_scene_info( + scene=scene, + scene_graph=scene_graph, + ) operations = _vlm_understand_scene_edit( scene=scene, edit_prompt=edit_prompt, @@ -192,7 +222,7 @@ def _apply_scene_edit_plan_to_scene_graph( """Apply the target graph updates implied by add and move operations.""" deleted_object_ids: set[str] = set() added_object_ids: list[str] = [] - on_parent_updates: list[tuple[str, str]] = [] + on_parent_updates: list[tuple[str, str, TableRegion | None]] = [] planar_relation_updates: list[tuple[str, PlanarRelationType, str]] = [] for operation in scene_edit_plan.operations: if operation.op == "delete": @@ -206,7 +236,13 @@ def _apply_scene_edit_plan_to_scene_graph( if operation.target_id is None or operation.relation is None: continue if operation.relation == "on": - on_parent_updates.append((operation.object_id, operation.target_id)) + on_parent_updates.append( + ( + operation.object_id, + operation.target_id, + operation.table_region, + ) + ) continue planar_relation_updates.append( (operation.object_id, operation.relation, operation.target_id) @@ -221,8 +257,15 @@ def _apply_scene_edit_plan_to_scene_graph( ) -def _simplify_scene_info(scene: Scene) -> dict[str, object]: +def _simplify_scene_info( + *, + scene: Scene, + scene_graph: SceneGraph, +) -> dict[str, object]: """Return the object metadata needed for edit instruction resolution.""" + table_regions_by_id = { + node.object_id: node.table_region for node in scene_graph.nodes + } return { "existing_object_ids": [scene_object.id for scene_object in scene.objects], "objects": [ @@ -232,6 +275,7 @@ def _simplify_scene_info(scene: Scene) -> dict[str, object]: "name": scene_object.name, "description": scene_object.description, "center_xy": scene_object.center_xy, + "table_region": table_regions_by_id.get(scene_object.id), } for scene_object in scene.objects ], @@ -302,6 +346,7 @@ def _parse_scene_edit_operations( "object_id", "target_id", "relation", + "table_region", "category", "name", "description", @@ -345,6 +390,7 @@ def _parse_scene_edit_operations( value.get("target_id"), field_name="target_id" ), relation=_optional_relation(value.get("relation")), + table_region=_optional_table_region(value.get("table_region")), category=category, name=_optional_string(value.get("name"), field_name="name"), description=_optional_string( @@ -387,3 +433,11 @@ def _optional_relation(value: object) -> str | None: if value not in {"on", "left_of", "right_of", "in_front_of", "behind"}: raise ValueError("Scene edit operation relation is invalid.") return value + + +def _optional_table_region(value: object) -> TableRegion | None: + if value is None: + return None + if value not in TABLE_REGIONS: + raise ValueError("Scene edit operation table_region is invalid.") + return value diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py index 1cfe3f592..d8d82b19b 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py @@ -48,8 +48,8 @@ quaternion_wxyz_to_euler_xyz_degrees, transform_matrix_to_layout_object, ) -from embodichain.gen_sim.scene_engine.pipeline.utils.simready_scene_processor import ( - SimReadySceneProcessor, +from embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor import ( + SimReadyProcessor, ) from embodichain.gen_sim.scene_engine.pipeline.utils.table_support_surface import ( TableSupportSurfaceDetector, @@ -105,7 +105,7 @@ def generate_scene_and_refine( coarse_layout_by_id = { layout_object["id"]: layout_object for layout_object in coarse_layout } - simready_processor = SimReadySceneProcessor( + simready_processor = SimReadyProcessor( scene=scene, coarse_layout_by_id=coarse_layout_by_id, coarse_geometry_root=coarse_geometry_output_root, diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py b/embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py index e0c3f772a..ccf9af949 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py @@ -94,6 +94,35 @@ def decode_rle_mask(mask_rle: dict[str, Any]) -> Image.Image: return Image.frombytes("L", (width, height), bytes(pixels)) +def invert_mask_if_foreground_is_off_center(candidate: MaskCandidate) -> MaskCandidate: + """Invert a mask when its foreground is less concentrated at image center. + + This heuristic is intended for generated single-object images, where the + object is expected near the center and SAM3 may return its background. + """ + mask = decode_rle_mask(candidate.mask_rle) + width, height = mask.size + left, top = width // 6, height // 6 + right, bottom = width - left, height - top + center_mask = mask.crop((left, top, right, bottom)) + + center_foreground_ratio = _foreground_ratio(center_mask) + total_foreground_pixels = _foreground_pixel_count(mask) + outside_foreground_pixels = total_foreground_pixels - _foreground_pixel_count( + center_mask + ) + outside_pixel_count = width * height - center_mask.width * center_mask.height + outside_foreground_ratio = outside_foreground_pixels / outside_pixel_count + if center_foreground_ratio >= outside_foreground_ratio: + return candidate + + inverted_mask = mask.point(lambda value: 0 if value else 255) + return MaskCandidate( + index=candidate.index, + mask_rle=_encode_binary_mask_rle(inverted_mask), + ) + + def union_overlapping_mask_candidates( candidates: list[MaskCandidate], *, @@ -361,6 +390,19 @@ def _mask_iou(first_mask: Image.Image, second_mask: Image.Image) -> float: return intersection.histogram()[255] / union_pixels +def _foreground_pixel_count(mask: Image.Image) -> int: + """Return the number of white pixels in one binary mask.""" + return mask.convert("L").histogram()[255] + + +def _foreground_ratio(mask: Image.Image) -> float: + """Return the white-pixel ratio in one non-empty image region.""" + pixel_count = mask.width * mask.height + if pixel_count == 0: + raise ValueError("Mask region must contain at least one pixel.") + return _foreground_pixel_count(mask) / pixel_count + + def _encode_binary_mask_rle(mask: Image.Image) -> dict[str, Any]: binary_mask = mask.convert("L").point( lambda value: 255 if value else 0 diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py index 1ae6cdb1b..a59bc7bdf 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py @@ -28,6 +28,7 @@ SceneGraph, SceneGraphNode, SceneGraphRelation, + TABLE_REGIONS, ) from embodichain.gen_sim.scene_engine.core.scene_object import ( ObjectPhysics, @@ -203,18 +204,7 @@ def _scene_graph_node_from_data(value: object) -> SceneGraphNode: raise ValueError("Scene graph node ids must be strings or null.") if parent_relation not in {None, "on"}: raise ValueError("Scene graph parent_relation must be 'on' or null.") - if table_region not in { - None, - "left_back", - "back_center", - "right_back", - "left_center", - "center", - "right_center", - "left_front", - "front_center", - "right_front", - }: + if table_region is not None and table_region not in TABLE_REGIONS: raise ValueError("Scene graph table_region is invalid.") return SceneGraphNode( object_id=object_id, diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_scene_processor.py b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py similarity index 97% rename from embodichain/gen_sim/scene_engine/pipeline/utils/simready_scene_processor.py rename to embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py index 40b4a6f94..6b5a09b0b 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_scene_processor.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py @@ -53,7 +53,7 @@ @dataclass(frozen=True) -class SimReadySceneProcessorConfig: +class SimReadyProcessorConfig: """Object-category policy for SimReady mesh canonicalization.""" upright_container_id_tokens: frozenset[str] = frozenset( @@ -61,8 +61,8 @@ class SimReadySceneProcessorConfig: ) # Object-id tokens that enable upright-container standardization. -class SimReadySceneProcessor: - """Create SimReady GLBs and layouts for one table and its scene assets.""" +class SimReadyProcessor: + """Create SimReady GLBs and layouts for scene objects.""" def __init__( self, @@ -71,7 +71,7 @@ def __init__( coarse_layout_by_id: dict[str, dict[str, object]], coarse_geometry_root: str | Path, simready_geometry_root: str | Path, - config: SimReadySceneProcessorConfig | None = None, + config: SimReadyProcessorConfig | None = None, ) -> None: self.scene = scene self.coarse_layout_by_id = coarse_layout_by_id @@ -81,7 +81,7 @@ def __init__( ) 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() + self.config = config if config is not None else SimReadyProcessorConfig() if not self.config.upright_container_id_tokens: raise ValueError("upright_container_id_tokens must not be empty.") @@ -317,8 +317,8 @@ def _standardize_bottle_z_up(mesh: trimesh.Trimesh) -> np.ndarray: 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) + upper_volume = SimReadyProcessor._convex_hull_volume(upper_points) + lower_volume = SimReadyProcessor._convex_hull_volume(lower_points) # Bottles usually have a smaller top (neck) than bottom; flip if necessary. if upper_volume > lower_volume: diff --git a/tests/gen_sim/scene_engine/test_scene_edit_plan.py b/tests/gen_sim/scene_engine/test_scene_edit_plan.py index 30c11c995..5579e8d55 100644 --- a/tests/gen_sim/scene_engine/test_scene_edit_plan.py +++ b/tests/gen_sim/scene_engine/test_scene_edit_plan.py @@ -17,8 +17,11 @@ from __future__ import annotations import json +from pathlib import Path import pytest +from PIL import Image +import trimesh from embodichain.gen_sim.scene_engine.core.scene import Scene from embodichain.gen_sim.scene_engine.core.scene_edit_plan import ( @@ -108,6 +111,7 @@ def test_scene_edit_plan_accepts_add_without_a_position() -> None: "object_id": "cup_001", "target_id": None, "relation": None, + "table_region": None, "category": "cup", "name": "green cup", "description": "A small green ceramic cup.", @@ -151,6 +155,7 @@ def test_scene_edit_parser_assigns_ids_to_same_category_adds_in_order() -> None: "object_id": None, "target_id": None, "relation": None, + "table_region": None, "category": "orange", "name": "small_orange", "description": "A small round orange with a textured peel.", @@ -160,6 +165,7 @@ def test_scene_edit_parser_assigns_ids_to_same_category_adds_in_order() -> None: "object_id": None, "target_id": None, "relation": None, + "table_region": None, "category": "orange", "name": "small_orange", "description": "A small round orange with a textured peel.", @@ -220,7 +226,9 @@ def test_scene_edit_plan_requires_a_position_for_move_operations() -> None: ) -def test_scene_edit_asset_preparation_skips_plans_without_adds() -> None: +def test_scene_edit_asset_preparation_skips_plans_without_adds( + tmp_path: Path, +) -> None: scene, scene_graph = _scene_and_graph() plan = SceneEditPlan( scene=scene, @@ -234,13 +242,141 @@ def test_scene_edit_asset_preparation_skips_plans_without_adds() -> None: ) ], ) + previous_asset_output = tmp_path / "scene_editing" / "asset_preparation" + previous_asset_output.mkdir(parents=True) + (previous_asset_output / "previous.txt").write_text("keep", encoding="utf-8") - prepared_scene = prepare_scene_edit_assets( + prepared_assets = prepare_scene_edit_assets( + scene_edit_plan=plan, + output_root=tmp_path, + image_generation_client=object(), # type: ignore[arg-type] + geometry_generation_client=object(), # type: ignore[arg-type] + image_segmentation_client=object(), # type: ignore[arg-type] + ) + + assert prepared_assets == [] + assert (previous_asset_output / "previous.txt").is_file() + + +def test_scene_edit_asset_preparation_generates_one_image_per_add( + tmp_path: Path, +) -> None: + class ImageGenerationClient: + def __init__(self) -> None: + self.requests: list[tuple[str, Path]] = [] + + def generate_image_by_prompt(self, *, prompt: str, output_path: Path) -> Path: + self.requests.append((prompt, output_path)) + Image.new("RGB", (6, 6), "white").save(output_path) + return output_path + + class ImageSegmentationClient: + def __init__(self) -> None: + self.requests: list[tuple[Path, str]] = [] + + def segment_single_object( + self, + *, + image_path: Path, + prompt: str, + ) -> list[dict[str, object]]: + self.requests.append((image_path, prompt)) + return [ + { + "size": [6, 6], + "counts": [7, 4, 2, 4, 2, 4, 2, 4, 7], + "starts_with": 1, + } + ] + + class GeometryGenerationClient: + def __init__(self) -> None: + self.requests: list[tuple[Path, list[tuple[str, Path]], Path]] = [] + + def generate_objects( + self, + *, + image_path: Path, + object_masks: list[tuple[str, Path]], + output_root: Path, + ) -> tuple[dict[str, object], list[dict[str, object]]]: + self.requests.append((image_path, object_masks, output_root)) + output_root.mkdir(parents=True, exist_ok=True) + for object_id, _ in object_masks: + trimesh.creation.box().export(output_root / f"{object_id}.glb") + return {}, [{"scale": [1.25, 1.5, 1.75]}] + + scene, scene_graph = _scene_and_graph() + plan = SceneEditPlan( scene=scene, + scene_graph=scene_graph, + operations=[ + SceneEditOperation( + op="add", + object_id="cup_001", + category="cup", + name="green cup", + description="A small green ceramic cup.", + ) + ], + ) + image_generation_client = ImageGenerationClient() + image_segmentation_client = ImageSegmentationClient() + geometry_generation_client = GeometryGenerationClient() + + prepared_assets = prepare_scene_edit_assets( scene_edit_plan=plan, + output_root=tmp_path, + image_generation_client=image_generation_client, # type: ignore[arg-type] + geometry_generation_client=geometry_generation_client, # type: ignore[arg-type] + image_segmentation_client=image_segmentation_client, # type: ignore[arg-type] ) - assert prepared_scene is scene + expected_image_path = ( + tmp_path + / "scene_editing" + / "asset_preparation" + / "generated_images" + / "cup_001.png" + ) + assert [asset.id for asset in prepared_assets] == ["cup_001"] + assert prepared_assets[0].simready_glb_path is not None + assert prepared_assets[0].rot == [0.0, 0.0, 0.0] + assert prepared_assets[0].pos == [0.0, 0.0, 0.0] + assert prepared_assets[0].scale == [1.0, 1.0, 1.0] + assert image_generation_client.requests == [ + ("A small green ceramic cup.", expected_image_path) + ] + assert expected_image_path.is_file() + assert image_segmentation_client.requests == [ + (expected_image_path, "A small green ceramic cup.") + ] + generated_mask_path = ( + tmp_path + / "scene_editing" + / "asset_preparation" + / "generated_masks" + / "cup_001_mask.png" + ) + assert generated_mask_path.is_file() + with Image.open(generated_mask_path) as mask: + assert mask.getpixel((3, 3)) == 255 + assert mask.getpixel((0, 0)) == 0 + generated_glb_path = ( + tmp_path + / "scene_editing" + / "asset_preparation" + / "coarse_geometry" + / "cup_001.glb" + ) + assert geometry_generation_client.requests == [ + ( + expected_image_path, + [("cup_001", generated_mask_path)], + generated_glb_path.parent, + ) + ] + assert generated_glb_path.read_bytes().startswith(b"glTF") def test_scene_edit_graph_builder_copies_the_pre_edit_graph() -> None: From 0bf17da9c6c7ae699a33cf35db65354ae90283b1 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Thu, 13 Aug 2026 20:27:18 +0800 Subject: [PATCH 12/29] finish basic simready real world scale + semantic-based rotation --- .../gen_sim/scene_engine/pipeline/edit.py | 1 + .../editing/scene_edit_asset_preparation.py | 11 + .../gen_sim/scene_engine/pipeline/generate.py | 1 + .../pipeline/generation/scene_generation.py | 11 + .../pipeline/utils/simready_processor.py | 96 ++++- .../utils/simready_processor_utils.py | 370 ++++++++++++++++++ 6 files changed, 486 insertions(+), 4 deletions(-) create mode 100644 embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor_utils.py diff --git a/embodichain/gen_sim/scene_engine/pipeline/edit.py b/embodichain/gen_sim/scene_engine/pipeline/edit.py index 6ba7898d1..fac2b6769 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/edit.py +++ b/embodichain/gen_sim/scene_engine/pipeline/edit.py @@ -87,6 +87,7 @@ def edit_scene( image_generation_client=image_generation_client, geometry_generation_client=geometry_generation_client, image_segmentation_client=image_segmentation_client, + vlm_client=vlm_client, ) finally: image_generation_client.close() diff --git a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py index 70c02d1a7..c6ec26fb3 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py @@ -35,6 +35,9 @@ from embodichain.gen_sim.scene_engine.core.scene import Scene from embodichain.gen_sim.scene_engine.core.scene_edit_plan import SceneEditPlan 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, @@ -44,6 +47,7 @@ ) from embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor import ( SimReadyProcessor, + SimReadyProcessorConfig, ) __all__ = ["prepare_scene_edit_assets"] @@ -66,6 +70,7 @@ def prepare_scene_edit_assets( image_generation_client: ImageGenerationClient, geometry_generation_client: GeometryGenerationClient, image_segmentation_client: ImageSegmentationClient, + vlm_client: OpenAICompatibleVLM | None = None, ) -> list[SceneObject]: """Prepare and return SimReady assets required by add operations.""" # Prepare descriptions for all newly added objects. @@ -112,6 +117,12 @@ def prepare_scene_edit_assets( coarse_layout_by_id=_coarse_layouts_by_id(generated_asset_glbs), coarse_geometry_root=stage_output_root / "coarse_geometry", simready_geometry_root=stage_output_root / "simready_geometry", + # Scene editing will later provide the VLM-selected scale and rotation. + config=SimReadyProcessorConfig( + use_vlm_scale=vlm_client is not None, + use_vlm_rotation=vlm_client is not None, + ), + vlm_client=vlm_client, ) # process_assets() validates and processes assets only; it does not require a table. simready_processor.process_assets() diff --git a/embodichain/gen_sim/scene_engine/pipeline/generate.py b/embodichain/gen_sim/scene_engine/pipeline/generate.py index 1df9027f9..144551c29 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generate.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generate.py @@ -81,6 +81,7 @@ def generate_scene_from_image( scene=scene, scene_graph=scene_graph, geometry_generation_client=geometry_generation_client, + vlm_client=vlm_client, ) finally: geometry_generation_client.close() # Kill the session to avoid resource leaks. diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py index d8d82b19b..0d736da76 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py @@ -30,6 +30,9 @@ from embodichain.gen_sim.scene_engine.core.scene import Scene from embodichain.gen_sim.scene_engine.core.scene_graph import SceneGraph 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.assets_group_support_clamp import ( AssetsGroupSupportClamp, ) @@ -50,6 +53,7 @@ ) from embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor import ( SimReadyProcessor, + SimReadyProcessorConfig, ) from embodichain.gen_sim.scene_engine.pipeline.utils.table_support_surface import ( TableSupportSurfaceDetector, @@ -66,6 +70,7 @@ def generate_scene_and_refine( scene_graph: SceneGraph, *, geometry_generation_client: GeometryGenerationClient, + vlm_client: OpenAICompatibleVLM, ) -> Scene: resolved_image_path = _validate_image_path(image_path) @@ -110,6 +115,12 @@ def generate_scene_and_refine( coarse_layout_by_id=coarse_layout_by_id, coarse_geometry_root=coarse_geometry_output_root, simready_geometry_root=simready_geometry_output_root, + # Image-to-scene uses the geometry service's coarse scale directly. + config=SimReadyProcessorConfig( + use_vlm_scale=False, + use_vlm_rotation=False, + ), + vlm_client=vlm_client, ) simready_assets_layout = simready_processor.process_assets() simready_table_layout = simready_processor.process_table() diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py index 6b5a09b0b..b22e24f6b 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py @@ -32,6 +32,15 @@ ObjectPhysics, SceneObject, ) +from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( + OpenAICompatibleVLM, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor_utils import ( + query_vlm_object_rotation_and_target_size, + compute_uniform_xy_scale_for_target, + render_object_front_top_views, + rotate_glb_about_x_axis, +) from embodichain.utils.logger import log_info _TABLE_PHYSICS_ATTRS = { @@ -56,6 +65,9 @@ class SimReadyProcessorConfig: """Object-category policy for SimReady mesh canonicalization.""" + use_vlm_scale: bool = False # Use the VLM-selected asset scale. + use_vlm_rotation: bool = False # Use the VLM-selected asset rotation. + upright_container_id_tokens: frozenset[str] = frozenset( {"bottle", "can", "jar", "flask", "thermos"} ) # Object-id tokens that enable upright-container standardization. @@ -72,6 +84,7 @@ def __init__( coarse_geometry_root: str | Path, simready_geometry_root: str | Path, config: SimReadyProcessorConfig | None = None, + vlm_client: OpenAICompatibleVLM | None = None, ) -> None: self.scene = scene self.coarse_layout_by_id = coarse_layout_by_id @@ -82,8 +95,13 @@ def __init__( 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 SimReadyProcessorConfig() + self.vlm_client = vlm_client if not self.config.upright_container_id_tokens: raise ValueError("upright_container_id_tokens must not be empty.") + if ( + self.config.use_vlm_scale or self.config.use_vlm_rotation + ) and vlm_client is None: + raise ValueError("vlm_client is required when VLM transforms are enabled.") def process_table(self) -> dict[str, object]: """Process the required scene table and return its SimReady layout.""" @@ -104,7 +122,13 @@ def process_assets(self) -> list[dict[str, object]]: self.simready_assets_layout = processed_assets return self.simready_assets_layout - def _process_object(self, scene_object: SceneObject) -> dict[str, object]: + def _process_object( + self, + scene_object: SceneObject, + *, + scale: object | None = None, + rot: object | None = None, + ) -> dict[str, object]: """Canonicalize one coarse object and write its SimReady GLB.""" object_id = scene_object.id object_role = scene_object.kind @@ -113,12 +137,18 @@ def _process_object(self, scene_object: SceneObject) -> dict[str, object]: 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}.") + prepared_glb_path, vlm_scale = self._prepare_vlm_rotated_glb(scene_object) + selected_scale = scale + if selected_scale is None: + selected_scale = vlm_scale or coarse_layout.get("scale") simready_mesh, simready_transform = self._canonicalize_object_mesh( - coarse_glb_path=self.coarse_geometry_root / f"{object_id}.glb", + coarse_glb_path=prepared_glb_path, object_id=object_id, - rot=coarse_layout.get("rot"), + # An enabled external rotation replaces the coarse-layout rotation. + rot=coarse_layout.get("rot") if rot is None else rot, pos=coarse_layout.get("pos"), - scale=coarse_layout.get("scale"), + # An enabled VLM scale replaces the coarse-layout scale. + scale=selected_scale, ) output_path = self.simready_geometry_root / f"{object_id}.glb" output_path.parent.mkdir(parents=True, exist_ok=True) @@ -132,6 +162,64 @@ def _process_object(self, scene_object: SceneObject) -> dict[str, object]: log_info(f"Created SimReady {object_role}: {object_id!r}.") return {"id": object_id, **simready_transform} + def _prepare_vlm_rotated_glb( + self, scene_object: SceneObject + ) -> tuple[Path, list[float] | None]: + """Render, query, and optionally bake the VLM-selected x-axis rotation.""" + coarse_path = self.coarse_geometry_root / f"{scene_object.id}.glb" + if not (self.config.use_vlm_scale or self.config.use_vlm_rotation): + return coarse_path, None + decision = self._vlm_transform_for_object( + scene_object, + use_scale=self.config.use_vlm_scale, + use_rotation=self.config.use_vlm_rotation, + ) + rotate_about_x = bool(decision["rotate_about_x"]) + vlm_scale = compute_uniform_xy_scale_for_target( + glb_path=coarse_path, + target_xy_size_cm=decision["target_xy_size_cm"], + rotate_about_x=rotate_about_x, + ) + rotated_path = rotate_glb_about_x_axis( + input_path=coarse_path, + output_path=self.simready_geometry_root + / "vlm_rotated" + / f"{scene_object.id}.glb", + rotate=rotate_about_x, + ) + # The scale flag controls whether this VLM-derived isotropic scale is used. + # Apply the same factor on x, y, and z to preserve the asset's proportions. + return ( + rotated_path, + [vlm_scale, vlm_scale, vlm_scale] if self.config.use_vlm_scale else None, + ) + + def _vlm_transform_for_object( + self, + scene_object: SceneObject, + *, + use_scale: bool, + use_rotation: bool, + ) -> dict[str, object]: + """Render the object and return the validated VLM pose decision.""" + del use_scale, use_rotation + assert self.vlm_client is not None + coarse_path = self.coarse_geometry_root / f"{scene_object.id}.glb" + needed_layout = "This asset needs to be place on the table that will not move a lot after simulation." + debug_root = self.simready_geometry_root.parent / "debug" + rendered_path = render_object_front_top_views( + glb_path=coarse_path, + output_path=debug_root / "vlm_views" / f"{scene_object.id}.png", + ) + # Both semantic questions are always answered in one multimodal call. + return query_vlm_object_rotation_and_target_size( + scene_object_description=scene_object.description, + needed_layout=needed_layout, + rendered_views_path=rendered_path, + vlm_client=self.vlm_client, + debug_output_path=debug_root / "vlm_outputs" / f"{scene_object.id}.json", + ) + @staticmethod def _fixed_physics_for_kind(kind: str) -> ObjectPhysics: """Create the fixed initial physics profile for one SimReady object.""" diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor_utils.py b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor_utils.py new file mode 100644 index 000000000..b7f718374 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor_utils.py @@ -0,0 +1,370 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF 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 PIL import Image, ImageDraw, ImageFont +from scipy.spatial.transform import Rotation +import trimesh + +from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( + OpenAICompatibleVLM, +) + +_VLM_SYSTEM_PROMPT = """You inspect one isolated 3D object from front and top views. +Use the object description and the rendered views together. + +Use the rendered views and the needed layout to decide whether the object should +be rotated around its own center by +90 degrees around the z-up world's x axis. +The z-up world is right-handed: x is left-right, y is front-back, and z is up. +In the composed image, FRONT VIEW is the left panel: x is horizontal and z is +vertical; the upper-right marker shows the positive z direction. TOP VIEW is +the right panel: x is horizontal and y is vertical; the upper-right markers +show the positive x and y directions. +Do not confuse the top view with looking at the object from above in the image +description: it is a projection along the z axis onto the x-y plane. +After deciding and applying that rotation, estimate the object's desired AABB +footprint on the x-y plane in real-world centimetres. The first value is the x +size and the second value is the y size. + +Return JSON only with exactly this schema: +{ + "rotate_about_x": false, + "target_xy_size_cm": [12.0, 5.0] +} + +Examples: +- Fork lying flat on a table: in FRONT VIEW the fork is mostly a thin + horizontal line; in TOP VIEW its length is visible. Keep it flat with + rotate_about_x=false, and use the tabletop footprint, for example + target_xy_size_cm=[15.0, 3.0]. +- Fork placed in a pen holder: the desired fork is upright, so its long axis is + approximately z. If the input coarse fork is lying in the x-y plane, set + rotate_about_x=true; if the input coarse fork is already upright, set it to + false. The target is the footprint inside the holder, not the fork's full + length, for example target_xy_size_cm=[3.0, 3.0]. +- Fork requested to lie flat on a table even when the input coarse fork is + upright: set rotate_about_x=true and estimate the final flat footprint, for + example target_xy_size_cm=[15.0, 3.0]. +- Bottle already standing on its flat base: keep it upright with + rotate_about_x=false and use target_xy_size_cm=[8.0, 8.0]. +""" + + +def render_object_front_top_views( + *, + glb_path: str | Path, + output_path: str | Path, + resolution: int = 512, +) -> Path: + """Render fixed z-up front/top views and compose them horizontally.""" + if resolution <= 0: + raise ValueError("resolution must be positive.") + source_path = Path(glb_path).expanduser().resolve() + if not source_path.is_file(): + raise FileNotFoundError(f"GLB for VLM rendering not found: {source_path}") + output_path = Path(output_path).expanduser().resolve() + output_path.parent.mkdir(parents=True, exist_ok=True) + front_path = output_path.with_name(f"{output_path.stem}_front.png") + top_path = output_path.with_name(f"{output_path.stem}_top.png") + try: + import bpy + from mathutils import Vector + except ImportError as exc: + raise RuntimeError( + "Blender's bpy is required for SimReady VLM view rendering." + ) from exc + + bpy.ops.wm.read_factory_settings(use_empty=True) + bpy.ops.import_scene.gltf(filepath=str(source_path)) + if not any(obj.type == "MESH" for obj in bpy.context.scene.objects): + raise ValueError(f"GLB contains no mesh objects: {source_path}") + scene = bpy.context.scene + # Eevee renders imported GLB materials and textures instead of Workbench previews. + try: + scene.render.engine = "BLENDER_EEVEE_NEXT" + except TypeError: + scene.render.engine = "BLENDER_EEVEE" + scene.render.resolution_x = resolution + scene.render.resolution_y = resolution + scene.render.resolution_percentage = 100 + scene.render.image_settings.file_format = "PNG" + scene.render.film_transparent = False + if scene.world is None: + scene.world = bpy.data.worlds.new("VLM_World") + scene.world.color = (0.08, 0.08, 0.08) + for name, location, energy in ( + ("VLM_Key", (2.0, -2.0, 3.0), 700.0), + ("VLM_Fill", (-2.0, 1.0, 2.0), 400.0), + ): + light_data = bpy.data.lights.new(name, type="AREA") + light_data.energy = energy + light_data.shape = "DISK" + light_data.size = 4.0 + light = bpy.data.objects.new(name, light_data) + light.location = location + light.rotation_euler = ( + (Vector((0.0, 0.0, 0.0)) - light.location) + .to_track_quat("-Z", "Y") + .to_euler() + ) + scene.collection.objects.link(light) + camera_data = bpy.data.cameras.new("VLM_Camera") + camera = bpy.data.objects.new("VLM_Camera", camera_data) + scene.collection.objects.link(camera) + scene.camera = camera + camera.data.type = "ORTHO" + camera.data.ortho_scale = 1.25 + + def render_view(path: Path, location: tuple[float, float, float]) -> None: + camera.location = location + camera.rotation_euler = ( + (Vector((0.0, 0.0, 0.0)) - camera.location) + .to_track_quat("-Z", "Y") + .to_euler() + ) + scene.render.filepath = str(path) + bpy.ops.render.render(write_still=True) + + # Blender uses a right-handed z-up world; front is viewed along +y. + render_view(front_path, (0.0, -3.0, 0.0)) + render_view(top_path, (0.0, 0.0, 3.0)) + with Image.open(front_path) as front, Image.open(top_path) as top: + composed = Image.new("RGB", (resolution * 2, resolution), "white") + composed.paste(front.convert("RGB"), (0, 0)) + composed.paste(top.convert("RGB"), (resolution, 0)) + draw = ImageDraw.Draw(composed) + # Use a readable scaled font for the panel labels when available. + try: + font = ImageFont.truetype( + "/usr/share/fonts/truetype/ubuntu/Ubuntu-B.ttf", + max(24, resolution // 16), + ) + except OSError: + font = ImageFont.load_default() + # Label each panel so the VLM and manual debugging can distinguish views. + for label, origin in (("FRONT VIEW", (0, 0)), ("TOP VIEW", (resolution, 0))): + x, y = origin + text_box = draw.textbbox((x + 16, y + 16), label, font=font) + draw.rectangle( + (text_box[0] - 8, text_box[1] - 6, text_box[2] + 8, text_box[3] + 6), + fill="white", + ) + draw.text((x + 16, y + 16), label, fill="black", font=font) + # Mark the positive axes used by each projection for VLM interpretation. + _draw_arrow( + draw, + (resolution - 62, 62), + (resolution - 62, 20), + "+Z", + font, + color="blue", + ) + _draw_arrow( + draw, + (2 * resolution - 92, 62), + (2 * resolution - 42, 62), + "+X", + font, + color="red", + ) + _draw_arrow( + draw, + (2 * resolution - 92, 62), + (2 * resolution - 92, 20), + "+Y", + font, + color="green", + ) + composed.save(output_path) + return output_path + + +def _draw_arrow( + draw: ImageDraw.ImageDraw, + start: tuple[int, int], + end: tuple[int, int], + label: str, + font: ImageFont.FreeTypeFont | ImageFont.ImageFont, + color: str, +) -> None: + """Draw one labeled positive-axis arrow on a rendered view.""" + dx, dy = end[0] - start[0], end[1] - start[1] + length = max(abs(dx), abs(dy)) + if length == 0: + raise ValueError("Axis arrow start and end must differ.") + unit_x, unit_y = dx / length, dy / length + perpendicular_x, perpendicular_y = -unit_y, unit_x + head_length = 14.0 + head_width = 8.0 + tip_x, tip_y = end + base_x = tip_x - unit_x * head_length + base_y = tip_y - unit_y * head_length + arrowhead = ( + (tip_x, tip_y), + ( + base_x + perpendicular_x * head_width, + base_y + perpendicular_y * head_width, + ), + ( + base_x - perpendicular_x * head_width, + base_y - perpendicular_y * head_width, + ), + ) + draw.line((*start, *end), fill=color, width=4) + draw.polygon(arrowhead, fill=color) + # Put each axis label beside its arrowhead so it does not cover the arrow. + draw.text((int(tip_x + 8), int(tip_y - 8)), label, fill=color, font=font) + + +def query_vlm_object_rotation_and_target_size( + *, + scene_object_description: str, + needed_layout: str, + rendered_views_path: str | Path, + vlm_client: OpenAICompatibleVLM, + debug_output_path: str | Path | None = None, +) -> dict[str, object]: + """Ask the VLM for rotation and post-rotation tabletop footprint.""" + response_text = vlm_client.complete( + system_prompt=_VLM_SYSTEM_PROMPT, + user_prompt=( + f"Object description:\n{scene_object_description}\n\n" + f"Needed layout:\n{needed_layout}\n\n" + "The image contains front view on the left and top view on the right." + ), + image_path=rendered_views_path, + ) + try: + value = json.loads(_strip_json_code_fence(response_text)) + except json.JSONDecodeError as exc: + raise ValueError(f"VLM transform response is not valid JSON: {exc}") from exc + if not isinstance(value, dict) or set(value) != { + "rotate_about_x", + "target_xy_size_cm", + }: + raise ValueError( + "VLM transform response must contain exactly rotate_about_x and " + "target_xy_size_cm." + ) + if not isinstance(value["rotate_about_x"], bool): + raise ValueError("VLM rotate_about_x must be boolean.") + target_size = value["target_xy_size_cm"] + if ( + not isinstance(target_size, list) + or len(target_size) != 2 + or not all(isinstance(item, (int, float)) for item in target_size) + or not all(np.isfinite(item) and item > 0 for item in target_size) + ): + raise ValueError("VLM target_xy_size_cm must contain two positive numbers.") + if debug_output_path is not None: + output_path = Path(debug_output_path).expanduser().resolve() + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text( + json.dumps( + { + "description": scene_object_description, + "needed_layout": needed_layout, + "rendered_views_path": str( + Path(rendered_views_path).expanduser().resolve() + ), + "vlm_output": value, + }, + indent=2, + ensure_ascii=False, + ) + + "\n", + encoding="utf-8", + ) + return value + + +def compute_uniform_xy_scale_for_target( + *, + glb_path: str | Path, + target_xy_size_cm: list[float], + rotate_about_x: bool, +) -> float: + """Compute an isotropic scale from the rotated mesh XY AABB and target size.""" + loaded = trimesh.load(Path(glb_path).expanduser().resolve(), process=False) + mesh = ( + loaded.dump(concatenate=True) if isinstance(loaded, trimesh.Scene) else loaded + ) + if not isinstance(mesh, trimesh.Trimesh): + raise ValueError(f"GLB is not a mesh: {glb_path}") + if len(target_xy_size_cm) != 2 or any(value <= 0 for value in target_xy_size_cm): + raise ValueError("target_xy_size_cm must contain two positive values.") + if rotate_about_x: + center = mesh.bounds.mean(axis=0) + mesh.apply_translation(-center) + transform = np.eye(4) + transform[:3, :3] = Rotation.from_euler("x", 90.0, degrees=True).as_matrix() + mesh.apply_transform(transform) + mesh.apply_translation(center) + actual_xy_size = mesh.bounds[1, :2] - mesh.bounds[0, :2] + if np.any(actual_xy_size <= 0): + raise ValueError("Rotated mesh must have a positive XY AABB.") + # Convert the VLM's centimetres to metres before comparing with the GLB AABB. + target_xy_size_m = np.asarray(target_xy_size_cm, dtype=float) / 100.0 + axis_scales = target_xy_size_m / actual_xy_size + # Use sqrt(target XY area / actual XY area) as one uniform scale on all axes. + return float(np.sqrt(axis_scales[0] * axis_scales[1])) + + +def rotate_glb_about_x_axis( + *, + input_path: str | Path, + output_path: str | Path, + rotate: bool, +) -> Path: + """Bake an optional +90-degree x-axis rotation around the mesh centre.""" + # Current coarse layouts are either flat on xy with possible random z rotation, + # or upright with almost no random y rotation, so this x-axis toggle is enough. + source_path = Path(input_path).expanduser().resolve() + destination_path = Path(output_path).expanduser().resolve() + destination_path.parent.mkdir(parents=True, exist_ok=True) + loaded = trimesh.load(source_path, process=False) + mesh = ( + loaded.dump(concatenate=True) if isinstance(loaded, trimesh.Scene) else loaded + ) + if not isinstance(mesh, trimesh.Trimesh): + raise ValueError(f"GLB is not a mesh: {source_path}") + if rotate: + center = mesh.bounds.mean(axis=0) + mesh.apply_translation(-center) + transform = np.eye(4) + transform[:3, :3] = Rotation.from_euler("x", 90.0, degrees=True).as_matrix() + mesh.apply_transform(transform) + mesh.apply_translation(center) + mesh.export(destination_path, file_type="glb") + return destination_path + + +def _strip_json_code_fence(response_text: str) -> str: + """Remove one optional Markdown JSON fence from a VLM response.""" + stripped = response_text.strip() + if not stripped.startswith("```"): + return stripped + lines = stripped.splitlines() + if lines and lines[0].startswith("```"): + lines = lines[1:] + if lines and lines[-1].strip() == "```": + lines = lines[:-1] + return "\n".join(lines).strip() From 36c98e430c159b49697c1508017306c6ac74f1fe Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Fri, 14 Aug 2026 15:34:00 +0800 Subject: [PATCH 13/29] Moved the table support infos into simready. --- .../gen_sim/scene_engine/core/scene_object.py | 6 ++ .../pipeline/generation/scene_generation.py | 67 +++++++++++-------- .../pipeline/utils/scene_exporter.py | 3 + .../pipeline/utils/scene_importer.py | 27 ++++++++ .../pipeline/utils/simready_processor.py | 39 +++++++++++ .../pipeline/utils/table_support_surface.py | 53 ++++++++++++++- 6 files changed, 167 insertions(+), 28 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/core/scene_object.py b/embodichain/gen_sim/scene_engine/core/scene_object.py index 0a8ed3aea..2b868e3c3 100644 --- a/embodichain/gen_sim/scene_engine/core/scene_object.py +++ b/embodichain/gen_sim/scene_engine/core/scene_object.py @@ -67,6 +67,9 @@ class SceneObject: pos: list[float] | None = None # Final y-up world position in metres. scale: list[float] | None = None # Final y-up object scale. center_xy: list[float] | None = None # Z-up table-frame XY AABB center. + support_surface_z: float | None = None # Detected tabletop height in z-up. + support_contour_xy: list[list[float]] | None = None # Outer support contour. + support_optimization_rect_xy: list[list[float]] | None = None # Safe XY rectangle. physics: ObjectPhysics | None = None # Assigned when SimReady processing succeeds. def to_dict(self) -> dict[str, object]: @@ -83,5 +86,8 @@ def to_dict(self) -> dict[str, object]: "pos": self.pos, "scale": self.scale, "center_xy": self.center_xy, + "support_surface_z": self.support_surface_z, + "support_contour_xy": self.support_contour_xy, + "support_optimization_rect_xy": self.support_optimization_rect_xy, "physics": self.physics.to_dict() if self.physics is not None else None, } diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py index 0d736da76..60f91ca81 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py @@ -23,6 +23,7 @@ import numpy as np import trimesh +from shapely.geometry import Polygon from embodichain.gen_sim.scene_engine.clients.geometry_generation import ( GeometryGenerationClient, @@ -55,9 +56,6 @@ SimReadyProcessor, SimReadyProcessorConfig, ) -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"} @@ -115,6 +113,7 @@ def generate_scene_and_refine( coarse_layout_by_id=coarse_layout_by_id, coarse_geometry_root=coarse_geometry_output_root, simready_geometry_root=simready_geometry_output_root, + debug_output_root=debug_output_root, # Image-to-scene uses the geometry service's coarse scale directly. config=SimReadyProcessorConfig( use_vlm_scale=False, @@ -248,6 +247,18 @@ def _update_scene_final_y_up_layout_and_z_up_centers( ) # Persist AABB centers for future scene-edit object disambiguation. scene.table.center_xy = table_mesh.bounds[:, :2].mean(axis=0).tolist() + if scene.table.support_contour_xy is not None: + # Move SimReady-local support geometry into the final table-frame position. + table_center_xy = np.asarray(scene.table.center_xy, dtype=float) + scene.table.support_contour_xy = [ + (np.asarray(point, dtype=float) + table_center_xy).tolist() + for point in scene.table.support_contour_xy + ] + if scene.table.support_optimization_rect_xy is not None: + scene.table.support_optimization_rect_xy = [ + (np.asarray(point, dtype=float) + table_center_xy).tolist() + for point in scene.table.support_optimization_rect_xy + ] for asset in scene.assets: asset.center_xy = assets_aabb_corners_by_id[asset.id].mean(axis=0).tolist() @@ -360,31 +371,36 @@ def _layout_refinement( log_info("Scene has no movable assets; skipping support-region clamping.") return refined_table_layout, [] - # 4. Detect the actual upward support triangles instead of projecting the - # entire table mesh to one convex hull. The result retains concavities - # (for example, an L-shaped tabletop) and is the only boundary used for - # placement below. - ( - table_world_mesh_z_up, - assets_aabb_2d_z_up_world_corners_by_id, - ) = _measure_table_and_assets_in_z_up_world( - table_layout=refined_table_layout, - assets_layout=refined_assets_layout, - geometry_root=simready_geometry_output_root, - ) - support_detector = TableSupportSurfaceDetector( - table_world_mesh=table_world_mesh_z_up, - debug_output_root=debug_output_root, + # 4. Reuse support geometry detected during SimReady processing. + if ( + scene.table is None + or scene.table.support_contour_xy is None + or scene.table.support_optimization_rect_xy is None + ): + raise ValueError("Scene table has no persisted support geometry.") + table_support_polygon = Polygon(scene.table.support_contour_xy) + table_optimization_rectangle = Polygon(scene.table.support_optimization_rect_xy) + if not table_support_polygon.is_valid or table_support_polygon.is_empty: + raise ValueError("Scene table support contour is not a valid polygon.") + if ( + not table_optimization_rectangle.is_valid + or table_optimization_rectangle.is_empty + ): + raise ValueError("Scene table optimization rectangle is not valid.") + _, 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, + ) ) - 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, + support_region=table_support_polygon, assets_aabb_2d_z_up_world_corners_by_id=( assets_aabb_2d_z_up_world_corners_by_id ), @@ -405,13 +421,10 @@ def _layout_refinement( ) ) - # 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. + # 6. Optimize independent asset positions inside the conservative rectangle. + # The clamp above already used the exact outer contour for the shared shift. overlap_optimizer = AssetsSupportLayoutOptimizer( - support_region=table_support_region.support_polygon, + support_region=table_optimization_rectangle, assets_aabb_2d_z_up_world_corners_by_id=( clamped_assets_aabb_2d_z_up_world_corners_by_id ), 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 391ccd156..4063bf92c 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py @@ -193,6 +193,9 @@ def _scene_object_config( # which SimulationManager itself converts to z-up. "body_scale": scale_y_up, "center_xy": scene_object.center_xy, + "support_surface_z": scene_object.support_surface_z, + "support_contour_xy": scene_object.support_contour_xy, + "support_optimization_rect_xy": scene_object.support_optimization_rect_xy, "max_convex_hull_num": scene_object.physics.max_convex_hull_num, } diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py index a59bc7bdf..eefeee6b6 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py @@ -265,6 +265,16 @@ def _scene_object_from_export_entry( center_xy = entry.get("center_xy") if center_xy is not None: center_xy = self._vector2(center_xy, field_name=f"{uid}.center_xy") + support_surface_z = entry.get("support_surface_z") + if support_surface_z is not None: + support_surface_z = float(support_surface_z) + support_contour_xy = self._points2( + entry.get("support_contour_xy"), field_name=f"{uid}.support_contour_xy" + ) + support_optimization_rect_xy = self._points2( + entry.get("support_optimization_rect_xy"), + field_name=f"{uid}.support_optimization_rect_xy", + ) pos_y_up = _Z_UP_TO_Y_UP_ROTATION @ np.asarray(pos_z_up, dtype=float) rotation_z_up = Rotation.from_euler("XYZ", rot_z_up, degrees=True).as_matrix() @@ -284,6 +294,9 @@ def _scene_object_from_export_entry( pos=pos_y_up.tolist(), scale=scale, center_xy=center_xy, + support_surface_z=support_surface_z, + support_contour_xy=support_contour_xy, + support_optimization_rect_xy=support_optimization_rect_xy, physics=ObjectPhysics( body_type=str(entry.get("body_type", "dynamic")), # type: ignore[arg-type] attrs=self._physics_attrs(entry.get("attrs", {"mass": 1.0})), @@ -345,6 +358,20 @@ def _vector2(value: object, *, field_name: str) -> list[float]: raise ValueError(f"Scene config field {field_name!r} must be finite.") return vector + @classmethod + def _points2(cls, value: object, *, field_name: str) -> list[list[float]] | None: + """Validate an optional list of XY points from the scene export.""" + if value is None: + return None + if not isinstance(value, list) or len(value) < 3: + raise ValueError( + f"Scene config field {field_name!r} must contain 3 points." + ) + return [ + cls._vector2(point, field_name=f"{field_name}[{index}]") + for index, point in enumerate(value) + ] + @staticmethod def _physics_attrs(value: object) -> dict[str, float | int]: """Validate exported physics attributes.""" diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py index b22e24f6b..55860be7f 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py @@ -41,6 +41,9 @@ render_object_front_top_views, rotate_glb_about_x_axis, ) +from embodichain.gen_sim.scene_engine.pipeline.utils.table_support_surface import ( + TableSupportSurfaceDetector, +) from embodichain.utils.logger import log_info _TABLE_PHYSICS_ATTRS = { @@ -83,6 +86,7 @@ def __init__( coarse_layout_by_id: dict[str, dict[str, object]], coarse_geometry_root: str | Path, simready_geometry_root: str | Path, + debug_output_root: str | Path | None = None, config: SimReadyProcessorConfig | None = None, vlm_client: OpenAICompatibleVLM | None = None, ) -> None: @@ -92,6 +96,12 @@ def __init__( self.simready_geometry_root = ( Path(simready_geometry_root).expanduser().resolve() ) + # Save rendered debug images. + self.debug_output_root = ( + Path(debug_output_root).expanduser().resolve() + if debug_output_root is not None + else None + ) 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 SimReadyProcessorConfig() @@ -159,9 +169,38 @@ def _process_object( ) scene_object.simready_glb_path = str(output_path) scene_object.physics = self._fixed_physics_for_kind(object_role) + # For table. (currently the id is fixed into table) + if object_role == "table": + # Detect and persist all reusable tabletop support geometry at SimReady time. + support_detector = TableSupportSurfaceDetector( + table_world_mesh=self._z_up_table_mesh(simready_mesh), + debug_output_root=self.debug_output_root, + ) + support_region = support_detector.detect() + scene_object.support_surface_z = support_region.top_z + scene_object.support_contour_xy = [ + [float(x), float(y)] + for x, y in support_region.support_polygon.exterior.coords[:-1] + ] + scene_object.support_optimization_rect_xy = [ + [float(x), float(y)] + for x, y in support_region.optimization_rectangle.exterior.coords[:-1] + ] + if self.debug_output_root is not None: + # Keep the 3D selected surface and 2D contour diagnostics beside SimReady output. + support_detector.save_support_surface_debug_images() log_info(f"Created SimReady {object_role}: {object_id!r}.") return {"id": object_id, **simready_transform} + @staticmethod + def _z_up_table_mesh(mesh: trimesh.Trimesh) -> trimesh.Trimesh: + """Convert one canonical y-up GLB mesh into the detector's z-up frame.""" + y_up_to_z_up = np.eye(4) + y_up_to_z_up[:3, :3] = Rotation.from_euler("x", 90.0, degrees=True).as_matrix() + z_up_mesh = mesh.copy() + z_up_mesh.apply_transform(y_up_to_z_up) + return z_up_mesh + def _prepare_vlm_rotated_glb( self, scene_object: SceneObject ) -> tuple[Path, list[float] | None]: 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 index 6c3ad3e94..983c16098 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/table_support_surface.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/table_support_surface.py @@ -55,6 +55,7 @@ class TableSupportRegion: 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. + optimization_rectangle: Polygon # Axis-aligned rectangle fully inside the contour. class TableSupportSurfaceDetector: @@ -131,11 +132,13 @@ def detect(self) -> TableSupportRegion: selected_vertices = face_vertices[selected_faces] vertices = mesh.vertices.copy() faces = mesh.faces[selected_faces].copy() + support_polygon = self._extract_largest_support_polygon(vertices[faces, :2]) 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]), + support_polygon=support_polygon, + optimization_rectangle=self._largest_inscribed_rectangle(support_polygon), ) return self.support_region @@ -367,6 +370,46 @@ def _extract_largest_support_polygon(cls, triangles_xy: np.ndarray) -> Polygon: raise ValueError("The merged 2D support contour is degenerate.") return Polygon(boundary_xy) + @staticmethod + def _largest_inscribed_rectangle(polygon: Polygon) -> Polygon: + """Find a conservative axis-aligned rectangle contained by the support contour.""" + coordinates = np.asarray(polygon.exterior.coords[:-1], dtype=float) + x_values = np.unique(coordinates[:, 0]) + y_values = np.unique(coordinates[:, 1]) + # Keep the search bounded for highly tessellated support contours. + if len(x_values) > 48: + x_values = x_values[np.linspace(0, len(x_values) - 1, 48, dtype=int)] + if len(y_values) > 48: + y_values = y_values[np.linspace(0, len(y_values) - 1, 48, dtype=int)] + + best_rectangle: Polygon | None = None + best_area = 0.0 + for x_index, minimum_x in enumerate(x_values[:-1]): + for maximum_x in x_values[x_index + 1 :]: + if maximum_x <= minimum_x: + continue + for y_index, minimum_y in enumerate(y_values[:-1]): + for maximum_y in y_values[y_index + 1 :]: + if maximum_y <= minimum_y: + continue + rectangle = Polygon( + [ + (minimum_x, minimum_y), + (maximum_x, minimum_y), + (maximum_x, maximum_y), + (minimum_x, maximum_y), + ] + ) + area = rectangle.area + if area > best_area and polygon.covers(rectangle): + best_rectangle = rectangle + best_area = area + if best_rectangle is None: + raise ValueError( + "Support contour has no non-degenerate inscribed rectangle." + ) + return best_rectangle + @staticmethod def _face_adjacency(mesh: trimesh.Trimesh) -> dict[int, set[int]]: """Build a face adjacency dictionary for the mesh.""" @@ -476,6 +519,14 @@ def _save_support_region_2d_image( linewidth=2.0, label="outer support contour", ) + rectangle_xy = np.asarray(support_region.optimization_rectangle.exterior.coords) + axis.plot( + rectangle_xy[:, 0], + rectangle_xy[:, 1], + color="seagreen", + linewidth=2.0, + label="optimization rectangle", + ) axis.autoscale_view() axis.set_aspect("equal", adjustable="box") axis.set_xlabel("x (z-up world)") From c6b0cc43f61f014c59c92ad1ad23766afca96c91 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Fri, 14 Aug 2026 19:16:56 +0800 Subject: [PATCH 14/29] Finished scene edit --- .../scene_engine/cli/test_text_to_simready.py | 183 ++++ .../gen_sim/scene_engine/pipeline/edit.py | 33 +- .../editing/scene_edit_layout_generation.py | 77 ++ .../pipeline/utils/scene_exporter.py | 35 +- .../pipeline/utils/scene_importer.py | 26 +- .../utils/scene_layout_constructor.py | 399 +++++++++ .../pipeline/utils/scene_layout_optimizer.py | 785 ++++++++++++++++++ .../test_scene_core_and_export.py | 69 ++ 8 files changed, 1593 insertions(+), 14 deletions(-) create mode 100644 embodichain/gen_sim/scene_engine/cli/test_text_to_simready.py create mode 100644 embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_layout_generation.py create mode 100644 embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_constructor.py create mode 100644 embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_optimizer.py diff --git a/embodichain/gen_sim/scene_engine/cli/test_text_to_simready.py b/embodichain/gen_sim/scene_engine/cli/test_text_to_simready.py new file mode 100644 index 000000000..e3db15b16 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/cli/test_text_to_simready.py @@ -0,0 +1,183 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF 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 +from collections.abc import Sequence +from pathlib import Path +import shutil + +from PIL import Image + +from embodichain.gen_sim.scene_engine.clients.geometry_generation import ( + GeometryGenerationClient, +) +from embodichain.gen_sim.scene_engine.clients.image_generation import ( + ImageGenerationClient, +) +from embodichain.gen_sim.scene_engine.clients.image_segmentation import ( + ImageSegmentationClient, +) +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( + OpenAICompatibleVLM, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.image_segmentation_utils import ( + build_mask_candidates, + invert_mask_if_foreground_is_off_center, + save_binary_mask, + union_overlapping_mask_candidates, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor import ( + SimReadyProcessor, + SimReadyProcessorConfig, +) + +__all__ = ["main", "run_text_to_simready"] + + +def run_text_to_simready(*, text: str, output_root: str | Path) -> SceneObject: + """Run one manual text-to-SimReady asset pipeline for debugging.""" + text = text.strip() + if not text: + raise ValueError("Text prompt must not be empty.") + + root = Path(output_root).expanduser().resolve() + if root.exists(): + shutil.rmtree(root) + debug_root = root / "debug" + image_root = root / "generated_images" + mask_root = root / "masks" + coarse_root = root / "coarse_geometry" + simready_root = root / "simready_geometry" + for directory in (debug_root, image_root, mask_root, coarse_root, simready_root): + directory.mkdir(parents=True, exist_ok=True) + + object_id = "asset_001" + scene_object = SceneObject( + id=object_id, + kind="asset", + category="asset", + name=text, + description=text, + ) + + image_generation_client = ImageGenerationClient.from_dotenv() + image_segmentation_client = ImageSegmentationClient.from_dotenv() + geometry_generation_client = GeometryGenerationClient.from_dotenv() + vlm_client = OpenAICompatibleVLM.from_dotenv() + try: + image_generation_client.check_health() + image_segmentation_client.check_health() + geometry_generation_client.check_health() + + # Generate a centered single-object image from the semantic text prompt. + image_path = image_generation_client.generate_image_by_prompt( + prompt=text, + output_path=image_root / f"{object_id}.png", + ) + with Image.open(image_path) as image: + image_size = image.size + + # Segment the generated object and apply the single-object foreground heuristic. + candidates = union_overlapping_mask_candidates( + build_mask_candidates( + image_segmentation_client.segment_single_object( + image_path=image_path, + prompt=text, + ) + ), + min_iou=0.8, + ) + if not candidates: + raise ValueError("Image segmentation returned no mask candidates.") + mask_path = save_binary_mask( + invert_mask_if_foreground_is_off_center(candidates[0]), + image_size=image_size, + output_path=mask_root / f"{object_id}.png", + ) + + # Generate one coarse GLB using the generated image and its binary mask. + geometry_generation_client.generate_objects( + image_path=image_path, + object_masks=[(object_id, mask_path)], + output_root=coarse_root, + ) + coarse_glb_path = coarse_root / f"{object_id}.glb" + if not coarse_glb_path.is_file(): + raise FileNotFoundError(f"Coarse GLB was not generated: {coarse_glb_path}") + + # Use identity coarse layout; VLM determines rotation and real-world size. + processor = SimReadyProcessor( + scene=Scene(objects=[scene_object]), + coarse_layout_by_id={ + object_id: { + "rot": [0.0, 0.0, 0.0], + "pos": [0.0, 0.0, 0.0], + "scale": [1.0, 1.0, 1.0], + } + }, + coarse_geometry_root=coarse_root, + simready_geometry_root=simready_root, + config=SimReadyProcessorConfig( + use_vlm_scale=True, + use_vlm_rotation=True, + ), + vlm_client=vlm_client, + ) + simready_layout = processor.process_assets() + (root / "result.json").write_text( + json.dumps( + { + "input_text": text, + "scene_object": scene_object.to_dict(), + "simready_layout": simready_layout, + }, + indent=2, + ensure_ascii=False, + ) + + "\n", + encoding="utf-8", + ) + return scene_object + finally: + image_generation_client.close() + image_segmentation_client.close() + geometry_generation_client.close() + + +def main(argv: Sequence[str] | None = None) -> None: + """Run the manual text-to-SimReady CLI.""" + parser = argparse.ArgumentParser( + prog="embodichain test-text-to-simready", + description="Debug text-to-image-to-segmentation-to-SimReady generation.", + ) + parser.add_argument("--text", required=True, help="Description of one object.") + parser.add_argument( + "--output_root", + required=True, + help="Directory for all intermediate and final artifacts.", + ) + args = parser.parse_args(argv) + scene_object = run_text_to_simready(text=args.text, output_root=args.output_root) + print(f"Generated SimReady asset: {scene_object.simready_glb_path}") + + +if __name__ == "__main__": + main() diff --git a/embodichain/gen_sim/scene_engine/pipeline/edit.py b/embodichain/gen_sim/scene_engine/pipeline/edit.py index fac2b6769..5b71af556 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/edit.py +++ b/embodichain/gen_sim/scene_engine/pipeline/edit.py @@ -33,12 +33,18 @@ from embodichain.gen_sim.scene_engine.pipeline.utils.scene_importer import ( SceneExportImporter, ) +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_exporter import ( + SceneExporter, +) from embodichain.gen_sim.scene_engine.pipeline.editing.scene_edit_understanding import ( understand_scene_edit, ) from embodichain.gen_sim.scene_engine.pipeline.editing.scene_edit_asset_preparation import ( prepare_scene_edit_assets, ) +from embodichain.gen_sim.scene_engine.pipeline.editing.scene_edit_layout_generation import ( + edit_layout, +) from embodichain.utils.logger import log_info @@ -95,20 +101,25 @@ def edit_scene( image_segmentation_client.close() log_info("Completed Objects Preparation") - # 3. Layout Editing - log_info("Starting Layout Editing") - # scene = edit_layout( - # scene=scene, - # edit_plan=edit_plan, - # scene_graph=updated_scene_graph, - # output_root=output_root, - # ) - log_info("Completed Layout Editing") + # 3. Layout Generation + log_info("Starting Layout Generation") + post_edit_scene = edit_layout( + scene=scene, + scene_edit_plan=scene_edit_plan, + updated_scene_graph=updated_scene_graph, + added_assets=added_assets, + output_root=resolved_output_root, + ) + log_info("Completed Layout Generation") # 4. Scene Export - # Re export the scene to the same output format, - # and delete some temporary files or folders. log_info("Starting Scene Export") + scene_exporter = SceneExporter( + scene=post_edit_scene, + scene_graph=updated_scene_graph, + output_root=resolved_output_root, + ) + scene_exporter.export() log_info("Completed Scene Export") return None diff --git a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_layout_generation.py b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_layout_generation.py new file mode 100644 index 000000000..9c6b06410 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_layout_generation.py @@ -0,0 +1,77 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +import shutil +from pathlib import Path + +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_edit_plan import SceneEditPlan +from embodichain.gen_sim.scene_engine.core.scene_graph import SceneGraph +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_constructor import ( + SceneLayoutConstructor, +) + + +def edit_layout( + *, + scene: Scene, + scene_edit_plan: SceneEditPlan, + updated_scene_graph: SceneGraph, + added_assets: list[SceneObject], + output_root: str | Path, +) -> Scene: + """Dispatch one edit-layout optimization from the goal scene graph.""" + formal_scene = scene + goal_scene_graph = updated_scene_graph + generated_scene_objects = added_assets + # Recreate this stage only when new assets need image, segmentation, and geometry outputs. + stage_output_root = ( + Path(output_root).expanduser().resolve() + / "scene_editing" + / "layout_optimization" + ) + if stage_output_root.exists(): + shutil.rmtree(stage_output_root) + stage_output_root.mkdir(parents=True, exist_ok=True) + + # Add and move operations are the only layout variables in this edit pass. + layout_variable_ids = { + operation.object_id + for operation in scene_edit_plan.operations + if operation.op in {"add", "move"} + } + if None in layout_variable_ids: + raise ValueError("Add and move operations must identify an object.") + layout_variable_ids = { + object_id for object_id in layout_variable_ids if object_id is not None + } + + # Optimize the layout constrained by the goal scene graph. + layout_constructor = SceneLayoutConstructor( + formal_scene=formal_scene, + goal_scene_graph=goal_scene_graph, + layout_variable_ids=layout_variable_ids, + generated_scene_objects=generated_scene_objects, + output_root=stage_output_root, + ) + # Optimize. + post_edit_scene = layout_constructor.construct() + + return post_edit_scene 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 4063bf92c..7451296b4 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py @@ -56,6 +56,7 @@ def __init__( self.export_root = self.output_root / "scene_export" self.scene_config_path: Path | None = None self.scene_graph_path: Path | None = None + self.scene_json_path: Path | None = None def export(self) -> Path: """Write a scene-only config and copy SimReady GLBs into ``mesh_assets``. @@ -119,6 +120,17 @@ def export(self) -> Path: encoding="utf-8", ) log_info(f"Exported scene graph: {self.scene_graph_path}") + self.scene_json_path = self.export_root / "scene.json" + self.scene_json_path.write_text( + json.dumps(self.scene.to_dict(), indent=2, ensure_ascii=False) + "\n", + encoding="utf-8", + ) + log_info(f"Exported scene JSON: {self.scene_json_path}") + # Remove only assets absent from the completed scene export. + self._remove_stale_mesh_assets( + mesh_assets_root=mesh_assets_root, + object_ids=set(object_ids), + ) return self.scene_config_path @staticmethod @@ -148,9 +160,28 @@ def _copy_scene_object_to_assets( ) 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) + # Imported assets already live at their export destination. + if not destination_glb_path.is_file() or not source_glb_path.samefile( + destination_glb_path + ): + shutil.copy2(source_glb_path, destination_glb_path) return destination_glb_path.relative_to(mesh_assets_root.parent).as_posix() + @staticmethod + def _remove_stale_mesh_assets( + *, + mesh_assets_root: Path, + object_ids: set[str], + ) -> None: + """Remove copied asset directories that no longer belong to the scene.""" + for asset_root in mesh_assets_root.iterdir(): + if asset_root.name in object_ids: + continue + if asset_root.is_dir(): + shutil.rmtree(asset_root) + else: + asset_root.unlink() + @staticmethod def _scene_object_config( *, @@ -179,6 +210,8 @@ def _scene_object_config( return { "uid": scene_object.id, + "category": scene_object.category, + "name": scene_object.name, "description": scene_object.description, "shape": { "shape_type": "Mesh", diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py index eefeee6b6..9ea558a2a 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py @@ -286,8 +286,16 @@ def _scene_object_from_export_entry( return SceneObject( id=uid, kind=kind, # type: ignore[arg-type] - category=uid, - name=uid, + category=self._semantic_text( + entry.get("category"), + field_name=f"{uid}.category", + default=uid, + ), + name=self._semantic_text( + entry.get("name"), + field_name=f"{uid}.name", + default=uid, + ), description=str(entry.get("description") or uid), simready_glb_path=str(glb_path), rot=rot_y_up.tolist(), @@ -346,6 +354,20 @@ def _vector3(value: object, *, field_name: str) -> list[float]: raise ValueError(f"Scene config field {field_name!r} must be finite.") return vector + @staticmethod + def _semantic_text( + value: object, + *, + field_name: str, + default: str, + ) -> str: + """Read one non-empty semantic label with a legacy-export fallback.""" + if value is None: + return default + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"Scene config field {field_name!r} must be non-empty.") + return value + @staticmethod def _vector2(value: object, *, field_name: str) -> list[float]: """Validate one length-2 numeric vector.""" diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_constructor.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_constructor.py new file mode 100644 index 000000000..38706ac47 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_constructor.py @@ -0,0 +1,399 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF 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 embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + TABLE_OBJECT_ID, + SceneGraph, +) +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_optimizer import ( + SceneLayoutOptimizerConfig, + SceneLayoutOptimizer, +) + + +@dataclass(frozen=True) +class SceneLayoutGroup: + """One parent and its direct on-children handled in one layout pass.""" + + parent_id: str + child_ids: list[str] + + +@dataclass(frozen=True) +class SceneLayoutProblem: + """Prepared graph-constrained inputs for one scene-layout construction.""" + + post_edit_scene: Scene + goal_scene_graph: SceneGraph + layout_variable_ids: set[str] + initial_xy_by_id: dict[str, list[float] | None] + groups: list[SceneLayoutGroup] + + +class SceneLayoutConstructor: + """Construct a scene layout from its goal graph. + + ``formal_scene`` may be empty for text-to-scene. In that case every table + and asset object must be supplied through ``generated_scene_objects``. + """ + + def __init__( + self, + *, + formal_scene: Scene, + goal_scene_graph: SceneGraph, + layout_variable_ids: set[str], + generated_scene_objects: list[SceneObject], + output_root: str | Path, + config: SceneLayoutOptimizerConfig | None = None, + ) -> None: + self.formal_scene = formal_scene + self.goal_scene_graph = goal_scene_graph + self.layout_variable_ids = layout_variable_ids + self.generated_scene_objects = generated_scene_objects + self.output_root = Path(output_root).expanduser().resolve() + self.layout_optimizer = SceneLayoutOptimizer(config=config) + self._current_xy_by_id: dict[str, list[float] | None] = {} + self._solved_delta_xy_by_id: dict[str, list[float]] = {} + self._updated_object_ids: set[str] = set() + + def construct(self) -> Scene: + """Construct table-root layouts before later stacked-group refinement.""" + layout_problem = self._build_problem() + self._current_xy_by_id = { + object_id: list(initial_xy) if initial_xy is not None else None + for object_id, initial_xy in layout_problem.initial_xy_by_id.items() + } + self._solved_delta_xy_by_id = {} + self._updated_object_ids = set() + if ( + layout_problem.groups + and layout_problem.groups[0].parent_id != TABLE_OBJECT_ID + ): + raise ValueError("The first layout group must be rooted at the table.") + + for group in layout_problem.groups: + if group.parent_id == TABLE_OBJECT_ID: + self._optimize_table_group( + layout_problem=layout_problem, + group=group, + ) + continue + self._optimize_parent_group( + layout_problem=layout_problem, + group=group, + ) + + return layout_problem.post_edit_scene + + def _optimize_table_group( + self, + *, + layout_problem: SceneLayoutProblem, + group: SceneLayoutGroup, + ) -> None: + """Optimize all direct on-table children before any stacked child groups.""" + table = layout_problem.post_edit_scene.table + if table is None: + raise ValueError("Table group optimization requires a table.") + if table.support_optimization_rect_xy is None: + raise ValueError( + "Table group optimization requires a table support optimization rectangle." + ) + + root_ids = set(group.child_ids) + root_relations = [ + relation + for relation in layout_problem.goal_scene_graph.relations + if relation.source_id in root_ids and relation.target_id in root_ids + ] + root_seed_xy_by_id: dict[str, list[float]] = {} + for root_id in group.child_ids: + inherited_xy = self._current_xy_by_id[root_id] + # New roots start from the table-local origin; imported roots keep their pose. + root_seed_xy_by_id[root_id] = ( + [0.0, 0.0] if inherited_xy is None else list(inherited_xy) + ) + self._current_xy_by_id[root_id] = root_seed_xy_by_id[root_id] + + nodes_by_id = layout_problem.goal_scene_graph.node_by_id() + solved_root_xy_by_id = self.layout_optimizer.optimize_table_root_xy( + assets_by_id={ + asset.id: asset for asset in layout_problem.post_edit_scene.assets + }, + root_ids=group.child_ids, + root_seed_xy_by_id=root_seed_xy_by_id, + imported_root_ids={ + root_id + for root_id in group.child_ids + if layout_problem.initial_xy_by_id[root_id] is not None + }, + fixed_root_xy_by_id={ + root_id: ( + None + if root_id in layout_problem.layout_variable_ids + else self._current_xy_by_id[root_id] + ) + for root_id in group.child_ids + }, + root_table_regions_by_id={ + root_id: nodes_by_id[root_id].table_region + for root_id in group.child_ids + }, + table_optimization_rect_xy=table.support_optimization_rect_xy, + root_relations=root_relations, + ) + if table.support_surface_z is None and any( + root_id in layout_problem.layout_variable_ids for root_id in group.child_ids + ): + raise ValueError("Table group optimization requires support_surface_z.") + assets_by_id = { + asset.id: asset for asset in layout_problem.post_edit_scene.assets + } + for root_id, solved_xy in solved_root_xy_by_id.items(): + seed_xy = root_seed_xy_by_id[root_id] + delta_xy = [ + solved_xy[0] - seed_xy[0], + solved_xy[1] - seed_xy[1], + ] + self._current_xy_by_id[root_id] = list(solved_xy) + self._solved_delta_xy_by_id[root_id] = delta_xy + if root_id in layout_problem.layout_variable_ids: + # Direct add/move roots receive a new pose on the table support. + assert table.support_surface_z is not None + self.layout_optimizer.update_scene_object_y_up_pose_from_z_up_support( + scene_object=assets_by_id[root_id], + support_region_z=table.support_surface_z, + center_xy=solved_xy, + ) + self._updated_object_ids.add(root_id) + self._propagate_descendant_delta( + scene=layout_problem.post_edit_scene, + root_id=root_id, + delta_xy=delta_xy, + ) + + def _propagate_descendant_delta( + self, + *, + scene: Scene, + root_id: str, + delta_xy: list[float], + ) -> None: + """Move every positioned descendant by one solved ancestor XY delta.""" + if delta_xy == [0.0, 0.0]: + return + assets_by_id = {asset.id: asset for asset in scene.assets} + children_by_parent: dict[str, list[str]] = {} + for node in self.goal_scene_graph.nodes: + if node.parent_id is not None: + children_by_parent.setdefault(node.parent_id, []).append(node.object_id) + + pending = list(children_by_parent.get(root_id, [])) + while pending: + descendant_id = pending.pop(0) + descendant_xy = self._current_xy_by_id[descendant_id] + if descendant_xy is not None: + self._current_xy_by_id[descendant_id] = [ + descendant_xy[0] + delta_xy[0], + descendant_xy[1] + delta_xy[1], + ] + self.layout_optimizer.translate_scene_object_y_up_by_z_up_delta( + scene_object=assets_by_id[descendant_id], + delta_xy=delta_xy, + ) + self._updated_object_ids.add(descendant_id) + pending.extend(children_by_parent.get(descendant_id, [])) + + def _optimize_parent_group( + self, + *, + layout_problem: SceneLayoutProblem, + group: SceneLayoutGroup, + ) -> None: + """Optimize one settled parent's direct on-children in local XY coordinates.""" + assets_by_id = { + asset.id: asset for asset in layout_problem.post_edit_scene.assets + } + parent = assets_by_id.get(group.parent_id) + if parent is None: + raise ValueError(f"Parent {group.parent_id!r} is not an asset.") + parent_aabb = self.layout_optimizer.scene_object_z_up_world_aabb( + scene_object=parent + ) + parent_aabb_xy = [ + [parent_aabb[0][0], parent_aabb[0][1]], + [parent_aabb[1][0], parent_aabb[1][1]], + ] + parent_center_xy = [ + (parent_aabb[0][0] + parent_aabb[1][0]) / 2.0, + (parent_aabb[0][1] + parent_aabb[1][1]) / 2.0, + ] + child_seed_xy_by_id: dict[str, list[float]] = {} + for child_id in group.child_ids: + inherited_xy = self._current_xy_by_id[child_id] + # New children start at their parent's current AABB center. + child_seed_xy_by_id[child_id] = ( + parent_center_xy if inherited_xy is None else list(inherited_xy) + ) + self._current_xy_by_id[child_id] = child_seed_xy_by_id[child_id] + + solved_child_xy_by_id = self.layout_optimizer.optimize_parent_child_xy( + assets_by_id=assets_by_id, + child_ids=group.child_ids, + child_seed_xy_by_id=child_seed_xy_by_id, + imported_child_ids={ + child_id + for child_id in group.child_ids + if layout_problem.initial_xy_by_id[child_id] is not None + }, + fixed_child_xy_by_id={ + child_id: ( + None + if child_id in layout_problem.layout_variable_ids + else self._current_xy_by_id[child_id] + ) + for child_id in group.child_ids + }, + parent_aabb_xy=parent_aabb_xy, + ) + parent_top_z = parent_aabb[1][2] + for child_id, solved_xy in solved_child_xy_by_id.items(): + seed_xy = child_seed_xy_by_id[child_id] + delta_xy = [ + solved_xy[0] - seed_xy[0], + solved_xy[1] - seed_xy[1], + ] + self._current_xy_by_id[child_id] = list(solved_xy) + self._solved_delta_xy_by_id[child_id] = delta_xy + if child_id in layout_problem.layout_variable_ids: + # Variable children are placed directly above the parent's current top. + self.layout_optimizer.update_scene_object_y_up_pose_from_z_up_support( + scene_object=assets_by_id[child_id], + support_region_z=parent_top_z, + center_xy=solved_xy, + ) + self._updated_object_ids.add(child_id) + self._propagate_descendant_delta( + scene=layout_problem.post_edit_scene, + root_id=child_id, + delta_xy=delta_xy, + ) + + def _build_problem(self) -> SceneLayoutProblem: + """Build post-edit objects and preserve formal-scene centers as seeds.""" + self.goal_scene_graph.validate() + graph_object_ids = set(self.goal_scene_graph.node_by_id()) + generated_objects_by_id = self._generated_scene_objects_by_id() + + # The goal graph removes deleted formal-scene objects from the layout input. + post_edit_objects = [ + scene_object + for scene_object in self.formal_scene.objects + if scene_object.id in graph_object_ids + ] + imported_object_ids = {scene_object.id for scene_object in post_edit_objects} + if imported_object_ids.intersection(generated_objects_by_id): + raise ValueError( + "Generated scene objects must not reuse formal scene object ids." + ) + post_edit_objects.extend(generated_objects_by_id.values()) + + post_edit_scene = Scene(objects=post_edit_objects) + post_edit_object_ids = { + scene_object.id for scene_object in post_edit_scene.objects + } + if post_edit_object_ids != graph_object_ids: + raise ValueError("Goal scene graph and post-edit scene have different ids.") + if not self.layout_variable_ids.issubset(post_edit_object_ids - {"table"}): + raise ValueError( + "Only post-edit assets may participate in layout optimization." + ) + + initial_xy_by_id = { + asset.id: self._initial_xy( + asset, + is_generated=asset.id in generated_objects_by_id, + ) + for asset in post_edit_scene.assets + } + for object_id, initial_xy in initial_xy_by_id.items(): + if initial_xy is None and object_id not in self.layout_variable_ids: + raise ValueError( + f"New asset {object_id!r} must participate in layout optimization." + ) + + return SceneLayoutProblem( + post_edit_scene=post_edit_scene, + goal_scene_graph=self.goal_scene_graph, + layout_variable_ids=set(self.layout_variable_ids), + initial_xy_by_id=initial_xy_by_id, + groups=self._build_groups(), + ) + + def _build_groups(self) -> list[SceneLayoutGroup]: + """Build table-rooted BFS groups of direct on-children.""" + children_by_parent: dict[str, list[str]] = {} + for node in self.goal_scene_graph.nodes: + if node.parent_id is None: + continue + if node.parent_relation != "on": + raise ValueError(f"Node {node.object_id!r} must be on its parent.") + children_by_parent.setdefault(node.parent_id, []).append(node.object_id) + + groups: list[SceneLayoutGroup] = [] + pending = [TABLE_OBJECT_ID] + while pending: + parent_id = pending.pop(0) + child_ids = children_by_parent.get(parent_id, []) + if not child_ids: + continue + groups.append(SceneLayoutGroup(parent_id=parent_id, child_ids=child_ids)) + pending.extend(child_ids) + return groups + + def _generated_scene_objects_by_id(self) -> dict[str, SceneObject]: + """Index generated scene objects before merging them into the formal scene.""" + generated_objects_by_id = { + scene_object.id: scene_object + for scene_object in self.generated_scene_objects + } + if len(generated_objects_by_id) != len(self.generated_scene_objects): + raise ValueError("Generated scene objects must use unique object ids.") + return generated_objects_by_id + + @staticmethod + def _initial_xy( + asset: SceneObject, + *, + is_generated: bool, + ) -> list[float] | None: + """Retain formal-scene centers while generated assets await initialization.""" + if is_generated: + return None + if asset.center_xy is None or len(asset.center_xy) != 2: + raise ValueError( + f"Formal-scene asset {asset.id!r} must have a 2D center_xy." + ) + return [float(value) for value in asset.center_xy] diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_optimizer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_optimizer.py new file mode 100644 index 000000000..85e98f0e8 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_optimizer.py @@ -0,0 +1,785 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF 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 numpy as np +from scipy.optimize import minimize + +from embodichain.gen_sim.scene_engine.core.scene_graph import SceneGraphRelation +from embodichain.gen_sim.scene_engine.core.scene_object import 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, +) + + +@dataclass(frozen=True) +class SceneLayoutOptimizerConfig: + """Numerical controls shared by each graph-layout solve.""" + + relation_clearance_m: float = 0.03 + collision_margin_m: float = 0.02 + max_slsqp_iterations: int = 500 + slsqp_ftol: float = 1e-6 + max_collision_rounds: int = 8 + max_added_collision_pairs: int = 64 + imported_seed_weight: float = 5.0 + min_center_distance_m: float = 0.01 + min_center_distance_weight: float = 0.05 + + def __post_init__(self) -> None: + """Reject invalid numerical controls before assembling a layout problem.""" + if self.relation_clearance_m < 0.0: + raise ValueError("relation_clearance_m must be non-negative.") + if self.collision_margin_m < 0.0: + raise ValueError("collision_margin_m must be non-negative.") + if self.max_slsqp_iterations <= 0: + raise ValueError("max_slsqp_iterations must be positive.") + if self.slsqp_ftol <= 0.0: + raise ValueError("slsqp_ftol must be positive.") + if self.max_collision_rounds <= 0: + raise ValueError("max_collision_rounds must be positive.") + if self.max_added_collision_pairs <= 0: + raise ValueError("max_added_collision_pairs must be positive.") + + +class SceneLayoutOptimizer: + """Solve graph-constrained XY layouts and apply resulting poses.""" + + def __init__(self, *, config: SceneLayoutOptimizerConfig | None = None) -> None: + self.config = config if config is not None else SceneLayoutOptimizerConfig() + + def optimize_table_root_xy( + self, + *, + assets_by_id: dict[str, SceneObject], + root_ids: list[str], + root_seed_xy_by_id: dict[str, list[float]], + imported_root_ids: set[str], + fixed_root_xy_by_id: dict[str, list[float] | None], + root_table_regions_by_id: dict[str, str | None], + table_optimization_rect_xy: list[list[float]], + root_relations: list[SceneGraphRelation], + ) -> dict[str, list[float]]: + """Solve direct table-child centers with graph and AABB constraints.""" + return _optimize_table_root_xy( + assets_by_id=assets_by_id, + root_ids=root_ids, + root_seed_xy_by_id=root_seed_xy_by_id, + imported_root_ids=imported_root_ids, + fixed_root_xy_by_id=fixed_root_xy_by_id, + root_table_regions_by_id=root_table_regions_by_id, + table_optimization_rect_xy=table_optimization_rect_xy, + root_relations=root_relations, + config=self.config, + ) + + def optimize_parent_child_xy( + self, + *, + assets_by_id: dict[str, SceneObject], + child_ids: list[str], + child_seed_xy_by_id: dict[str, list[float]], + imported_child_ids: set[str], + fixed_child_xy_by_id: dict[str, list[float] | None], + parent_aabb_xy: list[list[float]], + ) -> dict[str, list[float]]: + """Solve direct on-children inside one parent's current XY AABB.""" + child_half_extents_xy = _asset_half_extents_xy( + assets_by_id=assets_by_id, + object_ids=child_ids, + ) + inequality_constraints: list[tuple[np.ndarray, float]] = [] + equality_constraints: list[tuple[np.ndarray, float]] = [] + child_index = {child_id: index for index, child_id in enumerate(child_ids)} + parent_bounds = _bounds_from_points(parent_aabb_xy) + for child_id in child_ids: + _append_aabb_center_bounds( + constraints=inequality_constraints, + root_index=child_index, + root_id=child_id, + bounds=parent_bounds, + half_extents_xy=child_half_extents_xy[child_id], + ) + fixed_xy = fixed_child_xy_by_id[child_id] + if fixed_xy is not None: + _append_fixed_root_constraints( + constraints=equality_constraints, + root_index=child_index, + root_id=child_id, + fixed_xy=fixed_xy, + ) + + solved_child_xy_by_id = _solve_root_xy( + root_ids=child_ids, + root_seed_xy_by_id=child_seed_xy_by_id, + imported_root_ids=imported_child_ids, + inequality_constraints=inequality_constraints, + equality_constraints=equality_constraints, + config=self.config, + ) + return _refine_root_collisions( + root_ids=child_ids, + root_seed_xy_by_id=child_seed_xy_by_id, + imported_root_ids=imported_child_ids, + root_half_extents_xy=child_half_extents_xy, + inequality_constraints=inequality_constraints, + equality_constraints=equality_constraints, + solved_root_xy_by_id=solved_child_xy_by_id, + config=self.config, + ) + + @staticmethod + def scene_object_z_up_world_aabb( + *, + scene_object: SceneObject, + ) -> list[list[float]]: + """Return one object's current z-up world AABB as [min, max].""" + return _scene_object_z_up_world_aabb(scene_object=scene_object) + + @staticmethod + def update_scene_object_y_up_pose_from_z_up_support( + *, + scene_object: SceneObject, + support_region_z: float, + center_xy: list[float], + clearance_m: float = 0.02, + ) -> None: + """Place one SimReady asset on a horizontal z-up support region.""" + _update_scene_object_y_up_pose_from_z_up_support( + scene_object=scene_object, + support_region_z=support_region_z, + center_xy=center_xy, + clearance_m=clearance_m, + ) + + @staticmethod + def translate_scene_object_y_up_by_z_up_delta( + *, + scene_object: SceneObject, + delta_xy: list[float], + ) -> None: + """Translate one existing y-up pose by a solved z-up XY delta.""" + _translate_scene_object_y_up_by_z_up_delta( + scene_object=scene_object, + delta_xy=delta_xy, + ) + + +def _optimize_table_root_xy( + *, + assets_by_id: dict[str, SceneObject], + root_ids: list[str], + root_seed_xy_by_id: dict[str, list[float]], + imported_root_ids: set[str], + fixed_root_xy_by_id: dict[str, list[float] | None], + root_table_regions_by_id: dict[str, str | None], + table_optimization_rect_xy: list[list[float]], + root_relations: list[SceneGraphRelation], + config: SceneLayoutOptimizerConfig, +) -> dict[str, list[float]]: + """Solve direct table-child centers with graph and AABB constraints.""" + root_half_extents_xy = _asset_half_extents_xy( + assets_by_id=assets_by_id, + object_ids=root_ids, + ) + inequality_constraints, equality_constraints = _build_table_root_constraints( + root_ids=root_ids, + root_half_extents_xy=root_half_extents_xy, + root_relations=root_relations, + root_table_regions_by_id=root_table_regions_by_id, + table_optimization_rect_xy=table_optimization_rect_xy, + fixed_root_xy_by_id=fixed_root_xy_by_id, + config=config, + ) + solved_root_xy_by_id = _solve_root_xy( + root_ids=root_ids, + root_seed_xy_by_id=root_seed_xy_by_id, + imported_root_ids=imported_root_ids, + inequality_constraints=inequality_constraints, + equality_constraints=equality_constraints, + config=config, + ) + return _refine_root_collisions( + root_ids=root_ids, + root_seed_xy_by_id=root_seed_xy_by_id, + imported_root_ids=imported_root_ids, + root_half_extents_xy=root_half_extents_xy, + inequality_constraints=inequality_constraints, + equality_constraints=equality_constraints, + solved_root_xy_by_id=solved_root_xy_by_id, + config=config, + ) + + +def _update_scene_object_y_up_pose_from_z_up_support( + *, + scene_object: SceneObject, + support_region_z: float, + center_xy: list[float], + clearance_m: float = 0.02, +) -> None: + """Place one SimReady asset on a horizontal z-up support region. + + ``SceneObject`` stores poses in y-up before export. The target center and + support height are z-up values because layout optimization uses that frame. + """ + if not np.isfinite(support_region_z): + raise ValueError("support_region_z must be finite.") + if clearance_m < 0.0 or not np.isfinite(clearance_m): + raise ValueError("clearance_m must be finite and non-negative.") + target_xy = _two_floats(center_xy, field_name="center_xy") + rotation_y_up = _three_floats_or_default( + scene_object.rot, + field_name="rot", + default=[0.0, 0.0, 0.0], + ) + mesh = _asset_z_up_mesh_at_zero_translation( + scene_object=scene_object, + rotation_y_up=rotation_y_up, + ) + target_position_z_up = np.array( + [ + target_xy[0] - float(mesh.bounds[:, 0].mean()), + target_xy[1] - float(mesh.bounds[:, 1].mean()), + float(support_region_z) + clearance_m - float(mesh.bounds[0, 2]), + ] + ) + z_up_to_y_up = np.linalg.inv(_y_up_to_z_up_matrix()) + # Persist the y-up pose that SceneExporter later converts back to z-up. + scene_object.pos = (z_up_to_y_up[:3, :3] @ target_position_z_up).tolist() + scene_object.rot = rotation_y_up + scene_object.center_xy = target_xy + + +def _translate_scene_object_y_up_by_z_up_delta( + *, + scene_object: SceneObject, + delta_xy: list[float], +) -> None: + """Translate one existing y-up pose by a solved z-up XY delta.""" + dx, dy = _two_floats(delta_xy, field_name="delta_xy") + current_pos = _three_floats_or_default( + scene_object.pos, + field_name="pos", + default=None, + ) + # z-up x maps to y-up x, while z-up y maps to negative y-up z. + scene_object.pos = [ + current_pos[0] + dx, + current_pos[1], + current_pos[2] - dy, + ] + if scene_object.center_xy is not None: + scene_object.center_xy = [ + scene_object.center_xy[0] + dx, + scene_object.center_xy[1] + dy, + ] + + +def _scene_object_z_up_world_aabb( + *, + scene_object: SceneObject, +) -> list[list[float]]: + """Measure one current SceneObject pose in z-up world coordinates.""" + position_y_up = _three_floats_or_default( + scene_object.pos, + field_name="pos", + default=None, + ) + mesh = _asset_z_up_mesh_at_zero_translation(scene_object=scene_object) + position_z_up = _y_up_to_z_up_matrix()[:3, :3] @ np.asarray( + position_y_up, + dtype=float, + ) + mesh.apply_translation(position_z_up) + return mesh.bounds.tolist() + + +def _build_table_root_constraints( + *, + root_ids: list[str], + root_half_extents_xy: dict[str, np.ndarray], + root_relations: list[SceneGraphRelation], + root_table_regions_by_id: dict[str, str | None], + table_optimization_rect_xy: list[list[float]], + fixed_root_xy_by_id: dict[str, list[float] | None], + config: SceneLayoutOptimizerConfig, +) -> tuple[list[tuple[np.ndarray, float]], list[tuple[np.ndarray, float]]]: + """Build hard table, region, planar, and fixed-root constraints.""" + root_index = {root_id: index for index, root_id in enumerate(root_ids)} + table_bounds = _bounds_from_points(table_optimization_rect_xy) + inequality_constraints: list[tuple[np.ndarray, float]] = [] + equality_constraints: list[tuple[np.ndarray, float]] = [] + + for root_id in root_ids: + region_bounds = _table_region_bounds( + table_bounds=table_bounds, + table_region=root_table_regions_by_id[root_id], + ) + _append_aabb_center_bounds( + constraints=inequality_constraints, + root_index=root_index, + root_id=root_id, + bounds=region_bounds, + half_extents_xy=root_half_extents_xy[root_id], + ) + fixed_xy = fixed_root_xy_by_id[root_id] + if fixed_xy is not None: + _append_fixed_root_constraints( + constraints=equality_constraints, + root_index=root_index, + root_id=root_id, + fixed_xy=fixed_xy, + ) + + for relation in root_relations: + _append_planar_relation_constraint( + constraints=inequality_constraints, + root_index=root_index, + source_id=relation.source_id, + relation=relation.relation, + target_id=relation.target_id, + source_half_extents_xy=root_half_extents_xy[relation.source_id], + target_half_extents_xy=root_half_extents_xy[relation.target_id], + relation_clearance_m=config.relation_clearance_m, + ) + return inequality_constraints, equality_constraints + + +def _solve_root_xy( + *, + root_ids: list[str], + root_seed_xy_by_id: dict[str, list[float]], + imported_root_ids: set[str], + inequality_constraints: list[tuple[np.ndarray, float]], + equality_constraints: list[tuple[np.ndarray, float]], + config: SceneLayoutOptimizerConfig, +) -> dict[str, list[float]]: + """Solve one root-group center model with the legacy SLSQP settings.""" + root_index = {root_id: index for index, root_id in enumerate(root_ids)} + initial_xy = np.asarray( + [root_seed_xy_by_id[root_id] for root_id in root_ids], dtype=float + ) + x0 = initial_xy.reshape(-1) + + def unpack(values: np.ndarray) -> dict[str, list[float]]: + return { + root_id: [float(values[2 * index]), float(values[2 * index + 1])] + for root_id, index in root_index.items() + } + + def objective(values: np.ndarray) -> float: + coordinates = values.reshape(-1, 2) + loss = 0.0 + for root_id, index in root_index.items(): + if root_id in imported_root_ids: + delta = coordinates[index] - initial_xy[index] + loss += config.imported_seed_weight * float(delta @ delta) + for first_index in range(len(root_ids)): + for second_index in range(first_index + 1, len(root_ids)): + distance = float( + np.linalg.norm(coordinates[first_index] - coordinates[second_index]) + ) + shortfall = max(0.0, config.min_center_distance_m - distance) + loss += config.min_center_distance_weight * shortfall**2 + return loss + + constraints: list[dict[str, object]] = [] + for row, bound in inequality_constraints: + constraints.append( + { + "type": "ineq", + "fun": lambda values, row=row, bound=bound: bound - float(row @ values), + } + ) + for row, bound in equality_constraints: + constraints.append( + { + "type": "eq", + "fun": lambda values, row=row, bound=bound: float(row @ values) - bound, + } + ) + + result = minimize( + objective, + x0, + method="SLSQP", + constraints=constraints, + options={ + "maxiter": config.max_slsqp_iterations, + "ftol": config.slsqp_ftol, + "disp": False, + }, + ) + if not result.success: + raise ValueError(f"Table layout optimization failed: {result.message}") + return unpack(np.asarray(result.x, dtype=float)) + + +def _refine_root_collisions( + *, + root_ids: list[str], + root_seed_xy_by_id: dict[str, list[float]], + imported_root_ids: set[str], + root_half_extents_xy: dict[str, np.ndarray], + inequality_constraints: list[tuple[np.ndarray, float]], + equality_constraints: list[tuple[np.ndarray, float]], + solved_root_xy_by_id: dict[str, list[float]], + config: SceneLayoutOptimizerConfig, +) -> dict[str, list[float]]: + """Add AABB separation constraints until the table roots no longer overlap.""" + seen_pairs: set[tuple[str, str]] = set() + current_xy_by_id = solved_root_xy_by_id + for _ in range(config.max_collision_rounds): + overlaps = _root_aabb_overlaps( + root_ids=root_ids, + root_half_extents_xy=root_half_extents_xy, + xy_by_id=current_xy_by_id, + ) + if not overlaps: + return current_xy_by_id + added_constraint_count = 0 + for _, first_id, second_id in overlaps[: config.max_added_collision_pairs]: + pair_key = tuple(sorted((first_id, second_id))) + if pair_key in seen_pairs: + continue + inequality_constraints.append( + _aabb_separation_constraint( + root_ids=root_ids, + first_id=first_id, + second_id=second_id, + first_half_extents_xy=root_half_extents_xy[first_id], + second_half_extents_xy=root_half_extents_xy[second_id], + first_xy=current_xy_by_id[first_id], + second_xy=current_xy_by_id[second_id], + collision_margin_m=config.collision_margin_m, + ) + ) + seen_pairs.add(pair_key) + added_constraint_count += 1 + if added_constraint_count == 0: + break + current_xy_by_id = _solve_root_xy( + root_ids=root_ids, + root_seed_xy_by_id=current_xy_by_id, + imported_root_ids=imported_root_ids, + inequality_constraints=inequality_constraints, + equality_constraints=equality_constraints, + config=config, + ) + + remaining_pairs = [ + f"{first_id}/{second_id}" + for _, first_id, second_id in _root_aabb_overlaps( + root_ids=root_ids, + root_half_extents_xy=root_half_extents_xy, + xy_by_id=current_xy_by_id, + ) + ] + raise ValueError( + "Table-root AABB collisions remain after layout refinement: " + f"{remaining_pairs}." + ) + + +def _asset_half_extents_xy( + *, + assets_by_id: dict[str, SceneObject], + object_ids: list[str], +) -> dict[str, np.ndarray]: + """Measure each asset's oriented z-up footprint around its XY center.""" + half_extents_xy: dict[str, np.ndarray] = {} + for object_id in object_ids: + asset = assets_by_id.get(object_id) + if asset is None: + raise ValueError(f"Table root {object_id!r} is not an asset.") + half_extents_xy[object_id] = _asset_half_extent_xy(asset) + return half_extents_xy + + +def _asset_half_extent_xy(asset: SceneObject) -> np.ndarray: + """Measure one SimReady GLB with its current orientation and scale.""" + mesh = _asset_z_up_mesh_at_zero_translation(scene_object=asset) + return (mesh.bounds[1, :2] - mesh.bounds[0, :2]) / 2.0 + + +def _asset_z_up_mesh_at_zero_translation( + *, + scene_object: SceneObject, + rotation_y_up: list[float] | None = None, +): + """Load one SimReady GLB in z-up with orientation and scale but no position.""" + asset = scene_object + if asset.simready_glb_path is None: + raise ValueError(f"Asset {asset.id!r} has no SimReady GLB path.") + y_up_layout = { + "id": asset.id, + "rot": ( + rotation_y_up + if rotation_y_up is not None + else _three_floats_or_default( + asset.rot, + field_name="rot", + default=[0.0, 0.0, 0.0], + ) + ), + "pos": [0.0, 0.0, 0.0], + "scale": _three_floats_or_default( + asset.scale, + field_name="scale", + default=[1.0, 1.0, 1.0], + ), + } + y_up_to_z_up = _y_up_to_z_up_matrix() + z_up_layout = transform_matrix_to_layout_object( + asset.id, + y_up_to_z_up + @ layout_object_to_transform_matrix(y_up_layout) + @ np.linalg.inv(y_up_to_z_up), + ) + mesh = load_glb_mesh(asset.simready_glb_path) + mesh.apply_transform(y_up_to_z_up) + mesh.apply_transform(layout_object_to_transform_matrix(z_up_layout)) + return mesh + + +def _y_up_to_z_up_matrix() -> np.ndarray: + """Return the coordinate transform used by SceneExporter and layout stages.""" + matrix = np.eye(4) + matrix[:3, :3] = np.array([[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]]) + return matrix + + +def _two_floats(value: object, *, field_name: str) -> list[float]: + """Validate one finite two-value vector.""" + if not isinstance(value, (list, tuple)) or len(value) != 2: + raise ValueError(f"{field_name} must contain two values.") + vector = [float(component) for component in value] + if not np.all(np.isfinite(vector)): + raise ValueError(f"{field_name} must contain finite values.") + return vector + + +def _three_floats_or_default( + value: object, + *, + field_name: str, + default: list[float] | None, +) -> list[float]: + """Return a finite three-value vector or the canonical SimReady default.""" + if value is None: + if default is None: + raise ValueError(f"{field_name} must contain three values.") + return list(default) + if not isinstance(value, (list, tuple)) or len(value) != 3: + raise ValueError(f"{field_name} must contain three values.") + vector = [float(component) for component in value] + if not np.all(np.isfinite(vector)): + raise ValueError(f"{field_name} must contain finite values.") + return vector + + +def _bounds_from_points(points: list[list[float]]) -> np.ndarray: + """Return [[min_x, min_y], [max_x, max_y]] from finite XY points.""" + coordinates = np.asarray(points, dtype=float) + if coordinates.ndim != 2 or coordinates.shape[1] != 2 or len(coordinates) < 2: + raise ValueError("XY bounds must contain at least two points.") + if not np.all(np.isfinite(coordinates)): + raise ValueError("XY bounds must contain finite values.") + return np.stack([coordinates.min(axis=0), coordinates.max(axis=0)]) + + +def _table_region_bounds( + *, + table_bounds: np.ndarray, + table_region: str | None, +) -> np.ndarray: + """Return the requested 3x3 table region, with y increasing toward front.""" + if table_region is None: + return table_bounds.copy() + column_by_region = { + "left_back": 0, + "left_center": 0, + "left_front": 0, + "back_center": 1, + "center": 1, + "front_center": 1, + "right_back": 2, + "right_center": 2, + "right_front": 2, + } + row_by_region = { + "left_front": 0, + "front_center": 0, + "right_front": 0, + "left_center": 1, + "center": 1, + "right_center": 1, + "left_back": 2, + "back_center": 2, + "right_back": 2, + } + if table_region not in column_by_region: + raise ValueError(f"Unsupported table region {table_region!r}.") + minimum, maximum = table_bounds + cell_size = (maximum - minimum) / 3.0 + region_minimum = minimum + cell_size * np.array( + [column_by_region[table_region], row_by_region[table_region]] + ) + return np.stack([region_minimum, region_minimum + cell_size]) + + +def _append_aabb_center_bounds( + *, + constraints: list[tuple[np.ndarray, float]], + root_index: dict[str, int], + root_id: str, + bounds: np.ndarray, + half_extents_xy: np.ndarray, +) -> None: + """Keep one root's complete AABB inside the given rectangular bounds.""" + minimum = bounds[0] + half_extents_xy + maximum = bounds[1] - half_extents_xy + if np.any(minimum > maximum): + raise ValueError( + f"Asset {root_id!r} cannot fit inside its assigned table region." + ) + variable_count = 2 * len(root_index) + root_offset = 2 * root_index[root_id] + for axis in range(2): + upper_row = np.zeros(variable_count) + upper_row[root_offset + axis] = 1.0 + constraints.append((upper_row, float(maximum[axis]))) + lower_row = np.zeros(variable_count) + lower_row[root_offset + axis] = -1.0 + constraints.append((lower_row, -float(minimum[axis]))) + + +def _append_fixed_root_constraints( + *, + constraints: list[tuple[np.ndarray, float]], + root_index: dict[str, int], + root_id: str, + fixed_xy: list[float], +) -> None: + """Use equality constraints so unchanged formal objects remain fixed.""" + variable_count = 2 * len(root_index) + root_offset = 2 * root_index[root_id] + for axis, coordinate in enumerate(fixed_xy): + row = np.zeros(variable_count) + row[root_offset + axis] = 1.0 + constraints.append((row, float(coordinate))) + + +def _append_planar_relation_constraint( + *, + constraints: list[tuple[np.ndarray, float]], + root_index: dict[str, int], + source_id: str, + relation: str, + target_id: str, + source_half_extents_xy: np.ndarray, + target_half_extents_xy: np.ndarray, + relation_clearance_m: float, +) -> None: + """Require directional relations to clear both sibling AABB footprints.""" + if source_id not in root_index or target_id not in root_index: + raise ValueError("Table-root planar relations must reference table roots.") + axis, source_sign = { + "left_of": (0, 1.0), + "right_of": (0, -1.0), + "behind": (1, 1.0), + "in_front_of": (1, -1.0), + }.get(relation, (None, None)) + if axis is None or source_sign is None: + raise ValueError(f"Unsupported planar relation {relation!r}.") + row = np.zeros(2 * len(root_index)) + row[2 * root_index[source_id] + axis] = source_sign + row[2 * root_index[target_id] + axis] = -source_sign + required_distance = ( + source_half_extents_xy[axis] + + target_half_extents_xy[axis] + + relation_clearance_m + ) + constraints.append((row, -float(required_distance))) + + +def _root_aabb_overlaps( + *, + root_ids: list[str], + root_half_extents_xy: dict[str, np.ndarray], + xy_by_id: dict[str, list[float]], +) -> list[tuple[float, str, str]]: + """Return root pairs whose current XY AABBs overlap without a margin.""" + overlaps: list[tuple[float, str, str]] = [] + for first_index, first_id in enumerate(root_ids): + first_xy = np.asarray(xy_by_id[first_id], dtype=float) + first_half_extents = root_half_extents_xy[first_id] + for second_id in root_ids[first_index + 1 :]: + second_xy = np.asarray(xy_by_id[second_id], dtype=float) + second_half_extents = root_half_extents_xy[second_id] + overlap_xy = np.minimum( + first_xy + first_half_extents, + second_xy + second_half_extents, + ) - np.maximum( + first_xy - first_half_extents, + second_xy - second_half_extents, + ) + if np.all(overlap_xy > 1e-9): + overlaps.append((float(np.min(overlap_xy)), first_id, second_id)) + return sorted(overlaps, reverse=True) + + +def _aabb_separation_constraint( + *, + root_ids: list[str], + first_id: str, + second_id: str, + first_half_extents_xy: np.ndarray, + second_half_extents_xy: np.ndarray, + first_xy: list[float], + second_xy: list[float], + collision_margin_m: float, +) -> tuple[np.ndarray, float]: + """Separate one overlapping pair along its shallowest penetration axis.""" + root_index = {root_id: index for index, root_id in enumerate(root_ids)} + first_xy_array = np.asarray(first_xy, dtype=float) + second_xy_array = np.asarray(second_xy, dtype=float) + overlap_xy = np.minimum( + first_xy_array + first_half_extents_xy, + second_xy_array + second_half_extents_xy, + ) - np.maximum( + first_xy_array - first_half_extents_xy, + second_xy_array - second_half_extents_xy, + ) + axis = int(np.argmin(overlap_xy)) + first_is_lower = first_xy_array[axis] < second_xy_array[axis] or ( + first_xy_array[axis] == second_xy_array[axis] and first_id < second_id + ) + row = np.zeros(2 * len(root_ids)) + first_coefficient = 1.0 if first_is_lower else -1.0 + row[2 * root_index[first_id] + axis] = first_coefficient + row[2 * root_index[second_id] + axis] = -first_coefficient + required_distance = ( + first_half_extents_xy[axis] + second_half_extents_xy[axis] + collision_margin_m + ) + return row, -float(required_distance) 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 b7b614d49..b50126751 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 @@ -161,6 +161,8 @@ def test_scene_export_copies_meshes_and_converts_y_up_pose(tmp_path: Path) -> No 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["category"] == "asset" + assert entry["name"] == "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] @@ -188,9 +190,76 @@ def test_scene_export_copies_meshes_and_converts_y_up_pose(tmp_path: Path) -> No output_root=tmp_path / "output" ).import_scene_and_graph() assert [asset.id for asset in imported_scene.assets] == ["cup"] + assert imported_scene.assets[0].category == "asset" + assert imported_scene.assets[0].name == "cup" assert imported_graph.to_dict() == _scene_graph(scene).to_dict() +def test_scene_export_overwrites_an_existing_scene_export(tmp_path: Path) -> None: + table_glb = tmp_path / "table.glb" + cup_glb = tmp_path / "cup.glb" + banana_glb = tmp_path / "banana.glb" + table_glb.write_bytes(b"glTF-table") + cup_glb.write_bytes(b"glTF-cup") + banana_glb.write_bytes(b"glTF-banana") + output_root = tmp_path / "output" + + initial_table = _scene_object( + object_id="table", + kind="table", + glb_path=table_glb, + physics=_physics("kinematic"), + ) + initial_cup = _scene_object( + object_id="cup", + kind="asset", + glb_path=cup_glb, + physics=_physics("dynamic"), + ) + initial_scene = Scene(objects=[initial_table, initial_cup]) + SceneExporter( + scene=initial_scene, + scene_graph=_scene_graph(initial_scene), + output_root=output_root, + ).export() + + # The imported table mesh already occupies its final export location. + exported_table_glb = ( + output_root / "scene_export" / "mesh_assets" / "table" / "table.glb" + ) + updated_table = _scene_object( + object_id="table", + kind="table", + glb_path=exported_table_glb, + physics=_physics("kinematic"), + ) + banana = _scene_object( + object_id="banana", + kind="asset", + glb_path=banana_glb, + physics=_physics("dynamic"), + ) + updated_scene = Scene(objects=[updated_table, banana]) + SceneExporter( + scene=updated_scene, + scene_graph=_scene_graph(updated_scene), + output_root=output_root, + ).export() + + scene_export_root = output_root / "scene_export" + assert exported_table_glb.read_bytes() == b"glTF-table" + assert ( + scene_export_root / "mesh_assets" / "banana" / "banana.glb" + ).read_bytes() == b"glTF-banana" + assert not (scene_export_root / "mesh_assets" / "cup").exists() + assert ( + json.loads((scene_export_root / "scene.json").read_text(encoding="utf-8"))[ + "objects" + ][1]["id"] + == "banana" + ) + + def test_scene_export_requires_final_physics(tmp_path: Path) -> None: glb_path = tmp_path / "table.glb" glb_path.write_bytes(b"glTF") From 8550f14bd4e9f41e134ba3b2e56b1c57ac2bc94a Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Fri, 14 Aug 2026 19:28:27 +0800 Subject: [PATCH 15/29] fix some bug before pr --- .../features/generative_sim/scene_engine.md | 22 ++- .../scene_engine/cli/test_text_to_simready.py | 183 ------------------ .../pipeline/utils/scene_layout_optimizer.py | 12 +- .../test_scene_layout_optimizer.py | 138 +++++++++++++ 4 files changed, 165 insertions(+), 190 deletions(-) delete mode 100644 embodichain/gen_sim/scene_engine/cli/test_text_to_simready.py create mode 100644 tests/gen_sim/scene_engine/test_scene_layout_optimizer.py diff --git a/docs/source/features/generative_sim/scene_engine.md b/docs/source/features/generative_sim/scene_engine.md index c9f569ee2..e4255163e 100644 --- a/docs/source/features/generative_sim/scene_engine.md +++ b/docs/source/features/generative_sim/scene_engine.md @@ -27,6 +27,23 @@ python -m embodichain scene-engine \ --output_root /path/to/scene_output ``` +## Scene Editing + +Edit an existing valid Scene Engine output with an instruction: + +```bash +embodichain scene-engine \ + --output_root /path/to/scene_output \ + --edit_prompt "add a red cup to the front-center of the tabletop" +``` + +`--image` and `--edit_prompt` may also be provided together. Scene Engine then +generates the image-based scene first and applies the edit to that export. An +edit-only invocation requires an existing `scene_export` directory. The edit +overwrites its `scene_config.json`, `scene_graph.json`, `scene.json`, and final +`mesh_assets`; intermediate generation and edit artifacts remain available for +debugging. + ## Configuration Scene Engine reads the LLM, segmentation, image-generation, and @@ -72,9 +89,12 @@ The important final outputs are: scene_output/ |-- scene_understanding/ # Object analysis, masks, and stage JSON |-- scene_generation/ # Generated, SimReady, and layout-debug artifacts +|-- scene_editing/ # Present after edits; generated asset/debug artifacts `-- scene_export/ |-- mesh_assets/ # Final GLBs - `-- scene_config.json # Exported scene description + |-- scene_config.json # Exported z-up scene description + |-- scene_graph.json # Table support and planar relation graph + `-- scene.json # Scene Engine object metadata and y-up poses ``` Validate the export without opening a window: diff --git a/embodichain/gen_sim/scene_engine/cli/test_text_to_simready.py b/embodichain/gen_sim/scene_engine/cli/test_text_to_simready.py deleted file mode 100644 index e3db15b16..000000000 --- a/embodichain/gen_sim/scene_engine/cli/test_text_to_simready.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 - -import argparse -import json -from collections.abc import Sequence -from pathlib import Path -import shutil - -from PIL import Image - -from embodichain.gen_sim.scene_engine.clients.geometry_generation import ( - GeometryGenerationClient, -) -from embodichain.gen_sim.scene_engine.clients.image_generation import ( - ImageGenerationClient, -) -from embodichain.gen_sim.scene_engine.clients.image_segmentation import ( - ImageSegmentationClient, -) -from embodichain.gen_sim.scene_engine.core.scene import Scene -from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject -from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( - OpenAICompatibleVLM, -) -from embodichain.gen_sim.scene_engine.pipeline.utils.image_segmentation_utils import ( - build_mask_candidates, - invert_mask_if_foreground_is_off_center, - save_binary_mask, - union_overlapping_mask_candidates, -) -from embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor import ( - SimReadyProcessor, - SimReadyProcessorConfig, -) - -__all__ = ["main", "run_text_to_simready"] - - -def run_text_to_simready(*, text: str, output_root: str | Path) -> SceneObject: - """Run one manual text-to-SimReady asset pipeline for debugging.""" - text = text.strip() - if not text: - raise ValueError("Text prompt must not be empty.") - - root = Path(output_root).expanduser().resolve() - if root.exists(): - shutil.rmtree(root) - debug_root = root / "debug" - image_root = root / "generated_images" - mask_root = root / "masks" - coarse_root = root / "coarse_geometry" - simready_root = root / "simready_geometry" - for directory in (debug_root, image_root, mask_root, coarse_root, simready_root): - directory.mkdir(parents=True, exist_ok=True) - - object_id = "asset_001" - scene_object = SceneObject( - id=object_id, - kind="asset", - category="asset", - name=text, - description=text, - ) - - image_generation_client = ImageGenerationClient.from_dotenv() - image_segmentation_client = ImageSegmentationClient.from_dotenv() - geometry_generation_client = GeometryGenerationClient.from_dotenv() - vlm_client = OpenAICompatibleVLM.from_dotenv() - try: - image_generation_client.check_health() - image_segmentation_client.check_health() - geometry_generation_client.check_health() - - # Generate a centered single-object image from the semantic text prompt. - image_path = image_generation_client.generate_image_by_prompt( - prompt=text, - output_path=image_root / f"{object_id}.png", - ) - with Image.open(image_path) as image: - image_size = image.size - - # Segment the generated object and apply the single-object foreground heuristic. - candidates = union_overlapping_mask_candidates( - build_mask_candidates( - image_segmentation_client.segment_single_object( - image_path=image_path, - prompt=text, - ) - ), - min_iou=0.8, - ) - if not candidates: - raise ValueError("Image segmentation returned no mask candidates.") - mask_path = save_binary_mask( - invert_mask_if_foreground_is_off_center(candidates[0]), - image_size=image_size, - output_path=mask_root / f"{object_id}.png", - ) - - # Generate one coarse GLB using the generated image and its binary mask. - geometry_generation_client.generate_objects( - image_path=image_path, - object_masks=[(object_id, mask_path)], - output_root=coarse_root, - ) - coarse_glb_path = coarse_root / f"{object_id}.glb" - if not coarse_glb_path.is_file(): - raise FileNotFoundError(f"Coarse GLB was not generated: {coarse_glb_path}") - - # Use identity coarse layout; VLM determines rotation and real-world size. - processor = SimReadyProcessor( - scene=Scene(objects=[scene_object]), - coarse_layout_by_id={ - object_id: { - "rot": [0.0, 0.0, 0.0], - "pos": [0.0, 0.0, 0.0], - "scale": [1.0, 1.0, 1.0], - } - }, - coarse_geometry_root=coarse_root, - simready_geometry_root=simready_root, - config=SimReadyProcessorConfig( - use_vlm_scale=True, - use_vlm_rotation=True, - ), - vlm_client=vlm_client, - ) - simready_layout = processor.process_assets() - (root / "result.json").write_text( - json.dumps( - { - "input_text": text, - "scene_object": scene_object.to_dict(), - "simready_layout": simready_layout, - }, - indent=2, - ensure_ascii=False, - ) - + "\n", - encoding="utf-8", - ) - return scene_object - finally: - image_generation_client.close() - image_segmentation_client.close() - geometry_generation_client.close() - - -def main(argv: Sequence[str] | None = None) -> None: - """Run the manual text-to-SimReady CLI.""" - parser = argparse.ArgumentParser( - prog="embodichain test-text-to-simready", - description="Debug text-to-image-to-segmentation-to-SimReady generation.", - ) - parser.add_argument("--text", required=True, help="Description of one object.") - parser.add_argument( - "--output_root", - required=True, - help="Directory for all intermediate and final artifacts.", - ) - args = parser.parse_args(argv) - scene_object = run_text_to_simready(text=args.text, output_root=args.output_root) - print(f"Generated SimReady asset: {scene_object.simready_glb_path}") - - -if __name__ == "__main__": - main() diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_optimizer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_optimizer.py index 85e98f0e8..ea6387cd2 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_optimizer.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_optimizer.py @@ -628,15 +628,15 @@ def _table_region_bounds( "right_front": 2, } row_by_region = { - "left_front": 0, - "front_center": 0, - "right_front": 0, + "left_back": 0, + "back_center": 0, + "right_back": 0, "left_center": 1, "center": 1, "right_center": 1, - "left_back": 2, - "back_center": 2, - "right_back": 2, + "left_front": 2, + "front_center": 2, + "right_front": 2, } if table_region not in column_by_region: raise ValueError(f"Unsupported table region {table_region!r}.") diff --git a/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py b/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py new file mode 100644 index 000000000..36b4bfa79 --- /dev/null +++ b/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py @@ -0,0 +1,138 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF 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 numpy as np +import trimesh + +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + SceneGraph, + SceneGraphNode, +) +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_constructor import ( + SceneLayoutConstructor, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_optimizer import ( + _table_region_bounds, +) + + +def _asset( + *, + object_id: str, + glb_path: Path, + center_xy: list[float] | None = None, + pos: list[float] | None = None, +) -> SceneObject: + return SceneObject( + id=object_id, + kind="asset", + category=object_id, + name=object_id, + description=object_id, + simready_glb_path=str(glb_path), + rot=[0.0, 0.0, 0.0], + pos=pos or [0.0, 0.0, 0.0], + scale=[1.0, 1.0, 1.0], + center_xy=center_xy, + ) + + +def test_table_regions_put_front_at_larger_y() -> None: + table_bounds = np.asarray([[0.0, 0.0], [3.0, 3.0]]) + + assert np.allclose( + _table_region_bounds( + table_bounds=table_bounds, + table_region="back_center", + ), + [[1.0, 0.0], [2.0, 1.0]], + ) + assert np.allclose( + _table_region_bounds( + table_bounds=table_bounds, + table_region="front_center", + ), + [[1.0, 2.0], [2.0, 3.0]], + ) + + +def test_layout_constructor_places_new_child_on_parent_top( + tmp_path: Path, +) -> None: + book_glb = tmp_path / "book.glb" + cup_glb = tmp_path / "cup.glb" + # SimReady GLBs are y-up, so the book's short vertical axis is y. + trimesh.creation.box(extents=[1.0, 0.2, 1.0]).export(book_glb) + trimesh.creation.box(extents=[0.2, 0.2, 0.2]).export(cup_glb) + + table = SceneObject( + id="table", + kind="table", + category="table", + name="table", + description="table", + support_surface_z=0.0, + support_optimization_rect_xy=[ + [-2.0, -2.0], + [2.0, -2.0], + [2.0, 2.0], + [-2.0, 2.0], + ], + ) + book = _asset( + object_id="book_001", + glb_path=book_glb, + center_xy=[0.0, 0.0], + # This y-up position maps to a z-up center at z=0.52 m. + pos=[0.0, 0.52, 0.0], + ) + cup = _asset(object_id="cup_001", glb_path=cup_glb) + graph = SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="book_001", + parent_id="table", + parent_relation="on", + ), + SceneGraphNode( + object_id="cup_001", + parent_id="book_001", + parent_relation="on", + ), + ] + ) + + post_edit_scene = SceneLayoutConstructor( + formal_scene=Scene(objects=[table, book]), + goal_scene_graph=graph, + layout_variable_ids={"cup_001"}, + generated_scene_objects=[cup], + output_root=tmp_path, + ).construct() + + placed_cup = next( + asset for asset in post_edit_scene.assets if asset.id == "cup_001" + ) + assert placed_cup.center_xy == [0.0, 0.0] + # book top is z=0.62 m; cup half-height is 0.1 m and clearance is 0.02 m. + assert np.allclose(placed_cup.pos, [0.0, 0.74, 0.0]) From 455d585647b867ae35e7e7a953845aceb921e301 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Fri, 14 Aug 2026 20:05:55 +0800 Subject: [PATCH 16/29] fix a real-world size bug, using z-up mesh --- .../utils/simready_processor_utils.py | 36 ++++++++++++++-- .../test_simready_processor_utils.py | 41 +++++++++++++++++++ 2 files changed, 74 insertions(+), 3 deletions(-) create mode 100644 tests/gen_sim/scene_engine/test_simready_processor_utils.py diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor_utils.py b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor_utils.py index b7f718374..539bace70 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor_utils.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor_utils.py @@ -17,7 +17,10 @@ from __future__ import annotations import json +import os from pathlib import Path +import sys +from typing import Callable import numpy as np from PIL import Image, ImageDraw, ImageFont @@ -92,8 +95,12 @@ def render_object_front_top_views( "Blender's bpy is required for SimReady VLM view rendering." ) from exc - bpy.ops.wm.read_factory_settings(use_empty=True) - bpy.ops.import_scene.gltf(filepath=str(source_path)) + _run_blender_operation_silently( + lambda: bpy.ops.wm.read_factory_settings(use_empty=True) + ) + _run_blender_operation_silently( + lambda: bpy.ops.import_scene.gltf(filepath=str(source_path)) + ) if not any(obj.type == "MESH" for obj in bpy.context.scene.objects): raise ValueError(f"GLB contains no mesh objects: {source_path}") scene = bpy.context.scene @@ -141,7 +148,7 @@ def render_view(path: Path, location: tuple[float, float, float]) -> None: .to_euler() ) scene.render.filepath = str(path) - bpy.ops.render.render(write_still=True) + _run_blender_operation_silently(lambda: bpy.ops.render.render(write_still=True)) # Blender uses a right-handed z-up world; front is viewed along +y. render_view(front_path, (0.0, -3.0, 0.0)) @@ -197,6 +204,25 @@ def render_view(path: Path, location: tuple[float, float, float]) -> None: return output_path +def _run_blender_operation_silently(operation: Callable[[], object]) -> object: + """Run one bpy operation without forwarding Blender-native console output.""" + # bpy writes render progress directly to process file descriptors, not Python streams. + sys.stdout.flush() + sys.stderr.flush() + saved_stdout_fd = os.dup(1) + saved_stderr_fd = os.dup(2) + try: + with open(os.devnull, "w", encoding="utf-8") as null_output: + os.dup2(null_output.fileno(), 1) + os.dup2(null_output.fileno(), 2) + return operation() + finally: + os.dup2(saved_stdout_fd, 1) + os.dup2(saved_stderr_fd, 2) + os.close(saved_stdout_fd) + os.close(saved_stderr_fd) + + def _draw_arrow( draw: ImageDraw.ImageDraw, start: tuple[int, int], @@ -311,6 +337,10 @@ def compute_uniform_xy_scale_for_target( raise ValueError(f"GLB is not a mesh: {glb_path}") if len(target_xy_size_cm) != 2 or any(value <= 0 for value in target_xy_size_cm): raise ValueError("target_xy_size_cm must contain two positive values.") + # GLB geometry is y-up, while the target footprint is defined on z-up table XY. + y_up_to_z_up = np.eye(4) + y_up_to_z_up[:3, :3] = Rotation.from_euler("x", 90.0, degrees=True).as_matrix() + mesh.apply_transform(y_up_to_z_up) if rotate_about_x: center = mesh.bounds.mean(axis=0) mesh.apply_translation(-center) diff --git a/tests/gen_sim/scene_engine/test_simready_processor_utils.py b/tests/gen_sim/scene_engine/test_simready_processor_utils.py new file mode 100644 index 000000000..51459b56b --- /dev/null +++ b/tests/gen_sim/scene_engine/test_simready_processor_utils.py @@ -0,0 +1,41 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF 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.simready_processor_utils import ( + compute_uniform_xy_scale_for_target, +) + + +def test_uniform_scale_uses_the_z_up_tabletop_footprint(tmp_path: Path) -> None: + """Measure y-up GLBs against the VLM's z-up XY target footprint.""" + glb_path = tmp_path / "flat_fork.glb" + # In y-up, the thin vertical axis is y; in z-up it becomes the z axis. + trimesh.creation.box(extents=[2.0, 0.01, 0.5]).export(glb_path) + + scale = compute_uniform_xy_scale_for_target( + glb_path=glb_path, + target_xy_size_cm=[200.0, 50.0], + rotate_about_x=False, + ) + + assert scale == pytest.approx(1.0) From 4976b471c3b791913eadf2959f86cfeba9eafee7 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Sat, 15 Aug 2026 17:41:36 +0800 Subject: [PATCH 17/29] update image-to-scene scene understanding: let vlm give orientation state for assets who need to be calibrated --- .../gen_sim/scene_engine/core/scene_graph.py | 14 +- .../editing/scene_edit_understanding.py | 1 + .../generation/scene_understanding.py | 153 +++++++++++++++- .../utils/image_segmentation_utils.py | 61 +++++++ .../pipeline/utils/scene_importer.py | 5 + .../test_scene_core_and_export.py | 28 +++ .../gen_sim/scene_engine/test_scene_graph.py | 4 + .../scene_engine/test_scene_understanding.py | 170 +++++++++++++++++- 8 files changed, 427 insertions(+), 9 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/core/scene_graph.py b/embodichain/gen_sim/scene_engine/core/scene_graph.py index c5c49c625..e13f65aba 100644 --- a/embodichain/gen_sim/scene_engine/core/scene_graph.py +++ b/embodichain/gen_sim/scene_engine/core/scene_graph.py @@ -54,16 +54,23 @@ # A PlanarRelation with B, then A and B must have the same parent node. PlanarRelationType = Literal["left_of", "right_of", "in_front_of", "behind"] SceneConstraintType = SupportRelationType | PlanarRelationType +OrientationState = Literal["standing", "lying"] @dataclass class SceneGraphNode: - """One object node in the edit-time scene hierarchy.""" + """One object node in the edit-time scene hierarchy. + + ``orientation_state`` is an image-derived placement semantic, rather than + an edge to the node itself or an exact three-dimensional transform. + """ object_id: str parent_id: str | None parent_relation: SupportRelationType | None = None table_region: TableRegion | None = None + # Preserves image-observed placement semantics for later pose refinement. + orientation_state: OrientationState | None = None def __post_init__(self) -> None: """Validate local node fields before graph-level checks.""" @@ -71,12 +78,16 @@ def __post_init__(self) -> None: raise ValueError("object_id must be non-empty.") if self.table_region not in {None, *TABLE_REGIONS}: raise ValueError("table_region is invalid.") + if self.orientation_state not in {None, "standing", "lying"}: + raise ValueError("orientation_state is invalid.") # If the node is the table. if self.object_id == TABLE_OBJECT_ID: if self.parent_id is not None: raise ValueError("table must not have a parent.") if self.parent_relation is not None: raise ValueError("table must not have a parent relation.") + if self.orientation_state is not None: + raise ValueError("table must not have an orientation state.") # If the node is not the table. elif self.parent_id is None: raise ValueError("non-table nodes must have a parent.") @@ -90,6 +101,7 @@ def to_dict(self) -> dict[str, object]: "parent_id": self.parent_id, "parent_relation": self.parent_relation, "table_region": self.table_region, + "orientation_state": self.orientation_state, } diff --git a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py index e66fd8653..a648981e5 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py @@ -194,6 +194,7 @@ def _build_updated_scene_graph( parent_id=node.parent_id, parent_relation=node.parent_relation, table_region=node.table_region, + orientation_state=node.orientation_state, ) for node in scene_graph.nodes ], diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py index f2103a945..c5f340714 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py @@ -41,6 +41,7 @@ from embodichain.gen_sim.scene_engine.pipeline.utils.image_segmentation_utils import ( MaskCandidate, build_mask_candidates, + render_asset_mask_id_overlay, render_image_without_masks, render_numbered_mask_candidates, save_binary_mask, @@ -147,6 +148,24 @@ 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.""" +_ORIENTATION_STATE_SYSTEM_PROMPT = """You inspect an outlined tabletop-scene image. +Each visible asset has an outline and an ID label. Determine whether each listed +object is standing, lying, or unknown in the image. + +Use "standing" only when an upright container, such as a bottle, can, jar, +flask, or thermos, is resting vertically on its base. Use "lying" only when +such a container rests on its side. Use null for the table, every other object +type, or any uncertain case. + +Return JSON only, with exactly this schema. Include every supplied object ID +exactly once and do not add IDs: +{ + "orientation_states": [ + {"object_id": "bottle_001", "orientation_state": "standing"}, + {"object_id": "table", "orientation_state": null} + ] +}""" +_UPRIGHT_CONTAINER_ID_TOKENS = frozenset({"bottle", "can", "jar", "flask", "thermos"}) def understand_scene( @@ -174,7 +193,8 @@ def understand_scene( json_max_attempts=json_max_attempts, ) - _segment_scene( + # Receive the validated whole-scene mask, for VLM output the scene graph. + asset_mask_id_overlay_path = _segment_scene( image_path=resolved_image_path, stage_output_root=stage_output_root, scene=scene, @@ -185,7 +205,11 @@ def understand_scene( # Use the segmented image to initialize the scene graph # with the help of the VLM client. # But at here, we do with the simplest way (hard code). - scene_graph = _initialize_scene_graph_from_segmented_scene(scene) + scene_graph = _initialize_scene_graph_from_segmented_scene( + scene, + asset_mask_id_overlay_path=asset_mask_id_overlay_path, + vlm_client=vlm_client, + ) # Write the Updated scene JSON for debugging. (stage_output_root / "scene.json").write_text( @@ -199,18 +223,38 @@ def understand_scene( return scene, scene_graph -def _initialize_scene_graph_from_segmented_scene(scene: Scene) -> SceneGraph: +def _initialize_scene_graph_from_segmented_scene( + scene: Scene, + *, + asset_mask_id_overlay_path: str | Path, + vlm_client: OpenAICompatibleVLM, +) -> SceneGraph: """Build the initial graph assuming every segmented asset rests on the table.""" + # Get simplified scene info for VLM. + scene_info = _simplify_scene_info_for_graph_initialization(scene=scene) + resolved_asset_mask_id_overlay_path = _validate_image_path( + asset_mask_id_overlay_path + ) if scene.table is None: raise ValueError("Cannot initialize a scene graph without a table.") + orientation_states_by_id = _query_orientation_states( + scene_info=scene_info, + asset_mask_id_overlay_path=resolved_asset_mask_id_overlay_path, + vlm_client=vlm_client, + ) return SceneGraph( nodes=[ SceneGraphNode(object_id=TABLE_OBJECT_ID, parent_id=None), *[ - SceneGraphNode( + SceneGraphNode( # semi-hard code. object_id=asset.id, parent_id=TABLE_OBJECT_ID, parent_relation="on", + orientation_state=( + orientation_states_by_id[asset.id] + if _is_upright_container_id(asset.id) + else None + ), ) for asset in scene.assets ], @@ -218,6 +262,93 @@ def _initialize_scene_graph_from_segmented_scene(scene: Scene) -> SceneGraph: ) +def _simplify_scene_info_for_graph_initialization( + *, + scene: Scene, +) -> dict[str, object]: + """Return the object metadata needed to initialize an image-based graph.""" + return { + "existing_object_ids": [scene_object.id for scene_object in scene.objects], + } + + +def _query_orientation_states( + *, + scene_info: dict[str, object], + asset_mask_id_overlay_path: Path, + vlm_client: OpenAICompatibleVLM, +) -> dict[str, str | None]: + """Return validated image-observed orientation states keyed by object ID.""" + response_text = vlm_client.complete( + image_path=asset_mask_id_overlay_path, + system_prompt=_ORIENTATION_STATE_SYSTEM_PROMPT, + user_prompt=json.dumps(scene_info, ensure_ascii=False), + ) + return _parse_orientation_states_response( + response_text=response_text, + existing_object_ids=scene_info["existing_object_ids"], + ) + + +def _parse_orientation_states_response( + *, + response_text: str, + existing_object_ids: object, +) -> dict[str, str | None]: + """Parse a complete VLM orientation-state response for known object IDs.""" + if not isinstance(existing_object_ids, list) or not all( + isinstance(object_id, str) for object_id in existing_object_ids + ): + raise ValueError("Scene graph initialization requires string object 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) != {"orientation_states"}: + raise ValueError("VLM JSON must contain exactly the key: orientation_states.") + states_value = payload["orientation_states"] + if not isinstance(states_value, list): + raise ValueError("VLM JSON key orientation_states must be an array.") + + orientation_states_by_id: dict[str, str | None] = {} + for index, state_value in enumerate(states_value): + if not isinstance(state_value, dict) or set(state_value) != { + "object_id", + "orientation_state", + }: + raise ValueError( + "VLM JSON orientation_states[" + f"{index}] must contain exactly object_id and orientation_state." + ) + object_id = state_value["object_id"] + orientation_state = state_value["orientation_state"] + if not isinstance(object_id, str) or not object_id: + raise ValueError( + f"VLM JSON orientation_states[{index}].object_id is invalid." + ) + if orientation_state not in {None, "standing", "lying"}: + raise ValueError( + f"VLM JSON orientation_states[{index}].orientation_state is invalid." + ) + if object_id in orientation_states_by_id: + raise ValueError(f"VLM JSON repeats orientation state for {object_id!r}.") + orientation_states_by_id[object_id] = orientation_state + + if set(orientation_states_by_id) != set(existing_object_ids): + raise ValueError( + "VLM JSON orientation states must match all existing object IDs." + ) + return orientation_states_by_id + + +def _is_upright_container_id(object_id: str) -> bool: + """Return whether an object ID identifies a standardized upright container.""" + return bool( + set(re.findall(r"[a-z0-9]+", object_id.lower())) & _UPRIGHT_CONTAINER_ID_TOKENS + ) + + def _analyze_image_objects( *, scene: Scene, @@ -386,8 +517,8 @@ def _segment_scene( scene: Scene, vlm_client: OpenAICompatibleVLM, image_segmentation_client: ImageSegmentationClient, -) -> None: - """Add validated table and asset mask paths to a semantic scene.""" +) -> Path: + """Add validated masks and return an asset-only ID overlay image.""" debug_output_root = ( Path(stage_output_root) / "debug" ) # Keeps the mask debug images. @@ -429,6 +560,16 @@ def _segment_scene( vlm_client=vlm_client, image_segmentation_client=image_segmentation_client, ) + asset_masks: list[tuple[str, 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_masks.append((asset.id, asset.mask_path)) + return render_asset_mask_id_overlay( + image_path=image_path, + asset_masks=asset_masks, + output_path=Path(masks_output_root) / "asset_masks_with_ids.png", + ) def _segment_table( diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py b/embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py index ccf9af949..29defd4d6 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py @@ -306,6 +306,67 @@ def render_numbered_mask_candidates( return resolved_output_path +def render_asset_mask_id_overlay( + *, + image_path: str | Path, + asset_masks: list[tuple[str, str | Path]], + output_path: str | Path, +) -> Path: + """Overlay outlined asset masks and stable asset IDs on a scene image. + + The table mask is intentionally omitted so its large contour does not + obscure the asset labels or their visual context in the source image. + """ + asset_ids = [asset_id for asset_id, _ in asset_masks] + if any(not asset_id for asset_id in asset_ids): + raise ValueError("Every asset mask must have a non-empty asset id.") + if len(asset_ids) != len(set(asset_ids)): + raise ValueError("Asset mask ids must be unique.") + + image = Image.open(image_path).convert("RGBA") + overlay = Image.new("RGBA", image.size, (0, 0, 0, 0)) + colors = ( + (239, 83, 80, 255), + (66, 165, 245, 255), + (102, 187, 106, 255), + (255, 202, 40, 255), + (171, 71, 188, 255), + (38, 198, 218, 255), + ) + decoded_masks: list[tuple[str, Image.Image]] = [] + for index, (asset_id, mask_path) in enumerate(asset_masks): + mask = Image.open(mask_path).convert("L") + _require_image_size(mask, image.size) + decoded_masks.append((asset_id, mask)) + color_layer = Image.new("RGBA", image.size, colors[index % len(colors)]) + transparent_layer = Image.new("RGBA", image.size, (0, 0, 0, 0)) + overlay.alpha_composite( + Image.composite( + color_layer, + transparent_layer, + _mask_outer_outline(mask, image.size), + ) + ) + + draw = ImageDraw.Draw(overlay) + font = _load_label_font(image.size) + for asset_id, mask in decoded_masks: + bbox = mask.getbbox() + if bbox is None: + raise ValueError(f"Asset mask {asset_id!r} is empty.") + _draw_number_label( + draw=draw, + label=asset_id, + 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( diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py index 9ea558a2a..e730b88cb 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py @@ -192,12 +192,14 @@ def _scene_graph_node_from_data(value: object) -> SceneGraphNode: "parent_id", "parent_relation", "table_region", + "orientation_state", }: raise ValueError("Scene graph nodes must use the serialized node schema.") object_id = value["object_id"] parent_id = value["parent_id"] parent_relation = value["parent_relation"] table_region = value["table_region"] + orientation_state = value["orientation_state"] if not isinstance(object_id, str) or not isinstance( parent_id, (str, type(None)) ): @@ -206,11 +208,14 @@ def _scene_graph_node_from_data(value: object) -> SceneGraphNode: raise ValueError("Scene graph parent_relation must be 'on' or null.") if table_region is not None and table_region not in TABLE_REGIONS: raise ValueError("Scene graph table_region is invalid.") + if orientation_state not in {None, "standing", "lying"}: + raise ValueError("Scene graph orientation_state is invalid.") return SceneGraphNode( object_id=object_id, parent_id=parent_id, parent_relation=parent_relation, table_region=table_region, + orientation_state=orientation_state, ) @staticmethod 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 b50126751..95658cf16 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 @@ -175,12 +175,14 @@ def test_scene_export_copies_meshes_and_converts_y_up_pose(tmp_path: Path) -> No "parent_id": None, "parent_relation": None, "table_region": None, + "orientation_state": None, }, { "object_id": "cup", "parent_id": "table", "parent_relation": "on", "table_region": None, + "orientation_state": None, }, ], "relations": [], @@ -195,6 +197,32 @@ def test_scene_export_copies_meshes_and_converts_y_up_pose(tmp_path: Path) -> No assert imported_graph.to_dict() == _scene_graph(scene).to_dict() +def test_scene_graph_importer_restores_node_orientation_state() -> None: + imported_graph = SceneExportImporter._scene_graph_from_data( + { + "nodes": [ + { + "object_id": "table", + "parent_id": None, + "parent_relation": None, + "table_region": None, + "orientation_state": None, + }, + { + "object_id": "bottle_001", + "parent_id": "table", + "parent_relation": "on", + "table_region": None, + "orientation_state": "standing", + }, + ], + "relations": [], + } + ) + + assert imported_graph.node_by_id()["bottle_001"].orientation_state == "standing" + + def test_scene_export_overwrites_an_existing_scene_export(tmp_path: Path) -> None: table_glb = tmp_path / "table.glb" cup_glb = tmp_path / "cup.glb" diff --git a/tests/gen_sim/scene_engine/test_scene_graph.py b/tests/gen_sim/scene_engine/test_scene_graph.py index 970a3fc68..d5b6adaf7 100644 --- a/tests/gen_sim/scene_engine/test_scene_graph.py +++ b/tests/gen_sim/scene_engine/test_scene_graph.py @@ -34,6 +34,7 @@ def test_scene_graph_accepts_layered_on_relations() -> None: parent_id="table", parent_relation="on", table_region="center", + orientation_state="standing", ), SceneGraphNode( object_id="cup", @@ -284,6 +285,7 @@ def test_scene_graph_to_dict_serializes_graph_state() -> None: parent_id="table", parent_relation="on", table_region="center", + orientation_state="standing", ), ], ) @@ -297,12 +299,14 @@ def test_scene_graph_to_dict_serializes_graph_state() -> None: "parent_id": None, "parent_relation": None, "table_region": None, + "orientation_state": None, }, { "object_id": "plate", "parent_id": "table", "parent_relation": "on", "table_region": "center", + "orientation_state": "standing", }, ], "relations": [], diff --git a/tests/gen_sim/scene_engine/test_scene_understanding.py b/tests/gen_sim/scene_engine/test_scene_understanding.py index 8b9357f33..5313e75fa 100644 --- a/tests/gen_sim/scene_engine/test_scene_understanding.py +++ b/tests/gen_sim/scene_engine/test_scene_understanding.py @@ -20,11 +20,15 @@ import json from pathlib import Path +from PIL import Image, ImageDraw import pytest from embodichain.gen_sim.scene_engine.core.scene import Scene from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject from embodichain.gen_sim.scene_engine.pipeline.generation import scene_understanding +from embodichain.gen_sim.scene_engine.pipeline.utils.image_segmentation_utils import ( + render_asset_mask_id_overlay, +) def _response(*, asset_name: str = "cup") -> str: @@ -94,7 +98,44 @@ def complete(self, **_: object) -> str: assert [asset.id for asset in scene.assets] == ["cup_001"] -def test_initial_scene_graph_places_every_asset_on_table() -> None: +def test_asset_mask_id_overlay_excludes_the_table_mask(tmp_path: Path) -> None: + image_path = tmp_path / "scene.png" + table_mask_path = tmp_path / "table_mask.png" + asset_mask_path = tmp_path / "bottle_mask.png" + output_path = tmp_path / "asset_masks_with_ids.png" + image_size = (512, 512) + Image.new("RGB", image_size, "black").save(image_path) + + table_mask = Image.new("L", image_size, 0) + ImageDraw.Draw(table_mask).rectangle((10, 10, 100, 100), fill=255) + table_mask.save(table_mask_path) + asset_mask = Image.new("L", image_size, 0) + ImageDraw.Draw(asset_mask).rectangle((380, 180, 450, 360), fill=255) + asset_mask.save(asset_mask_path) + + rendered_path = render_asset_mask_id_overlay( + image_path=image_path, + asset_masks=[("bottle_001", asset_mask_path)], + output_path=output_path, + ) + + with Image.open(rendered_path) as overlay: + assert overlay.getpixel((10, 10)) == (0, 0, 0) + assert overlay.getpixel((377, 180)) != (0, 0, 0) + + +def test_initial_scene_graph_places_every_asset_on_table(tmp_path: Path) -> None: + class VLM: + def complete(self, **_: object) -> str: + return json.dumps( + { + "orientation_states": [ + {"object_id": "table", "orientation_state": None}, + {"object_id": "cup_001", "orientation_state": "lying"}, + ] + } + ) + scene = Scene( objects=[ SceneObject( @@ -114,8 +155,12 @@ def test_initial_scene_graph_places_every_asset_on_table() -> None: ], ) + overlay_path = tmp_path / "asset_masks_with_ids.png" + overlay_path.write_bytes(b"png") scene_graph = scene_understanding._initialize_scene_graph_from_segmented_scene( - scene + scene, + asset_mask_id_overlay_path=overlay_path, + vlm_client=VLM(), # type: ignore[arg-type] ) assert scene_graph.to_dict() == { @@ -125,13 +170,134 @@ def test_initial_scene_graph_places_every_asset_on_table() -> None: "parent_id": None, "parent_relation": None, "table_region": None, + "orientation_state": None, }, { "object_id": "cup_001", "parent_id": "table", "parent_relation": "on", "table_region": None, + "orientation_state": None, }, ], "relations": [], } + + +def test_scene_graph_initialization_uses_container_orientation_states( + tmp_path: Path, +) -> None: + class VLM: + def complete(self, **_: object) -> str: + return json.dumps( + { + "orientation_states": [ + {"object_id": "table", "orientation_state": None}, + { + "object_id": "bottle_001", + "orientation_state": "standing", + }, + {"object_id": "book_001", "orientation_state": "lying"}, + ] + } + ) + + overlay_path = tmp_path / "asset_masks_with_ids.png" + overlay_path.write_bytes(b"png") + scene = Scene( + objects=[ + SceneObject( + id="table", + kind="table", + category="table", + name="wooden table", + description="A wooden table.", + ), + SceneObject( + id="bottle_001", + kind="asset", + category="bottle", + name="blue bottle", + description="A blue bottle.", + ), + SceneObject( + id="book_001", + kind="asset", + category="book", + name="red book", + description="A red book.", + ), + ] + ) + + scene_graph = scene_understanding._initialize_scene_graph_from_segmented_scene( + scene, + asset_mask_id_overlay_path=overlay_path, + vlm_client=VLM(), # type: ignore[arg-type] + ) + + assert scene_graph.node_by_id()["bottle_001"].orientation_state == "standing" + assert scene_graph.node_by_id()["book_001"].orientation_state is None + + +def test_scene_graph_initialization_requires_asset_mask_id_overlay( + tmp_path: Path, +) -> None: + scene = Scene( + objects=[ + SceneObject( + id="table", + kind="table", + category="table", + name="wooden table", + description="A wooden table.", + ) + ] + ) + + with pytest.raises(FileNotFoundError, match="Image input not found"): + scene_understanding._initialize_scene_graph_from_segmented_scene( + scene, + asset_mask_id_overlay_path=tmp_path / "missing.png", + vlm_client=object(), # type: ignore[arg-type] + ) + + +def test_scene_graph_initialization_info_lists_existing_object_ids() -> None: + scene = Scene( + objects=[ + SceneObject( + id="table", + kind="table", + category="table", + name="wooden table", + description="A wooden table.", + ), + SceneObject( + id="bottle_001", + kind="asset", + category="bottle", + name="blue bottle", + description="A blue bottle.", + center_xy=[0.2, -0.1], + ), + SceneObject( + id="book_001", + kind="asset", + category="book", + name="red book", + description="A red book.", + center_xy=[-0.1, 0.2], + ), + ] + ) + + simplified_scene_info = ( + scene_understanding._simplify_scene_info_for_graph_initialization( + scene=scene + ) + ) + + assert simplified_scene_info == { + "existing_object_ids": ["table", "bottle_001", "book_001"], + } From 7fdffdcb2d97e4746a8c1a6405ba6395e7b82753 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Sat, 15 Aug 2026 18:03:48 +0800 Subject: [PATCH 18/29] fix big-font problem in asset_masks_with_ids; add scene graph based clibration in image-conditioned scene engine pipeline (but only calibrated the upright bottle-like assets currently) --- .../pipeline/generation/scene_generation.py | 119 +++++++++++++++++- .../utils/image_segmentation_utils.py | 60 ++++++++- .../scene_engine/test_scene_generation.py | 95 ++++++++++++++ .../scene_engine/test_scene_understanding.py | 21 ++++ 4 files changed, 287 insertions(+), 8 deletions(-) create mode 100644 tests/gen_sim/scene_engine/test_scene_generation.py diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py index 60f91ca81..6103f8cf4 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py @@ -22,6 +22,7 @@ import shutil import numpy as np +from scipy.spatial.transform import Rotation import trimesh from shapely.geometry import Polygon @@ -133,6 +134,7 @@ def generate_scene_and_refine( # Layout refinement will start with the table. refined_table_layout, refined_assets_layout = _layout_refinement( scene=scene, # Update this data structure internally. + scene_graph=scene_graph, 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. ) @@ -291,6 +293,7 @@ def _copy_y_up_layout_to_scene_object( def _layout_refinement( *, scene: Scene, + scene_graph: SceneGraph, simready_geometry_output_root: str | Path, debug_output_root: str | Path, ) -> tuple[dict[str, object], list[dict[str, object]]]: @@ -357,7 +360,14 @@ def _layout_refinement( ) ) - # 3. Move all assets as one rigid group so its lowest AABB point is 2cm above + # 3. Correct image-observed standing containers before every geometry-based + # layout stage measures their footprint. + refined_assets_layout = _scene_graph_based_calibration( + scene_graph=scene_graph, + assets_layout=refined_assets_layout, + ) + + # 4. 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. @@ -371,7 +381,7 @@ def _layout_refinement( log_info("Scene has no movable assets; skipping support-region clamping.") return refined_table_layout, [] - # 4. Reuse support geometry detected during SimReady processing. + # 5. Reuse support geometry detected during SimReady processing. if ( scene.table is None or scene.table.support_contour_xy is None @@ -395,7 +405,7 @@ def _layout_refinement( ) ) - # 5. Keep the complete clutter rigid in the table plane. A successful + # 6. 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. @@ -421,7 +431,7 @@ def _layout_refinement( ) ) - # 6. Optimize independent asset positions inside the conservative rectangle. + # 7. Optimize independent asset positions inside the conservative rectangle. # The clamp above already used the exact outer contour for the shared shift. overlap_optimizer = AssetsSupportLayoutOptimizer( support_region=table_optimization_rectangle, @@ -437,7 +447,7 @@ def _layout_refinement( 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. + # 8. Gravity simulation, to let all the assets to be stable and placed well on the table's support surface. # Notice that: we do not consider the assets like a bottle, which should be standing on the table but laid down # after the simulation. gravity_settler = AssetsGravitySettler( @@ -458,6 +468,105 @@ def _layout_refinement( return refined_table_layout, refined_assets_layout +def _scene_graph_based_calibration( + *, + scene_graph: SceneGraph, + assets_layout: list[dict[str, object]], +) -> list[dict[str, object]]: + """Minimally align graph-marked standing assets with the z-up table frame.""" + # This is the extension point for future image-conditioned scene generation + # calibration. The scene graph may later provide richer image-grounded + # constraints, but the current implementation deliberately consumes only + # ``orientation_state`` to correct standing container axes before 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) + nodes_by_id = scene_graph.node_by_id() + calibrated_assets_layout: list[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.") + node = nodes_by_id.get(asset_id) + if node is None: + raise ValueError(f"Scene graph does not contain asset {asset_id!r}.") + if node.orientation_state != "standing": + calibrated_assets_layout.append(asset_layout) + continue + + # Conjugate the y-up pose so the SimReady container axis is local z. + z_up_asset_to_table_matrix = ( + y_up_to_z_up_matrix + @ layout_object_to_transform_matrix(asset_layout) + @ z_up_to_y_up_matrix + ) + linear_matrix = z_up_asset_to_table_matrix[:3, :3] + # Layout transforms store rotation and per-axis scale in the same matrix. + scale = np.linalg.norm(linear_matrix, axis=0) + if np.any(scale <= 1e-8): + raise ValueError(f"Asset {asset_id!r} 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(f"Asset {asset_id!r} layout contains shear.") + + local_z_axis_in_table = rotation_matrix[:, 2] + # Treat the long axis as unsigned to avoid an unnecessary 180-degree flip. + target_z_axis = np.array( + [0.0, 0.0, 1.0 if local_z_axis_in_table[2] >= 0.0 else -1.0] + ) + # Left multiplication applies the correction in the table/world frame. + z_up_asset_to_table_matrix[:3, :3] = ( + _minimum_axis_alignment_rotation( + source_axis=local_z_axis_in_table, + target_axis=target_z_axis, + ) + @ rotation_matrix + @ np.diag(scale) + ) + calibrated_assets_layout.append( + transform_matrix_to_layout_object( + asset_id, + z_up_to_y_up_matrix + @ z_up_asset_to_table_matrix + @ y_up_to_z_up_matrix, + ) + ) + return calibrated_assets_layout + + +def _minimum_axis_alignment_rotation( + *, + source_axis: np.ndarray, + target_axis: np.ndarray, +) -> np.ndarray: + """Return the smallest proper rotation mapping one nonzero axis to another.""" + source = np.asarray(source_axis, dtype=float) + target = np.asarray(target_axis, dtype=float) + source_norm = np.linalg.norm(source) + target_norm = np.linalg.norm(target) + if source_norm <= 1e-8 or target_norm <= 1e-8: + raise ValueError("Axis alignment requires nonzero axes.") + source /= source_norm + target /= target_norm + + cross_product = np.cross(source, target) + sine = np.linalg.norm(cross_product) + cosine = float(np.clip(np.dot(source, target), -1.0, 1.0)) + if sine <= 1e-8: + if cosine > 0.0: + return np.eye(3) + basis_axis = np.eye(3)[np.argmin(np.abs(source))] + rotation_axis = np.cross(source, basis_axis) + rotation_axis /= np.linalg.norm(rotation_axis) + return Rotation.from_rotvec(np.pi * rotation_axis).as_matrix() + + rotation_axis = cross_product / sine + return Rotation.from_rotvec(np.arctan2(sine, cosine) * rotation_axis).as_matrix() + + def _measure_table_and_assets_in_z_up_world( *, table_layout: dict[str, object], diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py b/embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py index 29defd4d6..302ec2cbf 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py @@ -349,16 +349,21 @@ def render_asset_mask_id_overlay( ) draw = ImageDraw.Draw(overlay) - font = _load_label_font(image.size) for asset_id, mask in decoded_masks: bbox = mask.getbbox() if bbox is None: raise ValueError(f"Asset mask {asset_id!r} is empty.") + font = _load_asset_id_label_font( + image_size=image.size, + mask_bbox=bbox, + label=asset_id, + ) _draw_number_label( draw=draw, label=asset_id, center=((bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2), font=font, + minimum_padding=2, ) resolved_output_path = Path(output_path).expanduser().resolve() @@ -504,6 +509,47 @@ def _union_parent(parents: list[int], first_index: int, second_index: int) -> No def _load_label_font(image_size: tuple[int, int]) -> ImageFont.ImageFont: font_size = max(16, round(min(image_size) / 32)) + return _load_label_font_at_size(font_size) + + +def _load_asset_id_label_font( + *, + image_size: tuple[int, int], + mask_bbox: tuple[int, int, int, int], + label: str, +) -> ImageFont.ImageFont: + """Choose an ID-label font constrained by both image and mask dimensions.""" + # The image sets the readable upper bound; the individual mask then caps it. + image_font_size = min(32, max(8, round(min(image_size) / 48))) + mask_width = mask_bbox[2] - mask_bbox[0] + mask_height = mask_bbox[3] - mask_bbox[1] + maximum_label_width = max(24, round(mask_width * 0.9)) + maximum_label_height = max(16, round(mask_height * 0.75)) + # Measure the complete text-and-background rectangle, not glyphs alone. + probe_draw = ImageDraw.Draw(Image.new("RGBA", image_size)) + smallest_font = _load_label_font_at_size(6) + for font_size in range(image_font_size, 5, -1): + font = _load_label_font_at_size(font_size) + label_bounds = _number_label_bounds( + draw=probe_draw, + label=label, + center=(0.0, 0.0), + font=font, + minimum_padding=2, + ) + if ( + label_bounds[2] - label_bounds[0] <= maximum_label_width + and label_bounds[3] - label_bounds[1] <= maximum_label_height + ): + return font + smallest_font = font + return smallest_font + + +def _load_label_font_at_size(font_size: int) -> ImageFont.ImageFont: + """Load the shared bold label font at one validated pixel size.""" + if font_size < 1: + raise ValueError("Label font size must be positive.") try: return ImageFont.truetype("DejaVuSans-Bold.ttf", font_size) except OSError: @@ -516,10 +562,15 @@ def _draw_number_label( label: str, center: tuple[float, float], font: ImageFont.ImageFont, + minimum_padding: int = 4, ) -> 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 + draw=draw, + label=label, + center=center, + font=font, + minimum_padding=minimum_padding, ) label_box = draw.textbbox((0, 0), label, font=font) label_width = label_box[2] - label_box[0] @@ -541,12 +592,15 @@ def _number_label_bounds( label: str, center: tuple[float, float], font: ImageFont.ImageFont, + minimum_padding: int = 4, ) -> 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)) + if minimum_padding < 0: + raise ValueError("Label minimum padding must be non-negative.") + padding = max(minimum_padding, round(max(label_width, label_height) / 4)) x = center[0] - label_width / 2 y = center[1] - label_height / 2 return ( diff --git a/tests/gen_sim/scene_engine/test_scene_generation.py b/tests/gen_sim/scene_engine/test_scene_generation.py new file mode 100644 index 000000000..6a61a4d1b --- /dev/null +++ b/tests/gen_sim/scene_engine/test_scene_generation.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 + +import numpy as np +from scipy.spatial.transform import Rotation + +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + SceneGraph, + SceneGraphNode, +) +from embodichain.gen_sim.scene_engine.pipeline.generation.scene_generation import ( + _scene_graph_based_calibration, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_generation_utils import ( + layout_object_to_transform_matrix, + transform_matrix_to_layout_object, +) + + +def _y_up_layout_from_z_up_rotation( + object_id: str, + rotation_matrix: np.ndarray, +) -> dict[str, object]: + 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_transform = np.eye(4) + z_up_transform[:3, :3] = rotation_matrix + return transform_matrix_to_layout_object( + object_id, + z_up_to_y_up_matrix @ z_up_transform @ y_up_to_z_up_matrix, + ) + + +def _z_up_rotation_from_y_up_layout(layout: dict[str, object]) -> np.ndarray: + 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]] + ) + return ( + y_up_to_z_up_matrix + @ layout_object_to_transform_matrix(layout) + @ np.linalg.inv(y_up_to_z_up_matrix) + )[:3, :3] + + +def test_scene_graph_calibration_makes_standing_asset_vertical() -> None: + scene_graph = SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="bottle_001", + parent_id="table", + parent_relation="on", + orientation_state="standing", + ), + SceneGraphNode( + object_id="book_001", + parent_id="table", + parent_relation="on", + ), + ] + ) + lying_rotation = Rotation.from_euler("x", 90.0, degrees=True).as_matrix() + bottle_layout = _y_up_layout_from_z_up_rotation("bottle_001", lying_rotation) + book_layout = _y_up_layout_from_z_up_rotation("book_001", lying_rotation) + + calibrated_layouts = _scene_graph_based_calibration( + scene_graph=scene_graph, + assets_layout=[bottle_layout, book_layout], + ) + + bottle_axis = _z_up_rotation_from_y_up_layout(calibrated_layouts[0])[:, 2] + assert np.isclose(abs(bottle_axis[2]), 1.0) + assert np.allclose( + _z_up_rotation_from_y_up_layout(calibrated_layouts[1]), + lying_rotation, + ) diff --git a/tests/gen_sim/scene_engine/test_scene_understanding.py b/tests/gen_sim/scene_engine/test_scene_understanding.py index 5313e75fa..2a9af65b2 100644 --- a/tests/gen_sim/scene_engine/test_scene_understanding.py +++ b/tests/gen_sim/scene_engine/test_scene_understanding.py @@ -26,6 +26,7 @@ from embodichain.gen_sim.scene_engine.core.scene import Scene from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject from embodichain.gen_sim.scene_engine.pipeline.generation import scene_understanding +from embodichain.gen_sim.scene_engine.pipeline.utils import image_segmentation_utils from embodichain.gen_sim.scene_engine.pipeline.utils.image_segmentation_utils import ( render_asset_mask_id_overlay, ) @@ -124,6 +125,26 @@ def test_asset_mask_id_overlay_excludes_the_table_mask(tmp_path: Path) -> None: assert overlay.getpixel((377, 180)) != (0, 0, 0) +def test_asset_mask_id_label_font_fits_the_mask_bbox() -> None: + image_size = (512, 512) + mask_bbox = (380, 180, 450, 360) + label = "bottle_001" + font = image_segmentation_utils._load_asset_id_label_font( + image_size=image_size, + mask_bbox=mask_bbox, + label=label, + ) + label_bounds = image_segmentation_utils._number_label_bounds( + draw=ImageDraw.Draw(Image.new("RGBA", image_size)), + label=label, + center=(0.0, 0.0), + font=font, + minimum_padding=2, + ) + + assert label_bounds[2] - label_bounds[0] <= round((mask_bbox[2] - mask_bbox[0]) * 0.9) + + def test_initial_scene_graph_places_every_asset_on_table(tmp_path: Path) -> None: class VLM: def complete(self, **_: object) -> str: From a7ff6b17c0e64a99d5ef804bfca6953992c07698 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Sat, 15 Aug 2026 18:04:23 +0800 Subject: [PATCH 19/29] run black --- .../scene_engine/pipeline/generation/scene_generation.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py index 6103f8cf4..db7ed8ee2 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py @@ -529,9 +529,7 @@ def _scene_graph_based_calibration( calibrated_assets_layout.append( transform_matrix_to_layout_object( asset_id, - z_up_to_y_up_matrix - @ z_up_asset_to_table_matrix - @ y_up_to_z_up_matrix, + z_up_to_y_up_matrix @ z_up_asset_to_table_matrix @ y_up_to_z_up_matrix, ) ) return calibrated_assets_layout From 387646e6ae1f754e64a7becdae2ba90c58455d6f Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Sat, 15 Aug 2026 18:10:49 +0800 Subject: [PATCH 20/29] run black to tests/ --- tests/gen_sim/scene_engine/test_scene_understanding.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/gen_sim/scene_engine/test_scene_understanding.py b/tests/gen_sim/scene_engine/test_scene_understanding.py index 2a9af65b2..cd63bd2b0 100644 --- a/tests/gen_sim/scene_engine/test_scene_understanding.py +++ b/tests/gen_sim/scene_engine/test_scene_understanding.py @@ -142,7 +142,9 @@ def test_asset_mask_id_label_font_fits_the_mask_bbox() -> None: minimum_padding=2, ) - assert label_bounds[2] - label_bounds[0] <= round((mask_bbox[2] - mask_bbox[0]) * 0.9) + assert label_bounds[2] - label_bounds[0] <= round( + (mask_bbox[2] - mask_bbox[0]) * 0.9 + ) def test_initial_scene_graph_places_every_asset_on_table(tmp_path: Path) -> None: @@ -314,9 +316,7 @@ def test_scene_graph_initialization_info_lists_existing_object_ids() -> None: ) simplified_scene_info = ( - scene_understanding._simplify_scene_info_for_graph_initialization( - scene=scene - ) + scene_understanding._simplify_scene_info_for_graph_initialization(scene=scene) ) assert simplified_scene_info == { From 06f0b77af9e97c078a36e063d0ce93f8614a1461 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Tue, 18 Aug 2026 10:32:30 +0800 Subject: [PATCH 21/29] replace hard-code bottle-z-up with VLM auto-analyze --- .../pipeline/generation/scene_generation.py | 5 + .../generation/scene_understanding.py | 86 +++++++------- .../pipeline/utils/simready_processor.py | 106 +++++------------- .../scene_engine/test_scene_understanding.py | 83 ++++++++++++-- .../test_simready_processor_utils.py | 22 ++++ 5 files changed, 177 insertions(+), 125 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py index db7ed8ee2..7a0e19f35 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py @@ -119,6 +119,11 @@ def generate_scene_and_refine( config=SimReadyProcessorConfig( use_vlm_scale=False, use_vlm_rotation=False, + long_axis_object_ids=frozenset( + node.object_id + for node in scene_graph.nodes + if node.orientation_state is not None + ), ), vlm_client=vlm_client, ) diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py index c5f340714..9507c1d3a 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py @@ -150,22 +150,21 @@ text.""" _ORIENTATION_STATE_SYSTEM_PROMPT = """You inspect an outlined tabletop-scene image. Each visible asset has an outline and an ID label. Determine whether each listed -object is standing, lying, or unknown in the image. +asset is standing, lying, or unknown in the image. -Use "standing" only when an upright container, such as a bottle, can, jar, -flask, or thermos, is resting vertically on its base. Use "lying" only when -such a container rests on its side. Use null for the table, every other object -type, or any uncertain case. +Use a non-null state only for an elongated object with a clear primary long axis. +Use "standing" when its primary axis is approximately vertical to the tabletop. +Use "lying" when its primary axis is approximately parallel to the tabletop. +Use null for every object without a clear primary long axis or when uncertain. -Return JSON only, with exactly this schema. Include every supplied object ID -exactly once and do not add IDs: +Return JSON only, with exactly this schema. Include every supplied asset ID +exactly once. Never include the table or any ID that was not supplied: { "orientation_states": [ {"object_id": "bottle_001", "orientation_state": "standing"}, - {"object_id": "table", "orientation_state": null} + {"object_id": "book_001", "orientation_state": null} ] }""" -_UPRIGHT_CONTAINER_ID_TOKENS = frozenset({"bottle", "can", "jar", "flask", "thermos"}) def understand_scene( @@ -209,6 +208,7 @@ def understand_scene( scene, asset_mask_id_overlay_path=asset_mask_id_overlay_path, vlm_client=vlm_client, + json_max_attempts=json_max_attempts, ) # Write the Updated scene JSON for debugging. @@ -228,6 +228,7 @@ def _initialize_scene_graph_from_segmented_scene( *, asset_mask_id_overlay_path: str | Path, vlm_client: OpenAICompatibleVLM, + json_max_attempts: int = 3, ) -> SceneGraph: """Build the initial graph assuming every segmented asset rests on the table.""" # Get simplified scene info for VLM. @@ -241,20 +242,17 @@ def _initialize_scene_graph_from_segmented_scene( scene_info=scene_info, asset_mask_id_overlay_path=resolved_asset_mask_id_overlay_path, vlm_client=vlm_client, + json_max_attempts=json_max_attempts, ) return SceneGraph( nodes=[ SceneGraphNode(object_id=TABLE_OBJECT_ID, parent_id=None), *[ - SceneGraphNode( # semi-hard code. + SceneGraphNode( object_id=asset.id, parent_id=TABLE_OBJECT_ID, - parent_relation="on", - orientation_state=( - orientation_states_by_id[asset.id] - if _is_upright_container_id(asset.id) - else None - ), + parent_relation="on", # semi-hard-code. + orientation_state=orientation_states_by_id[asset.id], ) for asset in scene.assets ], @@ -268,7 +266,7 @@ def _simplify_scene_info_for_graph_initialization( ) -> dict[str, object]: """Return the object metadata needed to initialize an image-based graph.""" return { - "existing_object_ids": [scene_object.id for scene_object in scene.objects], + "asset_ids": [asset.id for asset in scene.assets], } @@ -277,29 +275,42 @@ def _query_orientation_states( scene_info: dict[str, object], asset_mask_id_overlay_path: Path, vlm_client: OpenAICompatibleVLM, + json_max_attempts: int, ) -> dict[str, str | None]: - """Return validated image-observed orientation states keyed by object ID.""" - response_text = vlm_client.complete( - image_path=asset_mask_id_overlay_path, - system_prompt=_ORIENTATION_STATE_SYSTEM_PROMPT, - user_prompt=json.dumps(scene_info, ensure_ascii=False), - ) - return _parse_orientation_states_response( - response_text=response_text, - existing_object_ids=scene_info["existing_object_ids"], - ) + """Return validated image-observed orientation states keyed by asset ID.""" + if json_max_attempts < 1: + raise ValueError("json_max_attempts must be at least 1.") + last_validation_error: ValueError | None = None + for _ in range(json_max_attempts): + response_text = vlm_client.complete( + image_path=asset_mask_id_overlay_path, + system_prompt=_ORIENTATION_STATE_SYSTEM_PROMPT, + user_prompt=json.dumps(scene_info, ensure_ascii=False), + ) + try: + return _parse_orientation_states_response( + response_text=response_text, + asset_ids=scene_info["asset_ids"], + ) + except ValueError as exc: + last_validation_error = exc + assert last_validation_error is not None + raise ValueError( + "VLM returned invalid orientation-state JSON after " + f"{json_max_attempts} attempts: {last_validation_error}" + ) from last_validation_error def _parse_orientation_states_response( *, response_text: str, - existing_object_ids: object, + asset_ids: object, ) -> dict[str, str | None]: - """Parse a complete VLM orientation-state response for known object IDs.""" - if not isinstance(existing_object_ids, list) or not all( - isinstance(object_id, str) for object_id in existing_object_ids + """Parse a complete VLM orientation-state response for known asset IDs.""" + if not isinstance(asset_ids, list) or not all( + isinstance(object_id, str) for object_id in asset_ids ): - raise ValueError("Scene graph initialization requires string object IDs.") + raise ValueError("Scene graph initialization requires string asset IDs.") json_text = _strip_json_code_fence(response_text) try: payload = json.loads(json_text) @@ -335,20 +346,13 @@ def _parse_orientation_states_response( raise ValueError(f"VLM JSON repeats orientation state for {object_id!r}.") orientation_states_by_id[object_id] = orientation_state - if set(orientation_states_by_id) != set(existing_object_ids): + if set(orientation_states_by_id) != set(asset_ids): raise ValueError( - "VLM JSON orientation states must match all existing object IDs." + "VLM JSON orientation states must match all supplied asset IDs." ) return orientation_states_by_id -def _is_upright_container_id(object_id: str) -> bool: - """Return whether an object ID identifies a standardized upright container.""" - return bool( - set(re.findall(r"[a-z0-9]+", object_id.lower())) & _UPRIGHT_CONTAINER_ID_TOKENS - ) - - def _analyze_image_objects( *, scene: Scene, diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py index 55860be7f..7bc041abf 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py @@ -19,11 +19,9 @@ 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 @@ -66,14 +64,12 @@ @dataclass(frozen=True) class SimReadyProcessorConfig: - """Object-category policy for SimReady mesh canonicalization.""" + """SceneGraph-conditioned policy for SimReady mesh canonicalization.""" use_vlm_scale: bool = False # Use the VLM-selected asset scale. use_vlm_rotation: bool = False # Use the VLM-selected asset rotation. - upright_container_id_tokens: frozenset[str] = frozenset( - {"bottle", "can", "jar", "flask", "thermos"} - ) # Object-id tokens that enable upright-container standardization. + long_axis_object_ids: frozenset[str] = frozenset() class SimReadyProcessor: @@ -106,8 +102,6 @@ def __init__( self.simready_assets_layout: list[dict[str, object]] | None = None self.config = config if config is not None else SimReadyProcessorConfig() self.vlm_client = vlm_client - if not self.config.upright_container_id_tokens: - raise ValueError("upright_container_id_tokens must not be empty.") if ( self.config.use_vlm_scale or self.config.use_vlm_rotation ) and vlm_client is None: @@ -312,8 +306,6 @@ def _canonicalize_object_mesh( ) 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.") @@ -324,15 +316,15 @@ def _canonicalize_object_mesh( 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) + # Standardize graph-marked elongated assets before shared mesh processing. + # This makes local z their primary axis, so later scene-graph calibration + # can reliably recover the image-observed standing or lying orientation. + long_axis_alignment_matrix = np.eye(3) + if self._requires_long_axis_standardization(object_id): + long_axis_alignment_matrix = self._standardize_long_axis_z_up(mesh) + long_axis_alignment_transform = np.eye(4) + long_axis_alignment_transform[:3, :3] = long_axis_alignment_matrix + mesh.apply_transform(long_axis_alignment_transform) # First make the object's AABB center at the origin. original_aabb_center = mesh.bounds.mean(axis=0) @@ -343,11 +335,11 @@ def _canonicalize_object_mesh( 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 + long_axis_alignment_matrix @ y_up_to_z_up_matrix @ np.diag(coarse_scale) @ y_up_to_z_up_matrix.T - @ bottle_alignment_matrix.T + @ long_axis_alignment_matrix.T ) mesh.apply_transform(scale_transform) @@ -367,16 +359,16 @@ def _canonicalize_object_mesh( 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 + # Compensate the local canonicalization so its coarse world pose does not + # change until layout refinement applies the image-observed correction. + local_long_axis_rotation = Rotation.from_matrix( + y_up_to_z_up_matrix.T @ long_axis_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() + coarse_rotation_matrix @ local_long_axis_rotation.inv().as_matrix() ) # Update the pos. position_offset = y_up_to_z_up_matrix.T @ ( @@ -388,24 +380,18 @@ def _canonicalize_object_mesh( "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) + def _requires_long_axis_standardization(self, object_id: str) -> bool: + """Return whether graph semantics identified one asset with a long axis.""" + return object_id in self.config.long_axis_object_ids @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! + def _standardize_long_axis_z_up(mesh: trimesh.Trimesh) -> np.ndarray: + """Return a proper rotation that maps a mesh's primary axis to z-up. + Thanks to chanjian's idea. """ if len(mesh.vertices) < 4 or len(mesh.faces) < 4: raise ValueError( - "Bottle standardization requires a non-degenerate triangle mesh." + "Long-axis standardization requires a non-degenerate triangle mesh." ) open3d_mesh = o3d.geometry.TriangleMesh( vertices=o3d.utility.Vector3dVector(mesh.vertices), @@ -419,7 +405,7 @@ def _standardize_bottle_z_up(mesh: trimesh.Trimesh) -> np.ndarray: # 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." + "Long-axis standardization could not sample valid mesh points." ) centered_points = sampled_points - sampled_points.mean(axis=0) @@ -428,46 +414,12 @@ def _standardize_bottle_z_up(mesh: trimesh.Trimesh) -> np.ndarray: if np.linalg.det(principal_axes) < 0: principal_axes[2, :] *= -1 # in case the SVD returns a reflection. - bottle_rotation = Rotation.from_euler( + long_axis_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 = SimReadyProcessor._convex_hull_volume(upper_points) - lower_volume = SimReadyProcessor._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 + long_axis_rotation = long_axis_rotation @ principal_axes + return long_axis_rotation @staticmethod def _three_floats(value: object, *, field_name: str) -> list[float]: diff --git a/tests/gen_sim/scene_engine/test_scene_understanding.py b/tests/gen_sim/scene_engine/test_scene_understanding.py index cd63bd2b0..d2b3898da 100644 --- a/tests/gen_sim/scene_engine/test_scene_understanding.py +++ b/tests/gen_sim/scene_engine/test_scene_understanding.py @@ -153,8 +153,7 @@ def complete(self, **_: object) -> str: return json.dumps( { "orientation_states": [ - {"object_id": "table", "orientation_state": None}, - {"object_id": "cup_001", "orientation_state": "lying"}, + {"object_id": "cup_001", "orientation_state": None}, ] } ) @@ -207,15 +206,18 @@ def complete(self, **_: object) -> str: } -def test_scene_graph_initialization_uses_container_orientation_states( +def test_scene_graph_initialization_uses_image_orientation_states( tmp_path: Path, ) -> None: class VLM: + def __init__(self) -> None: + self.user_prompt: str | None = None + def complete(self, **_: object) -> str: + self.user_prompt = _["user_prompt"] # type: ignore[assignment,index] return json.dumps( { "orientation_states": [ - {"object_id": "table", "orientation_state": None}, { "object_id": "bottle_001", "orientation_state": "standing", @@ -253,14 +255,81 @@ def complete(self, **_: object) -> str: ] ) + vlm = VLM() + scene_graph = scene_understanding._initialize_scene_graph_from_segmented_scene( + scene, + asset_mask_id_overlay_path=overlay_path, + vlm_client=vlm, # type: ignore[arg-type] + ) + + assert json.loads(vlm.user_prompt or "{}") == { + "asset_ids": ["bottle_001", "book_001"], + } + assert scene_graph.node_by_id()["bottle_001"].orientation_state == "standing" + assert scene_graph.node_by_id()["book_001"].orientation_state == "lying" + + +def test_scene_graph_initialization_retries_a_response_containing_table( + tmp_path: Path, +) -> None: + class VLM: + def __init__(self) -> None: + self.responses = [ + json.dumps( + { + "orientation_states": [ + {"object_id": "table", "orientation_state": None}, + { + "object_id": "bottle_001", + "orientation_state": "standing", + }, + ] + } + ), + json.dumps( + { + "orientation_states": [ + { + "object_id": "bottle_001", + "orientation_state": "standing", + }, + ] + } + ), + ] + + def complete(self, **_: object) -> str: + return self.responses.pop(0) + + overlay_path = tmp_path / "asset_masks_with_ids.png" + overlay_path.write_bytes(b"png") + scene = Scene( + objects=[ + SceneObject( + id="table", + kind="table", + category="table", + name="wooden table", + description="A wooden table.", + ), + SceneObject( + id="bottle_001", + kind="asset", + category="bottle", + name="blue bottle", + description="A blue bottle.", + ), + ] + ) + scene_graph = scene_understanding._initialize_scene_graph_from_segmented_scene( scene, asset_mask_id_overlay_path=overlay_path, vlm_client=VLM(), # type: ignore[arg-type] + json_max_attempts=2, ) assert scene_graph.node_by_id()["bottle_001"].orientation_state == "standing" - assert scene_graph.node_by_id()["book_001"].orientation_state is None def test_scene_graph_initialization_requires_asset_mask_id_overlay( @@ -286,7 +355,7 @@ def test_scene_graph_initialization_requires_asset_mask_id_overlay( ) -def test_scene_graph_initialization_info_lists_existing_object_ids() -> None: +def test_scene_graph_initialization_info_lists_asset_ids() -> None: scene = Scene( objects=[ SceneObject( @@ -320,5 +389,5 @@ def test_scene_graph_initialization_info_lists_existing_object_ids() -> None: ) assert simplified_scene_info == { - "existing_object_ids": ["table", "bottle_001", "book_001"], + "asset_ids": ["bottle_001", "book_001"], } diff --git a/tests/gen_sim/scene_engine/test_simready_processor_utils.py b/tests/gen_sim/scene_engine/test_simready_processor_utils.py index 51459b56b..13cf73e10 100644 --- a/tests/gen_sim/scene_engine/test_simready_processor_utils.py +++ b/tests/gen_sim/scene_engine/test_simready_processor_utils.py @@ -21,11 +21,33 @@ import pytest import trimesh +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor import ( + SimReadyProcessor, + SimReadyProcessorConfig, +) from embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor_utils import ( compute_uniform_xy_scale_for_target, ) +def test_simready_long_axis_standardization_uses_graph_selected_ids( + tmp_path: Path, +) -> None: + processor = SimReadyProcessor( + scene=Scene(), + coarse_layout_by_id={}, + coarse_geometry_root=tmp_path / "coarse", + simready_geometry_root=tmp_path / "simready", + config=SimReadyProcessorConfig( + long_axis_object_ids=frozenset({"rolling_pin_001"}), + ), + ) + + assert processor._requires_long_axis_standardization("rolling_pin_001") + assert not processor._requires_long_axis_standardization("bottle_001") + + def test_uniform_scale_uses_the_z_up_tabletop_footprint(tmp_path: Path) -> None: """Measure y-up GLBs against the VLM's z-up XY target footprint.""" glb_path = tmp_path / "flat_fork.glb" From b27cdb6b83027fc4ff49d412b0ecd4aa4187fba5 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Tue, 18 Aug 2026 12:20:28 +0800 Subject: [PATCH 22/29] make assets group layout optimizer more rubust --- .../utils/assets_group_layout_optimizer.py | 58 ++++++++++++++++--- .../scene_engine/test_support_and_layout.py | 35 +++++++++++ 2 files changed, 84 insertions(+), 9 deletions(-) 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 index d4021fede..8b3933b09 100644 --- 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 @@ -120,24 +120,20 @@ def optimize(self) -> list[dict[str, object]]: 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." - ) + # Independently project each AABB into the rectangular optimization region. + offsets = self._project_aabbs_inside_rectangle(safe_support, base_aabbs) initial_overlaps = self._overlaps(base_aabbs, offsets) + projected_asset_count = int(np.count_nonzero(np.any(offsets != 0.0, axis=1))) log_info( "Support-constrained AABB overlap optimization started: " f"assets={len(asset_ids)}, initial_overlaps={len(initial_overlaps)}, " + f"initial_projections={projected_asset_count}, " 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.") + log_info("AABB overlap optimization succeeded without pair separation.") self.refined_assets_layout = self._apply_offsets_to_y_up_layouts( asset_ids=asset_ids, offsets=offsets, @@ -213,6 +209,50 @@ def optimize(self) -> list[dict[str, object]]: "inside the detected table support region." ) + @staticmethod + def _project_aabbs_inside_rectangle( + support: Polygon | MultiPolygon, + base_aabbs: np.ndarray, + ) -> np.ndarray: + """Return minimum per-AABB offsets that place AABBs in a rectangle.""" + if not isinstance(support, Polygon) or support.interiors: + raise ValueError( + "Initial AABB projection requires an axis-aligned rectangular " + "support region." + ) + minimum_x, minimum_y, maximum_x, maximum_y = support.bounds + rectangle = Polygon( + [ + (minimum_x, minimum_y), + (maximum_x, minimum_y), + (maximum_x, maximum_y), + (minimum_x, maximum_y), + ] + ) + if not support.equals(rectangle): + raise ValueError( + "Initial AABB projection requires an axis-aligned rectangular " + "support region." + ) + + aabb_minimums, aabb_maximums = base_aabbs.min(axis=1), base_aabbs.max(axis=1) + half_extents = (aabb_maximums - aabb_minimums) / 2.0 + support_minimum = np.array([minimum_x, minimum_y], dtype=float) + support_maximum = np.array([maximum_x, maximum_y], dtype=float) + valid_center_minimums = support_minimum + half_extents + valid_center_maximums = support_maximum - half_extents + if np.any(valid_center_minimums > valid_center_maximums + 1e-9): + raise ValueError( + "An asset AABB is larger than the rectangular support region." + ) + + centers = (aabb_minimums + aabb_maximums) / 2.0 + # A center must stay inset from each boundary by its AABB half extent. + projected_centers = np.clip( + centers, valid_center_minimums, valid_center_maximums + ) + return projected_centers - centers + def _apply_offsets_to_y_up_layouts( self, *, asset_ids: list[str], offsets: np.ndarray ) -> list[dict[str, object]]: 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 98ed30491..fda7f9e08 100644 --- a/tests/gen_sim/scene_engine/test_support_and_layout.py +++ b/tests/gen_sim/scene_engine/test_support_and_layout.py @@ -166,6 +166,41 @@ def test_layout_optimizer_resolves_a_simple_pair_overlap() -> None: assert not optimizer._overlaps(base_aabbs, refined_offsets) +def test_layout_optimizer_projects_out_of_bounds_aabb_into_rectangle() -> None: + support = Polygon([(0, 0), (3, 0), (3, 3), (0, 3)]) + layout = _layout("cup", 0.0, 1.5) + aabb = _aabb(-0.5, 1.0, 0.5, 2.0) + optimizer = AssetsSupportLayoutOptimizer( + support_region=support, + assets_aabb_2d_z_up_world_corners_by_id={"cup": aabb}, + assets_layout=[layout], + ) + + refined = optimizer.optimize() + + offset = np.array( + [ + refined[0]["pos"][0] - layout["pos"][0], # type: ignore[index] + layout["pos"][2] - refined[0]["pos"][2], # type: ignore[index] + ] + ) + assert offset[0] == pytest.approx(0.5) + assert optimizer._all_contained(support, np.stack([aabb]), np.stack([offset])) + + +def test_layout_optimizer_rejects_aabb_larger_than_rectangle() -> None: + optimizer = AssetsSupportLayoutOptimizer( + support_region=Polygon([(0, 0), (2, 0), (2, 2), (0, 2)]), + assets_aabb_2d_z_up_world_corners_by_id={ + "large": _aabb(-0.5, 0.5, 2.5, 1.5) + }, + assets_layout=[_layout("large", 1.0, 1.0)], + ) + + with pytest.raises(ValueError, match="larger than the rectangular support"): + optimizer.optimize() + + def test_layout_optimizer_rejects_unresolvable_overlap() -> None: optimizer = AssetsSupportLayoutOptimizer( support_region=Polygon([(0, 0), (2, 0), (2, 2), (0, 2)]), From 4365da4b65837135333f7c5eb40e398686fd8ace Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Tue, 18 Aug 2026 12:21:19 +0800 Subject: [PATCH 23/29] delete the hard-coded spatial check in assets' names and descriptionsa --- .../generation/scene_understanding.py | 20 +++-------------- .../scene_engine/test_scene_understanding.py | 22 +++++++++++-------- 2 files changed, 16 insertions(+), 26 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py index 9507c1d3a..d93ada019 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py @@ -50,12 +50,6 @@ _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. @@ -82,6 +76,9 @@ 8. For assets, description contains only visible category, material, color, texture, shape, and structural details. Do not mention location, the table, or any relationship to another object. + Structural direction words are allowed when they describe the object itself: + "bottle with a black cap on top" is valid, while "bottle on the left of the + table" is not. Return JSON only: no Markdown, comments, or prose outside this exact schema: { @@ -494,17 +491,6 @@ def _parse_scene_object_fields( f"VLM JSON key {field_name}.category must be a lower-case snake_case " "class name." ) - # Check whether the name and description contain location or relationship words. - if _LOCATION_WORD_PATTERN.search(fields["name"]): - raise ValueError( - f"VLM JSON key {field_name}.name must not contain location or " - "relationship words." - ) - if _LOCATION_WORD_PATTERN.search(fields["description"]): - raise ValueError( - f"VLM JSON key {field_name}.description must not contain location or " - "relationship words." - ) return fields diff --git a/tests/gen_sim/scene_engine/test_scene_understanding.py b/tests/gen_sim/scene_engine/test_scene_understanding.py index d2b3898da..382e4933d 100644 --- a/tests/gen_sim/scene_engine/test_scene_understanding.py +++ b/tests/gen_sim/scene_engine/test_scene_understanding.py @@ -61,19 +61,23 @@ def test_image_object_analysis_parses_code_fence_and_assigns_stable_ids() -> Non 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_accepts_name_with_spatial_words() -> None: + scene = scene_understanding._parse_image_object_analysis_response( + _response(asset_name="left cup") + ) + + assert scene.assets[0].name == "left cup" -def test_image_object_analysis_rejects_location_words_in_object_descriptions() -> None: +def test_image_object_analysis_accepts_description_with_structural_words() -> None: response = json.loads(_response()) - response["assets"][0]["description"] = "A small ceramic cup on the table." + response["assets"][0]["description"] = "A small ceramic cup with a lid on top." + + scene = scene_understanding._parse_image_object_analysis_response( + json.dumps(response) + ) - with pytest.raises(ValueError, match="description must not contain location"): - scene_understanding._parse_image_object_analysis_response(json.dumps(response)) + assert scene.assets[0].description == "A small ceramic cup with a lid on top." def test_image_object_analysis_retries_then_updates_scene(tmp_path: Path) -> None: From b3f6e63fd0a4d385215e3cffa8a25e39a17a122b Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Tue, 18 Aug 2026 16:44:34 +0800 Subject: [PATCH 24/29] replace jiange's pca heuristic method with VLM operate rotation tools follow the assets' orientation_states --- .../scene_engine/core/scene_edit_plan.py | 20 ++- .../gen_sim/scene_engine/core/scene_graph.py | 4 + .../editing/scene_edit_asset_preparation.py | 10 +- .../editing/scene_edit_understanding.py | 72 +++++++-- .../pipeline/generation/scene_generation.py | 16 +- .../pipeline/utils/simready_processor.py | 140 +++++++----------- .../utils/simready_processor_utils.py | 103 +++++++++---- .../scene_engine/test_scene_edit_plan.py | 32 ++++ .../test_simready_processor_utils.py | 44 +++++- .../scene_engine/test_support_and_layout.py | 4 +- 10 files changed, 300 insertions(+), 145 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py b/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py index 83ad5bd79..0f78b9699 100644 --- a/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py +++ b/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py @@ -21,6 +21,7 @@ from embodichain.gen_sim.scene_engine.core.scene import Scene from embodichain.gen_sim.scene_engine.core.scene_graph import ( + OrientationState, SceneConstraintType, SceneGraph, TableRegion, @@ -44,6 +45,7 @@ class SceneEditOperation: category: str | None = None name: str | None = None description: str | None = None + orientation_state: OrientationState | None = None def to_dict(self) -> dict[str, object]: """Serialize one normalized edit operation.""" @@ -56,6 +58,7 @@ def to_dict(self) -> dict[str, object]: "category": self.category, "name": self.name, "description": self.description, + "orientation_state": self.orientation_state, } @@ -84,6 +87,7 @@ def validate(self) -> None: # Edit-plan rules: # - move and delete identify one existing non-table object with object_id. # - add carries generated object_id plus non-empty category, name, and description. + # - add may preserve an explicit standing or lying user placement intent. # - move always supplies target_id and relation; add may omit both. # - table_region is only valid with target_id=table and relation=on. # - target_id and relation are otherwise supplied together or both absent. @@ -164,6 +168,7 @@ def _validate_operation( operation.category, operation.name, operation.description, + operation.orientation_state, ) ): raise ValueError("Delete operations may only specify object_id.") @@ -171,6 +176,13 @@ def _validate_operation( if operation.target_id is None or operation.relation is None: raise ValueError("Move operations must specify target_id and relation.") + existing_orientation_state = self.scene_graph.node_by_id()[ + operation.object_id + ].orientation_state + if operation.orientation_state not in {None, existing_orientation_state}: + raise ValueError( + "Move operations may only preserve the existing orientation_state." + ) self._validate_position_reference( operation=operation, existing_object_ids=existing_object_ids, @@ -178,7 +190,11 @@ def _validate_operation( ) if any( value is not None - for value in (operation.category, operation.name, operation.description) + for value in ( + operation.category, + operation.name, + operation.description, + ) ): raise ValueError("Move operations must not declare a new object.") @@ -203,6 +219,8 @@ def _validate_add_operation( for value in (operation.category, operation.name, operation.description) ): raise ValueError("Add operations require category, name, and description.") + if operation.orientation_state not in {None, "standing", "lying"}: + raise ValueError("Add operation orientation_state is invalid.") SceneEditPlan._validate_position_reference( operation=operation, existing_object_ids=existing_object_ids, diff --git a/embodichain/gen_sim/scene_engine/core/scene_graph.py b/embodichain/gen_sim/scene_engine/core/scene_graph.py index e13f65aba..500980b53 100644 --- a/embodichain/gen_sim/scene_engine/core/scene_graph.py +++ b/embodichain/gen_sim/scene_engine/core/scene_graph.py @@ -193,6 +193,7 @@ def apply_updates( *, deleted_object_ids: set[str], added_object_ids: list[str], + added_orientation_states_by_id: dict[str, OrientationState | None], on_parent_updates: list[tuple[str, str, TableRegion | None]], planar_relation_updates: list[tuple[str, PlanarRelationType, str]], ) -> None: @@ -226,6 +227,8 @@ def apply_updates( raise ValueError( f"Duplicate scene graph nodes: {sorted(duplicate_object_ids)}" ) + if set(added_orientation_states_by_id) != set(added_object_ids): + raise ValueError("Added orientation states must match added node ids.") # New nodes default to the table; later updates replace that parent when needed. self.nodes.extend( @@ -233,6 +236,7 @@ def apply_updates( object_id=object_id, parent_id=TABLE_OBJECT_ID, parent_relation="on", + orientation_state=added_orientation_states_by_id[object_id], ) for object_id in added_object_ids ) diff --git a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py index c6ec26fb3..220a91014 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py @@ -117,10 +117,18 @@ def prepare_scene_edit_assets( coarse_layout_by_id=_coarse_layouts_by_id(generated_asset_glbs), coarse_geometry_root=stage_output_root / "coarse_geometry", simready_geometry_root=stage_output_root / "simready_geometry", - # Scene editing will later provide the VLM-selected scale and rotation. + # Every added asset uses the VLM's pose and post-pose XY footprint scale. config=SimReadyProcessorConfig( use_vlm_scale=vlm_client is not None, use_vlm_rotation=vlm_client is not None, + # An explicit edit state overrides the default stable tabletop pose. + orientation_states_by_id={ + operation.object_id: operation.orientation_state + for operation in scene_edit_plan.operations + if operation.op == "add" + and operation.object_id is not None + and operation.orientation_state is not None + }, ), vlm_client=vlm_client, ) diff --git a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py index a648981e5..92720c7ec 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py @@ -25,6 +25,7 @@ ) from embodichain.gen_sim.scene_engine.core.scene import Scene from embodichain.gen_sim.scene_engine.core.scene_graph import ( + OrientationState, PlanarRelationType, SceneGraph, SceneGraphNode, @@ -50,7 +51,9 @@ singular snake_case category, name, and description. Multiple add operations may have the same category and name; their final IDs are assigned by the program in operation order. target_id and relation are either both provided - or both null. + or both null. Set orientation_state to standing or lying only when the user + explicitly asks for that placement; otherwise set it to null so the object + uses its natural, physically stable tabletop pose. For every move and every positioned add, target_id must be an Existing object ID and relation must be one of on, left_of, right_of, in_front_of, or behind. @@ -77,11 +80,14 @@ class. name contains only color, material, texture, shape, and object details. description contains only visible category, material, color, texture, shape, and structural details. name and description must not mention position, the -table, or relations to any object. +table, relations to any object, or orientation. orientation_state must be null +unless the user explicitly requests standing/upright/vertical or lying/flat/ +horizontal placement. Follow that explicit user intent even if it is not the +object's natural stable pose. Return JSON only: no Markdown, comments, or prose. Every operation must contain exactly these fields: op, object_id, target_id, relation, table_region, category, -name, and description. Use null for every field that does not apply: +name, description, and orientation_state. Use null for every field that does not apply: { "operations": [ { @@ -92,7 +98,8 @@ "table_region": null, "category": null, "name": null, - "description": null + "description": null, + "orientation_state": null }, { "op": "delete", @@ -102,7 +109,8 @@ "table_region": null, "category": null, "name": null, - "description": null + "description": null, + "orientation_state": null }, { "op": "add", @@ -112,7 +120,8 @@ "table_region": "back_center", "category": "orange", "name": "small orange", - "description": "small round orange with a textured peel" + "description": "small round orange with a textured peel", + "orientation_state": null }, { "op": "add", @@ -122,7 +131,8 @@ "table_region": null, "category": "orange", "name": "small orange", - "description": "small round orange with a textured peel" + "description": "small round orange with a textured peel", + "orientation_state": null }, { "op": "add", @@ -130,14 +140,31 @@ "target_id": null, "relation": null, "table_region": null, - "category": "banana", - "name": "yellow banana", - "description": "curved yellow banana with a green stem" + "category": "bottle", + "name": "blue glass bottle", + "description": "tall transparent blue glass bottle with a narrow neck", + "orientation_state": "standing" + }, + { + "op": "add", + "object_id": null, + "target_id": null, + "relation": null, + "table_region": null, + "category": "fork", + "name": "silver metal fork", + "description": "four-tined silver stainless-steel fork with a plain handle", + "orientation_state": "lying" } ] } -The two orange additions intentionally share category and name. Do not add -fields beyond the required schema.""" +The two orange additions intentionally share category and name. The bottle +example represents an explicit user request to stand it upright, and the fork +example represents an explicit user request to lay it flat. Only add operations +may introduce a new non-null orientation_state. A move may use null or repeat +its existing orientation_state from the supplied scene metadata, but it must not +change that state. Delete operations must use null. Do not add fields beyond the +required schema.""" def understand_scene_edit( @@ -223,6 +250,7 @@ def _apply_scene_edit_plan_to_scene_graph( """Apply the target graph updates implied by add and move operations.""" deleted_object_ids: set[str] = set() added_object_ids: list[str] = [] + added_orientation_states_by_id: dict[str, OrientationState | None] = {} on_parent_updates: list[tuple[str, str, TableRegion | None]] = [] planar_relation_updates: list[tuple[str, PlanarRelationType, str]] = [] for operation in scene_edit_plan.operations: @@ -234,6 +262,9 @@ def _apply_scene_edit_plan_to_scene_graph( raise ValueError("Add and move operations must have an object_id.") if operation.op == "add": added_object_ids.append(operation.object_id) + added_orientation_states_by_id[operation.object_id] = ( + operation.orientation_state + ) if operation.target_id is None or operation.relation is None: continue if operation.relation == "on": @@ -253,6 +284,7 @@ def _apply_scene_edit_plan_to_scene_graph( scene_graph.apply_updates( deleted_object_ids=deleted_object_ids, added_object_ids=added_object_ids, + added_orientation_states_by_id=added_orientation_states_by_id, on_parent_updates=on_parent_updates, planar_relation_updates=planar_relation_updates, ) @@ -267,6 +299,9 @@ def _simplify_scene_info( table_regions_by_id = { node.object_id: node.table_region for node in scene_graph.nodes } + orientation_states_by_id = { + node.object_id: node.orientation_state for node in scene_graph.nodes + } return { "existing_object_ids": [scene_object.id for scene_object in scene.objects], "objects": [ @@ -277,6 +312,7 @@ def _simplify_scene_info( "description": scene_object.description, "center_xy": scene_object.center_xy, "table_region": table_regions_by_id.get(scene_object.id), + "orientation_state": orientation_states_by_id.get(scene_object.id), } for scene_object in scene.objects ], @@ -351,6 +387,7 @@ def _parse_scene_edit_operations( "category", "name", "description", + "orientation_state", } # Get ids and counts of existing objects to assign new add IDs. assigned_object_ids = {scene_object.id for scene_object in scene.objects} @@ -371,6 +408,7 @@ def _parse_scene_edit_operations( raise ValueError("Scene edit operations must use the required schema.") object_id = _optional_string(value.get("object_id"), field_name="object_id") category = _optional_string(value.get("category"), field_name="category") + orientation_state = _optional_orientation_state(value.get("orientation_state")) if op == "add": if object_id is not None: raise ValueError("VLM add operations must set object_id to null.") @@ -397,6 +435,7 @@ def _parse_scene_edit_operations( description=_optional_string( value.get("description"), field_name="description" ), + orientation_state=orientation_state, ) ) return operations @@ -442,3 +481,12 @@ def _optional_table_region(value: object) -> TableRegion | None: if value not in TABLE_REGIONS: raise ValueError("Scene edit operation table_region is invalid.") return value + + +def _optional_orientation_state(value: object) -> OrientationState | None: + """Validate an optional explicit upright or lying edit intent.""" + if value is None: + return None + if value not in {"standing", "lying"}: + raise ValueError("Scene edit operation orientation_state is invalid.") + return value diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py index 7a0e19f35..0e1b17f95 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py @@ -109,21 +109,24 @@ def generate_scene_and_refine( coarse_layout_by_id = { layout_object["id"]: layout_object for layout_object in coarse_layout } + # Coarse poses already preserve lying and unconstrained assets; only standing + # assets need a VLM semantic-axis correction before later z-up calibration. + standing_orientation_states_by_id = { + node.object_id: node.orientation_state + for node in scene_graph.nodes + if node.orientation_state == "standing" + } simready_processor = SimReadyProcessor( scene=scene, coarse_layout_by_id=coarse_layout_by_id, coarse_geometry_root=coarse_geometry_output_root, simready_geometry_root=simready_geometry_output_root, debug_output_root=debug_output_root, - # Image-to-scene uses the geometry service's coarse scale directly. + # Keep the geometry-server scale and only correct unstable standing poses. config=SimReadyProcessorConfig( use_vlm_scale=False, use_vlm_rotation=False, - long_axis_object_ids=frozenset( - node.object_id - for node in scene_graph.nodes - if node.orientation_state is not None - ), + orientation_states_by_id=standing_orientation_states_by_id, ), vlm_client=vlm_client, ) @@ -498,6 +501,7 @@ def _scene_graph_based_calibration( node = nodes_by_id.get(asset_id) if node is None: raise ValueError(f"Scene graph does not contain asset {asset_id!r}.") + # Only correct the standing assets. if node.orientation_state != "standing": calibrated_assets_layout.append(asset_layout) continue diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py index 7bc041abf..864e39a59 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py @@ -17,15 +17,15 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path import numpy as np -import open3d as o3d 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_graph import OrientationState from embodichain.gen_sim.scene_engine.core.scene_object import ( ObjectPhysics, SceneObject, @@ -34,8 +34,11 @@ OpenAICompatibleVLM, ) from embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor_utils import ( - query_vlm_object_rotation_and_target_size, + DEFAULT_NEEDED_LAYOUT, + LYING_NEEDED_LAYOUT, + STANDING_NEEDED_LAYOUT, compute_uniform_xy_scale_for_target, + query_vlm_object_rotation_and_target_size, render_object_front_top_views, rotate_glb_about_x_axis, ) @@ -68,8 +71,8 @@ class SimReadyProcessorConfig: use_vlm_scale: bool = False # Use the VLM-selected asset scale. use_vlm_rotation: bool = False # Use the VLM-selected asset rotation. - - long_axis_object_ids: frozenset[str] = frozenset() + # Explicit graph orientation overrides the default stable tabletop pose. + orientation_states_by_id: dict[str, OrientationState] = field(default_factory=dict) class SimReadyProcessor: @@ -103,7 +106,9 @@ def __init__( self.config = config if config is not None else SimReadyProcessorConfig() self.vlm_client = vlm_client if ( - self.config.use_vlm_scale or self.config.use_vlm_rotation + self.config.use_vlm_scale + or self.config.use_vlm_rotation + or self.config.orientation_states_by_id ) and vlm_client is None: raise ValueError("vlm_client is required when VLM transforms are enabled.") @@ -200,25 +205,34 @@ def _prepare_vlm_rotated_glb( ) -> tuple[Path, list[float] | None]: """Render, query, and optionally bake the VLM-selected x-axis rotation.""" coarse_path = self.coarse_geometry_root / f"{scene_object.id}.glb" - if not (self.config.use_vlm_scale or self.config.use_vlm_rotation): + orientation_state = self._orientation_state_for_object(scene_object.id) + orientation_pose_required = orientation_state is not None + if not ( + self.config.use_vlm_scale + or self.config.use_vlm_rotation + or orientation_pose_required + ): return coarse_path, None decision = self._vlm_transform_for_object( scene_object, - use_scale=self.config.use_vlm_scale, - use_rotation=self.config.use_vlm_rotation, + needed_layout=self._needed_layout_for_object(scene_object.id), ) rotate_about_x = bool(decision["rotate_about_x"]) - vlm_scale = compute_uniform_xy_scale_for_target( - glb_path=coarse_path, - target_xy_size_cm=decision["target_xy_size_cm"], - rotate_about_x=rotate_about_x, - ) + vlm_scale = None + if self.config.use_vlm_scale: + # The VLM target describes the final, post-rotation z-up XY footprint. + vlm_scale = compute_uniform_xy_scale_for_target( + glb_path=coarse_path, + target_xy_size_cm=decision["target_xy_size_cm"], + rotate_about_x=rotate_about_x, + ) rotated_path = rotate_glb_about_x_axis( input_path=coarse_path, output_path=self.simready_geometry_root / "vlm_rotated" / f"{scene_object.id}.glb", - rotate=rotate_about_x, + rotate=rotate_about_x + and (orientation_pose_required or self.config.use_vlm_rotation), ) # The scale flag controls whether this VLM-derived isotropic scale is used. # Apply the same factor on x, y, and z to preserve the asset's proportions. @@ -227,19 +241,34 @@ def _prepare_vlm_rotated_glb( [vlm_scale, vlm_scale, vlm_scale] if self.config.use_vlm_scale else None, ) + def _orientation_state_for_object(self, object_id: str) -> OrientationState | None: + """Return the explicit graph orientation requested for one object.""" + return self.config.orientation_states_by_id.get(object_id) + + def _needed_layout_for_object(self, object_id: str) -> str: + """Return the VLM layout instruction for one object's graph semantics.""" + return ( + STANDING_NEEDED_LAYOUT + if self._orientation_state_for_object(object_id) == "standing" + else ( + LYING_NEEDED_LAYOUT + if self._orientation_state_for_object(object_id) == "lying" + else DEFAULT_NEEDED_LAYOUT + ) + ) + def _vlm_transform_for_object( self, scene_object: SceneObject, *, - use_scale: bool, - use_rotation: bool, + needed_layout: str, ) -> dict[str, object]: """Render the object and return the validated VLM pose decision.""" - del use_scale, use_rotation assert self.vlm_client is not None coarse_path = self.coarse_geometry_root / f"{scene_object.id}.glb" - needed_layout = "This asset needs to be place on the table that will not move a lot after simulation." - debug_root = self.simready_geometry_root.parent / "debug" + debug_root = ( + self.debug_output_root or self.simready_geometry_root.parent / "debug" + ) rendered_path = render_object_front_top_views( glb_path=coarse_path, output_path=debug_root / "vlm_views" / f"{scene_object.id}.png", @@ -316,16 +345,6 @@ def _canonicalize_object_mesh( y_up_to_z_up_transform[:3, :3] = y_up_to_z_up_matrix mesh.apply_transform(y_up_to_z_up_transform) - # Standardize graph-marked elongated assets before shared mesh processing. - # This makes local z their primary axis, so later scene-graph calibration - # can reliably recover the image-observed standing or lying orientation. - long_axis_alignment_matrix = np.eye(3) - if self._requires_long_axis_standardization(object_id): - long_axis_alignment_matrix = self._standardize_long_axis_z_up(mesh) - long_axis_alignment_transform = np.eye(4) - long_axis_alignment_transform[:3, :3] = long_axis_alignment_matrix - mesh.apply_transform(long_axis_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) @@ -333,13 +352,7 @@ def _canonicalize_object_mesh( # 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. - long_axis_alignment_matrix - @ y_up_to_z_up_matrix - @ np.diag(coarse_scale) - @ y_up_to_z_up_matrix.T - @ long_axis_alignment_matrix.T + y_up_to_z_up_matrix @ np.diag(coarse_scale) @ y_up_to_z_up_matrix.T ) mesh.apply_transform(scale_transform) @@ -359,17 +372,7 @@ def _canonicalize_object_mesh( 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 local canonicalization so its coarse world pose does not - # change until layout refinement applies the image-observed correction. - local_long_axis_rotation = Rotation.from_matrix( - y_up_to_z_up_matrix.T @ long_axis_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_long_axis_rotation.inv().as_matrix() - ) + rotation = Rotation.from_euler("xyz", coarse_rot, degrees=True) # Update the pos. position_offset = y_up_to_z_up_matrix.T @ ( scale_transform[:3, :3] @ original_aabb_center + scaled_aabb_bottom_center @@ -380,47 +383,6 @@ def _canonicalize_object_mesh( "scale": [1.0, 1.0, 1.0], } - def _requires_long_axis_standardization(self, object_id: str) -> bool: - """Return whether graph semantics identified one asset with a long axis.""" - return object_id in self.config.long_axis_object_ids - - @staticmethod - def _standardize_long_axis_z_up(mesh: trimesh.Trimesh) -> np.ndarray: - """Return a proper rotation that maps a mesh's primary axis to z-up. - Thanks to chanjian's idea. - """ - if len(mesh.vertices) < 4 or len(mesh.faces) < 4: - raise ValueError( - "Long-axis 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( - "Long-axis 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. - - long_axis_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. - long_axis_rotation = long_axis_rotation @ principal_axes - return long_axis_rotation - @staticmethod def _three_floats(value: object, *, field_name: str) -> list[float]: """Validate and convert a three-value layout field to floats.""" diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor_utils.py b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor_utils.py index 539bace70..4523aa3e6 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor_utils.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor_utils.py @@ -70,6 +70,27 @@ rotate_about_x=false and use target_xy_size_cm=[8.0, 8.0]. """ +DEFAULT_NEEDED_LAYOUT = ( + "Place this asset on the table in its natural, physically stable resting " + "orientation. For example, a fork should lie flat on the table rather " + "than stand on an edge." +) +STANDING_NEEDED_LAYOUT = ( + "The scene graph requires this asset to stand vertically on the table, " + "even when its natural stable pose would be lying down. For example, a " + "bottle should stand on its base and a fork should stand upright. If the " + "coarse GLB is lying flat, set rotate_about_x=true so its semantic vertical " + "axis aligns with the z-up world's z axis; if it is already upright, set " + "it to false." +) +LYING_NEEDED_LAYOUT = ( + "The scene graph requires this asset to lie flat on the table, even when " + "its natural stable pose would be standing. For example, a bottle should " + "lie on its side and a fork should lie flat. Choose rotate_about_x so the " + "asset's semantic long axis remains in the tabletop x-y plane rather than " + "along the z-up world's z axis." +) + def render_object_front_top_views( *, @@ -267,17 +288,60 @@ def query_vlm_object_rotation_and_target_size( rendered_views_path: str | Path, vlm_client: OpenAICompatibleVLM, debug_output_path: str | Path | None = None, + json_max_attempts: int = 3, ) -> dict[str, object]: - """Ask the VLM for rotation and post-rotation tabletop footprint.""" - response_text = vlm_client.complete( - system_prompt=_VLM_SYSTEM_PROMPT, - user_prompt=( - f"Object description:\n{scene_object_description}\n\n" - f"Needed layout:\n{needed_layout}\n\n" - "The image contains front view on the left and top view on the right." - ), - image_path=rendered_views_path, - ) + """Ask the VLM for a valid rotation and post-rotation tabletop footprint.""" + if json_max_attempts < 1: + raise ValueError("json_max_attempts must be at least 1.") + last_validation_error: ValueError | None = None + for _ in range(json_max_attempts): + response_text = vlm_client.complete( + system_prompt=_VLM_SYSTEM_PROMPT, + user_prompt=( + f"Object description:\n{scene_object_description}\n\n" + f"Needed layout:\n{needed_layout}\n\n" + "The image contains front view on the left and top view on the right." + ), + image_path=rendered_views_path, + ) + try: + value = _parse_vlm_rotation_and_target_size_response(response_text) + break + except ValueError as exc: + last_validation_error = exc + else: + assert last_validation_error is not None + raise ValueError( + "VLM transform response is invalid after " + f"{json_max_attempts} attempts: {last_validation_error}" + ) from last_validation_error + + if debug_output_path is not None: + output_path = Path(debug_output_path).expanduser().resolve() + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text( + json.dumps( + { + "description": scene_object_description, + "needed_layout": needed_layout, + "rendered_views_path": str( + Path(rendered_views_path).expanduser().resolve() + ), + "vlm_output": value, + }, + indent=2, + ensure_ascii=False, + ) + + "\n", + encoding="utf-8", + ) + return value + + +def _parse_vlm_rotation_and_target_size_response( + response_text: str, +) -> dict[str, object]: + """Validate one VLM rotation-and-scale JSON response.""" try: value = json.loads(_strip_json_code_fence(response_text)) except json.JSONDecodeError as exc: @@ -300,25 +364,6 @@ def query_vlm_object_rotation_and_target_size( or not all(np.isfinite(item) and item > 0 for item in target_size) ): raise ValueError("VLM target_xy_size_cm must contain two positive numbers.") - if debug_output_path is not None: - output_path = Path(debug_output_path).expanduser().resolve() - output_path.parent.mkdir(parents=True, exist_ok=True) - output_path.write_text( - json.dumps( - { - "description": scene_object_description, - "needed_layout": needed_layout, - "rendered_views_path": str( - Path(rendered_views_path).expanduser().resolve() - ), - "vlm_output": value, - }, - indent=2, - ensure_ascii=False, - ) - + "\n", - encoding="utf-8", - ) return value diff --git a/tests/gen_sim/scene_engine/test_scene_edit_plan.py b/tests/gen_sim/scene_engine/test_scene_edit_plan.py index 5579e8d55..a99e04384 100644 --- a/tests/gen_sim/scene_engine/test_scene_edit_plan.py +++ b/tests/gen_sim/scene_engine/test_scene_edit_plan.py @@ -115,6 +115,7 @@ def test_scene_edit_plan_accepts_add_without_a_position() -> None: "category": "cup", "name": "green cup", "description": "A small green ceramic cup.", + "orientation_state": None, } ] @@ -159,6 +160,7 @@ def test_scene_edit_parser_assigns_ids_to_same_category_adds_in_order() -> None: "category": "orange", "name": "small_orange", "description": "A small round orange with a textured peel.", + "orientation_state": None, }, { "op": "add", @@ -169,6 +171,7 @@ def test_scene_edit_parser_assigns_ids_to_same_category_adds_in_order() -> None: "category": "orange", "name": "small_orange", "description": "A small round orange with a textured peel.", + "orientation_state": "lying", }, ] } @@ -181,6 +184,30 @@ def test_scene_edit_parser_assigns_ids_to_same_category_adds_in_order() -> None: "orange_002", "orange_003", ] + assert [operation.orientation_state for operation in operations] == [ + None, + "lying", + ] + + +def test_scene_edit_plan_rejects_a_changed_move_orientation_state() -> None: + scene, scene_graph = _scene_and_graph() + scene_graph.node_by_id()["book_001"].orientation_state = "lying" + + with pytest.raises(ValueError, match="may only preserve"): + SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[ + SceneEditOperation( + op="move", + object_id="book_001", + target_id="table", + relation="on", + orientation_state="standing", + ) + ], + ) def test_scene_edit_plan_rejects_targets_outside_the_input_scene() -> None: @@ -426,6 +453,7 @@ def test_scene_edit_graph_builder_adds_unpositioned_objects_on_the_table() -> No category="cup", name="green cup", description="A small green ceramic cup.", + orientation_state="standing", ) ], ) @@ -438,10 +466,12 @@ def test_scene_edit_graph_builder_adds_unpositioned_objects_on_the_table() -> No added_node = updated_scene_graph.node_by_id()["cup_001"] assert added_node.parent_id == "table" assert added_node.parent_relation == "on" + assert added_node.orientation_state == "standing" def test_scene_edit_graph_builder_updates_move_on_parent() -> None: scene, scene_graph = _scene_and_graph() + scene_graph.node_by_id()["orange_001"].orientation_state = "lying" plan = SceneEditPlan( scene=scene, scene_graph=scene_graph, @@ -451,6 +481,7 @@ def test_scene_edit_graph_builder_updates_move_on_parent() -> None: object_id="orange_001", target_id="table", relation="on", + orientation_state="lying", ) ], ) @@ -461,6 +492,7 @@ def test_scene_edit_graph_builder_updates_move_on_parent() -> None: ) assert updated_scene_graph.node_by_id()["orange_001"].parent_id == "table" + assert updated_scene_graph.node_by_id()["orange_001"].orientation_state == "lying" def test_scene_edit_graph_builder_adds_planar_relation_with_target_parent() -> None: diff --git a/tests/gen_sim/scene_engine/test_simready_processor_utils.py b/tests/gen_sim/scene_engine/test_simready_processor_utils.py index 13cf73e10..dd872f913 100644 --- a/tests/gen_sim/scene_engine/test_simready_processor_utils.py +++ b/tests/gen_sim/scene_engine/test_simready_processor_utils.py @@ -27,25 +27,61 @@ SimReadyProcessorConfig, ) from embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor_utils import ( + DEFAULT_NEEDED_LAYOUT, + LYING_NEEDED_LAYOUT, + STANDING_NEEDED_LAYOUT, compute_uniform_xy_scale_for_target, + query_vlm_object_rotation_and_target_size, ) -def test_simready_long_axis_standardization_uses_graph_selected_ids( +def test_simready_pose_layout_uses_graph_orientation_states( tmp_path: Path, ) -> None: + class VLM: + def complete(self, **_: object) -> str: + raise AssertionError("This selection test must not call the VLM.") + processor = SimReadyProcessor( scene=Scene(), coarse_layout_by_id={}, coarse_geometry_root=tmp_path / "coarse", simready_geometry_root=tmp_path / "simready", config=SimReadyProcessorConfig( - long_axis_object_ids=frozenset({"rolling_pin_001"}), + orientation_states_by_id={"bottle_001": "standing", "fork_001": "lying"}, ), + vlm_client=VLM(), # type: ignore[arg-type] + ) + + assert processor._orientation_state_for_object("bottle_001") == "standing" + assert processor._orientation_state_for_object("fork_001") == "lying" + assert processor._orientation_state_for_object("knife_001") is None + assert processor._needed_layout_for_object("bottle_001") == STANDING_NEEDED_LAYOUT + assert processor._needed_layout_for_object("fork_001") == LYING_NEEDED_LAYOUT + assert processor._needed_layout_for_object("knife_001") == DEFAULT_NEEDED_LAYOUT + + +def test_vlm_transform_query_retries_an_empty_response(tmp_path: Path) -> None: + class VLM: + def __init__(self) -> None: + self.responses = [ + "", + '{"rotate_about_x": false, "target_xy_size_cm": [8.0, 8.0]}', + ] + + def complete(self, **_: object) -> str: + return self.responses.pop(0) + + vlm_client = VLM() + decision = query_vlm_object_rotation_and_target_size( + scene_object_description="small blue bottle", + needed_layout=STANDING_NEEDED_LAYOUT, + rendered_views_path=tmp_path / "views.png", + vlm_client=vlm_client, # type: ignore[arg-type] ) - assert processor._requires_long_axis_standardization("rolling_pin_001") - assert not processor._requires_long_axis_standardization("bottle_001") + assert decision == {"rotate_about_x": False, "target_xy_size_cm": [8.0, 8.0]} + assert vlm_client.responses == [] def test_uniform_scale_uses_the_z_up_tabletop_footprint(tmp_path: Path) -> None: 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 fda7f9e08..b6c9df3ff 100644 --- a/tests/gen_sim/scene_engine/test_support_and_layout.py +++ b/tests/gen_sim/scene_engine/test_support_and_layout.py @@ -191,9 +191,7 @@ def test_layout_optimizer_projects_out_of_bounds_aabb_into_rectangle() -> None: def test_layout_optimizer_rejects_aabb_larger_than_rectangle() -> None: optimizer = AssetsSupportLayoutOptimizer( support_region=Polygon([(0, 0), (2, 0), (2, 2), (0, 2)]), - assets_aabb_2d_z_up_world_corners_by_id={ - "large": _aabb(-0.5, 0.5, 2.5, 1.5) - }, + assets_aabb_2d_z_up_world_corners_by_id={"large": _aabb(-0.5, 0.5, 2.5, 1.5)}, assets_layout=[_layout("large", 1.0, 1.0)], ) From bc92aa50df7b06fc6060ec779049027384ccb30f Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Wed, 19 Aug 2026 11:08:42 +0800 Subject: [PATCH 25/29] ignore the fixed 2d aabb overlap when doing the on-table optimization --- .../pipeline/utils/scene_layout_optimizer.py | 16 +++- .../test_scene_layout_optimizer.py | 86 +++++++++++++++++++ 2 files changed, 100 insertions(+), 2 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_optimizer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_optimizer.py index ea6387cd2..ed877c734 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_optimizer.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_optimizer.py @@ -143,6 +143,7 @@ def optimize_parent_child_xy( root_half_extents_xy=child_half_extents_xy, inequality_constraints=inequality_constraints, equality_constraints=equality_constraints, + fixed_root_xy_by_id=fixed_child_xy_by_id, solved_root_xy_by_id=solved_child_xy_by_id, config=self.config, ) @@ -225,6 +226,7 @@ def _optimize_table_root_xy( root_half_extents_xy=root_half_extents_xy, inequality_constraints=inequality_constraints, equality_constraints=equality_constraints, + fixed_root_xy_by_id=fixed_root_xy_by_id, solved_root_xy_by_id=solved_root_xy_by_id, config=config, ) @@ -443,18 +445,26 @@ def _refine_root_collisions( root_half_extents_xy: dict[str, np.ndarray], inequality_constraints: list[tuple[np.ndarray, float]], equality_constraints: list[tuple[np.ndarray, float]], + fixed_root_xy_by_id: dict[str, list[float] | None], solved_root_xy_by_id: dict[str, list[float]], config: SceneLayoutOptimizerConfig, ) -> dict[str, list[float]]: - """Add AABB separation constraints until the table roots no longer overlap.""" + """Separate only sibling pairs that include a layout-variable object.""" seen_pairs: set[tuple[str, str]] = set() current_xy_by_id = solved_root_xy_by_id for _ in range(config.max_collision_rounds): - overlaps = _root_aabb_overlaps( + all_overlaps = _root_aabb_overlaps( root_ids=root_ids, root_half_extents_xy=root_half_extents_xy, xy_by_id=current_xy_by_id, ) + # Fixed/fixed siblings are outside this edit and therefore cannot be solved here. + overlaps = [ + overlap + for overlap in all_overlaps + if fixed_root_xy_by_id[overlap[1]] is None + or fixed_root_xy_by_id[overlap[2]] is None + ] if not overlaps: return current_xy_by_id added_constraint_count = 0 @@ -494,6 +504,8 @@ def _refine_root_collisions( root_half_extents_xy=root_half_extents_xy, xy_by_id=current_xy_by_id, ) + if fixed_root_xy_by_id[first_id] is None + or fixed_root_xy_by_id[second_id] is None ] raise ValueError( "Table-root AABB collisions remain after layout refinement: " diff --git a/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py b/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py index 36b4bfa79..bca20da8e 100644 --- a/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py +++ b/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py @@ -31,9 +31,19 @@ SceneLayoutConstructor, ) from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_optimizer import ( + SceneLayoutOptimizer, _table_region_bounds, ) +_TABLE_BOUNDS = [ + [-2.0, -2.0], + [2.0, -2.0], + [2.0, 2.0], + [-2.0, 2.0], +] +_OVERLAPPING_CENTER_XY = [0.0, 0.0] +_ASSET_SIDE_LENGTH_M = 0.2 + def _asset( *, @@ -136,3 +146,79 @@ def test_layout_constructor_places_new_child_on_parent_top( assert placed_cup.center_xy == [0.0, 0.0] # book top is z=0.62 m; cup half-height is 0.1 m and clearance is 0.02 m. assert np.allclose(placed_cup.pos, [0.0, 0.74, 0.0]) + + +def test_table_optimizer_ignores_fixed_sibling_overlap(tmp_path: Path) -> None: + asset_glb = tmp_path / "asset.glb" + trimesh.creation.box(extents=[_ASSET_SIDE_LENGTH_M] * 3).export(asset_glb) + first_id, second_id = "first_001", "second_001" + optimizer = SceneLayoutOptimizer() + + solved_xy_by_id = optimizer.optimize_table_root_xy( + assets_by_id={ + first_id: _asset( + object_id=first_id, + glb_path=asset_glb, + center_xy=_OVERLAPPING_CENTER_XY, + ), + second_id: _asset( + object_id=second_id, + glb_path=asset_glb, + center_xy=_OVERLAPPING_CENTER_XY, + ), + }, + root_ids=[first_id, second_id], + root_seed_xy_by_id={ + first_id: _OVERLAPPING_CENTER_XY, + second_id: _OVERLAPPING_CENTER_XY, + }, + imported_root_ids={first_id, second_id}, + fixed_root_xy_by_id={ + first_id: _OVERLAPPING_CENTER_XY, + second_id: _OVERLAPPING_CENTER_XY, + }, + root_table_regions_by_id={first_id: None, second_id: None}, + table_optimization_rect_xy=_TABLE_BOUNDS, + root_relations=[], + ) + + assert solved_xy_by_id == { + first_id: _OVERLAPPING_CENTER_XY, + second_id: _OVERLAPPING_CENTER_XY, + } + + +def test_table_optimizer_separates_variable_sibling_from_fixed_sibling( + tmp_path: Path, +) -> None: + asset_glb = tmp_path / "asset.glb" + trimesh.creation.box(extents=[_ASSET_SIDE_LENGTH_M] * 3).export(asset_glb) + fixed_id, variable_id = "fixed_001", "variable_001" + optimizer = SceneLayoutOptimizer() + + solved_xy_by_id = optimizer.optimize_table_root_xy( + assets_by_id={ + fixed_id: _asset( + object_id=fixed_id, + glb_path=asset_glb, + center_xy=_OVERLAPPING_CENTER_XY, + ), + variable_id: _asset( + object_id=variable_id, + glb_path=asset_glb, + ), + }, + root_ids=[fixed_id, variable_id], + root_seed_xy_by_id={ + fixed_id: _OVERLAPPING_CENTER_XY, + variable_id: _OVERLAPPING_CENTER_XY, + }, + imported_root_ids={fixed_id}, + fixed_root_xy_by_id={fixed_id: _OVERLAPPING_CENTER_XY, variable_id: None}, + root_table_regions_by_id={fixed_id: None, variable_id: None}, + table_optimization_rect_xy=_TABLE_BOUNDS, + root_relations=[], + ) + + assert solved_xy_by_id[fixed_id] == _OVERLAPPING_CENTER_XY + assert np.max(np.abs(solved_xy_by_id[variable_id])) >= _ASSET_SIDE_LENGTH_M - 1e-6 From 361b69df4c2455bc438c0eca6be3d964d7147371 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Wed, 19 Aug 2026 15:06:04 +0800 Subject: [PATCH 26/29] finished the table as root layout optimization (but still needs to be tested) --- .../utils/scene_layout_constructor.py | 165 ++-- .../pipeline/utils/scene_layout_optimizer.py | 797 ------------------ .../pipeline/utils/scene_layout_utils.py | 155 ++++ .../utils/table_surface_layout_optimizer.py | 542 ++++++++++++ .../test_scene_layout_optimizer.py | 113 +-- 5 files changed, 815 insertions(+), 957 deletions(-) delete mode 100644 embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_optimizer.py create mode 100644 embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_utils.py create mode 100644 embodichain/gen_sim/scene_engine/pipeline/utils/table_surface_layout_optimizer.py diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_constructor.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_constructor.py index 38706ac47..9279bd1ae 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_constructor.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_constructor.py @@ -26,9 +26,19 @@ SceneGraph, ) from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject -from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_optimizer import ( - SceneLayoutOptimizerConfig, - SceneLayoutOptimizer, +from embodichain.gen_sim.scene_engine.pipeline.utils.parent_surface_layout_optimizer import ( + ParentSurfaceLayoutOptimizer, + ParentSurfaceLayoutOptimizerConfig, + ParentSurfaceLayoutProblem, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_utils import ( + translate_scene_object_y_up_by_z_up_delta, + update_scene_object_y_up_pose_from_z_up_support, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.table_surface_layout_optimizer import ( + TableSurfaceLayoutOptimizer, + TableSurfaceLayoutOptimizerConfig, + TableSurfaceLayoutProblem, ) @@ -66,33 +76,45 @@ def __init__( layout_variable_ids: set[str], generated_scene_objects: list[SceneObject], output_root: str | Path, - config: SceneLayoutOptimizerConfig | None = None, + table_surface_config: TableSurfaceLayoutOptimizerConfig | None = None, + parent_surface_config: ParentSurfaceLayoutOptimizerConfig | None = None, ) -> None: self.formal_scene = formal_scene self.goal_scene_graph = goal_scene_graph self.layout_variable_ids = layout_variable_ids self.generated_scene_objects = generated_scene_objects self.output_root = Path(output_root).expanduser().resolve() - self.layout_optimizer = SceneLayoutOptimizer(config=config) + # Table surface optimizer. + self.table_surface_layout_optimizer = TableSurfaceLayoutOptimizer( + config=table_surface_config + ) + # Parent surface (on) optimizer. + self.parent_surface_layout_optimizer = ParentSurfaceLayoutOptimizer( + config=parent_surface_config + ) self._current_xy_by_id: dict[str, list[float] | None] = {} self._solved_delta_xy_by_id: dict[str, list[float]] = {} self._updated_object_ids: set[str] = set() def construct(self) -> Scene: """Construct table-root layouts before later stacked-group refinement.""" + # Build layout problem. layout_problem = self._build_problem() + # Get current XY centers. self._current_xy_by_id = { object_id: list(initial_xy) if initial_xy is not None else None for object_id, initial_xy in layout_problem.initial_xy_by_id.items() } self._solved_delta_xy_by_id = {} self._updated_object_ids = set() + # Check the group. if ( layout_problem.groups and layout_problem.groups[0].parent_id != TABLE_OBJECT_ID ): raise ValueError("The first layout group must be rooted at the table.") + # Optimize each group in BFS order, propagating solved deltas to descendants. for group in layout_problem.groups: if group.parent_id == TABLE_OBJECT_ID: self._optimize_table_group( @@ -114,56 +136,18 @@ def _optimize_table_group( group: SceneLayoutGroup, ) -> None: """Optimize all direct on-table children before any stacked child groups.""" + table_surface_problem = TableSurfaceLayoutProblem.from_layout_problem( + layout_problem=layout_problem, + group=group, + current_xy_by_id=self._current_xy_by_id, + ) + solved_root_xy_by_id = self.table_surface_layout_optimizer.optimize( + table_surface_problem + ) table = layout_problem.post_edit_scene.table if table is None: raise ValueError("Table group optimization requires a table.") - if table.support_optimization_rect_xy is None: - raise ValueError( - "Table group optimization requires a table support optimization rectangle." - ) - - root_ids = set(group.child_ids) - root_relations = [ - relation - for relation in layout_problem.goal_scene_graph.relations - if relation.source_id in root_ids and relation.target_id in root_ids - ] - root_seed_xy_by_id: dict[str, list[float]] = {} - for root_id in group.child_ids: - inherited_xy = self._current_xy_by_id[root_id] - # New roots start from the table-local origin; imported roots keep their pose. - root_seed_xy_by_id[root_id] = ( - [0.0, 0.0] if inherited_xy is None else list(inherited_xy) - ) - self._current_xy_by_id[root_id] = root_seed_xy_by_id[root_id] - - nodes_by_id = layout_problem.goal_scene_graph.node_by_id() - solved_root_xy_by_id = self.layout_optimizer.optimize_table_root_xy( - assets_by_id={ - asset.id: asset for asset in layout_problem.post_edit_scene.assets - }, - root_ids=group.child_ids, - root_seed_xy_by_id=root_seed_xy_by_id, - imported_root_ids={ - root_id - for root_id in group.child_ids - if layout_problem.initial_xy_by_id[root_id] is not None - }, - fixed_root_xy_by_id={ - root_id: ( - None - if root_id in layout_problem.layout_variable_ids - else self._current_xy_by_id[root_id] - ) - for root_id in group.child_ids - }, - root_table_regions_by_id={ - root_id: nodes_by_id[root_id].table_region - for root_id in group.child_ids - }, - table_optimization_rect_xy=table.support_optimization_rect_xy, - root_relations=root_relations, - ) + # Check the table's z. if table.support_surface_z is None and any( root_id in layout_problem.layout_variable_ids for root_id in group.child_ids ): @@ -172,7 +156,7 @@ def _optimize_table_group( asset.id: asset for asset in layout_problem.post_edit_scene.assets } for root_id, solved_xy in solved_root_xy_by_id.items(): - seed_xy = root_seed_xy_by_id[root_id] + seed_xy = table_surface_problem.root_seed_xy_by_id[root_id] delta_xy = [ solved_xy[0] - seed_xy[0], solved_xy[1] - seed_xy[1], @@ -182,10 +166,11 @@ def _optimize_table_group( if root_id in layout_problem.layout_variable_ids: # Direct add/move roots receive a new pose on the table support. assert table.support_surface_z is not None - self.layout_optimizer.update_scene_object_y_up_pose_from_z_up_support( + update_scene_object_y_up_pose_from_z_up_support( scene_object=assets_by_id[root_id], support_region_z=table.support_surface_z, center_xy=solved_xy, + clearance_m=0.00, # Directly place on the support surface. ) self._updated_object_ids.add(root_id) self._propagate_descendant_delta( @@ -202,6 +187,7 @@ def _propagate_descendant_delta( delta_xy: list[float], ) -> None: """Move every positioned descendant by one solved ancestor XY delta.""" + # A zero root delta cannot change any descendant pose, so skip the subtree walk. if delta_xy == [0.0, 0.0]: return assets_by_id = {asset.id: asset for asset in scene.assets} @@ -219,7 +205,7 @@ def _propagate_descendant_delta( descendant_xy[0] + delta_xy[0], descendant_xy[1] + delta_xy[1], ] - self.layout_optimizer.translate_scene_object_y_up_by_z_up_delta( + translate_scene_object_y_up_by_z_up_delta( scene_object=assets_by_id[descendant_id], delta_xy=delta_xy, ) @@ -233,54 +219,17 @@ def _optimize_parent_group( group: SceneLayoutGroup, ) -> None: """Optimize one settled parent's direct on-children in local XY coordinates.""" - assets_by_id = { - asset.id: asset for asset in layout_problem.post_edit_scene.assets - } - parent = assets_by_id.get(group.parent_id) - if parent is None: - raise ValueError(f"Parent {group.parent_id!r} is not an asset.") - parent_aabb = self.layout_optimizer.scene_object_z_up_world_aabb( - scene_object=parent + parent_surface_problem = ParentSurfaceLayoutProblem.from_layout_problem( + layout_problem=layout_problem, + group=group, + current_xy_by_id=self._current_xy_by_id, ) - parent_aabb_xy = [ - [parent_aabb[0][0], parent_aabb[0][1]], - [parent_aabb[1][0], parent_aabb[1][1]], - ] - parent_center_xy = [ - (parent_aabb[0][0] + parent_aabb[1][0]) / 2.0, - (parent_aabb[0][1] + parent_aabb[1][1]) / 2.0, - ] - child_seed_xy_by_id: dict[str, list[float]] = {} - for child_id in group.child_ids: - inherited_xy = self._current_xy_by_id[child_id] - # New children start at their parent's current AABB center. - child_seed_xy_by_id[child_id] = ( - parent_center_xy if inherited_xy is None else list(inherited_xy) - ) - self._current_xy_by_id[child_id] = child_seed_xy_by_id[child_id] - - solved_child_xy_by_id = self.layout_optimizer.optimize_parent_child_xy( - assets_by_id=assets_by_id, - child_ids=group.child_ids, - child_seed_xy_by_id=child_seed_xy_by_id, - imported_child_ids={ - child_id - for child_id in group.child_ids - if layout_problem.initial_xy_by_id[child_id] is not None - }, - fixed_child_xy_by_id={ - child_id: ( - None - if child_id in layout_problem.layout_variable_ids - else self._current_xy_by_id[child_id] - ) - for child_id in group.child_ids - }, - parent_aabb_xy=parent_aabb_xy, + # Get results. + solved_child_xy_by_id = self.parent_surface_layout_optimizer.optimize( + parent_surface_problem ) - parent_top_z = parent_aabb[1][2] for child_id, solved_xy in solved_child_xy_by_id.items(): - seed_xy = child_seed_xy_by_id[child_id] + seed_xy = parent_surface_problem.child_seed_xy_by_id[child_id] delta_xy = [ solved_xy[0] - seed_xy[0], solved_xy[1] - seed_xy[1], @@ -289,9 +238,9 @@ def _optimize_parent_group( self._solved_delta_xy_by_id[child_id] = delta_xy if child_id in layout_problem.layout_variable_ids: # Variable children are placed directly above the parent's current top. - self.layout_optimizer.update_scene_object_y_up_pose_from_z_up_support( - scene_object=assets_by_id[child_id], - support_region_z=parent_top_z, + update_scene_object_y_up_pose_from_z_up_support( + scene_object=parent_surface_problem.assets_by_id[child_id], + support_region_z=parent_surface_problem.parent_top_z, center_xy=solved_xy, ) self._updated_object_ids.add(child_id) @@ -303,6 +252,7 @@ def _optimize_parent_group( def _build_problem(self) -> SceneLayoutProblem: """Build post-edit objects and preserve formal-scene centers as seeds.""" + # Validate the graph first. self.goal_scene_graph.validate() graph_object_ids = set(self.goal_scene_graph.node_by_id()) generated_objects_by_id = self._generated_scene_objects_by_id() @@ -326,13 +276,14 @@ def _build_problem(self) -> SceneLayoutProblem: } if post_edit_object_ids != graph_object_ids: raise ValueError("Goal scene graph and post-edit scene have different ids.") + # Get the movable asset ids. if not self.layout_variable_ids.issubset(post_edit_object_ids - {"table"}): raise ValueError( "Only post-edit assets may participate in layout optimization." ) - + # Get initial XY centers. initial_xy_by_id = { - asset.id: self._initial_xy( + asset.id: self._initial_xy( # The assets' center XY should always be updated whenever changes are made. asset, is_generated=asset.id in generated_objects_by_id, ) @@ -343,13 +294,15 @@ def _build_problem(self) -> SceneLayoutProblem: raise ValueError( f"New asset {object_id!r} must participate in layout optimization." ) + # Build the table-rooted BFS groups. + groups = self._build_groups() return SceneLayoutProblem( post_edit_scene=post_edit_scene, goal_scene_graph=self.goal_scene_graph, layout_variable_ids=set(self.layout_variable_ids), initial_xy_by_id=initial_xy_by_id, - groups=self._build_groups(), + groups=groups, ) def _build_groups(self) -> list[SceneLayoutGroup]: diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_optimizer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_optimizer.py deleted file mode 100644 index ed877c734..000000000 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_optimizer.py +++ /dev/null @@ -1,797 +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 - -import numpy as np -from scipy.optimize import minimize - -from embodichain.gen_sim.scene_engine.core.scene_graph import SceneGraphRelation -from embodichain.gen_sim.scene_engine.core.scene_object import 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, -) - - -@dataclass(frozen=True) -class SceneLayoutOptimizerConfig: - """Numerical controls shared by each graph-layout solve.""" - - relation_clearance_m: float = 0.03 - collision_margin_m: float = 0.02 - max_slsqp_iterations: int = 500 - slsqp_ftol: float = 1e-6 - max_collision_rounds: int = 8 - max_added_collision_pairs: int = 64 - imported_seed_weight: float = 5.0 - min_center_distance_m: float = 0.01 - min_center_distance_weight: float = 0.05 - - def __post_init__(self) -> None: - """Reject invalid numerical controls before assembling a layout problem.""" - if self.relation_clearance_m < 0.0: - raise ValueError("relation_clearance_m must be non-negative.") - if self.collision_margin_m < 0.0: - raise ValueError("collision_margin_m must be non-negative.") - if self.max_slsqp_iterations <= 0: - raise ValueError("max_slsqp_iterations must be positive.") - if self.slsqp_ftol <= 0.0: - raise ValueError("slsqp_ftol must be positive.") - if self.max_collision_rounds <= 0: - raise ValueError("max_collision_rounds must be positive.") - if self.max_added_collision_pairs <= 0: - raise ValueError("max_added_collision_pairs must be positive.") - - -class SceneLayoutOptimizer: - """Solve graph-constrained XY layouts and apply resulting poses.""" - - def __init__(self, *, config: SceneLayoutOptimizerConfig | None = None) -> None: - self.config = config if config is not None else SceneLayoutOptimizerConfig() - - def optimize_table_root_xy( - self, - *, - assets_by_id: dict[str, SceneObject], - root_ids: list[str], - root_seed_xy_by_id: dict[str, list[float]], - imported_root_ids: set[str], - fixed_root_xy_by_id: dict[str, list[float] | None], - root_table_regions_by_id: dict[str, str | None], - table_optimization_rect_xy: list[list[float]], - root_relations: list[SceneGraphRelation], - ) -> dict[str, list[float]]: - """Solve direct table-child centers with graph and AABB constraints.""" - return _optimize_table_root_xy( - assets_by_id=assets_by_id, - root_ids=root_ids, - root_seed_xy_by_id=root_seed_xy_by_id, - imported_root_ids=imported_root_ids, - fixed_root_xy_by_id=fixed_root_xy_by_id, - root_table_regions_by_id=root_table_regions_by_id, - table_optimization_rect_xy=table_optimization_rect_xy, - root_relations=root_relations, - config=self.config, - ) - - def optimize_parent_child_xy( - self, - *, - assets_by_id: dict[str, SceneObject], - child_ids: list[str], - child_seed_xy_by_id: dict[str, list[float]], - imported_child_ids: set[str], - fixed_child_xy_by_id: dict[str, list[float] | None], - parent_aabb_xy: list[list[float]], - ) -> dict[str, list[float]]: - """Solve direct on-children inside one parent's current XY AABB.""" - child_half_extents_xy = _asset_half_extents_xy( - assets_by_id=assets_by_id, - object_ids=child_ids, - ) - inequality_constraints: list[tuple[np.ndarray, float]] = [] - equality_constraints: list[tuple[np.ndarray, float]] = [] - child_index = {child_id: index for index, child_id in enumerate(child_ids)} - parent_bounds = _bounds_from_points(parent_aabb_xy) - for child_id in child_ids: - _append_aabb_center_bounds( - constraints=inequality_constraints, - root_index=child_index, - root_id=child_id, - bounds=parent_bounds, - half_extents_xy=child_half_extents_xy[child_id], - ) - fixed_xy = fixed_child_xy_by_id[child_id] - if fixed_xy is not None: - _append_fixed_root_constraints( - constraints=equality_constraints, - root_index=child_index, - root_id=child_id, - fixed_xy=fixed_xy, - ) - - solved_child_xy_by_id = _solve_root_xy( - root_ids=child_ids, - root_seed_xy_by_id=child_seed_xy_by_id, - imported_root_ids=imported_child_ids, - inequality_constraints=inequality_constraints, - equality_constraints=equality_constraints, - config=self.config, - ) - return _refine_root_collisions( - root_ids=child_ids, - root_seed_xy_by_id=child_seed_xy_by_id, - imported_root_ids=imported_child_ids, - root_half_extents_xy=child_half_extents_xy, - inequality_constraints=inequality_constraints, - equality_constraints=equality_constraints, - fixed_root_xy_by_id=fixed_child_xy_by_id, - solved_root_xy_by_id=solved_child_xy_by_id, - config=self.config, - ) - - @staticmethod - def scene_object_z_up_world_aabb( - *, - scene_object: SceneObject, - ) -> list[list[float]]: - """Return one object's current z-up world AABB as [min, max].""" - return _scene_object_z_up_world_aabb(scene_object=scene_object) - - @staticmethod - def update_scene_object_y_up_pose_from_z_up_support( - *, - scene_object: SceneObject, - support_region_z: float, - center_xy: list[float], - clearance_m: float = 0.02, - ) -> None: - """Place one SimReady asset on a horizontal z-up support region.""" - _update_scene_object_y_up_pose_from_z_up_support( - scene_object=scene_object, - support_region_z=support_region_z, - center_xy=center_xy, - clearance_m=clearance_m, - ) - - @staticmethod - def translate_scene_object_y_up_by_z_up_delta( - *, - scene_object: SceneObject, - delta_xy: list[float], - ) -> None: - """Translate one existing y-up pose by a solved z-up XY delta.""" - _translate_scene_object_y_up_by_z_up_delta( - scene_object=scene_object, - delta_xy=delta_xy, - ) - - -def _optimize_table_root_xy( - *, - assets_by_id: dict[str, SceneObject], - root_ids: list[str], - root_seed_xy_by_id: dict[str, list[float]], - imported_root_ids: set[str], - fixed_root_xy_by_id: dict[str, list[float] | None], - root_table_regions_by_id: dict[str, str | None], - table_optimization_rect_xy: list[list[float]], - root_relations: list[SceneGraphRelation], - config: SceneLayoutOptimizerConfig, -) -> dict[str, list[float]]: - """Solve direct table-child centers with graph and AABB constraints.""" - root_half_extents_xy = _asset_half_extents_xy( - assets_by_id=assets_by_id, - object_ids=root_ids, - ) - inequality_constraints, equality_constraints = _build_table_root_constraints( - root_ids=root_ids, - root_half_extents_xy=root_half_extents_xy, - root_relations=root_relations, - root_table_regions_by_id=root_table_regions_by_id, - table_optimization_rect_xy=table_optimization_rect_xy, - fixed_root_xy_by_id=fixed_root_xy_by_id, - config=config, - ) - solved_root_xy_by_id = _solve_root_xy( - root_ids=root_ids, - root_seed_xy_by_id=root_seed_xy_by_id, - imported_root_ids=imported_root_ids, - inequality_constraints=inequality_constraints, - equality_constraints=equality_constraints, - config=config, - ) - return _refine_root_collisions( - root_ids=root_ids, - root_seed_xy_by_id=root_seed_xy_by_id, - imported_root_ids=imported_root_ids, - root_half_extents_xy=root_half_extents_xy, - inequality_constraints=inequality_constraints, - equality_constraints=equality_constraints, - fixed_root_xy_by_id=fixed_root_xy_by_id, - solved_root_xy_by_id=solved_root_xy_by_id, - config=config, - ) - - -def _update_scene_object_y_up_pose_from_z_up_support( - *, - scene_object: SceneObject, - support_region_z: float, - center_xy: list[float], - clearance_m: float = 0.02, -) -> None: - """Place one SimReady asset on a horizontal z-up support region. - - ``SceneObject`` stores poses in y-up before export. The target center and - support height are z-up values because layout optimization uses that frame. - """ - if not np.isfinite(support_region_z): - raise ValueError("support_region_z must be finite.") - if clearance_m < 0.0 or not np.isfinite(clearance_m): - raise ValueError("clearance_m must be finite and non-negative.") - target_xy = _two_floats(center_xy, field_name="center_xy") - rotation_y_up = _three_floats_or_default( - scene_object.rot, - field_name="rot", - default=[0.0, 0.0, 0.0], - ) - mesh = _asset_z_up_mesh_at_zero_translation( - scene_object=scene_object, - rotation_y_up=rotation_y_up, - ) - target_position_z_up = np.array( - [ - target_xy[0] - float(mesh.bounds[:, 0].mean()), - target_xy[1] - float(mesh.bounds[:, 1].mean()), - float(support_region_z) + clearance_m - float(mesh.bounds[0, 2]), - ] - ) - z_up_to_y_up = np.linalg.inv(_y_up_to_z_up_matrix()) - # Persist the y-up pose that SceneExporter later converts back to z-up. - scene_object.pos = (z_up_to_y_up[:3, :3] @ target_position_z_up).tolist() - scene_object.rot = rotation_y_up - scene_object.center_xy = target_xy - - -def _translate_scene_object_y_up_by_z_up_delta( - *, - scene_object: SceneObject, - delta_xy: list[float], -) -> None: - """Translate one existing y-up pose by a solved z-up XY delta.""" - dx, dy = _two_floats(delta_xy, field_name="delta_xy") - current_pos = _three_floats_or_default( - scene_object.pos, - field_name="pos", - default=None, - ) - # z-up x maps to y-up x, while z-up y maps to negative y-up z. - scene_object.pos = [ - current_pos[0] + dx, - current_pos[1], - current_pos[2] - dy, - ] - if scene_object.center_xy is not None: - scene_object.center_xy = [ - scene_object.center_xy[0] + dx, - scene_object.center_xy[1] + dy, - ] - - -def _scene_object_z_up_world_aabb( - *, - scene_object: SceneObject, -) -> list[list[float]]: - """Measure one current SceneObject pose in z-up world coordinates.""" - position_y_up = _three_floats_or_default( - scene_object.pos, - field_name="pos", - default=None, - ) - mesh = _asset_z_up_mesh_at_zero_translation(scene_object=scene_object) - position_z_up = _y_up_to_z_up_matrix()[:3, :3] @ np.asarray( - position_y_up, - dtype=float, - ) - mesh.apply_translation(position_z_up) - return mesh.bounds.tolist() - - -def _build_table_root_constraints( - *, - root_ids: list[str], - root_half_extents_xy: dict[str, np.ndarray], - root_relations: list[SceneGraphRelation], - root_table_regions_by_id: dict[str, str | None], - table_optimization_rect_xy: list[list[float]], - fixed_root_xy_by_id: dict[str, list[float] | None], - config: SceneLayoutOptimizerConfig, -) -> tuple[list[tuple[np.ndarray, float]], list[tuple[np.ndarray, float]]]: - """Build hard table, region, planar, and fixed-root constraints.""" - root_index = {root_id: index for index, root_id in enumerate(root_ids)} - table_bounds = _bounds_from_points(table_optimization_rect_xy) - inequality_constraints: list[tuple[np.ndarray, float]] = [] - equality_constraints: list[tuple[np.ndarray, float]] = [] - - for root_id in root_ids: - region_bounds = _table_region_bounds( - table_bounds=table_bounds, - table_region=root_table_regions_by_id[root_id], - ) - _append_aabb_center_bounds( - constraints=inequality_constraints, - root_index=root_index, - root_id=root_id, - bounds=region_bounds, - half_extents_xy=root_half_extents_xy[root_id], - ) - fixed_xy = fixed_root_xy_by_id[root_id] - if fixed_xy is not None: - _append_fixed_root_constraints( - constraints=equality_constraints, - root_index=root_index, - root_id=root_id, - fixed_xy=fixed_xy, - ) - - for relation in root_relations: - _append_planar_relation_constraint( - constraints=inequality_constraints, - root_index=root_index, - source_id=relation.source_id, - relation=relation.relation, - target_id=relation.target_id, - source_half_extents_xy=root_half_extents_xy[relation.source_id], - target_half_extents_xy=root_half_extents_xy[relation.target_id], - relation_clearance_m=config.relation_clearance_m, - ) - return inequality_constraints, equality_constraints - - -def _solve_root_xy( - *, - root_ids: list[str], - root_seed_xy_by_id: dict[str, list[float]], - imported_root_ids: set[str], - inequality_constraints: list[tuple[np.ndarray, float]], - equality_constraints: list[tuple[np.ndarray, float]], - config: SceneLayoutOptimizerConfig, -) -> dict[str, list[float]]: - """Solve one root-group center model with the legacy SLSQP settings.""" - root_index = {root_id: index for index, root_id in enumerate(root_ids)} - initial_xy = np.asarray( - [root_seed_xy_by_id[root_id] for root_id in root_ids], dtype=float - ) - x0 = initial_xy.reshape(-1) - - def unpack(values: np.ndarray) -> dict[str, list[float]]: - return { - root_id: [float(values[2 * index]), float(values[2 * index + 1])] - for root_id, index in root_index.items() - } - - def objective(values: np.ndarray) -> float: - coordinates = values.reshape(-1, 2) - loss = 0.0 - for root_id, index in root_index.items(): - if root_id in imported_root_ids: - delta = coordinates[index] - initial_xy[index] - loss += config.imported_seed_weight * float(delta @ delta) - for first_index in range(len(root_ids)): - for second_index in range(first_index + 1, len(root_ids)): - distance = float( - np.linalg.norm(coordinates[first_index] - coordinates[second_index]) - ) - shortfall = max(0.0, config.min_center_distance_m - distance) - loss += config.min_center_distance_weight * shortfall**2 - return loss - - constraints: list[dict[str, object]] = [] - for row, bound in inequality_constraints: - constraints.append( - { - "type": "ineq", - "fun": lambda values, row=row, bound=bound: bound - float(row @ values), - } - ) - for row, bound in equality_constraints: - constraints.append( - { - "type": "eq", - "fun": lambda values, row=row, bound=bound: float(row @ values) - bound, - } - ) - - result = minimize( - objective, - x0, - method="SLSQP", - constraints=constraints, - options={ - "maxiter": config.max_slsqp_iterations, - "ftol": config.slsqp_ftol, - "disp": False, - }, - ) - if not result.success: - raise ValueError(f"Table layout optimization failed: {result.message}") - return unpack(np.asarray(result.x, dtype=float)) - - -def _refine_root_collisions( - *, - root_ids: list[str], - root_seed_xy_by_id: dict[str, list[float]], - imported_root_ids: set[str], - root_half_extents_xy: dict[str, np.ndarray], - inequality_constraints: list[tuple[np.ndarray, float]], - equality_constraints: list[tuple[np.ndarray, float]], - fixed_root_xy_by_id: dict[str, list[float] | None], - solved_root_xy_by_id: dict[str, list[float]], - config: SceneLayoutOptimizerConfig, -) -> dict[str, list[float]]: - """Separate only sibling pairs that include a layout-variable object.""" - seen_pairs: set[tuple[str, str]] = set() - current_xy_by_id = solved_root_xy_by_id - for _ in range(config.max_collision_rounds): - all_overlaps = _root_aabb_overlaps( - root_ids=root_ids, - root_half_extents_xy=root_half_extents_xy, - xy_by_id=current_xy_by_id, - ) - # Fixed/fixed siblings are outside this edit and therefore cannot be solved here. - overlaps = [ - overlap - for overlap in all_overlaps - if fixed_root_xy_by_id[overlap[1]] is None - or fixed_root_xy_by_id[overlap[2]] is None - ] - if not overlaps: - return current_xy_by_id - added_constraint_count = 0 - for _, first_id, second_id in overlaps[: config.max_added_collision_pairs]: - pair_key = tuple(sorted((first_id, second_id))) - if pair_key in seen_pairs: - continue - inequality_constraints.append( - _aabb_separation_constraint( - root_ids=root_ids, - first_id=first_id, - second_id=second_id, - first_half_extents_xy=root_half_extents_xy[first_id], - second_half_extents_xy=root_half_extents_xy[second_id], - first_xy=current_xy_by_id[first_id], - second_xy=current_xy_by_id[second_id], - collision_margin_m=config.collision_margin_m, - ) - ) - seen_pairs.add(pair_key) - added_constraint_count += 1 - if added_constraint_count == 0: - break - current_xy_by_id = _solve_root_xy( - root_ids=root_ids, - root_seed_xy_by_id=current_xy_by_id, - imported_root_ids=imported_root_ids, - inequality_constraints=inequality_constraints, - equality_constraints=equality_constraints, - config=config, - ) - - remaining_pairs = [ - f"{first_id}/{second_id}" - for _, first_id, second_id in _root_aabb_overlaps( - root_ids=root_ids, - root_half_extents_xy=root_half_extents_xy, - xy_by_id=current_xy_by_id, - ) - if fixed_root_xy_by_id[first_id] is None - or fixed_root_xy_by_id[second_id] is None - ] - raise ValueError( - "Table-root AABB collisions remain after layout refinement: " - f"{remaining_pairs}." - ) - - -def _asset_half_extents_xy( - *, - assets_by_id: dict[str, SceneObject], - object_ids: list[str], -) -> dict[str, np.ndarray]: - """Measure each asset's oriented z-up footprint around its XY center.""" - half_extents_xy: dict[str, np.ndarray] = {} - for object_id in object_ids: - asset = assets_by_id.get(object_id) - if asset is None: - raise ValueError(f"Table root {object_id!r} is not an asset.") - half_extents_xy[object_id] = _asset_half_extent_xy(asset) - return half_extents_xy - - -def _asset_half_extent_xy(asset: SceneObject) -> np.ndarray: - """Measure one SimReady GLB with its current orientation and scale.""" - mesh = _asset_z_up_mesh_at_zero_translation(scene_object=asset) - return (mesh.bounds[1, :2] - mesh.bounds[0, :2]) / 2.0 - - -def _asset_z_up_mesh_at_zero_translation( - *, - scene_object: SceneObject, - rotation_y_up: list[float] | None = None, -): - """Load one SimReady GLB in z-up with orientation and scale but no position.""" - asset = scene_object - if asset.simready_glb_path is None: - raise ValueError(f"Asset {asset.id!r} has no SimReady GLB path.") - y_up_layout = { - "id": asset.id, - "rot": ( - rotation_y_up - if rotation_y_up is not None - else _three_floats_or_default( - asset.rot, - field_name="rot", - default=[0.0, 0.0, 0.0], - ) - ), - "pos": [0.0, 0.0, 0.0], - "scale": _three_floats_or_default( - asset.scale, - field_name="scale", - default=[1.0, 1.0, 1.0], - ), - } - y_up_to_z_up = _y_up_to_z_up_matrix() - z_up_layout = transform_matrix_to_layout_object( - asset.id, - y_up_to_z_up - @ layout_object_to_transform_matrix(y_up_layout) - @ np.linalg.inv(y_up_to_z_up), - ) - mesh = load_glb_mesh(asset.simready_glb_path) - mesh.apply_transform(y_up_to_z_up) - mesh.apply_transform(layout_object_to_transform_matrix(z_up_layout)) - return mesh - - -def _y_up_to_z_up_matrix() -> np.ndarray: - """Return the coordinate transform used by SceneExporter and layout stages.""" - matrix = np.eye(4) - matrix[:3, :3] = np.array([[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]]) - return matrix - - -def _two_floats(value: object, *, field_name: str) -> list[float]: - """Validate one finite two-value vector.""" - if not isinstance(value, (list, tuple)) or len(value) != 2: - raise ValueError(f"{field_name} must contain two values.") - vector = [float(component) for component in value] - if not np.all(np.isfinite(vector)): - raise ValueError(f"{field_name} must contain finite values.") - return vector - - -def _three_floats_or_default( - value: object, - *, - field_name: str, - default: list[float] | None, -) -> list[float]: - """Return a finite three-value vector or the canonical SimReady default.""" - if value is None: - if default is None: - raise ValueError(f"{field_name} must contain three values.") - return list(default) - if not isinstance(value, (list, tuple)) or len(value) != 3: - raise ValueError(f"{field_name} must contain three values.") - vector = [float(component) for component in value] - if not np.all(np.isfinite(vector)): - raise ValueError(f"{field_name} must contain finite values.") - return vector - - -def _bounds_from_points(points: list[list[float]]) -> np.ndarray: - """Return [[min_x, min_y], [max_x, max_y]] from finite XY points.""" - coordinates = np.asarray(points, dtype=float) - if coordinates.ndim != 2 or coordinates.shape[1] != 2 or len(coordinates) < 2: - raise ValueError("XY bounds must contain at least two points.") - if not np.all(np.isfinite(coordinates)): - raise ValueError("XY bounds must contain finite values.") - return np.stack([coordinates.min(axis=0), coordinates.max(axis=0)]) - - -def _table_region_bounds( - *, - table_bounds: np.ndarray, - table_region: str | None, -) -> np.ndarray: - """Return the requested 3x3 table region, with y increasing toward front.""" - if table_region is None: - return table_bounds.copy() - column_by_region = { - "left_back": 0, - "left_center": 0, - "left_front": 0, - "back_center": 1, - "center": 1, - "front_center": 1, - "right_back": 2, - "right_center": 2, - "right_front": 2, - } - row_by_region = { - "left_back": 0, - "back_center": 0, - "right_back": 0, - "left_center": 1, - "center": 1, - "right_center": 1, - "left_front": 2, - "front_center": 2, - "right_front": 2, - } - if table_region not in column_by_region: - raise ValueError(f"Unsupported table region {table_region!r}.") - minimum, maximum = table_bounds - cell_size = (maximum - minimum) / 3.0 - region_minimum = minimum + cell_size * np.array( - [column_by_region[table_region], row_by_region[table_region]] - ) - return np.stack([region_minimum, region_minimum + cell_size]) - - -def _append_aabb_center_bounds( - *, - constraints: list[tuple[np.ndarray, float]], - root_index: dict[str, int], - root_id: str, - bounds: np.ndarray, - half_extents_xy: np.ndarray, -) -> None: - """Keep one root's complete AABB inside the given rectangular bounds.""" - minimum = bounds[0] + half_extents_xy - maximum = bounds[1] - half_extents_xy - if np.any(minimum > maximum): - raise ValueError( - f"Asset {root_id!r} cannot fit inside its assigned table region." - ) - variable_count = 2 * len(root_index) - root_offset = 2 * root_index[root_id] - for axis in range(2): - upper_row = np.zeros(variable_count) - upper_row[root_offset + axis] = 1.0 - constraints.append((upper_row, float(maximum[axis]))) - lower_row = np.zeros(variable_count) - lower_row[root_offset + axis] = -1.0 - constraints.append((lower_row, -float(minimum[axis]))) - - -def _append_fixed_root_constraints( - *, - constraints: list[tuple[np.ndarray, float]], - root_index: dict[str, int], - root_id: str, - fixed_xy: list[float], -) -> None: - """Use equality constraints so unchanged formal objects remain fixed.""" - variable_count = 2 * len(root_index) - root_offset = 2 * root_index[root_id] - for axis, coordinate in enumerate(fixed_xy): - row = np.zeros(variable_count) - row[root_offset + axis] = 1.0 - constraints.append((row, float(coordinate))) - - -def _append_planar_relation_constraint( - *, - constraints: list[tuple[np.ndarray, float]], - root_index: dict[str, int], - source_id: str, - relation: str, - target_id: str, - source_half_extents_xy: np.ndarray, - target_half_extents_xy: np.ndarray, - relation_clearance_m: float, -) -> None: - """Require directional relations to clear both sibling AABB footprints.""" - if source_id not in root_index or target_id not in root_index: - raise ValueError("Table-root planar relations must reference table roots.") - axis, source_sign = { - "left_of": (0, 1.0), - "right_of": (0, -1.0), - "behind": (1, 1.0), - "in_front_of": (1, -1.0), - }.get(relation, (None, None)) - if axis is None or source_sign is None: - raise ValueError(f"Unsupported planar relation {relation!r}.") - row = np.zeros(2 * len(root_index)) - row[2 * root_index[source_id] + axis] = source_sign - row[2 * root_index[target_id] + axis] = -source_sign - required_distance = ( - source_half_extents_xy[axis] - + target_half_extents_xy[axis] - + relation_clearance_m - ) - constraints.append((row, -float(required_distance))) - - -def _root_aabb_overlaps( - *, - root_ids: list[str], - root_half_extents_xy: dict[str, np.ndarray], - xy_by_id: dict[str, list[float]], -) -> list[tuple[float, str, str]]: - """Return root pairs whose current XY AABBs overlap without a margin.""" - overlaps: list[tuple[float, str, str]] = [] - for first_index, first_id in enumerate(root_ids): - first_xy = np.asarray(xy_by_id[first_id], dtype=float) - first_half_extents = root_half_extents_xy[first_id] - for second_id in root_ids[first_index + 1 :]: - second_xy = np.asarray(xy_by_id[second_id], dtype=float) - second_half_extents = root_half_extents_xy[second_id] - overlap_xy = np.minimum( - first_xy + first_half_extents, - second_xy + second_half_extents, - ) - np.maximum( - first_xy - first_half_extents, - second_xy - second_half_extents, - ) - if np.all(overlap_xy > 1e-9): - overlaps.append((float(np.min(overlap_xy)), first_id, second_id)) - return sorted(overlaps, reverse=True) - - -def _aabb_separation_constraint( - *, - root_ids: list[str], - first_id: str, - second_id: str, - first_half_extents_xy: np.ndarray, - second_half_extents_xy: np.ndarray, - first_xy: list[float], - second_xy: list[float], - collision_margin_m: float, -) -> tuple[np.ndarray, float]: - """Separate one overlapping pair along its shallowest penetration axis.""" - root_index = {root_id: index for index, root_id in enumerate(root_ids)} - first_xy_array = np.asarray(first_xy, dtype=float) - second_xy_array = np.asarray(second_xy, dtype=float) - overlap_xy = np.minimum( - first_xy_array + first_half_extents_xy, - second_xy_array + second_half_extents_xy, - ) - np.maximum( - first_xy_array - first_half_extents_xy, - second_xy_array - second_half_extents_xy, - ) - axis = int(np.argmin(overlap_xy)) - first_is_lower = first_xy_array[axis] < second_xy_array[axis] or ( - first_xy_array[axis] == second_xy_array[axis] and first_id < second_id - ) - row = np.zeros(2 * len(root_ids)) - first_coefficient = 1.0 if first_is_lower else -1.0 - row[2 * root_index[first_id] + axis] = first_coefficient - row[2 * root_index[second_id] + axis] = -first_coefficient - required_distance = ( - first_half_extents_xy[axis] + second_half_extents_xy[axis] + collision_margin_m - ) - return row, -float(required_distance) diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_utils.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_utils.py new file mode 100644 index 000000000..100f5b269 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_utils.py @@ -0,0 +1,155 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF 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 + +from embodichain.gen_sim.scene_engine.core.scene_object import 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, +) + + +def update_scene_object_y_up_pose_from_z_up_support( + *, + scene_object: SceneObject, + support_region_z: float, + center_xy: list[float], + clearance_m: float = 0.02, +) -> None: + """Place a SimReady asset on a horizontal z-up support region.""" + if ( + not np.isfinite(support_region_z) + or clearance_m < 0.0 + or not np.isfinite(clearance_m) + ): + raise ValueError("support_region_z and clearance_m must be finite and valid.") + target_xy = two_floats(center_xy, field_name="center_xy") + rotation_y_up = three_floats_or_default( + scene_object.rot, field_name="rot", default=[0.0, 0.0, 0.0] + ) + mesh = load_scene_object_z_up_mesh( + scene_object=scene_object, rotation_y_up=rotation_y_up + ) + target_position_z_up = np.array( + [ + target_xy[0] - float(mesh.bounds[:, 0].mean()), + target_xy[1] - float(mesh.bounds[:, 1].mean()), + float(support_region_z) + clearance_m - float(mesh.bounds[0, 2]), + ] + ) + scene_object.pos = ( + np.linalg.inv(y_up_to_z_up_matrix())[:3, :3] @ target_position_z_up + ).tolist() + scene_object.rot = rotation_y_up + scene_object.center_xy = target_xy + + +def translate_scene_object_y_up_by_z_up_delta( + *, scene_object: SceneObject, delta_xy: list[float] +) -> None: + """Translate an existing y-up pose by a solved z-up XY delta.""" + dx, dy = two_floats(delta_xy, field_name="delta_xy") + position = three_floats_or_default(scene_object.pos, field_name="pos", default=None) + scene_object.pos = [position[0] + dx, position[1], position[2] - dy] + if scene_object.center_xy is not None: + scene_object.center_xy = [ + scene_object.center_xy[0] + dx, + scene_object.center_xy[1] + dy, + ] + + +def measure_scene_object_z_up_world_aabb( + *, scene_object: SceneObject +) -> list[list[float]]: + """Measure one current SceneObject pose in z-up world coordinates.""" + position_y_up = three_floats_or_default( + scene_object.pos, field_name="pos", default=None + ) + mesh = load_scene_object_z_up_mesh(scene_object=scene_object) + mesh.apply_translation( + y_up_to_z_up_matrix()[:3, :3] @ np.asarray(position_y_up, dtype=float) + ) + return mesh.bounds.tolist() + + +def load_scene_object_z_up_mesh( + *, scene_object: SceneObject, rotation_y_up: list[float] | None = None +): + """Load a SimReady mesh in z-up with orientation and scale but no translation.""" + if scene_object.simready_glb_path is None: + raise ValueError(f"Asset {scene_object.id!r} has no SimReady GLB path.") + y_up_layout = { + "id": scene_object.id, + "rot": ( + rotation_y_up + if rotation_y_up is not None + else three_floats_or_default( + scene_object.rot, field_name="rot", default=[0.0, 0.0, 0.0] + ) + ), + "pos": [0.0, 0.0, 0.0], + "scale": three_floats_or_default( + scene_object.scale, field_name="scale", default=[1.0, 1.0, 1.0] + ), + } + y_up_to_z_up = y_up_to_z_up_matrix() + z_up_layout = transform_matrix_to_layout_object( + scene_object.id, + y_up_to_z_up + @ layout_object_to_transform_matrix(y_up_layout) + @ np.linalg.inv(y_up_to_z_up), + ) + mesh = load_glb_mesh(scene_object.simready_glb_path) + mesh.apply_transform(y_up_to_z_up) + mesh.apply_transform(layout_object_to_transform_matrix(z_up_layout)) + return mesh + + +def y_up_to_z_up_matrix() -> np.ndarray: + """Return the coordinate transform used by layout and export stages.""" + matrix = np.eye(4) + matrix[:3, :3] = np.array([[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]]) + return matrix + + +def two_floats(value: object, *, field_name: str) -> list[float]: + """Validate and return one finite two-value vector.""" + if not isinstance(value, (list, tuple)) or len(value) != 2: + raise ValueError(f"{field_name} must contain two values.") + result = [float(component) for component in value] + if not np.all(np.isfinite(result)): + raise ValueError(f"{field_name} must contain finite values.") + return result + + +def three_floats_or_default( + value: object, *, field_name: str, default: list[float] | None +) -> list[float]: + """Validate three finite values, or return a canonical default.""" + if value is None: + if default is None: + raise ValueError(f"{field_name} must contain three values.") + return list(default) + if not isinstance(value, (list, tuple)) or len(value) != 3: + raise ValueError(f"{field_name} must contain three values.") + result = [float(component) for component in value] + if not np.all(np.isfinite(result)): + raise ValueError(f"{field_name} must contain finite values.") + return result diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/table_surface_layout_optimizer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/table_surface_layout_optimizer.py new file mode 100644 index 000000000..38fe7e6c4 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/table_surface_layout_optimizer.py @@ -0,0 +1,542 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF 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 TYPE_CHECKING + +import numpy as np +from scipy.optimize import minimize + +from embodichain.gen_sim.scene_engine.core.scene_graph import SceneGraphRelation +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_utils import ( + load_scene_object_z_up_mesh, +) + +if TYPE_CHECKING: + from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_constructor import ( + SceneLayoutGroup, + SceneLayoutProblem, + ) + + +@dataclass +class TableSurfaceLayoutProblem: + """All scene-graph and geometry inputs for one table-surface solve.""" + + assets_by_id: dict[str, SceneObject] + root_ids: list[str] + root_seed_xy_by_id: dict[str, list[float]] + imported_root_ids: set[str] + fixed_root_xy_by_id: dict[str, list[float] | None] + root_table_regions_by_id: dict[str, str | None] + table_optimization_rect_xy: list[list[float]] + root_relations: list[SceneGraphRelation] + + @classmethod + def from_layout_problem( + cls, + *, + layout_problem: SceneLayoutProblem, + group: SceneLayoutGroup, + current_xy_by_id: dict[str, list[float] | None], + ) -> TableSurfaceLayoutProblem: + """Build one table-surface problem without mutating layout state.""" + table = layout_problem.post_edit_scene.table + if table is None: + raise ValueError("Table group optimization requires a table.") + if table.support_optimization_rect_xy is None: + raise ValueError( + "Table group optimization requires a table support optimization rectangle." + ) + root_ids = set(group.child_ids) + nodes_by_id = layout_problem.goal_scene_graph.node_by_id() + root_seed_xy_by_id = {} + for root_id in group.child_ids: + inherited_xy = current_xy_by_id[root_id] + # New roots begin from the table origin; imported roots retain their seed. + root_seed_xy_by_id[root_id] = ( + [0.0, 0.0] if inherited_xy is None else list(inherited_xy) + ) + return cls( + assets_by_id={ + asset.id: asset for asset in layout_problem.post_edit_scene.assets + }, + root_ids=group.child_ids, + root_seed_xy_by_id=root_seed_xy_by_id, + imported_root_ids={ + root_id + for root_id in group.child_ids + if layout_problem.initial_xy_by_id[root_id] is not None + }, + fixed_root_xy_by_id={ + root_id: ( + None + if root_id in layout_problem.layout_variable_ids + else current_xy_by_id[root_id] + ) + for root_id in group.child_ids + }, + root_table_regions_by_id={ + root_id: nodes_by_id[root_id].table_region + for root_id in group.child_ids + }, + table_optimization_rect_xy=table.support_optimization_rect_xy, + root_relations=[ + relation + for relation in layout_problem.goal_scene_graph.relations + if relation.source_id in root_ids and relation.target_id in root_ids + ], + ) + + +@dataclass(frozen=True) +class TableSurfaceLayoutOptimizerConfig: + """Numerical controls for one direct-table sibling layout solve.""" + + relation_clearance_m: float = 0.03 + collision_margin_m: float = 0.02 + max_slsqp_iterations: int = 500 + slsqp_ftol: float = 1e-6 + max_collision_rounds: int = 8 + max_added_collision_pairs: int = 64 + imported_seed_weight: float = 5.0 + min_center_distance_m: float = 0.01 + min_center_distance_weight: float = 0.05 + + def __post_init__(self) -> None: + """Reject invalid controls before assembling table-surface constraints.""" + if self.relation_clearance_m < 0.0: + raise ValueError("relation_clearance_m must be non-negative.") + if self.collision_margin_m < 0.0: + raise ValueError("collision_margin_m must be non-negative.") + if self.max_slsqp_iterations <= 0: + raise ValueError("max_slsqp_iterations must be positive.") + if self.slsqp_ftol <= 0.0: + raise ValueError("slsqp_ftol must be positive.") + if self.max_collision_rounds <= 0: + raise ValueError("max_collision_rounds must be positive.") + if self.max_added_collision_pairs <= 0: + raise ValueError("max_added_collision_pairs must be positive.") + + +class TableSurfaceLayoutOptimizer: + """Solve direct table children with table, relation, and collision constraints.""" + + def __init__( + self, + *, + config: TableSurfaceLayoutOptimizerConfig | None = None, + ) -> None: + self.config = ( + config if config is not None else TableSurfaceLayoutOptimizerConfig() + ) + + def optimize( + self, + problem: TableSurfaceLayoutProblem, + ) -> dict[str, list[float]]: + """Return the table-frame XY centers satisfying this atomic problem.""" + # Measure only this sibling group from the complete scene-asset index. + root_half_extents_xy = _asset_half_extents_xy( + assets_by_id=problem.assets_by_id, + object_ids=problem.root_ids, + ) + # Equality constraints for fixed roots, and inequality constraints for table-region and planar-relation bounds. + inequality_constraints, equality_constraints = _build_constraints( + problem=problem, + root_half_extents_xy=root_half_extents_xy, + config=self.config, + ) + # Solve with the SLSQP optimizer. + solved_root_xy_by_id = _solve_root_xy( + root_ids=problem.root_ids, + root_seed_xy_by_id=problem.root_seed_xy_by_id, + imported_root_ids=problem.imported_root_ids, + inequality_constraints=inequality_constraints, + equality_constraints=equality_constraints, + config=self.config, + ) + return _refine_root_collisions( + root_ids=problem.root_ids, + root_seed_xy_by_id=problem.root_seed_xy_by_id, + imported_root_ids=problem.imported_root_ids, + root_half_extents_xy=root_half_extents_xy, + inequality_constraints=inequality_constraints, + equality_constraints=equality_constraints, + fixed_root_xy_by_id=problem.fixed_root_xy_by_id, + solved_root_xy_by_id=solved_root_xy_by_id, + config=self.config, + ) + + +def _build_constraints( + *, + problem: TableSurfaceLayoutProblem, + root_half_extents_xy: dict[str, np.ndarray], + config: TableSurfaceLayoutOptimizerConfig, +) -> tuple[list[tuple[np.ndarray, float]], list[tuple[np.ndarray, float]]]: + """Build hard table-region, planar-relation, and fixed-root constraints.""" + # Objects which need to be optimized. + root_index = {root_id: index for index, root_id in enumerate(problem.root_ids)} + table_bounds = _bounds_from_points(problem.table_optimization_rect_xy) + # Initi constraints. + inequality_constraints: list[tuple[np.ndarray, float]] = [] + equality_constraints: list[tuple[np.ndarray, float]] = [] + for root_id in problem.root_ids: + # Get the table region bound for this root asset. + region_bounds = _table_region_bounds( + table_bounds=table_bounds, + table_region=problem.root_table_regions_by_id[root_id], + ) + # Add AABB constraints for each root's center inside the table region. + _append_aabb_center_bounds( + constraints=inequality_constraints, + root_index=root_index, + root_id=root_id, + bounds=region_bounds, + half_extents_xy=root_half_extents_xy[root_id], + ) + fixed_xy = problem.fixed_root_xy_by_id[root_id] + if fixed_xy is not None: + # Add fixed-root constraints for each root with a fixed XY center. + _append_fixed_root_constraints( + constraints=equality_constraints, + root_index=root_index, + root_id=root_id, + fixed_xy=fixed_xy, + ) + for relation in problem.root_relations: + # Add planar-relation constraints for each sibling relation in this group. + _append_planar_relation_constraint( + constraints=inequality_constraints, + root_index=root_index, + source_id=relation.source_id, + relation=relation.relation, + target_id=relation.target_id, + source_half_extents_xy=root_half_extents_xy[relation.source_id], + target_half_extents_xy=root_half_extents_xy[relation.target_id], + relation_clearance_m=config.relation_clearance_m, + ) + return inequality_constraints, equality_constraints + + +def _table_region_bounds( + *, + table_bounds: np.ndarray, + table_region: str | None, +) -> np.ndarray: + """Return the requested 3x3 table region, with y increasing toward front.""" + if table_region is None: + return table_bounds.copy() + column_by_region = { + "left_back": 0, + "left_center": 0, + "left_front": 0, + "back_center": 1, + "center": 1, + "front_center": 1, + "right_back": 2, + "right_center": 2, + "right_front": 2, + } + row_by_region = { + "left_back": 0, + "back_center": 0, + "right_back": 0, + "left_center": 1, + "center": 1, + "right_center": 1, + "left_front": 2, + "front_center": 2, + "right_front": 2, + } + if table_region not in column_by_region: + raise ValueError(f"Unsupported table region {table_region!r}.") + minimum, maximum = table_bounds + # 9-grid. + cell_size = (maximum - minimum) / 3.0 + region_minimum = minimum + cell_size * np.array( + [column_by_region[table_region], row_by_region[table_region]] + ) + return np.stack([region_minimum, region_minimum + cell_size]) + + +def _asset_half_extents_xy( + *, assets_by_id: dict[str, SceneObject], object_ids: list[str] +) -> dict[str, np.ndarray]: + """Measure each optimized asset's oriented z-up XY half-extents.""" + result = {} + for object_id in object_ids: + asset = assets_by_id.get(object_id) + if asset is None: + raise ValueError(f"Table root {object_id!r} is not an asset.") + mesh = load_scene_object_z_up_mesh(scene_object=asset) + result[object_id] = (mesh.bounds[1, :2] - mesh.bounds[0, :2]) / 2.0 + return result + + +def _bounds_from_points(points: list[list[float]]) -> np.ndarray: + coordinates = np.asarray(points, dtype=float) + if ( + coordinates.ndim != 2 + or coordinates.shape[1] != 2 + or len(coordinates) < 2 + or not np.all(np.isfinite(coordinates)) + ): + raise ValueError("XY bounds must contain at least two finite points.") + return np.stack([coordinates.min(axis=0), coordinates.max(axis=0)]) + + +def _append_aabb_center_bounds( + *, + constraints: list[tuple[np.ndarray, float]], + root_index: dict[str, int], + root_id: str, + bounds: np.ndarray, + half_extents_xy: np.ndarray, +) -> None: + minimum, maximum = bounds[0] + half_extents_xy, bounds[1] - half_extents_xy + if np.any(minimum > maximum): + raise ValueError( + f"Asset {root_id!r} cannot fit inside its assigned table region." + ) + # root_id is the sibling whose center is constrained in this AABB bound. + offset, count = 2 * root_index[root_id], 2 * len(root_index) + # offset selects this root's XY pair; count is the full flattened XY vector size. + for axis in range(2): + upper, lower = np.zeros(count), np.zeros(count) + upper[offset + axis], lower[offset + axis] = 1.0, -1.0 + constraints.extend( + [(upper, float(maximum[axis])), (lower, -float(minimum[axis]))] + ) + + +def _append_fixed_root_constraints( + *, + constraints: list[tuple[np.ndarray, float]], + root_index: dict[str, int], + root_id: str, + fixed_xy: list[float], +) -> None: + offset, count = 2 * root_index[root_id], 2 * len(root_index) + for axis, coordinate in enumerate(fixed_xy): + row = np.zeros(count) + row[offset + axis] = 1.0 + constraints.append((row, float(coordinate))) + + +def _append_planar_relation_constraint( + *, + constraints: list[tuple[np.ndarray, float]], + root_index: dict[str, int], + source_id: str, + relation: str, + target_id: str, + source_half_extents_xy: np.ndarray, + target_half_extents_xy: np.ndarray, + relation_clearance_m: float, +) -> None: + axis, sign = { + "left_of": (0, 1.0), + "right_of": (0, -1.0), + "behind": (1, 1.0), + "in_front_of": (1, -1.0), + }.get(relation, (None, None)) + if axis is None or source_id not in root_index or target_id not in root_index: + raise ValueError(f"Unsupported table-root planar relation {relation!r}.") + row = np.zeros(2 * len(root_index)) + row[2 * root_index[source_id] + axis] = sign + row[2 * root_index[target_id] + axis] = -sign + constraints.append( + ( + row, + -float( + source_half_extents_xy[axis] + + target_half_extents_xy[axis] + + relation_clearance_m + ), + ) + ) + + +def _solve_root_xy( + *, + root_ids: list[str], + root_seed_xy_by_id: dict[str, list[float]], + imported_root_ids: set[str], + inequality_constraints: list[tuple[np.ndarray, float]], + equality_constraints: list[tuple[np.ndarray, float]], + config: TableSurfaceLayoutOptimizerConfig, +) -> dict[str, list[float]]: + # Init with XY-seeds. + initial = np.asarray( + [root_seed_xy_by_id[root_id] for root_id in root_ids], dtype=float + ) + + def objective(values: np.ndarray) -> float: + xy = values.reshape(-1, 2) + loss = 0.0 + for index, root_id in enumerate(root_ids): + if root_id in imported_root_ids: + delta = xy[index] - initial[index] + loss += config.imported_seed_weight * float(delta @ delta) + return loss + + constraints = [ + { + "type": "ineq", + "fun": lambda values, row=row, bound=bound: bound - float(row @ values), + } + for row, bound in inequality_constraints + ] + [ + { + "type": "eq", + "fun": lambda values, row=row, bound=bound: float(row @ values) - bound, + } + for row, bound in equality_constraints + ] + result = minimize( + objective, + initial.reshape(-1), + method="SLSQP", + constraints=constraints, + options={ + "maxiter": config.max_slsqp_iterations, + "ftol": config.slsqp_ftol, + "disp": False, + }, + ) + if not result.success: + raise ValueError(f"Table layout optimization failed: {result.message}") + return { + root_id: [float(result.x[2 * index]), float(result.x[2 * index + 1])] + for index, root_id in enumerate(root_ids) + } + + +def _refine_root_collisions( + *, + root_ids: list[str], + root_seed_xy_by_id: dict[str, list[float]], + imported_root_ids: set[str], + root_half_extents_xy: dict[str, np.ndarray], + inequality_constraints: list[tuple[np.ndarray, float]], + equality_constraints: list[tuple[np.ndarray, float]], + fixed_root_xy_by_id: dict[str, list[float] | None], + solved_root_xy_by_id: dict[str, list[float]], + config: TableSurfaceLayoutOptimizerConfig, +) -> dict[str, list[float]]: + # Get current SLSQP solution. + current = solved_root_xy_by_id + seen: set[tuple[str, str]] = set() + for _ in range(config.max_collision_rounds): + # Fine overlaps. + overlaps = [ + pair + for pair in _root_aabb_overlaps( + root_ids=root_ids, half_extents=root_half_extents_xy, xy_by_id=current + ) + if fixed_root_xy_by_id[pair[1]] is None + or fixed_root_xy_by_id[pair[2]] is None + ] + if not overlaps: + return current + added = 0 + for _, first_id, second_id in overlaps[: config.max_added_collision_pairs]: + key = tuple(sorted((first_id, second_id))) + if key in seen: + continue + # Add a new SLSQP constraint to separate this overlapping pair. + inequality_constraints.append( + _aabb_separation_constraint( + root_ids=root_ids, + first_id=first_id, + second_id=second_id, + half_extents=root_half_extents_xy, + xy_by_id=current, + margin=config.collision_margin_m, + ) + ) + seen.add(key) + added += 1 + if not added: + break + current = _solve_root_xy( + root_ids=root_ids, + root_seed_xy_by_id=current, + imported_root_ids=imported_root_ids, + inequality_constraints=inequality_constraints, + equality_constraints=equality_constraints, + config=config, + ) + raise ValueError("Table-root AABB collisions remain after layout refinement.") + + +def _root_aabb_overlaps( + *, + root_ids: list[str], + half_extents: dict[str, np.ndarray], + xy_by_id: dict[str, list[float]], +) -> list[tuple[float, str, str]]: + """Return all overlapping root pairs with their minimum XY overlap distance.""" + result = [] + for index, first_id in enumerate(root_ids): + for second_id in root_ids[index + 1 :]: + overlap = np.minimum( + np.asarray(xy_by_id[first_id]) + half_extents[first_id], + np.asarray(xy_by_id[second_id]) + half_extents[second_id], + ) - np.maximum( + np.asarray(xy_by_id[first_id]) - half_extents[first_id], + np.asarray(xy_by_id[second_id]) - half_extents[second_id], + ) + if np.all(overlap > 1e-9): + result.append((float(np.min(overlap)), first_id, second_id)) + return sorted(result, reverse=True) + + +def _aabb_separation_constraint( + *, + root_ids: list[str], + first_id: str, + second_id: str, + half_extents: dict[str, np.ndarray], + xy_by_id: dict[str, list[float]], + margin: float, +) -> tuple[np.ndarray, float]: + first, second = np.asarray(xy_by_id[first_id]), np.asarray(xy_by_id[second_id]) + # Positive overlap on both axes means these two center-based AABBs intersect. + overlap = np.minimum( + first + half_extents[first_id], second + half_extents[second_id] + ) - np.maximum(first - half_extents[first_id], second - half_extents[second_id]) + # Separate along the least-penetrating axis to require the smallest local shift. + axis = int(np.argmin(overlap)) + # Preserve the current order on that axis; object IDs break an exact tie deterministically. + lower = first[axis] < second[axis] or ( + first[axis] == second[axis] and first_id < second_id + ) + index = {root_id: i for i, root_id in enumerate(root_ids)} + # One row addresses the x/y variable pair of each root in the flattened solver vector. + row = np.zeros(2 * len(root_ids)) + sign = 1.0 if lower else -1.0 + row[2 * index[first_id] + axis], row[2 * index[second_id] + axis] = sign, -sign + # row @ values <= bound keeps the selected AABB faces apart by the requested margin. + return row, -float( + half_extents[first_id][axis] + half_extents[second_id][axis] + margin + ) diff --git a/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py b/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py index bca20da8e..9b650469c 100644 --- a/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py +++ b/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py @@ -30,8 +30,9 @@ from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_constructor import ( SceneLayoutConstructor, ) -from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_optimizer import ( - SceneLayoutOptimizer, +from embodichain.gen_sim.scene_engine.pipeline.utils.table_surface_layout_optimizer import ( + TableSurfaceLayoutOptimizer, + TableSurfaceLayoutProblem, _table_region_bounds, ) @@ -152,34 +153,36 @@ def test_table_optimizer_ignores_fixed_sibling_overlap(tmp_path: Path) -> None: asset_glb = tmp_path / "asset.glb" trimesh.creation.box(extents=[_ASSET_SIDE_LENGTH_M] * 3).export(asset_glb) first_id, second_id = "first_001", "second_001" - optimizer = SceneLayoutOptimizer() - - solved_xy_by_id = optimizer.optimize_table_root_xy( - assets_by_id={ - first_id: _asset( - object_id=first_id, - glb_path=asset_glb, - center_xy=_OVERLAPPING_CENTER_XY, - ), - second_id: _asset( - object_id=second_id, - glb_path=asset_glb, - center_xy=_OVERLAPPING_CENTER_XY, - ), - }, - root_ids=[first_id, second_id], - root_seed_xy_by_id={ - first_id: _OVERLAPPING_CENTER_XY, - second_id: _OVERLAPPING_CENTER_XY, - }, - imported_root_ids={first_id, second_id}, - fixed_root_xy_by_id={ - first_id: _OVERLAPPING_CENTER_XY, - second_id: _OVERLAPPING_CENTER_XY, - }, - root_table_regions_by_id={first_id: None, second_id: None}, - table_optimization_rect_xy=_TABLE_BOUNDS, - root_relations=[], + optimizer = TableSurfaceLayoutOptimizer() + + solved_xy_by_id = optimizer.optimize( + TableSurfaceLayoutProblem( + assets_by_id={ + first_id: _asset( + object_id=first_id, + glb_path=asset_glb, + center_xy=_OVERLAPPING_CENTER_XY, + ), + second_id: _asset( + object_id=second_id, + glb_path=asset_glb, + center_xy=_OVERLAPPING_CENTER_XY, + ), + }, + root_ids=[first_id, second_id], + root_seed_xy_by_id={ + first_id: _OVERLAPPING_CENTER_XY, + second_id: _OVERLAPPING_CENTER_XY, + }, + imported_root_ids={first_id, second_id}, + fixed_root_xy_by_id={ + first_id: _OVERLAPPING_CENTER_XY, + second_id: _OVERLAPPING_CENTER_XY, + }, + root_table_regions_by_id={first_id: None, second_id: None}, + table_optimization_rect_xy=_TABLE_BOUNDS, + root_relations=[], + ) ) assert solved_xy_by_id == { @@ -194,30 +197,32 @@ def test_table_optimizer_separates_variable_sibling_from_fixed_sibling( asset_glb = tmp_path / "asset.glb" trimesh.creation.box(extents=[_ASSET_SIDE_LENGTH_M] * 3).export(asset_glb) fixed_id, variable_id = "fixed_001", "variable_001" - optimizer = SceneLayoutOptimizer() - - solved_xy_by_id = optimizer.optimize_table_root_xy( - assets_by_id={ - fixed_id: _asset( - object_id=fixed_id, - glb_path=asset_glb, - center_xy=_OVERLAPPING_CENTER_XY, - ), - variable_id: _asset( - object_id=variable_id, - glb_path=asset_glb, - ), - }, - root_ids=[fixed_id, variable_id], - root_seed_xy_by_id={ - fixed_id: _OVERLAPPING_CENTER_XY, - variable_id: _OVERLAPPING_CENTER_XY, - }, - imported_root_ids={fixed_id}, - fixed_root_xy_by_id={fixed_id: _OVERLAPPING_CENTER_XY, variable_id: None}, - root_table_regions_by_id={fixed_id: None, variable_id: None}, - table_optimization_rect_xy=_TABLE_BOUNDS, - root_relations=[], + optimizer = TableSurfaceLayoutOptimizer() + + solved_xy_by_id = optimizer.optimize( + TableSurfaceLayoutProblem( + assets_by_id={ + fixed_id: _asset( + object_id=fixed_id, + glb_path=asset_glb, + center_xy=_OVERLAPPING_CENTER_XY, + ), + variable_id: _asset( + object_id=variable_id, + glb_path=asset_glb, + ), + }, + root_ids=[fixed_id, variable_id], + root_seed_xy_by_id={ + fixed_id: _OVERLAPPING_CENTER_XY, + variable_id: _OVERLAPPING_CENTER_XY, + }, + imported_root_ids={fixed_id}, + fixed_root_xy_by_id={fixed_id: _OVERLAPPING_CENTER_XY, variable_id: None}, + root_table_regions_by_id={fixed_id: None, variable_id: None}, + table_optimization_rect_xy=_TABLE_BOUNDS, + root_relations=[], + ) ) assert solved_xy_by_id[fixed_id] == _OVERLAPPING_CENTER_XY From 4f6aabb5fdbfdf32a458f1efe94c4201869833e8 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Wed, 19 Aug 2026 15:15:55 +0800 Subject: [PATCH 27/29] finish a simple on-relationship optimizer --- .../utils/parent_surface_layout_optimizer.py | 489 ++++++++++++++++++ .../utils/scene_layout_constructor.py | 3 +- .../test_scene_layout_optimizer.py | 46 +- 3 files changed, 535 insertions(+), 3 deletions(-) create mode 100644 embodichain/gen_sim/scene_engine/pipeline/utils/parent_surface_layout_optimizer.py diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/parent_surface_layout_optimizer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/parent_surface_layout_optimizer.py new file mode 100644 index 000000000..0032f1d66 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/parent_surface_layout_optimizer.py @@ -0,0 +1,489 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF 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 TYPE_CHECKING + +import numpy as np +from scipy.optimize import minimize + +from embodichain.gen_sim.scene_engine.core.scene_graph import SceneGraphRelation +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_utils import ( + load_scene_object_z_up_mesh, + measure_scene_object_z_up_world_aabb, +) + +if TYPE_CHECKING: + from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_constructor import ( + SceneLayoutGroup, + SceneLayoutProblem, + ) + + +@dataclass +class ParentSurfaceLayoutProblem: + """All geometry and layout-state inputs for one parent-surface solve.""" + + assets_by_id: dict[str, SceneObject] + child_ids: list[str] + child_seed_xy_by_id: dict[str, list[float]] + imported_child_ids: set[str] + fixed_child_xy_by_id: dict[str, list[float] | None] + parent_aabb_xy: list[list[float]] + parent_top_z: float + child_relations: list[SceneGraphRelation] + + @classmethod + def from_layout_problem( + cls, + *, + layout_problem: SceneLayoutProblem, + group: SceneLayoutGroup, + current_xy_by_id: dict[str, list[float] | None], + ) -> ParentSurfaceLayoutProblem: + """Build one parent-surface problem without mutating layout state.""" + assets_by_id = { + asset.id: asset for asset in layout_problem.post_edit_scene.assets + } + child_ids = set(group.child_ids) + parent = assets_by_id.get(group.parent_id) + if parent is None: + raise ValueError(f"Parent {group.parent_id!r} is not an asset.") + parent_aabb = measure_scene_object_z_up_world_aabb(scene_object=parent) + parent_aabb_xy = [ + [parent_aabb[0][0], parent_aabb[0][1]], + [parent_aabb[1][0], parent_aabb[1][1]], + ] + parent_center_xy = [ + (parent_aabb[0][0] + parent_aabb[1][0]) / 2.0, + (parent_aabb[0][1] + parent_aabb[1][1]) / 2.0, + ] + child_seed_xy_by_id = {} + for child_id in group.child_ids: + inherited_xy = current_xy_by_id[child_id] + # New children begin from the solved parent's AABB center. + child_seed_xy_by_id[child_id] = ( + parent_center_xy if inherited_xy is None else list(inherited_xy) + ) + return cls( + assets_by_id=assets_by_id, + child_ids=group.child_ids, + child_seed_xy_by_id=child_seed_xy_by_id, + imported_child_ids={ + child_id + for child_id in group.child_ids + if layout_problem.initial_xy_by_id[child_id] is not None + }, + fixed_child_xy_by_id={ + child_id: ( + None + if child_id in layout_problem.layout_variable_ids + else current_xy_by_id[child_id] + ) + for child_id in group.child_ids + }, + parent_aabb_xy=parent_aabb_xy, + parent_top_z=parent_aabb[1][2], + child_relations=[ + relation + for relation in layout_problem.goal_scene_graph.relations + if relation.source_id in child_ids and relation.target_id in child_ids + ], + ) + + +@dataclass(frozen=True) +class ParentSurfaceLayoutOptimizerConfig: + """Numerical controls for one non-table parent-surface sibling solve.""" + + relation_clearance_m: float = 0.03 + collision_margin_m: float = 0.02 + max_slsqp_iterations: int = 500 + slsqp_ftol: float = 1e-6 + max_collision_rounds: int = 8 + max_added_collision_pairs: int = 64 + imported_seed_weight: float = 5.0 + min_center_distance_m: float = 0.01 + min_center_distance_weight: float = 0.05 + + def __post_init__(self) -> None: + """Reject invalid controls before assembling parent-surface constraints.""" + if self.relation_clearance_m < 0.0: + raise ValueError("relation_clearance_m must be non-negative.") + if self.collision_margin_m < 0.0: + raise ValueError("collision_margin_m must be non-negative.") + if self.max_slsqp_iterations <= 0: + raise ValueError("max_slsqp_iterations must be positive.") + if self.slsqp_ftol <= 0.0: + raise ValueError("slsqp_ftol must be positive.") + if self.max_collision_rounds <= 0: + raise ValueError("max_collision_rounds must be positive.") + if self.max_added_collision_pairs <= 0: + raise ValueError("max_added_collision_pairs must be positive.") + + +class ParentSurfaceLayoutOptimizer: + """Solve direct ``on`` children inside one parent's current XY footprint.""" + + def __init__( + self, + *, + config: ParentSurfaceLayoutOptimizerConfig | None = None, + ) -> None: + self.config = ( + config if config is not None else ParentSurfaceLayoutOptimizerConfig() + ) + + def optimize( + self, + problem: ParentSurfaceLayoutProblem, + ) -> dict[str, list[float]]: + """Return sibling XY centers inside the parent AABB without overlap.""" + child_half_extents_xy = _asset_half_extents_xy( + assets_by_id=problem.assets_by_id, + object_ids=problem.child_ids, + ) + inequality_constraints, equality_constraints = _build_constraints( + problem=problem, + child_half_extents_xy=child_half_extents_xy, + config=self.config, + ) + solved_child_xy_by_id = _solve_root_xy( + root_ids=problem.child_ids, + root_seed_xy_by_id=problem.child_seed_xy_by_id, + imported_root_ids=problem.imported_child_ids, + inequality_constraints=inequality_constraints, + equality_constraints=equality_constraints, + config=self.config, + ) + return _refine_root_collisions( + root_ids=problem.child_ids, + root_seed_xy_by_id=problem.child_seed_xy_by_id, + imported_root_ids=problem.imported_child_ids, + root_half_extents_xy=child_half_extents_xy, + inequality_constraints=inequality_constraints, + equality_constraints=equality_constraints, + fixed_root_xy_by_id=problem.fixed_child_xy_by_id, + solved_root_xy_by_id=solved_child_xy_by_id, + config=self.config, + ) + + +def _build_constraints( + *, + problem: ParentSurfaceLayoutProblem, + child_half_extents_xy: dict[str, np.ndarray], + config: ParentSurfaceLayoutOptimizerConfig, +) -> tuple[list[tuple[np.ndarray, float]], list[tuple[np.ndarray, float]]]: + """Build hard parent-AABB, planar-relation, and fixed-child constraints.""" + root_index = {child_id: index for index, child_id in enumerate(problem.child_ids)} + parent_bounds = _bounds_from_points(problem.parent_aabb_xy) + inequality_constraints: list[tuple[np.ndarray, float]] = [] + equality_constraints: list[tuple[np.ndarray, float]] = [] + for child_id in problem.child_ids: + # Keep each child's 2D AABB inside the parent AABB support proxy. + _append_aabb_center_bounds( + constraints=inequality_constraints, + root_index=root_index, + root_id=child_id, + bounds=parent_bounds, + half_extents_xy=child_half_extents_xy[child_id], + ) + fixed_xy = problem.fixed_child_xy_by_id[child_id] + if fixed_xy is not None: + _append_fixed_root_constraints( + constraints=equality_constraints, + root_index=root_index, + root_id=child_id, + fixed_xy=fixed_xy, + ) + for relation in problem.child_relations: + # Apply planar relations between direct on-children of this parent. + _append_planar_relation_constraint( + constraints=inequality_constraints, + root_index=root_index, + source_id=relation.source_id, + relation=relation.relation, + target_id=relation.target_id, + source_half_extents_xy=child_half_extents_xy[relation.source_id], + target_half_extents_xy=child_half_extents_xy[relation.target_id], + relation_clearance_m=config.relation_clearance_m, + ) + return inequality_constraints, equality_constraints + + +def _asset_half_extents_xy( + *, assets_by_id: dict[str, SceneObject], object_ids: list[str] +) -> dict[str, np.ndarray]: + """Measure each optimized child asset's z-up XY half-extents.""" + result = {} + for object_id in object_ids: + asset = assets_by_id.get(object_id) + if asset is None: + raise ValueError(f"Parent child {object_id!r} is not an asset.") + mesh = load_scene_object_z_up_mesh(scene_object=asset) + result[object_id] = (mesh.bounds[1, :2] - mesh.bounds[0, :2]) / 2.0 + return result + + +def _bounds_from_points(points: list[list[float]]) -> np.ndarray: + """Return finite XY minimum and maximum bounds from polygon points.""" + coordinates = np.asarray(points, dtype=float) + if ( + coordinates.ndim != 2 + or coordinates.shape[1] != 2 + or len(coordinates) < 2 + or not np.all(np.isfinite(coordinates)) + ): + raise ValueError("XY bounds must contain at least two finite points.") + return np.stack([coordinates.min(axis=0), coordinates.max(axis=0)]) + + +def _append_aabb_center_bounds( + *, + constraints: list[tuple[np.ndarray, float]], + root_index: dict[str, int], + root_id: str, + bounds: np.ndarray, + half_extents_xy: np.ndarray, +) -> None: + """Constrain one child AABB center to lie completely inside XY bounds.""" + minimum, maximum = bounds[0] + half_extents_xy, bounds[1] - half_extents_xy + if np.any(minimum > maximum): + raise ValueError(f"Asset {root_id!r} cannot fit inside its parent AABB.") + # root_id is the child whose center is constrained in this AABB bound. + offset, count = 2 * root_index[root_id], 2 * len(root_index) + # offset selects this child's XY pair; count is the full flattened XY vector size. + for axis in range(2): + upper, lower = np.zeros(count), np.zeros(count) + upper[offset + axis], lower[offset + axis] = 1.0, -1.0 + constraints.extend( + [(upper, float(maximum[axis])), (lower, -float(minimum[axis]))] + ) + + +def _append_fixed_root_constraints( + *, + constraints: list[tuple[np.ndarray, float]], + root_index: dict[str, int], + root_id: str, + fixed_xy: list[float], +) -> None: + """Lock one fixed child's center to its imported XY coordinates.""" + offset, count = 2 * root_index[root_id], 2 * len(root_index) + for axis, coordinate in enumerate(fixed_xy): + row = np.zeros(count) + row[offset + axis] = 1.0 + constraints.append((row, float(coordinate))) + + +def _append_planar_relation_constraint( + *, + constraints: list[tuple[np.ndarray, float]], + root_index: dict[str, int], + source_id: str, + relation: str, + target_id: str, + source_half_extents_xy: np.ndarray, + target_half_extents_xy: np.ndarray, + relation_clearance_m: float, +) -> None: + """Append one world-XY separation constraint for sibling planar semantics.""" + axis, sign = { + "left_of": (0, 1.0), + "right_of": (0, -1.0), + "behind": (1, 1.0), + "in_front_of": (1, -1.0), + }.get(relation, (None, None)) + if axis is None or source_id not in root_index or target_id not in root_index: + raise ValueError(f"Unsupported parent-child planar relation {relation!r}.") + row = np.zeros(2 * len(root_index)) + row[2 * root_index[source_id] + axis] = sign + row[2 * root_index[target_id] + axis] = -sign + constraints.append( + ( + row, + -float( + source_half_extents_xy[axis] + + target_half_extents_xy[axis] + + relation_clearance_m + ), + ) + ) + + +def _solve_root_xy( + *, + root_ids: list[str], + root_seed_xy_by_id: dict[str, list[float]], + imported_root_ids: set[str], + inequality_constraints: list[tuple[np.ndarray, float]], + equality_constraints: list[tuple[np.ndarray, float]], + config: ParentSurfaceLayoutOptimizerConfig, +) -> dict[str, list[float]]: + """Solve one parent child-group's XY positions with SLSQP.""" + initial = np.asarray( + [root_seed_xy_by_id[root_id] for root_id in root_ids], dtype=float + ) + + def objective(values: np.ndarray) -> float: + xy = values.reshape(-1, 2) + loss = 0.0 + for index, root_id in enumerate(root_ids): + if root_id in imported_root_ids: + delta = xy[index] - initial[index] + loss += config.imported_seed_weight * float(delta @ delta) + return loss + + constraints = [ + { + "type": "ineq", + "fun": lambda values, row=row, bound=bound: bound - float(row @ values), + } + for row, bound in inequality_constraints + ] + [ + { + "type": "eq", + "fun": lambda values, row=row, bound=bound: float(row @ values) - bound, + } + for row, bound in equality_constraints + ] + result = minimize( + objective, + initial.reshape(-1), + method="SLSQP", + constraints=constraints, + options={ + "maxiter": config.max_slsqp_iterations, + "ftol": config.slsqp_ftol, + "disp": False, + }, + ) + if not result.success: + raise ValueError(f"Parent layout optimization failed: {result.message}") + return { + root_id: [float(result.x[2 * index]), float(result.x[2 * index + 1])] + for index, root_id in enumerate(root_ids) + } + + +def _refine_root_collisions( + *, + root_ids: list[str], + root_seed_xy_by_id: dict[str, list[float]], + imported_root_ids: set[str], + root_half_extents_xy: dict[str, np.ndarray], + inequality_constraints: list[tuple[np.ndarray, float]], + equality_constraints: list[tuple[np.ndarray, float]], + fixed_root_xy_by_id: dict[str, list[float] | None], + solved_root_xy_by_id: dict[str, list[float]], + config: ParentSurfaceLayoutOptimizerConfig, +) -> dict[str, list[float]]: + """Iteratively add separation constraints for overlapping child AABBs.""" + current = solved_root_xy_by_id + seen: set[tuple[str, str]] = set() + for _ in range(config.max_collision_rounds): + overlaps = [ + pair + for pair in _root_aabb_overlaps( + root_ids=root_ids, half_extents=root_half_extents_xy, xy_by_id=current + ) + if fixed_root_xy_by_id[pair[1]] is None + or fixed_root_xy_by_id[pair[2]] is None + ] + if not overlaps: + return current + added = 0 + for _, first_id, second_id in overlaps[: config.max_added_collision_pairs]: + key = tuple(sorted((first_id, second_id))) + if key in seen: + continue + inequality_constraints.append( + _aabb_separation_constraint( + root_ids=root_ids, + first_id=first_id, + second_id=second_id, + half_extents=root_half_extents_xy, + xy_by_id=current, + margin=config.collision_margin_m, + ) + ) + seen.add(key) + added += 1 + if not added: + break + current = _solve_root_xy( + root_ids=root_ids, + root_seed_xy_by_id=current, + imported_root_ids=imported_root_ids, + inequality_constraints=inequality_constraints, + equality_constraints=equality_constraints, + config=config, + ) + raise ValueError("Parent-child AABB collisions remain after layout refinement.") + + +def _root_aabb_overlaps( + *, + root_ids: list[str], + half_extents: dict[str, np.ndarray], + xy_by_id: dict[str, list[float]], +) -> list[tuple[float, str, str]]: + """Return overlapping child pairs with their minimum XY overlap distance.""" + result = [] + for index, first_id in enumerate(root_ids): + for second_id in root_ids[index + 1 :]: + overlap = np.minimum( + np.asarray(xy_by_id[first_id]) + half_extents[first_id], + np.asarray(xy_by_id[second_id]) + half_extents[second_id], + ) - np.maximum( + np.asarray(xy_by_id[first_id]) - half_extents[first_id], + np.asarray(xy_by_id[second_id]) - half_extents[second_id], + ) + if np.all(overlap > 1e-9): + result.append((float(np.min(overlap)), first_id, second_id)) + return sorted(result, reverse=True) + + +def _aabb_separation_constraint( + *, + root_ids: list[str], + first_id: str, + second_id: str, + half_extents: dict[str, np.ndarray], + xy_by_id: dict[str, list[float]], + margin: float, +) -> tuple[np.ndarray, float]: + """Return one least-penetration AABB separation inequality.""" + first, second = np.asarray(xy_by_id[first_id]), np.asarray(xy_by_id[second_id]) + overlap = np.minimum( + first + half_extents[first_id], second + half_extents[second_id] + ) - np.maximum(first - half_extents[first_id], second - half_extents[second_id]) + axis = int(np.argmin(overlap)) + lower = first[axis] < second[axis] or ( + first[axis] == second[axis] and first_id < second_id + ) + index = {root_id: i for i, root_id in enumerate(root_ids)} + row = np.zeros(2 * len(root_ids)) + sign = 1.0 if lower else -1.0 + row[2 * index[first_id] + axis], row[2 * index[second_id] + axis] = sign, -sign + return row, -float( + half_extents[first_id][axis] + half_extents[second_id][axis] + margin + ) diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_constructor.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_constructor.py index 9279bd1ae..95c77d712 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_constructor.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_constructor.py @@ -170,7 +170,7 @@ def _optimize_table_group( scene_object=assets_by_id[root_id], support_region_z=table.support_surface_z, center_xy=solved_xy, - clearance_m=0.00, # Directly place on the support surface. + clearance_m=0.00, # Directly place on the support surface. ) self._updated_object_ids.add(root_id) self._propagate_descendant_delta( @@ -242,6 +242,7 @@ def _optimize_parent_group( scene_object=parent_surface_problem.assets_by_id[child_id], support_region_z=parent_surface_problem.parent_top_z, center_xy=solved_xy, + clearance_m=0.00, # Directly place on the parent's top surface. ) self._updated_object_ids.add(child_id) self._propagate_descendant_delta( diff --git a/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py b/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py index 9b650469c..39a44ecc7 100644 --- a/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py +++ b/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py @@ -25,8 +25,13 @@ from embodichain.gen_sim.scene_engine.core.scene_graph import ( SceneGraph, SceneGraphNode, + SceneGraphRelation, ) from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.pipeline.utils.parent_surface_layout_optimizer import ( + ParentSurfaceLayoutOptimizer, + ParentSurfaceLayoutProblem, +) from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_constructor import ( SceneLayoutConstructor, ) @@ -44,6 +49,7 @@ ] _OVERLAPPING_CENTER_XY = [0.0, 0.0] _ASSET_SIDE_LENGTH_M = 0.2 +_RELATION_CLEARANCE_M = 0.03 def _asset( @@ -145,8 +151,8 @@ def test_layout_constructor_places_new_child_on_parent_top( asset for asset in post_edit_scene.assets if asset.id == "cup_001" ) assert placed_cup.center_xy == [0.0, 0.0] - # book top is z=0.62 m; cup half-height is 0.1 m and clearance is 0.02 m. - assert np.allclose(placed_cup.pos, [0.0, 0.74, 0.0]) + # Book top is z=0.62 m; cup half-height is 0.1 m with zero support clearance. + assert np.allclose(placed_cup.pos, [0.0, 0.72, 0.0]) def test_table_optimizer_ignores_fixed_sibling_overlap(tmp_path: Path) -> None: @@ -227,3 +233,39 @@ def test_table_optimizer_separates_variable_sibling_from_fixed_sibling( assert solved_xy_by_id[fixed_id] == _OVERLAPPING_CENTER_XY assert np.max(np.abs(solved_xy_by_id[variable_id])) >= _ASSET_SIDE_LENGTH_M - 1e-6 + + +def test_parent_optimizer_applies_sibling_planar_relation(tmp_path: Path) -> None: + asset_glb = tmp_path / "asset.glb" + trimesh.creation.box(extents=[_ASSET_SIDE_LENGTH_M] * 3).export(asset_glb) + left_id, right_id = "left_001", "right_001" + + solved_xy_by_id = ParentSurfaceLayoutOptimizer().optimize( + ParentSurfaceLayoutProblem( + assets_by_id={ + left_id: _asset(object_id=left_id, glb_path=asset_glb), + right_id: _asset(object_id=right_id, glb_path=asset_glb), + }, + child_ids=[left_id, right_id], + child_seed_xy_by_id={ + left_id: _OVERLAPPING_CENTER_XY, + right_id: _OVERLAPPING_CENTER_XY, + }, + imported_child_ids=set(), + fixed_child_xy_by_id={left_id: None, right_id: None}, + parent_aabb_xy=_TABLE_BOUNDS, + parent_top_z=0.0, + child_relations=[ + SceneGraphRelation( + source_id=left_id, + relation="left_of", + target_id=right_id, + ) + ], + ) + ) + + assert ( + solved_xy_by_id[right_id][0] - solved_xy_by_id[left_id][0] + >= _ASSET_SIDE_LENGTH_M + _RELATION_CLEARANCE_M - 1e-6 + ) From e78ba19b8d6c6b231f67c5a4888980b496a13cd5 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Wed, 19 Aug 2026 17:11:04 +0800 Subject: [PATCH 28/29] add the heuristic attaches in collision optimization --- .../utils/parent_surface_layout_optimizer.py | 107 +++++++++++++---- .../utils/table_surface_layout_optimizer.py | 110 +++++++++++++----- .../test_scene_layout_optimizer.py | 71 +++++++++++ 3 files changed, 236 insertions(+), 52 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/parent_surface_layout_optimizer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/parent_surface_layout_optimizer.py index 0032f1d66..6d2a3b0a5 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/parent_surface_layout_optimizer.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/parent_surface_layout_optimizer.py @@ -186,6 +186,10 @@ def optimize( ) +class _LayoutInfeasibleError(ValueError): + """Internal marker for an SLSQP failure while testing one collision direction.""" + + def _build_constraints( *, problem: ParentSurfaceLayoutProblem, @@ -377,7 +381,9 @@ def objective(values: np.ndarray) -> float: }, ) if not result.success: - raise ValueError(f"Parent layout optimization failed: {result.message}") + raise _LayoutInfeasibleError( + f"Parent layout optimization failed: {result.message}" + ) return { root_id: [float(result.x[2 * index]), float(result.x[2 * index + 1])] for index, root_id in enumerate(root_ids) @@ -415,28 +421,51 @@ def _refine_root_collisions( key = tuple(sorted((first_id, second_id))) if key in seen: continue - inequality_constraints.append( - _aabb_separation_constraint( + # Earlier pair updates may already have separated this stale overlap. + if key not in { + tuple(sorted((first, second))) + for _, first, second in _root_aabb_overlaps( root_ids=root_ids, - first_id=first_id, - second_id=second_id, half_extents=root_half_extents_xy, xy_by_id=current, - margin=config.collision_margin_m, ) - ) - seen.add(key) - added += 1 + }: + continue + for separation_constraint in _aabb_separation_constraints( + root_ids=root_ids, + first_id=first_id, + second_id=second_id, + half_extents=root_half_extents_xy, + xy_by_id=current, + margin=config.collision_margin_m, + ): + # Keep a candidate only when it is compatible with all hard constraints. + try: + solved_xy_by_id = _solve_root_xy( + root_ids=root_ids, + root_seed_xy_by_id=current, + imported_root_ids=imported_root_ids, + inequality_constraints=[ + *inequality_constraints, + separation_constraint, + ], + equality_constraints=equality_constraints, + config=config, + ) + except _LayoutInfeasibleError: + continue + inequality_constraints.append(separation_constraint) + current = solved_xy_by_id + seen.add(key) + added += 1 + break + else: + raise ValueError( + "Parent-child AABB pair has no feasible separation direction: " + f"{first_id!r}, {second_id!r}." + ) if not added: break - current = _solve_root_xy( - root_ids=root_ids, - root_seed_xy_by_id=current, - imported_root_ids=imported_root_ids, - inequality_constraints=inequality_constraints, - equality_constraints=equality_constraints, - config=config, - ) raise ValueError("Parent-child AABB collisions remain after layout refinement.") @@ -462,7 +491,7 @@ def _root_aabb_overlaps( return sorted(result, reverse=True) -def _aabb_separation_constraint( +def _aabb_separation_constraints( *, root_ids: list[str], first_id: str, @@ -470,19 +499,47 @@ def _aabb_separation_constraint( half_extents: dict[str, np.ndarray], xy_by_id: dict[str, list[float]], margin: float, -) -> tuple[np.ndarray, float]: - """Return one least-penetration AABB separation inequality.""" +) -> list[tuple[np.ndarray, float]]: + """Return ordered feasible-direction candidates for one overlapping AABB pair.""" first, second = np.asarray(xy_by_id[first_id]), np.asarray(xy_by_id[second_id]) overlap = np.minimum( first + half_extents[first_id], second + half_extents[second_id] ) - np.maximum(first - half_extents[first_id], second - half_extents[second_id]) - axis = int(np.argmin(overlap)) - lower = first[axis] < second[axis] or ( - first[axis] == second[axis] and first_id < second_id - ) + axes = np.argsort(overlap) + constraints = [] + for axis in axes: + current_order = first[axis] < second[axis] or ( + first[axis] == second[axis] and first_id < second_id + ) + for first_is_lower in (current_order, not current_order): + constraints.append( + _aabb_separation_constraint_for_direction( + root_ids=root_ids, + first_id=first_id, + second_id=second_id, + half_extents=half_extents, + axis=int(axis), + first_is_lower=first_is_lower, + margin=margin, + ) + ) + return constraints + + +def _aabb_separation_constraint_for_direction( + *, + root_ids: list[str], + first_id: str, + second_id: str, + half_extents: dict[str, np.ndarray], + axis: int, + first_is_lower: bool, + margin: float, +) -> tuple[np.ndarray, float]: + """Return one directed AABB separation inequality on a selected axis.""" index = {root_id: i for i, root_id in enumerate(root_ids)} row = np.zeros(2 * len(root_ids)) - sign = 1.0 if lower else -1.0 + sign = 1.0 if first_is_lower else -1.0 row[2 * index[first_id] + axis], row[2 * index[second_id] + axis] = sign, -sign return row, -float( half_extents[first_id][axis] + half_extents[second_id][axis] + margin diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/table_surface_layout_optimizer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/table_surface_layout_optimizer.py index 38fe7e6c4..9d874f239 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/table_surface_layout_optimizer.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/table_surface_layout_optimizer.py @@ -186,6 +186,10 @@ def optimize( ) +class _LayoutInfeasibleError(ValueError): + """Internal marker for an SLSQP failure while testing one collision direction.""" + + def _build_constraints( *, problem: TableSurfaceLayoutProblem, @@ -424,7 +428,9 @@ def objective(values: np.ndarray) -> float: }, ) if not result.success: - raise ValueError(f"Table layout optimization failed: {result.message}") + raise _LayoutInfeasibleError( + f"Table layout optimization failed: {result.message}" + ) return { root_id: [float(result.x[2 * index]), float(result.x[2 * index + 1])] for index, root_id in enumerate(root_ids) @@ -463,29 +469,51 @@ def _refine_root_collisions( key = tuple(sorted((first_id, second_id))) if key in seen: continue - # Add a new SLSQP constraint to separate this overlapping pair. - inequality_constraints.append( - _aabb_separation_constraint( + # Earlier pair updates may already have separated this stale overlap. + if key not in { + tuple(sorted((first, second))) + for _, first, second in _root_aabb_overlaps( root_ids=root_ids, - first_id=first_id, - second_id=second_id, half_extents=root_half_extents_xy, xy_by_id=current, - margin=config.collision_margin_m, ) - ) - seen.add(key) - added += 1 + }: + continue + for separation_constraint in _aabb_separation_constraints( + root_ids=root_ids, + first_id=first_id, + second_id=second_id, + half_extents=root_half_extents_xy, + xy_by_id=current, + margin=config.collision_margin_m, + ): + # Keep a candidate only when it is compatible with all hard constraints. + try: + solved_xy_by_id = _solve_root_xy( + root_ids=root_ids, + root_seed_xy_by_id=current, + imported_root_ids=imported_root_ids, + inequality_constraints=[ + *inequality_constraints, + separation_constraint, + ], + equality_constraints=equality_constraints, + config=config, + ) + except _LayoutInfeasibleError: + continue + inequality_constraints.append(separation_constraint) + current = solved_xy_by_id + seen.add(key) + added += 1 + break + else: + raise ValueError( + "Table-root AABB pair has no feasible separation direction: " + f"{first_id!r}, {second_id!r}." + ) if not added: break - current = _solve_root_xy( - root_ids=root_ids, - root_seed_xy_by_id=current, - imported_root_ids=imported_root_ids, - inequality_constraints=inequality_constraints, - equality_constraints=equality_constraints, - config=config, - ) raise ValueError("Table-root AABB collisions remain after layout refinement.") @@ -511,7 +539,7 @@ def _root_aabb_overlaps( return sorted(result, reverse=True) -def _aabb_separation_constraint( +def _aabb_separation_constraints( *, root_ids: list[str], first_id: str, @@ -519,22 +547,50 @@ def _aabb_separation_constraint( half_extents: dict[str, np.ndarray], xy_by_id: dict[str, list[float]], margin: float, -) -> tuple[np.ndarray, float]: +) -> list[tuple[np.ndarray, float]]: + """Return ordered feasible-direction candidates for one overlapping AABB pair.""" first, second = np.asarray(xy_by_id[first_id]), np.asarray(xy_by_id[second_id]) # Positive overlap on both axes means these two center-based AABBs intersect. overlap = np.minimum( first + half_extents[first_id], second + half_extents[second_id] ) - np.maximum(first - half_extents[first_id], second - half_extents[second_id]) - # Separate along the least-penetrating axis to require the smallest local shift. - axis = int(np.argmin(overlap)) - # Preserve the current order on that axis; object IDs break an exact tie deterministically. - lower = first[axis] < second[axis] or ( - first[axis] == second[axis] and first_id < second_id - ) + # Try the least-penetrating axis first, but permit order reversal if required. + axes = np.argsort(overlap) + constraints = [] + for axis in axes: + current_order = first[axis] < second[axis] or ( + first[axis] == second[axis] and first_id < second_id + ) + for first_is_lower in (current_order, not current_order): + constraints.append( + _aabb_separation_constraint_for_direction( + root_ids=root_ids, + first_id=first_id, + second_id=second_id, + half_extents=half_extents, + axis=int(axis), + first_is_lower=first_is_lower, + margin=margin, + ) + ) + return constraints + + +def _aabb_separation_constraint_for_direction( + *, + root_ids: list[str], + first_id: str, + second_id: str, + half_extents: dict[str, np.ndarray], + axis: int, + first_is_lower: bool, + margin: float, +) -> tuple[np.ndarray, float]: + """Return one directed AABB separation inequality on a selected axis.""" index = {root_id: i for i, root_id in enumerate(root_ids)} # One row addresses the x/y variable pair of each root in the flattened solver vector. row = np.zeros(2 * len(root_ids)) - sign = 1.0 if lower else -1.0 + sign = 1.0 if first_is_lower else -1.0 row[2 * index[first_id] + axis], row[2 * index[second_id] + axis] = sign, -sign # row @ values <= bound keeps the selected AABB faces apart by the requested margin. return row, -float( diff --git a/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py b/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py index 39a44ecc7..6698d5c83 100644 --- a/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py +++ b/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py @@ -50,6 +50,10 @@ _OVERLAPPING_CENTER_XY = [0.0, 0.0] _ASSET_SIDE_LENGTH_M = 0.2 _RELATION_CLEARANCE_M = 0.03 +_COLLISION_MARGIN_M = 0.02 +_BOARD_XY_SIZE_M = 0.6 +_CAN_XY_SIZE_M = 0.1 +_PENCIL_XY_SIZE_M = [0.04, 0.2] def _asset( @@ -235,6 +239,73 @@ def test_table_optimizer_separates_variable_sibling_from_fixed_sibling( assert np.max(np.abs(solved_xy_by_id[variable_id])) >= _ASSET_SIDE_LENGTH_M - 1e-6 +def test_table_optimizer_can_reverse_a_collision_order(tmp_path: Path) -> None: + board_glb = tmp_path / "board.glb" + can_glb = tmp_path / "can.glb" + pencil_glb = tmp_path / "pencil.glb" + # SimReady GLBs are y-up, so z-up XY uses the source XZ extents. + trimesh.creation.box( + extents=[_BOARD_XY_SIZE_M, _ASSET_SIDE_LENGTH_M, _BOARD_XY_SIZE_M] + ).export(board_glb) + trimesh.creation.box( + extents=[_CAN_XY_SIZE_M, _ASSET_SIDE_LENGTH_M, _CAN_XY_SIZE_M] + ).export(can_glb) + trimesh.creation.box( + extents=[ + _PENCIL_XY_SIZE_M[0], + _ASSET_SIDE_LENGTH_M, + _PENCIL_XY_SIZE_M[1], + ] + ).export(pencil_glb) + board_id, pencil_id, can_id = "board_001", "pencil_001", "can_001" + board_xy, pencil_xy, can_xy = [0.0, 0.0], [-0.2, 0.0], [-0.4, 0.0] + + solved_xy_by_id = TableSurfaceLayoutOptimizer().optimize( + TableSurfaceLayoutProblem( + assets_by_id={ + board_id: _asset( + object_id=board_id, + glb_path=board_glb, + center_xy=board_xy, + ), + pencil_id: _asset(object_id=pencil_id, glb_path=pencil_glb), + can_id: _asset( + object_id=can_id, + glb_path=can_glb, + center_xy=can_xy, + ), + }, + root_ids=[board_id, pencil_id, can_id], + root_seed_xy_by_id={ + board_id: board_xy, + pencil_id: pencil_xy, + can_id: can_xy, + }, + imported_root_ids={board_id, can_id}, + fixed_root_xy_by_id={ + board_id: board_xy, + pencil_id: None, + can_id: can_xy, + }, + root_table_regions_by_id={ + board_id: None, + pencil_id: None, + can_id: None, + }, + table_optimization_rect_xy=_TABLE_BOUNDS, + root_relations=[], + ) + ) + + expected_pencil_x_upper_bound = ( + can_xy[0] + - _CAN_XY_SIZE_M / 2.0 + - _PENCIL_XY_SIZE_M[0] / 2.0 + - _COLLISION_MARGIN_M + ) + assert solved_xy_by_id[pencil_id][0] <= expected_pencil_x_upper_bound + 1e-6 + + def test_parent_optimizer_applies_sibling_planar_relation(tmp_path: Path) -> None: asset_glb = tmp_path / "asset.glb" trimesh.creation.box(extents=[_ASSET_SIDE_LENGTH_M] * 3).export(asset_glb) From d55581c60eb4e8e3ddb6c1df3c675f97985aae6b Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Wed, 19 Aug 2026 18:13:15 +0800 Subject: [PATCH 29/29] replace the assets gravity settler with gravity settler --- .../pipeline/generation/scene_generation.py | 40 +- .../pipeline/utils/assets_gravity_settler.py | 349 ----------------- .../pipeline/utils/gravity_settler.py | 356 ++++++++++++++++++ .../scene_engine/test_gravity_settler.py | 81 ++++ 4 files changed, 465 insertions(+), 361 deletions(-) delete mode 100644 embodichain/gen_sim/scene_engine/pipeline/utils/assets_gravity_settler.py create mode 100644 embodichain/gen_sim/scene_engine/pipeline/utils/gravity_settler.py create mode 100644 tests/gen_sim/scene_engine/test_gravity_settler.py diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py index 0e1b17f95..7010e5f14 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py @@ -44,8 +44,9 @@ 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.gravity_settler import ( + GravitySettleBody, + GravitySettler, ) from embodichain.gen_sim.scene_engine.pipeline.utils.scene_generation_utils import ( layout_object_to_transform_matrix, @@ -455,16 +456,31 @@ def _layout_refinement( refined_assets_layout = overlap_optimizer.optimize() overlap_optimizer.save_overlap_optimization_debug_images() - # 8. Gravity simulation, to let all the assets to be stable and placed well on the table's support surface. - # Notice that: we do not consider the assets like a bottle, which should be standing on the table but laid down - # after the simulation. - gravity_settler = AssetsGravitySettler( - scene=scene, - table_layout=refined_table_layout, - assets_layout=refined_assets_layout, - geometry_root=simready_geometry_output_root, - ) - refined_assets_layout = gravity_settler.settle() + # 8. The initial image graph has one on-table level, so every asset settles + # dynamically against the table in this first generic gravity pass. + assets_by_id = {asset.id: asset for asset in scene.assets} + # All the assets are dynamic; the table is static. + settled_pose_by_id = GravitySettler( + table_body=GravitySettleBody( + scene_object=scene.table, + y_up_layout=refined_table_layout, + ), + participant_bodies=[ + GravitySettleBody( + scene_object=assets_by_id[str(asset_layout["id"])], + y_up_layout=asset_layout, + ) + for asset_layout in refined_assets_layout + ], + dynamic_asset_ids=set(assets_by_id), + static_asset_ids=set(), + ).settle() + # Update. + for asset_layout in refined_assets_layout: + asset_id = str(asset_layout["id"]) + settled_pose = settled_pose_by_id[asset_id] + asset_layout["pos"] = settled_pose["pos"] + asset_layout["rot"] = settled_pose["rot"] # Update the scene data structure with the final layout and spatial metadata. _update_scene_final_y_up_layout_and_z_up_centers( 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 deleted file mode 100644 index 31e9e8443..000000000 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/assets_gravity_settler.py +++ /dev/null @@ -1,349 +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 -from pathlib import Path -from typing import Sequence - -import numpy as np -from scipy.spatial.transform import Rotation -import trimesh - -from embodichain.gen_sim.scene_engine.core.scene import Scene -from embodichain.gen_sim.scene_engine.core.scene_object import ( - ObjectPhysics, - SceneObject, -) -from embodichain.gen_sim.scene_engine.pipeline.utils.scene_generation_utils import ( - layout_object_to_transform_matrix, - load_glb_mesh, - transform_matrix_to_layout_object, -) -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg -from embodichain.lab.sim.shapes import MeshCfg -from embodichain.utils.logger import log_info - - -@dataclass(frozen=True) -class AssetsGravitySettlerConfig: - """Physics controls for table-top asset settling.""" - - clearance_m: float = 0.02 # Initial gap between each asset and the table top. - settle_steps: int = 300 # Fixed number of simulator steps to execute. - physics_dt: float = 1.0 / 100.0 # Physics timestep in seconds. - sim_device: str = "cpu" # Simulation device requested from EmbodiChain Lab. - - -class AssetsGravitySettler: - """Settle all assets together on one kinematic table in a z-up simulation.""" - - def __init__( - self, - *, - scene: Scene, - table_layout: dict[str, object], - assets_layout: list[dict[str, object]], - geometry_root: str | Path, - config: AssetsGravitySettlerConfig | None = None, - ) -> None: - self.scene = scene - self.table_layout = table_layout - self.assets_layout = assets_layout - self.geometry_root = Path(geometry_root).expanduser().resolve() - self.settled_assets_layout: list[dict[str, object]] | None = None - self.config = config if config is not None else AssetsGravitySettlerConfig() - # Check. - if self.config.clearance_m < 0.0: - raise ValueError("Gravity-settle clearance_m must be non-negative.") - if self.config.settle_steps <= 0: - raise ValueError("Gravity-settle settle_steps must be positive.") - if self.config.physics_dt <= 0.0: - raise ValueError("Gravity-settle physics_dt must be positive.") - - def settle(self) -> list[dict[str, object]]: - """Run gravity settling and return the resulting y-up asset layouts.""" - self.settled_assets_layout = None - if not self.assets_layout: - self.settled_assets_layout = [] - log_info("Scene has no movable assets; skipping gravity settling.") - return self.settled_assets_layout - - table_id = self._require_layout_id(self.table_layout, name="Table") - table_object = self._require_scene_object(table_id, kind="table") - asset_ids: set[str] = set() - asset_objects_by_id: dict[str, SceneObject] = {} - for asset_layout in self.assets_layout: - asset_id = self._require_layout_id(asset_layout, name="Asset") - if asset_id in asset_ids: - raise ValueError(f"Asset layouts contain duplicate id {asset_id!r}.") - asset_ids.add(asset_id) - asset_objects_by_id[asset_id] = self._require_scene_object( - asset_id, kind="asset" - ) - expected_asset_ids = {asset.id for asset in self.scene.assets} - if asset_ids != expected_asset_ids: - raise ValueError( - "Gravity-settle layouts must contain exactly the scene asset ids." - ) - - y_up_to_z_up_matrix = np.eye(4) - y_up_to_z_up_matrix[:3, :3] = np.array( - [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]] - ) - z_up_to_y_up_matrix = np.linalg.inv(y_up_to_z_up_matrix) - table_info = self._prepare_sim_body( - layout_object=self.table_layout, - y_up_to_z_up_matrix=y_up_to_z_up_matrix, - ) - table_world_mesh = self._mesh_to_z_up_world_for_aabb( - y_up_mesh=table_info["mesh"], - z_up_rigid_layout=table_info["rigid_layout"], - z_up_scale=table_info["z_up_scale"], - y_up_to_z_up_matrix=y_up_to_z_up_matrix, - ) - table_top_z = float(table_world_mesh.bounds[1, 2]) - - prepared_assets: dict[str, dict[str, object]] = {} - for asset_layout in self.assets_layout: - asset_id = str(asset_layout["id"]) - asset_info = self._prepare_sim_body( - layout_object=asset_layout, - y_up_to_z_up_matrix=y_up_to_z_up_matrix, - ) - asset_world_mesh = self._mesh_to_z_up_world_for_aabb( - y_up_mesh=asset_info["mesh"], - z_up_rigid_layout=asset_info["rigid_layout"], - z_up_scale=asset_info["z_up_scale"], - y_up_to_z_up_matrix=y_up_to_z_up_matrix, - ) - asset_bottom_z = float(asset_world_mesh.bounds[0, 2]) - asset_info["rigid_layout"]["pos"][2] += ( - table_top_z + self.config.clearance_m - asset_bottom_z - ) - prepared_assets[asset_id] = asset_info - - log_info( - "Gravity settling started: " - f"assets={len(prepared_assets)}, steps={self.config.settle_steps}, " - f"physics_dt={self.config.physics_dt:.4f} s." - ) - sim = SimulationManager( - SimulationManagerCfg( - headless=True, - physics_dt=self.config.physics_dt, - sim_device=self.config.sim_device, - ) - ) - try: - # Add table. - sim.add_rigid_object( - RigidObjectCfg( - uid=table_id, - shape=MeshCfg(fpath=str(table_info["mesh_path"])), - init_pos=tuple(table_info["rigid_layout"]["pos"]), - init_rot=tuple( - self._simulation_euler_xyz_degrees(table_info["rigid_layout"]) - ), - body_scale=tuple(table_info["y_up_scale"]), - attrs=self._rigid_body_attrs(table_object.physics), - body_type=table_object.physics.body_type, - max_convex_hull_num=table_object.physics.max_convex_hull_num, - acd_method="vhacd", - ) - ) - # Add assets. - simulated_assets: dict[str, object] = {} - for asset_id, asset_info in prepared_assets.items(): - rigid_layout = asset_info["rigid_layout"] - simulated_assets[asset_id] = sim.add_rigid_object( - RigidObjectCfg( - uid=asset_id, - shape=MeshCfg(fpath=str(asset_info["mesh_path"])), - init_pos=tuple(rigid_layout["pos"]), - init_rot=tuple( - self._simulation_euler_xyz_degrees(rigid_layout) - ), - body_scale=tuple(asset_info["y_up_scale"]), - attrs=self._rigid_body_attrs( - asset_objects_by_id[asset_id].physics - ), - body_type=asset_objects_by_id[asset_id].physics.body_type, - max_convex_hull_num=( - asset_objects_by_id[asset_id].physics.max_convex_hull_num - ), - acd_method="vhacd", - ) - ) - # Run simulation to settle all assets. - sim.update(step=self.config.settle_steps) - - # Update the final layouts. - settled_layout_by_id: dict[str, dict[str, object]] = {} - for asset_id, simulated_asset in simulated_assets.items(): - final_rigid_pose_z_up = np.asarray( - simulated_asset.get_local_pose(to_matrix=True)[0] - .detach() - .cpu() - .numpy(), - dtype=float, - ) - scale_matrix = np.eye(4) - scale_matrix[:3, :3] = np.diag(prepared_assets[asset_id]["z_up_scale"]) - final_z_up_layout_matrix = final_rigid_pose_z_up @ scale_matrix - settled_layout_by_id[asset_id] = transform_matrix_to_layout_object( - asset_id, - z_up_to_y_up_matrix - @ final_z_up_layout_matrix - @ y_up_to_z_up_matrix, - ) - finally: - # Release resources. - sim.destroy(exit_process=False) - SimulationManager.flush_cleanup_queue() - - self.settled_assets_layout = [ - settled_layout_by_id[str(asset_layout["id"])] - for asset_layout in self.assets_layout - ] - log_info("Gravity settling completed for all assets.") - return self.settled_assets_layout - - def _prepare_sim_body( - self, - *, - layout_object: dict[str, object], - y_up_to_z_up_matrix: np.ndarray, - ) -> dict[str, object]: - """Load one y-up GLB and prepare its z-up simulation pose.""" - object_id = self._require_layout_id(layout_object, name="Layout object") - source_mesh_path = self.geometry_root / f"{object_id}.glb" - source_mesh = load_glb_mesh(source_mesh_path) - z_up_layout = self._convert_layout_coordinate_system( - layout_object, - source_to_target_matrix=y_up_to_z_up_matrix, - ) - return { - "mesh_path": source_mesh_path, - "mesh": source_mesh, - "rigid_layout": { - "id": object_id, - "rot": self._three_floats(z_up_layout.get("rot"), field_name="rot"), - "pos": self._three_floats(z_up_layout.get("pos"), field_name="pos"), - "scale": [1.0, 1.0, 1.0], - }, - "y_up_scale": self._three_floats( - layout_object.get("scale"), field_name="scale" - ), - "z_up_scale": self._three_floats( - z_up_layout.get("scale"), field_name="scale" - ), - } - - def _require_scene_object(self, object_id: str, *, kind: str) -> SceneObject: - """Return one physics-ready scene object with the expected semantic kind.""" - matching_objects = [ - scene_object - for scene_object in self.scene.objects - if scene_object.id == object_id - ] - if len(matching_objects) != 1: - raise ValueError( - f"Gravity settling requires exactly one scene object {object_id!r}." - ) - scene_object = matching_objects[0] - if scene_object.kind != kind: - raise ValueError( - f"Scene object {object_id!r} must have kind {kind!r} before " - "gravity settling." - ) - if scene_object.physics is None: - raise ValueError( - f"Scene object {object_id!r} has no SimReady physics settings." - ) - return scene_object - - @staticmethod - def _rigid_body_attrs(physics: ObjectPhysics | None) -> RigidBodyAttributesCfg: - """Convert persisted SceneObject physics attributes into Lab config.""" - if physics is None: - raise ValueError("Gravity settling requires SimReady physics settings.") - return RigidBodyAttributesCfg(**physics.attrs) - - @staticmethod - def _mesh_to_z_up_world_for_aabb( - *, - y_up_mesh: trimesh.Trimesh, - z_up_rigid_layout: dict[str, object], - z_up_scale: Sequence[float], - y_up_to_z_up_matrix: np.ndarray, - ) -> trimesh.Trimesh: - """Transform a y-up mesh into its z-up world pose for AABB measurement.""" - mesh = y_up_mesh.copy() - mesh.apply_transform(y_up_to_z_up_matrix) - scale_matrix = np.eye(4) - scale_matrix[:3, :3] = np.diag(z_up_scale) - mesh.apply_transform(scale_matrix) - mesh.apply_transform(layout_object_to_transform_matrix(z_up_rigid_layout)) - return mesh - - @staticmethod - def _simulation_euler_xyz_degrees(layout_object: dict[str, object]) -> list[float]: - """Convert lowercase-xyz layout rotation to SimulationManager's XYZ order.""" - layout_rotation = Rotation.from_euler( - "xyz", - AssetsGravitySettler._three_floats( - layout_object.get("rot"), field_name="rot" - ), - degrees=True, - ) - return layout_rotation.as_euler("XYZ", degrees=True).tolist() - - @staticmethod - def _convert_layout_coordinate_system( - layout_object: dict[str, object], - *, - source_to_target_matrix: np.ndarray, - ) -> dict[str, object]: - """Convert one layout object between coordinate frames through its matrix.""" - return transform_matrix_to_layout_object( - str(layout_object["id"]), - source_to_target_matrix - @ layout_object_to_transform_matrix(layout_object) - @ np.linalg.inv(source_to_target_matrix), - ) - - @staticmethod - def _require_layout_id(layout_object: dict[str, object], *, name: str) -> str: - """Check id.""" - object_id = layout_object.get("id") - if not isinstance(object_id, str) or not object_id: - raise ValueError(f"{name} layout must contain a non-empty string id.") - return object_id - - @staticmethod - def _three_floats(value: object, *, field_name: str) -> list[float]: - """Check three values.""" - if not isinstance(value, list) or len(value) != 3: - raise ValueError(f"Layout field {field_name} must contain three values.") - try: - return [float(item) for item in value] - except (TypeError, ValueError) as exc: - raise ValueError( - f"Layout field {field_name} must contain numeric values." - ) from exc diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/gravity_settler.py b/embodichain/gen_sim/scene_engine/pipeline/utils/gravity_settler.py new file mode 100644 index 000000000..b5078c75f --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/gravity_settler.py @@ -0,0 +1,356 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF 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 scipy.spatial.transform import Rotation + +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, + transform_matrix_to_layout_object, +) +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg +from embodichain.lab.sim.shapes import MeshCfg +from embodichain.utils.logger import log_info + + +@dataclass(frozen=True) +class GravitySettlerConfig: + """Physics controls for one caller-defined gravity-settlement pass.""" + + settle_steps: int = 300 + physics_dt: float = 1.0 / 100.0 + sim_device: str = "cpu" + + def __post_init__(self) -> None: + """Reject invalid numerical controls before starting a simulation.""" + if self.settle_steps <= 0: + raise ValueError("Gravity-settle settle_steps must be positive.") + if self.physics_dt <= 0.0: + raise ValueError("Gravity-settle physics_dt must be positive.") + + +@dataclass(frozen=True) +class GravitySettleBody: + """One scene object and its latest complete y-up pipeline layout.""" + + scene_object: SceneObject + y_up_layout: dict[str, object] + + +class GravitySettler: + """Settle caller-selected dynamic assets against a mandatory table body. + + All supplied layouts use Scene Engine's y-up pipeline convention. The + settler converts them to z-up only at the simulation boundary. The caller + explicitly classifies every participant as dynamic or static; static + participants and the table are kinematic collision bodies. + """ + + def __init__( + self, + *, + table_body: GravitySettleBody, + participant_bodies: list[GravitySettleBody], + dynamic_asset_ids: set[str], + static_asset_ids: set[str], + config: GravitySettlerConfig | None = None, + ) -> None: + self.table_body = table_body + self.participant_bodies = participant_bodies + self.dynamic_asset_ids = set(dynamic_asset_ids) + self.static_asset_ids = set(static_asset_ids) + self.config = config if config is not None else GravitySettlerConfig() + + def settle(self) -> dict[str, dict[str, list[float]]]: + """Return final y-up poses for dynamic participants only. + + Input layouts are used as-is. Placement clearance and support-surface + alignment remain the responsibility of the calling layout optimizer. + Static participants and every object's scale are unchanged, so they are + deliberately omitted from the result. + """ + # Check table. + table = self.table_body.scene_object + if table.kind != "table": + raise ValueError("Gravity settling requires a table body.") + table_id = self._require_body_layout_id(self.table_body, name="Table") + + participant_bodies_by_id: dict[str, GravitySettleBody] = {} + for participant_body in self.participant_bodies: + asset_id = self._require_body_layout_id( + participant_body, name="Participant asset" + ) + if asset_id == table_id: + raise ValueError( + "Gravity-settle participants cannot include the table." + ) + if participant_body.scene_object.kind != "asset": + raise ValueError( + f"Gravity-settle participant {asset_id!r} must be an asset body." + ) + if asset_id in participant_bodies_by_id: + raise ValueError( + f"Gravity-settle participant assets repeat id {asset_id!r}." + ) + participant_bodies_by_id[asset_id] = participant_body + + participant_ids = set(participant_bodies_by_id) + classified_ids = self.dynamic_asset_ids | self.static_asset_ids + if self.dynamic_asset_ids & self.static_asset_ids: + raise ValueError( + "Gravity-settle dynamic and static asset IDs must not overlap." + ) + if classified_ids != participant_ids: + raise ValueError( + "Gravity-settle dynamic and static asset IDs must exactly match " + f"participants; participants={sorted(participant_ids)}, " + f"classified={sorted(classified_ids)}." + ) + if not self.dynamic_asset_ids: + log_info("Gravity settle has no dynamic participants; skipping simulation.") + return {} + + y_up_to_z_up_matrix = self._y_up_to_z_up_matrix() + z_up_to_y_up_matrix = np.linalg.inv(y_up_to_z_up_matrix) + table_info = self._prepare_sim_body( + body=self.table_body, + y_up_to_z_up_matrix=y_up_to_z_up_matrix, + ) + participant_infos_by_id = { + asset_id: self._prepare_sim_body( + body=participant_body, + y_up_to_z_up_matrix=y_up_to_z_up_matrix, + ) + for asset_id, participant_body in participant_bodies_by_id.items() + } + + log_info( + "Gravity settling started: " + f"dynamic_assets={len(self.dynamic_asset_ids)}, " + f"kinematic_assets={len(participant_infos_by_id) - len(self.dynamic_asset_ids)}, " + f"steps={self.config.settle_steps}, " + f"physics_dt={self.config.physics_dt:.4f} s." + ) + sim = SimulationManager( + SimulationManagerCfg( + headless=True, + physics_dt=self.config.physics_dt, + sim_device=self.config.sim_device, + ) + ) + try: + self._add_sim_body( + sim=sim, + object_id=table_id, + body_info=table_info, + physics=table.physics, + body_type="kinematic", + ) + simulated_assets: dict[str, object] = {} + for asset_id, asset_info in participant_infos_by_id.items(): + simulated_assets[asset_id] = self._add_sim_body( + sim=sim, + object_id=asset_id, + body_info=asset_info, + physics=participant_bodies_by_id[asset_id].scene_object.physics, + body_type=( + "dynamic" if asset_id in self.dynamic_asset_ids else "kinematic" + ), + ) + sim.update(step=self.config.settle_steps) + + settled_pose_by_id: dict[str, dict[str, list[float]]] = {} + for asset_id in self.dynamic_asset_ids: + simulated_asset = simulated_assets[asset_id] + 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( + participant_infos_by_id[asset_id]["z_up_scale"] + ) + final_y_up_layout = transform_matrix_to_layout_object( + asset_id, + z_up_to_y_up_matrix + @ final_rigid_pose_z_up + @ scale_matrix + @ y_up_to_z_up_matrix, + ) + settled_pose_by_id[asset_id] = { + "pos": self._three_floats( + final_y_up_layout.get("pos"), field_name="pos" + ), + "rot": self._three_floats( + final_y_up_layout.get("rot"), field_name="rot" + ), + } + finally: + sim.destroy(exit_process=False) + SimulationManager.flush_cleanup_queue() + + log_info("Gravity settling completed for participating assets.") + return settled_pose_by_id + + def _prepare_sim_body( + self, + *, + body: GravitySettleBody, + y_up_to_z_up_matrix: np.ndarray, + ) -> dict[str, object]: + """Convert one supplied y-up layout into a simulator body pose.""" + scene_object = body.scene_object + if scene_object.simready_glb_path is None: + raise ValueError( + f"Gravity-settle object {scene_object.id!r} has no SimReady GLB path." + ) + mesh_path = Path(scene_object.simready_glb_path).expanduser().resolve() + if not mesh_path.is_file(): + raise FileNotFoundError( + f"Gravity-settle GLB for {scene_object.id!r} not found: {mesh_path}" + ) + y_up_layout = body.y_up_layout + z_up_layout = self._convert_layout_coordinate_system( + y_up_layout, + source_to_target_matrix=y_up_to_z_up_matrix, + ) + return { + "mesh_path": mesh_path, + "rigid_layout": { + "id": scene_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( + y_up_layout.get("scale"), field_name="scale" + ), + "z_up_scale": self._three_floats( + z_up_layout.get("scale"), field_name="scale" + ), + } + + def _add_sim_body( + self, + *, + sim: SimulationManager, + object_id: str, + body_info: dict[str, object], + physics: ObjectPhysics | None, + body_type: str, + ) -> object: + """Add one supplied body with a pass-specific dynamic or kinematic type.""" + rigid_layout = body_info["rigid_layout"] + if not isinstance(rigid_layout, dict): + raise ValueError("Gravity-settle body has invalid rigid layout.") + return sim.add_rigid_object( + RigidObjectCfg( + uid=object_id, + shape=MeshCfg(fpath=str(body_info["mesh_path"])), + init_pos=tuple( + self._three_floats(rigid_layout.get("pos"), field_name="pos") + ), + init_rot=tuple(self._simulation_euler_xyz_degrees(rigid_layout)), + body_scale=tuple( + self._three_floats(body_info["y_up_scale"], field_name="scale") + ), + attrs=self._rigid_body_attrs(physics), + body_type=body_type, + max_convex_hull_num=self._max_convex_hull_num(physics), + acd_method="vhacd", + ) + ) + + @staticmethod + def _rigid_body_attrs(physics: ObjectPhysics | None) -> RigidBodyAttributesCfg: + """Convert persisted collision material data into one Lab config.""" + if physics is None: + raise ValueError("Gravity settling requires SimReady physics settings.") + return RigidBodyAttributesCfg(**physics.attrs) + + @staticmethod + def _max_convex_hull_num(physics: ObjectPhysics | None) -> int: + """Read the persisted collision-hull budget after validating physics.""" + if physics is None: + raise ValueError("Gravity settling requires SimReady physics settings.") + return physics.max_convex_hull_num + + @staticmethod + def _require_body_layout_id(body: GravitySettleBody, *, name: str) -> str: + """Validate that a body layout belongs to its scene object.""" + object_id = body.y_up_layout.get("id") + if not isinstance(object_id, str) or not object_id: + raise ValueError(f"{name} layout requires a non-empty string id.") + if object_id != body.scene_object.id: + raise ValueError( + f"{name} layout id {object_id!r} does not match its scene object." + ) + return object_id + + @staticmethod + def _three_floats(value: object, *, field_name: str) -> list[float]: + """Return three finite layout values as Python floats.""" + if not isinstance(value, (list, tuple)) or len(value) != 3: + raise ValueError(f"Gravity-settle {field_name} must contain three values.") + result = [float(component) for component in value] + if not np.all(np.isfinite(result)): + raise ValueError(f"Gravity-settle {field_name} must contain finite values.") + return result + + @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", + GravitySettler._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 complete layout through the y-up/z-up basis change.""" + 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 _y_up_to_z_up_matrix() -> np.ndarray: + """Return the coordinate conversion used by Scene Engine layouts.""" + matrix = np.eye(4) + matrix[:3, :3] = np.array([[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]]) + return matrix diff --git a/tests/gen_sim/scene_engine/test_gravity_settler.py b/tests/gen_sim/scene_engine/test_gravity_settler.py new file mode 100644 index 000000000..5829fa39a --- /dev/null +++ b/tests/gen_sim/scene_engine/test_gravity_settler.py @@ -0,0 +1,81 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF 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 pytest + +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.pipeline.utils.gravity_settler import ( + GravitySettleBody, + GravitySettler, +) + +_TABLE_ID = "table" +_ASSET_ID = "cube_001" +_IDENTITY_LAYOUT = { + "rot": [0.0, 0.0, 0.0], + "pos": [0.0, 0.0, 0.0], + "scale": [1.0, 1.0, 1.0], +} + + +def _table_body() -> GravitySettleBody: + return GravitySettleBody( + scene_object=SceneObject( + id=_TABLE_ID, + kind="table", + category="table", + name="table", + description="table", + ), + y_up_layout={"id": _TABLE_ID, **_IDENTITY_LAYOUT}, + ) + + +def _asset_body() -> GravitySettleBody: + return GravitySettleBody( + scene_object=SceneObject( + id=_ASSET_ID, + kind="asset", + category="cube", + name="cube", + description="cube", + ), + y_up_layout={"id": _ASSET_ID, **_IDENTITY_LAYOUT}, + ) + + +def test_gravity_settler_returns_no_poses_without_dynamic_assets() -> None: + settled_pose_by_id = GravitySettler( + table_body=_table_body(), + participant_bodies=[_asset_body()], + dynamic_asset_ids=set(), + static_asset_ids={_ASSET_ID}, + ).settle() + + assert settled_pose_by_id == {} + + +def test_gravity_settler_rejects_dynamic_assets_outside_participants() -> None: + with pytest.raises(ValueError, match="exactly match participants"): + GravitySettler( + table_body=_table_body(), + participant_bodies=[], + dynamic_asset_ids={_ASSET_ID}, + static_asset_ids=set(), + ).settle()