carpet_detection.py

"""Entry point for carpet detection (side-view layer separation).

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

    # Run the pipeline on a saved cloud and print the separated layers:
    python carpet_detection.py --input captures/stack01.npy --stage layer

    # Live crease / pick-edge tuner on the camera (move the carpet, it re-detects):
    python carpet_detection.py --camera-type vzense --stage sideview

    # Same on a saved cloud (no camera):
    python carpet_detection.py --input captures/side01.npy --stage sideview
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np

from carpet.config import default_config, load_config, save_config
from carpet.pipeline import (
    crop_mask,
    rotate_xyz,
    run_crop_only,
    run_crop_then_plane,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Carpet detection: side-view layer separation")
    parser.add_argument(
        "--stage",
        choices=["crop", "plane", "layer", "sideview", "topcarpet"],
        default="topcarpet",
        help="Pipeline stage (default: topcarpet = isolate the top carpet by side curvature + pose)",
    )
    parser.add_argument(
        "--input",
        help="Run on a saved point cloud (.npy of shape (N,3), or .ply) instead of the camera",
    )
    parser.add_argument(
        "--save",
        help="Capture one frame from the camera, save the cloud to this path (.npy) "
             "plus the 2D camera view as a .png beside it, 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; just print metrics once.",
    )
    parser.add_argument(
        "--config",
        default=str(Path("config") / "carpet_params.json"),
        help="Path to JSON config for carpet 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:
    """Load an (N,3) point cloud from .npy or .ply."""
    p = Path(path)
    if p.suffix.lower() == ".npy":
        pts = np.load(p)
    elif p.suffix.lower() == ".ply":
        import open3d as o3d  # imported lazily; only needed for .ply
        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):
    """Build (but do not start) the shared camera source from carpet 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:
    """Tilt-correct then crop to the ROI — the shared first step that cuts the
    point count before the top-carpet / side-view detectors run."""
    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):
    if stage == "topcarpet":
        from carpet.topcarpet import extract_top_carpet, top_carpet_summary
        result = extract_top_carpet(_cropped_rotated(points_xyz, config), config.top_carpet)
        return {"metrics": top_carpet_summary(result)}
    if stage == "sideview":
        from carpet.sideview import detect_side_view, result_summary
        result = detect_side_view(_cropped_rotated(points_xyz, config), config.side_view)
        return {"metrics": result_summary(result)}
    if stage == "layer":
        return run_crop_then_plane(points_xyz, config)
    if stage == "plane":
        return run_crop_then_plane(points_xyz, config)
    return run_crop_only(points_xyz, config)


def launch_visualizer(points_xyz, config, config_path, stage, point_provider, color_provider=None) -> None:
    from carpet.visualizer import CarpetVisualizer
    CarpetVisualizer(
        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

    # The GUI tuner covers every stage, including crop/plane for tuning the ROI.
    visualize = not args.no_visualize and not args.save and args.stage in (
        "crop", "plane", "layer", "sideview", "topcarpet"
    )

    # Offline: run on a saved cloud (no camera needed).
    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

    # Live: open the camera and keep it open for streaming / capture.
    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()