box_detection.py

"""Entry point for box rim detection (top-down open-crate localization).

Examples:
    # Capture one frame from the camera and save it for offline tuning:
    python box_detection.py --camera-type vzense --save captures/box01.npy

    # Run the detector on a saved cloud and print the rim geometry:
    python box_detection.py --input captures/box01.npy --stage box --no-visualize

    # Live tuner on the camera:
    python box_detection.py --camera-type vzense --stage box
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np

from box.config import default_config, load_config, save_config
from box.pipeline import crop_mask, rotate_xyz


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Box rim detection (open-crate pickup)")
    parser.add_argument("--stage", choices=["crop", "box"], default="box",
                        help="Pipeline stage (default: box = full rim detection)")
    parser.add_argument("--input", help="Run on a saved point cloud (.npy of shape (N,3), or .ply)")
    parser.add_argument("--save", help="Capture one camera frame, save the cloud (.npy) + a .png, and exit")
    parser.add_argument("--camera-type",
                        choices=["realsense", "d435", "vzense", "nyx650", "tof"],
                        help="Camera backend. Overrides config.camera.camera_type.")
    parser.add_argument("--no-visualize", action="store_true",
                        help="Disable the interactive 3D tuner; print metrics once.")
    parser.add_argument("--config", default=str(Path("config") / "box_params.json"),
                        help="Path to JSON config for box detection tuning")
    parser.add_argument("--save-default-config", action="store_true",
                        help="Write default config JSON to --config and exit")
    return parser.parse_args()


def load_points(path: str) -> np.ndarray:
    p = Path(path)
    if p.suffix.lower() == ".npy":
        pts = np.load(p)
    elif p.suffix.lower() == ".ply":
        import open3d as o3d
        pts = np.asarray(o3d.io.read_point_cloud(str(p)).points)
    else:
        raise ValueError(f"Unsupported input format: {p.suffix} (use .npy or .ply)")
    pts = np.asarray(pts, dtype=np.float64)
    if pts.ndim != 2 or pts.shape[1] != 3:
        raise ValueError(f"Expected (N,3) points, got shape {pts.shape}")
    return pts


def build_camera(config):
    from segmentation.cameras import build_camera_source
    from segmentation.config import CameraParams as SegCameraParams

    seg_cam = SegCameraParams(
        camera_type=config.camera.camera_type,
        width=config.camera.width,
        height=config.camera.height,
        fps=config.camera.fps,
        decimation_magnitude=config.camera.decimation_magnitude,
        max_points=config.camera.max_points,
        wait_timeout_ms=config.camera.wait_timeout_ms,
        device_serial=config.camera.device_serial,
        scepter_sdk_path=config.camera.scepter_sdk_path,
        nyx650_config_path=config.camera.nyx650_config_path,
    )
    return build_camera_source(seg_cam)


def _cropped_rotated(points_xyz: np.ndarray, config) -> np.ndarray:
    rotated = rotate_xyz(points_xyz, config.crop.tilt_x_deg,
                         config.crop.tilt_y_deg, config.crop.tilt_z_deg)
    return rotated[crop_mask(points_xyz, config.crop)]


def run_stage(points_xyz: np.ndarray, stage: str, config) -> dict:
    from box.box import detect_box, box_summary
    from box.pipeline import metrics
    if stage == "box":
        result = detect_box(_cropped_rotated(points_xyz, config), config.box)
        return {"metrics": box_summary(result)}
    keep = crop_mask(points_xyz, config.crop)
    return {"metrics": metrics(points_xyz, keep)}


def launch_visualizer(points_xyz, config, config_path, stage, point_provider, color_provider=None) -> None:
    from box.visualizer import BoxVisualizer
    BoxVisualizer(
        points_xyz=points_xyz,
        config=config,
        config_path=config_path,
        point_provider=point_provider,
        color_provider=color_provider,
        stage=stage,
    ).run()


def main() -> None:
    args = parse_args()
    config_path = Path(args.config)

    if args.save_default_config:
        save_config(default_config(), config_path)
        print(f"Wrote default config to: {config_path}")
        return

    config = load_config(config_path)
    if args.camera_type:
        config.camera.camera_type = args.camera_type

    visualize = not args.no_visualize and not args.save and args.stage in ("crop", "box")

    if args.input:
        points_xyz = load_points(args.input)
        print(f"Loaded {points_xyz.shape[0]} points from: {args.input}")
        if visualize:
            launch_visualizer(points_xyz, config, config_path, args.stage, point_provider=None)
        else:
            print(json.dumps(run_stage(points_xyz, args.stage, config)["metrics"], indent=2))
        return

    with build_camera(config) as camera_source:
        points_xyz = camera_source.read_points()

        if args.save:
            from segmentation.cameras import save_color_png

            out = Path(args.save)
            out.parent.mkdir(parents=True, exist_ok=True)
            np.save(out, points_xyz)
            print(f"Saved {points_xyz.shape[0]} points to: {out}")
            png = out.with_suffix(".png")
            if save_color_png(camera_source.read_color(), png):
                print(f"Saved color image to: {png}")
            else:
                print("No color image available from camera; skipped PNG.")
            return

        if visualize:
            launch_visualizer(points_xyz, config, config_path, args.stage,
                              point_provider=camera_source.read_points,
                              color_provider=camera_source.read_color)
        else:
            print(json.dumps(run_stage(points_xyz, args.stage, config)["metrics"], indent=2))


if __name__ == "__main__":
    main()