diff --git a/docs/source/features/generative_sim/scene_engine.md b/docs/source/features/generative_sim/scene_engine.md index 03e4b3fa4..e4255163e 100644 --- a/docs/source/features/generative_sim/scene_engine.md +++ b/docs/source/features/generative_sim/scene_engine.md @@ -27,10 +27,27 @@ 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, 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" @@ -43,7 +60,13 @@ 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_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 @@ -66,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/start.py b/embodichain/gen_sim/scene_engine/cli/start.py index 59454b09f..6d46f32dc 100644 --- a/embodichain/gen_sim/scene_engine/cli/start.py +++ b/embodichain/gen_sim/scene_engine/cli/start.py @@ -21,15 +21,31 @@ 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"} 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}") @@ -40,27 +56,30 @@ def cli_scene_engine( "Image input must have one of these extensions: .jpg, .jpeg, .png" ) - resolved_output_root = Path(output_root).expanduser().resolve() resolved_output_root.mkdir(parents=True, exist_ok=True) - generate_scene_from_image( image_path=resolved_image_path, output_root=resolved_output_root, ) + if edit_prompt is not None: + edit_scene( + output_root=resolved_output_root, + edit_prompt=edit_prompt, + ) print("Successfully completed!") 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", @@ -68,9 +87,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="Text instruction for editing an existing or newly generated 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/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/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/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..0f78b9699 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py @@ -0,0 +1,280 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF 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 ( + OrientationState, + SceneConstraintType, + SceneGraph, + TableRegion, + 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 + table_region: TableRegion | None = None + 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.""" + return { + "op": self.op, + "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, + "orientation_state": self.orientation_state, + } + + +@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.""" + # 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. + # - 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): + 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 deletion intents first so move and add cannot target them regardless of order. + 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.table_region, + operation.category, + operation.name, + operation.description, + operation.orientation_state, + ) + ): + 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.") + 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, + 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.") + 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, + 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.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. + 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/core/scene_graph.py b/embodichain/gen_sim/scene_engine/core/scene_graph.py new file mode 100644 index 000000000..500980b53 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/core/scene_graph.py @@ -0,0 +1,546 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF 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" + +# Static type constraint for the nine regions of the tabletop 3x3 grid. +TableRegion = Literal[ + "left_back", + "back_center", + "right_back", + "left_center", + "center", + "right_center", + "left_front", + "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"] + +# 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. + + ``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.""" + 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 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.") + 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, + "orientation_state": self.orientation_state, + } + + +@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 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], + 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: + """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)}" + ) + 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( + SceneGraphNode( + 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 + ) + + # Apply support-parent changes before planar updates need the final parent. + 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) + 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, + 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: + 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 = table_region + 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() + 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/embodichain/gen_sim/scene_engine/core/scene_object.py b/embodichain/gen_sim/scene_engine/core/scene_object.py index d9a837e87..2b868e3c3 100644 --- a/embodichain/gen_sim/scene_engine/core/scene_object.py +++ b/embodichain/gen_sim/scene_engine/core/scene_object.py @@ -66,6 +66,10 @@ 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. + 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]: @@ -81,5 +85,9 @@ def to_dict(self) -> dict[str, object]: "rot": self.rot, "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/edit.py b/embodichain/gen_sim/scene_engine/pipeline/edit.py new file mode 100644 index 000000000..5b71af556 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/edit.py @@ -0,0 +1,125 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF 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.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, +) +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 + + +def edit_scene( + *, + output_root: str | Path, + 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() + + # 1. Edit Understanding + # Will return an already checked scene edit plan + # and a validated updated scene graph. + log_info("Starting Edit Understanding") + 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") + + # 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, + vlm_client=vlm_client, + ) + finally: + image_generation_client.close() + geometry_generation_client.close() + image_segmentation_client.close() + log_info("Completed Objects Preparation") + + # 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 + 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/__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_asset_preparation.py b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py new file mode 100644 index 000000000..220a91014 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py @@ -0,0 +1,328 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF 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.llms.openai_compatible_client import ( + OpenAICompatibleVLM, +) +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, + SimReadyProcessorConfig, +) + +__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, + vlm_client: OpenAICompatibleVLM | None = None, +) -> 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", + # 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, + ) + # 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_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/editing/scene_edit_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py new file mode 100644 index 000000000..92720c7ec --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py @@ -0,0 +1,492 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF 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 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 ( + OrientationState, + PlanarRelationType, + SceneGraph, + SceneGraphNode, + SceneGraphRelation, + TABLE_REGIONS, + TableRegion, +) +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, 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 + 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. 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. +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 +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, 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, description, and orientation_state. Use null for every field that does not apply: +{ + "operations": [ + { + "op": "move", + "object_id": "bottle_001", + "target_id": "book_001", + "relation": "right_of", + "table_region": null, + "category": null, + "name": null, + "description": null, + "orientation_state": null + }, + { + "op": "delete", + "object_id": "cup_001", + "target_id": null, + "relation": null, + "table_region": null, + "category": null, + "name": null, + "description": null, + "orientation_state": null + }, + { + "op": "add", + "object_id": null, + "target_id": "table", + "relation": "on", + "table_region": "back_center", + "category": "orange", + "name": "small orange", + "description": "small round orange with a textured peel", + "orientation_state": null + }, + { + "op": "add", + "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", + "orientation_state": null + }, + { + "op": "add", + "object_id": null, + "target_id": null, + "relation": null, + "table_region": null, + "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. 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( + *, + scene: Scene, + scene_graph: SceneGraph, + edit_prompt: str, + vlm_client: OpenAICompatibleVLM, + json_max_attempts: int = 3, +) -> tuple[SceneEditPlan, SceneGraph]: + """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, + scene_graph=scene_graph, + ) + 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. + 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, + orientation_state=node.orientation_state, + ) + 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] = [] + 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: + 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) + 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": + 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) + ) + + # 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, + added_orientation_states_by_id=added_orientation_states_by_id, + on_parent_updates=on_parent_updates, + planar_relation_updates=planar_relation_updates, + ) + + +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 + } + 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": [ + { + "id": scene_object.id, + "category": scene_object.category, + "name": scene_object.name, + "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 + ], + } + + +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", + "table_region", + "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} + 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") + 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.") + 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")), + table_region=_optional_table_region(value.get("table_region")), + category=category, + name=_optional_string(value.get("name"), field_name="name"), + description=_optional_string( + value.get("description"), field_name="description" + ), + orientation_state=orientation_state, + ) + ) + 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 + + +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 + + +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/generate.py b/embodichain/gen_sim/scene_engine/pipeline/generate.py index 773a897bf..144551c29 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generate.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generate.py @@ -25,13 +25,16 @@ 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.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 @@ -51,12 +54,19 @@ def generate_scene_from_image( # 1. Scene Understanding log_info("Starting Scene Understanding") - scene = understand_scene( - scene=scene, - image_path=image_path, - output_root=resolved_output_root, - vlm_client=vlm_client, - ) + # 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 @@ -69,7 +79,9 @@ 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, + vlm_client=vlm_client, ) finally: geometry_generation_client.close() # Kill the session to avoid resource leaks. @@ -79,6 +91,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/__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 65% rename from embodichain/gen_sim/scene_engine/pipeline/scene_generation.py rename to embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py index d14ce364f..7010e5f14 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/scene_generation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py @@ -22,13 +22,19 @@ import shutil import numpy as np +from scipy.spatial.transform import Rotation import trimesh +from shapely.geometry import Polygon from embodichain.gen_sim.scene_engine.clients.geometry_generation import ( 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.llms.openai_compatible_client import ( + OpenAICompatibleVLM, +) from embodichain.gen_sim.scene_engine.pipeline.utils.assets_group_support_clamp import ( AssetsGroupSupportClamp, ) @@ -38,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, @@ -47,11 +54,9 @@ 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.table_support_surface import ( - TableSupportSurfaceDetector, +from embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor import ( + SimReadyProcessor, + SimReadyProcessorConfig, ) from embodichain.utils.logger import log_info @@ -62,11 +67,15 @@ def generate_scene_and_refine( image_path: str | Path, output_root: str | Path, scene: Scene, + scene_graph: SceneGraph, *, geometry_generation_client: GeometryGenerationClient, + vlm_client: OpenAICompatibleVLM, ) -> 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(): @@ -101,11 +110,26 @@ def generate_scene_and_refine( coarse_layout_by_id = { layout_object["id"]: layout_object for layout_object in coarse_layout } - simready_processor = SimReadySceneProcessor( + # 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, + # Keep the geometry-server scale and only correct unstable standing poses. + config=SimReadyProcessorConfig( + use_vlm_scale=False, + use_vlm_rotation=False, + orientation_states_by_id=standing_orientation_states_by_id, + ), + vlm_client=vlm_client, ) simready_assets_layout = simready_processor.process_assets() simready_table_layout = simready_processor.process_table() @@ -119,6 +143,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. ) @@ -195,16 +220,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 +250,29 @@ 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() + 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() + def _copy_y_up_layout_to_scene_object( scene_object: SceneObject, @@ -252,6 +302,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]]]: @@ -318,7 +369,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. @@ -332,31 +390,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, + # 5. 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 + # 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. 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 ), @@ -377,13 +440,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. + # 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_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 ), @@ -396,24 +456,138 @@ 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. - # 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( + # 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( scene=scene, table_layout=refined_table_layout, assets_layout=refined_assets_layout, geometry_root=simready_geometry_output_root, ) - refined_assets_layout = gravity_settler.settle() + return refined_table_layout, refined_assets_layout - # Update the scene data structure with the final y-up layout values. - _update_scene_final_y_up_layout( - scene=scene, - table_layout=refined_table_layout, - assets_layout=refined_assets_layout, + +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]] ) - return refined_table_layout, refined_assets_layout + 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}.") + # Only correct the standing assets. + 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( 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 74% rename from embodichain/gen_sim/scene_engine/pipeline/scene_understanding.py rename to embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py index f91c665e6..d93ada019 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/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, @@ -36,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, @@ -44,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. @@ -62,19 +62,23 @@ 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. - 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 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. + 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: { @@ -87,14 +91,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." @@ -141,6 +145,23 @@ 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 +asset is standing, lying, or unknown in the image. + +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 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": "book_001", "orientation_state": null} + ] +}""" def understand_scene( @@ -149,8 +170,9 @@ def understand_scene( output_root: str | Path, *, vlm_client: OpenAICompatibleVLM, + image_segmentation_client: ImageSegmentationClient, 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 @@ -167,26 +189,165 @@ 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. + # 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, + 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. + # But at here, we do with the simplest way (hard code). + scene_graph = _initialize_scene_graph_from_segmented_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. (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, + *, + 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. + 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, + json_max_attempts=json_max_attempts, + ) + 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", # semi-hard-code. + orientation_state=orientation_states_by_id[asset.id], + ) + for asset in scene.assets + ], + ], + ) + + +def _simplify_scene_info_for_graph_initialization( + *, + scene: Scene, +) -> dict[str, object]: + """Return the object metadata needed to initialize an image-based graph.""" + return { + "asset_ids": [asset.id for asset in scene.assets], + } + + +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 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, + asset_ids: object, +) -> dict[str, str | None]: + """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 asset 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(asset_ids): + raise ValueError( + "VLM JSON orientation states must match all supplied asset IDs." + ) + return orientation_states_by_id def _analyze_image_objects( @@ -330,13 +491,6 @@ 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. - raise ValueError( - f"VLM JSON key {field_name}.name must not contain location or " - "relationship words." - ) return fields @@ -353,8 +507,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. @@ -396,6 +550,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/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/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/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/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..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 @@ -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], *, @@ -277,6 +306,72 @@ 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) + 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() + 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( @@ -361,6 +456,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 @@ -401,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: @@ -413,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] @@ -438,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/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..6d2a3b0a5 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/parent_surface_layout_optimizer.py @@ -0,0 +1,546 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF 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, + ) + + +class _LayoutInfeasibleError(ValueError): + """Internal marker for an SLSQP failure while testing one collision direction.""" + + +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 _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) + } + + +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 + # 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, + half_extents=root_half_extents_xy, + xy_by_id=current, + ) + }: + 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 + 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_constraints( + *, + 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, +) -> 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]) + 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 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/scene_exporter.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py index fb66c30c4..7451296b4 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,16 @@ 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 + self.scene_json_path: Path | None = None def export(self) -> Path: """Write a scene-only config and copy SimReady GLBs into ``mesh_assets``. @@ -72,6 +77,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 +114,23 @@ 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}") + 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 @@ -135,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( *, @@ -166,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", @@ -179,6 +225,10 @@ 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, + "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 new file mode 100644 index 000000000..e730b88cb --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py @@ -0,0 +1,417 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF 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_graph import ( + SceneGraph, + SceneGraphNode, + SceneGraphRelation, + TABLE_REGIONS, +) +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_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( + "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.") + + 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}") + + 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 + ], + ] + ) + + @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", + "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)) + ): + 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 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 + 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, + *, + 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", + ) + 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() + 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=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(), + 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})), + 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.") + 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( + 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 _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.""" + 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 + + @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.""" + 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/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..95c77d712 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_constructor.py @@ -0,0 +1,353 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF 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.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, +) + + +@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, + 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() + # 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( + 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_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.") + # 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 + ): + 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 = table_surface_problem.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 + 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( + 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.""" + # 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} + 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], + ] + 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.""" + parent_surface_problem = ParentSurfaceLayoutProblem.from_layout_problem( + layout_problem=layout_problem, + group=group, + current_xy_by_id=self._current_xy_by_id, + ) + # Get results. + solved_child_xy_by_id = self.parent_surface_layout_optimizer.optimize( + parent_surface_problem + ) + for child_id, solved_xy in solved_child_xy_by_id.items(): + 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], + ] + 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. + 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, + clearance_m=0.00, # Directly place on the parent's top surface. + ) + 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.""" + # 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() + + # 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.") + # 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( # The assets' center XY should always be updated whenever changes are made. + 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." + ) + # 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=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_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/simready_scene_processor.py b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py similarity index 56% 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..864e39a59 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_scene_processor.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py @@ -17,21 +17,34 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path -import re import numpy as np -import open3d as o3d -from scipy.spatial import ConvexHull, QhullError from scipy.spatial.transform import Rotation import trimesh from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_graph import OrientationState from embodichain.gen_sim.scene_engine.core.scene_object import ( 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 ( + 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, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.table_support_surface import ( + TableSupportSurfaceDetector, +) from embodichain.utils.logger import log_info _TABLE_PHYSICS_ATTRS = { @@ -53,16 +66,17 @@ @dataclass(frozen=True) -class SimReadySceneProcessorConfig: - """Object-category policy for SimReady mesh canonicalization.""" +class SimReadyProcessorConfig: + """SceneGraph-conditioned policy for SimReady mesh canonicalization.""" - upright_container_id_tokens: frozenset[str] = frozenset( - {"bottle", "can", "jar", "flask", "thermos"} - ) # Object-id tokens that enable upright-container standardization. + use_vlm_scale: bool = False # Use the VLM-selected asset scale. + use_vlm_rotation: bool = False # Use the VLM-selected asset rotation. + # Explicit graph orientation overrides the default stable tabletop pose. + orientation_states_by_id: dict[str, OrientationState] = field(default_factory=dict) -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 +85,9 @@ 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, + debug_output_root: str | Path | None = None, + config: SimReadyProcessorConfig | None = None, + vlm_client: OpenAICompatibleVLM | None = None, ) -> None: self.scene = scene self.coarse_layout_by_id = coarse_layout_by_id @@ -79,11 +95,22 @@ 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 SimReadySceneProcessorConfig() - if not self.config.upright_container_id_tokens: - raise ValueError("upright_container_id_tokens must not be empty.") + 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 + or self.config.orientation_states_by_id + ) 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 +131,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 +146,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) @@ -129,9 +168,120 @@ def _process_object(self, scene_object: SceneObject) -> dict[str, 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]: + """Render, query, and optionally bake the VLM-selected x-axis rotation.""" + coarse_path = self.coarse_geometry_root / f"{scene_object.id}.glb" + 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, + needed_layout=self._needed_layout_for_object(scene_object.id), + ) + rotate_about_x = bool(decision["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 + 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. + return ( + rotated_path, + [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, + *, + needed_layout: str, + ) -> dict[str, object]: + """Render the object and return the validated VLM pose decision.""" + assert self.vlm_client is not None + coarse_path = self.coarse_geometry_root / f"{scene_object.id}.glb" + 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", + ) + # 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.""" @@ -185,8 +335,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.") @@ -197,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 upright containers in temporary z-up coordinates before the - # shared center, scale, and bottom-center preprocessing. - # This is to ensure the action agent can pick up the bottle or can-like objects. - bottle_alignment_matrix = np.eye(3) - if self._is_upright_container_id(object_id): - bottle_alignment_matrix = self._standardize_bottle_z_up(mesh) - bottle_alignment_transform = np.eye(4) - bottle_alignment_transform[:3, :3] = bottle_alignment_matrix - mesh.apply_transform(bottle_alignment_transform) - # First make the object's AABB center at the origin. original_aabb_center = mesh.bounds.mean(axis=0) mesh.apply_translation(-original_aabb_center) @@ -214,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. - bottle_alignment_matrix - @ y_up_to_z_up_matrix - @ np.diag(coarse_scale) - @ y_up_to_z_up_matrix.T - @ bottle_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) @@ -240,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 bottle's local rotation so that its coarse world pose does - # not change. - local_bottle_rotation = Rotation.from_matrix( - y_up_to_z_up_matrix.T @ bottle_alignment_matrix @ y_up_to_z_up_matrix - ) - coarse_rotation_matrix = Rotation.from_euler( - "xyz", coarse_rot, degrees=True - ).as_matrix() - rotation = Rotation.from_matrix( - coarse_rotation_matrix @ local_bottle_rotation.inv().as_matrix() - ) + 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 @@ -261,87 +383,6 @@ 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) - - @staticmethod - def _standardize_bottle_z_up(mesh: trimesh.Trimesh) -> np.ndarray: - """Return a proper rotation that maps a bottle-like mesh's long axis to z-up. - - Thanks to chenjian for this idea! - """ - if len(mesh.vertices) < 4 or len(mesh.faces) < 4: - raise ValueError( - "Bottle standardization requires a non-degenerate triangle mesh." - ) - open3d_mesh = o3d.geometry.TriangleMesh( - vertices=o3d.utility.Vector3dVector(mesh.vertices), - triangles=o3d.utility.Vector3iVector(mesh.faces), - ) - sampled_points = np.asarray( - open3d_mesh.sample_points_uniformly(number_of_points=10_000).points - ) # (10000, 3) x (x, y, z) - - # Check the number of the points again, and check whether have some - # non-finite values. - if sampled_points.shape[0] < 4 or not np.all(np.isfinite(sampled_points)): - raise ValueError( - "Bottle standardization could not sample valid mesh points." - ) - - centered_points = sampled_points - sampled_points.mean(axis=0) - # SVD find the longest axis. - _, _, principal_axes = np.linalg.svd(centered_points, full_matrices=False) - if np.linalg.det(principal_axes) < 0: - principal_axes[2, :] *= -1 # in case the SVD returns a reflection. - - bottle_rotation = Rotation.from_euler( - "y", 90.0, degrees=True - ).as_matrix() # 3x3 matrix - # The first PCA axis is the longest axis; rotate it onto the temporary z axis. - bottle_rotation = bottle_rotation @ principal_axes - standardized_points = (bottle_rotation @ centered_points.T).T - - axis_min = standardized_points[:, 2].min() - axis_max = standardized_points[:, 2].max() - axis_range = axis_max - axis_min - upper_points = standardized_points[ - standardized_points[:, 2] > axis_min + axis_range * 0.8 - ] - lower_points = standardized_points[ - standardized_points[:, 2] < axis_min + axis_range * 0.2 - ] - upper_volume = SimReadySceneProcessor._convex_hull_volume(upper_points) - lower_volume = SimReadySceneProcessor._convex_hull_volume(lower_points) - - # Bottles usually have a smaller top (neck) than bottom; flip if necessary. - if upper_volume > lower_volume: - bottle_rotation = ( - Rotation.from_euler("x", 180.0, degrees=True).as_matrix() - @ bottle_rotation - ) - return bottle_rotation - - @staticmethod - def _convex_hull_volume(points: np.ndarray) -> float: - """Return the volume of a non-degenerate point set's convex hull.""" - if points.shape[0] < 4: - raise ValueError( - "Bottle standardization needs at least four points per end." - ) - try: - return float(ConvexHull(points).volume) - except QhullError as exc: - raise ValueError( - "Bottle standardization found a degenerate end volume." - ) from exc - @staticmethod def _three_floats(value: object, *, field_name: str) -> list[float]: """Validate and convert a three-value layout field to floats.""" 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..4523aa3e6 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor_utils.py @@ -0,0 +1,445 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF 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 os +from pathlib import Path +import sys +from typing import Callable + +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]. +""" + +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( + *, + 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 + + _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 + # 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) + _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)) + 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 _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], + 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, + json_max_attempts: int = 3, +) -> dict[str, object]: + """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: + 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.") + 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.") + # 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) + 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() 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)") 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..9d874f239 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/table_surface_layout_optimizer.py @@ -0,0 +1,598 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF 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, + ) + + +class _LayoutInfeasibleError(ValueError): + """Internal marker for an SLSQP failure while testing one collision direction.""" + + +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 _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) + } + + +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 + # 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, + half_extents=root_half_extents_xy, + xy_by_id=current, + ) + }: + 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 + 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_constraints( + *, + 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, +) -> 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]) + # 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 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( + half_extents[first_id][axis] + half_extents[second_id][axis] + margin + ) diff --git a/tests/gen_sim/scene_engine/test_clients.py b/tests/gen_sim/scene_engine/test_clients.py index 513f0c5c0..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 @@ -79,7 +87,14 @@ 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", + } + 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", @@ -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_single_object_path == "/predict" + 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" @@ -146,15 +172,93 @@ 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, ) + 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: @@ -170,7 +274,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 +282,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"} 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() 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..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 @@ -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")) @@ -133,10 +161,131 @@ 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] + 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, + "orientation_state": None, + }, + { + "object_id": "cup", + "parent_id": "table", + "parent_relation": "on", + "table_region": None, + "orientation_state": 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_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_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" + 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: @@ -145,7 +294,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 +320,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 new file mode 100644 index 000000000..fd1a9e933 --- /dev/null +++ b/tests/gen_sim/scene_engine/test_scene_edit.py @@ -0,0 +1,135 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF 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], + "center_xy": [1.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].center_xy == [1.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_edit_plan.py b/tests/gen_sim/scene_engine/test_scene_edit_plan.py new file mode 100644 index 000000000..a99e04384 --- /dev/null +++ b/tests/gen_sim/scene_engine/test_scene_edit_plan.py @@ -0,0 +1,555 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF 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 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 ( + 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 ( + _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]: + 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, + "table_region": None, + "category": "cup", + "name": "green cup", + "description": "A small green ceramic cup.", + "orientation_state": None, + } + ] + + +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, + "table_region": None, + "category": "orange", + "name": "small_orange", + "description": "A small round orange with a textured peel.", + "orientation_state": None, + }, + { + "op": "add", + "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.", + "orientation_state": "lying", + }, + ] + } + + operations = _parse_scene_edit_operations( + json.loads(json.dumps(draft)), scene=scene + ) + + assert [operation.object_id for operation in operations] == [ + "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: + 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")], + ) + + +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( + tmp_path: Path, +) -> 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", + ) + ], + ) + 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_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] + ) + + 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: + 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.", + orientation_state="standing", + ) + ], + ) + + 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" + 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, + operations=[ + SceneEditOperation( + op="move", + object_id="orange_001", + target_id="table", + relation="on", + orientation_state="lying", + ) + ], + ) + + 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" + 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: + 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" 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_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_graph.py b/tests/gen_sim/scene_engine/test_scene_graph.py new file mode 100644 index 000000000..d5b6adaf7 --- /dev/null +++ b/tests/gen_sim/scene_engine/test_scene_graph.py @@ -0,0 +1,313 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF 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", + orientation_state="standing", + ), + 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", + orientation_state="standing", + ), + ], + ) + + graph_dict = graph.to_dict() + + assert graph_dict == { + "nodes": [ + { + "object_id": "table", + "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_layout_optimizer.py b/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py new file mode 100644 index 000000000..6698d5c83 --- /dev/null +++ b/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py @@ -0,0 +1,342 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF 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, + 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, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.table_surface_layout_optimizer import ( + TableSurfaceLayoutOptimizer, + TableSurfaceLayoutProblem, + _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 +_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( + *, + 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 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: + 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 = 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 == { + 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 = 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 + 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) + 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 + ) diff --git a/tests/gen_sim/scene_engine/test_scene_understanding.py b/tests/gen_sim/scene_engine/test_scene_understanding.py index 1fdb51f7b..382e4933d 100644 --- a/tests/gen_sim/scene_engine/test_scene_understanding.py +++ b/tests/gen_sim/scene_engine/test_scene_understanding.py @@ -20,10 +20,16 @@ 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.pipeline import scene_understanding +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, +) def _response(*, asset_name: str = "cup") -> str: @@ -55,11 +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_accepts_description_with_structural_words() -> None: + response = json.loads(_response()) + response["assets"][0]["description"] = "A small ceramic cup with a lid on top." + + scene = 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: @@ -83,3 +101,297 @@ def complete(self, **_: object) -> str: assert scene.table is not None assert [asset.id for asset in scene.assets] == ["cup_001"] + + +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_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: + return json.dumps( + { + "orientation_states": [ + {"object_id": "cup_001", "orientation_state": 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.", + ), + ], + ) + + 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, + asset_mask_id_overlay_path=overlay_path, + vlm_client=VLM(), # type: ignore[arg-type] + ) + + assert scene_graph.to_dict() == { + "nodes": [ + { + "object_id": "table", + "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_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": "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.", + ), + ] + ) + + 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" + + +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_asset_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 == { + "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 new file mode 100644 index 000000000..dd872f913 --- /dev/null +++ b/tests/gen_sim/scene_engine/test_simready_processor_utils.py @@ -0,0 +1,99 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF 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.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 ( + 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_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( + 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 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: + """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) 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..b6c9df3ff 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,39 @@ 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)]),