embodied_gen/skills/spatial_computing/cli/main.py

# Project EmbodiedGen
#
# Copyright (c) 2025 Horizon Robotics. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied. See the License for the specific language governing
# permissions and limitations under the License.

from __future__ import annotations

import logging
import sys
import warnings

import tyro
from embodied_gen.skills.spatial_computing.api.floorplan_api import (
    FloorplanConfig,
    FloorplanManager,
)

warnings.filterwarnings("ignore", category=RuntimeWarning)

logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(message)s",
    level=logging.INFO,
    force=True,
)
logger = logging.getLogger(__name__)


def _get_gpt_client() -> object | None:
    """Lazy-import GPT_CLIENT for semantic --on_instance resolution."""
    try:
        from embodied_gen.utils.gpt_clients import GPT_CLIENT

        return GPT_CLIENT
    except Exception:
        return None


def entrypoint(cfg: FloorplanConfig) -> None:
    """Main entry point for floorplan visualization and scene manipulation.

    Args:
        cfg: Configuration object with all parameters.

    """
    manager = FloorplanManager(
        urdf_path=cfg.urdf_path,
        usd_path=cfg.usd_path,
        mesh_sample_num=cfg.mesh_sample_num,
        ignore_items=cfg.ignore_items,
    )

    # List instances/rooms and exit if requested
    if cfg.list_instances:
        names = manager.get_instance_names()
        rooms = manager.get_room_names()
        logger.info("instance_names:", names)
        logger.info("room_names:", rooms)
        return

    gpt_client = _get_gpt_client()
    on_instance = cfg.on_instance
    if on_instance is not None:
        resolved = manager.resolve_on_instance(
            on_instance, gpt_client=gpt_client
        )
        if resolved is None:
            logger.error(
                "No object matched \"%s\"。Current scene instance name: %s。",
                on_instance,
                manager.get_instance_names(),
            )
            sys.exit(1)
        on_instance = resolved
        if resolved != cfg.on_instance:
            logger.info("\"%s\" -> \"%s\"", cfg.on_instance, resolved)

    in_room = cfg.in_room
    if in_room is not None:
        resolved = manager.resolve_in_room(in_room, gpt_client=gpt_client)
        if resolved is None:
            logger.error(
                "No room matched \"%s\"。Current scene room names: %s。",
                in_room,
                manager.get_room_names(),
            )
            sys.exit(1)
        in_room = resolved
        if resolved != cfg.in_room:
            logger.info("\"%s\" -> \"%s\"", cfg.in_room, resolved)

    beside_instance = cfg.beside_instance
    if beside_instance is not None:
        resolved = manager.resolve_beside_instance(
            beside_instance, gpt_client=gpt_client, in_room=in_room
        )
        if resolved is None:
            logger.error(
                "No object matched \"%s\"。Current scene instance name: %s。",
                beside_instance,
                manager.get_instance_names(),
            )
            sys.exit(1)
        beside_instance = resolved
        if resolved != cfg.beside_instance:
            logger.info("\"%s\" -> \"%s\"", cfg.beside_instance, resolved)

    # Add asset instance if specified
    center = None
    if cfg.asset_path is not None:
        center = manager.insert_object(
            asset_path=cfg.asset_path,
            instance_key=cfg.instance_key,
            in_room=in_room,
            on_instance=on_instance,
            beside_instance=beside_instance,
            beside_distance=cfg.beside_distance,
            rotation_rpy=cfg.rotation_rpy,
            n_max_attempt=cfg.max_placement_attempts,
            place_strategy=cfg.place_strategy,
        )

        if center is not None:
            logger.info(
                f"Successfully placed '{cfg.instance_key}' at "
                f"({center[0]:.3f}, {center[1]:.3f}, {center[2]:.3f})"
            )
        else:
            logger.error(
                f"❌ Failed to place '{cfg.instance_key}' in the scene."
            )
            sys.exit(1)

    # Generate floorplan visualization
    if cfg.output_path is not None:
        manager.visualize(output_path=cfg.output_path)


if __name__ == "__main__":
    config = tyro.cli(FloorplanConfig)
    entrypoint(config)