File size: 3,912 Bytes
d93804e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | 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()
|