bundle/sim/build_scene.py

from __future__ import annotations

import argparse
from pathlib import Path

import cv2
import numpy as np
import pandas as pd
from PIL import Image

from scene_tools import (
    create_background_inpaint,
    load_scene_calibration,
    render_scene,
    save_scene_json,
)


def _save_mask_overlay(rgb: np.ndarray, mask: np.ndarray, color: tuple[int, int, int], path: Path) -> None:
    overlay = rgb.copy()
    overlay[mask] = color
    Image.fromarray(overlay).save(path)


def _save_box_outline(rgb: np.ndarray, polygon_world: np.ndarray, calibration, path: Path) -> None:
    canvas = rgb.copy()
    image_points, _ = cv2.projectPoints(
        polygon_world.astype(np.float64), calibration.rvec, calibration.tvec, calibration.camera_matrix, None
    )
    pts = np.round(image_points[:, 0, :]).astype(np.int32)
    cv2.polylines(canvas, [pts], True, (255, 0, 0), 4)
    Image.fromarray(canvas).save(path)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--session-root",
        default="/workspace/data/teddybear_raw/session_20260327_165944_bear",
        help="Raw session root used for calibration.",
    )
    parser.add_argument(
        "--sync-row-index",
        type=int,
        default=0,
        help="Row of sync_index.csv to use as the start frame.",
    )
    parser.add_argument(
        "--output-dir",
        default="/workspace/kinova_scene_sim/outputs",
        help="Directory for the scene JSON and visual checks.",
    )
    args = parser.parse_args()

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    calibration = load_scene_calibration(args.session_root, args.sync_row_index)
    sync = pd.read_csv(Path(args.session_root) / "sync_index.csv")
    start_q = sync.iloc[args.sync_row_index][[f"q{i}_deg" for i in range(1, 8)]].to_numpy(np.float64)
    motion_q = sync.iloc[[0, 56, 70, 84, 98, 112]][[f"q{i}_deg" for i in range(1, 8)]].to_numpy(np.float64)

    background = create_background_inpaint(calibration.rgb)
    scene_json = output_dir / f"{Path(args.session_root).name}_scene_fit.json"
    save_scene_json(calibration, scene_json)

    Image.fromarray(calibration.rgb).save(output_dir / "real_start.png")
    Image.fromarray(background).save(output_dir / "background_inpaint.png")

    _save_mask_overlay(calibration.rgb, calibration.box_mask, (255, 0, 0), output_dir / "box_mask_overlay.png")
    _save_mask_overlay(calibration.rgb, calibration.teddy_mask, (0, 255, 0), output_dir / "teddy_mask_overlay.png")
    _save_box_outline(calibration.rgb, calibration.box_world_polygon, calibration, output_dir / "box_world_outline.png")

    start_render = render_scene(calibration, start_q, background=background)
    Image.fromarray(start_render).save(output_dir / "scene_start_render.png")

    colors = [
        (255, 64, 64),
        (64, 255, 64),
        (64, 64, 255),
        (255, 255, 64),
        (255, 64, 255),
        (64, 255, 255),
    ]
    motion_render = background.copy()
    for q_deg, color in zip(motion_q, colors):
        motion_render = render_scene(calibration, q_deg, background=motion_render, color=color, alpha=0.95)
    Image.fromarray(motion_render).save(output_dir / "scene_motion_fit.png")

    side_by_side = np.concatenate([calibration.rgb, start_render], axis=1)
    Image.fromarray(side_by_side).save(output_dir / "real_vs_scene_start.png")

    print(f"saved scene json to {scene_json}")
    print(f"start frame: azure_rgb/{calibration.azure_rgb_seq:06d}.jpg")
    print(f"depth frame: azure_depth/{calibration.azure_depth_seq:06d}.png")
    print(f"robot seq: {calibration.robot_seq}")
    print(f"box height estimate: {calibration.box_height_m:.3f} m")
    print(f"teddy height estimate: {calibration.teddy_height_m:.3f} m")
    print(f"teddy center world: {calibration.teddy_world_center.tolist()}")


if __name__ == "__main__":
    main()