Presa-Mano/src/presa_mano/api/perception_routes.py

"""Perception endpoints: YOLO detect + object point cloud in base frame."""

from __future__ import annotations

from typing import Any, Optional

import numpy as np
from fastapi import APIRouter, HTTPException, Query
from pydantic import BaseModel

from ..config import get_config
from ..geometry import transform_points
from ..perception import ObjectDetector, filter_outliers, nuwa_to_cv_frame
from ..timing import Stopwatch
from . import _deps

router = APIRouter()


class DetectReq(BaseModel):
    conf: Optional[float] = None
    class_name: Optional[str] = None
    max_points: int = 20000


class ScanReq(BaseModel):
    conf: Optional[float] = None
    class_name: Optional[str] = None  # if set, point cloud is computed for this class
    max_points: int = 20000
    include_points: bool = True       # set false to skip the ROI pointcloud (faster)


@router.post("/detect")
def detect_objects(req: DetectReq) -> dict[str, Any]:
    sw = Stopwatch()
    yolo = _deps.yolo()
    nuwa = _deps.nuwa()
    detector = ObjectDetector(yolo, nuwa)
    try:
        with sw.step("yolo"):
            dets = detector.detect(conf=req.conf)
    except Exception as e:
        raise HTTPException(status_code=502, detail=f"YOLO detect failed: {e}")
    if req.class_name:
        dets = [d for d in dets if (d.get("class") or "").lower() == req.class_name.lower()]
    return {"detections": dets, "timings_ms": sw.as_dict()}


@router.get("/classes")
def list_classes() -> dict[str, Any]:
    """Class names exposed by the YOLO model. Empty list if yolo-server is older
    and does not serve /classes yet — the UI should still work by deriving the
    classes from /scan detections."""
    try:
        items = _deps.yolo().classes()
    except Exception as e:
        raise HTTPException(status_code=502, detail=f"YOLO classes failed: {e}")
    return {"classes": items}


@router.post("/scan")
def scan(req: ScanReq) -> dict[str, Any]:
    """One YOLO call → all detections (sorted by confidence) + optional point
    cloud for the selected class. If class_name is provided the point cloud is
    computed for its highest-confidence detection; otherwise for the top overall
    detection."""
    sw = Stopwatch()
    yolo = _deps.yolo()
    nuwa = _deps.nuwa()
    detector = ObjectDetector(yolo, nuwa)
    try:
        with sw.step("yolo"):
            all_dets = detector.detect(conf=req.conf)
    except Exception as e:
        raise HTTPException(status_code=502, detail=f"YOLO detect failed: {e}")

    detections = sorted(
        (d for d in all_dets if (d.get("class") or "").strip()),
        key=lambda d: d["confidence"],
        reverse=True,
    )

    out: dict[str, Any] = {
        "detections": detections,
        "selected": None,
        "n_points": 0,
        "timings_ms": sw.as_dict(),
    }
    if not req.include_points or not detections:
        out["timings_ms"] = sw.as_dict()
        return out

    target: dict[str, Any] | None = None
    if req.class_name:
        target = next(
            (d for d in detections if (d.get("class") or "").lower() == req.class_name.lower()),
            None,
        )
    if target is None and not req.class_name:
        target = detections[0]
    if target is None:
        # requested class not detected in this frame
        out["timings_ms"] = sw.as_dict()
        return out

    bbox = target["bbox"]
    cfg = get_config()
    step = int(cfg.get("grasp", {}).get("pointcloud_step", 1))
    max_depth = int(cfg.get("grasp", {}).get("pointcloud_max_depth_mm", 2000))
    u_min, v_min, u_max, v_max = _bbox_to_roi(bbox)
    try:
        with sw.step("pointcloud_roi"):
            pc, _meta = nuwa.pointcloud_roi(u_min, v_min, u_max, v_max,
                                            step=step, max_depth_mm=max_depth)
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=502, detail=f"Nuwa pointcloud ROI failed: {e}")

    with sw.step("filter_outliers"):
        cropped = filter_outliers(pc)
    if len(cropped) > req.max_points:
        idx = np.random.choice(len(cropped), req.max_points, replace=False)
        cropped = cropped[idx]

    out["selected"] = target
    out["n_points"] = int(len(cropped))
    out["points_camera"] = cropped[:, :3].tolist() if len(cropped) else []

    T_base_cam = _load_calibration()
    if T_base_cam is not None and len(cropped) > 0:
        with sw.step("to_base_frame"):
            pts_cv = nuwa_to_cv_frame(cropped[:, :3])
            pts_base = transform_points(T_base_cam, pts_cv)
        out["points_base"] = pts_base.tolist()
        out["T_base_camera"] = T_base_cam.tolist()
    out["timings_ms"] = sw.as_dict()
    return out


@router.post("/object_points")
def object_points(req: DetectReq) -> dict[str, Any]:
    """Runs YOLO, takes the top detection (or class-filtered), returns its point cloud
    both in camera frame and in base frame (if calibration is available)."""
    sw = Stopwatch()
    yolo = _deps.yolo()
    nuwa = _deps.nuwa()
    detector = ObjectDetector(yolo, nuwa)
    try:
        with sw.step("yolo"):
            dets = detector.detect(conf=req.conf)
    except Exception as e:
        raise HTTPException(status_code=502, detail=f"YOLO detect failed: {e}")
    if req.class_name:
        dets = [d for d in dets if (d.get("class") or "").lower() == req.class_name.lower()]
    if not dets:
        raise HTTPException(status_code=404, detail="no object detected")
    dets.sort(key=lambda d: d["confidence"], reverse=True)
    top = dets[0]
    bbox = top["bbox"]

    cfg = get_config()
    step = int(cfg.get("grasp", {}).get("pointcloud_step", 1))
    max_depth = int(cfg.get("grasp", {}).get("pointcloud_max_depth_mm", 2000))
    u_min, v_min, u_max, v_max = _bbox_to_roi(bbox)
    try:
        with sw.step("pointcloud_roi"):
            pc, _meta = nuwa.pointcloud_roi(u_min, v_min, u_max, v_max,
                                            step=step, max_depth_mm=max_depth)
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=502, detail=f"Nuwa pointcloud ROI failed: {e}")

    with sw.step("filter_outliers"):
        cropped = filter_outliers(pc)
    if len(cropped) > req.max_points:
        idx = np.random.choice(len(cropped), req.max_points, replace=False)
        cropped = cropped[idx]

    out: dict[str, Any] = {
        "detection": top,
        "n_points": int(len(cropped)),
        "points_camera": cropped[:, :3].tolist(),
    }

    T_base_cam = _load_calibration()
    if T_base_cam is not None and len(cropped) > 0:
        with sw.step("to_base_frame"):
            pts_cv = nuwa_to_cv_frame(cropped[:, :3])
            pts_base = transform_points(T_base_cam, pts_cv)
        out["points_base"] = pts_base.tolist()
        out["T_base_camera"] = T_base_cam.tolist()
    out["timings_ms"] = sw.as_dict()
    return out


def _load_calibration() -> np.ndarray | None:
    from ..calibration import CalibrationStore
    return CalibrationStore(_deps.calibration_file_path(_deps.default_side())).T_base_camera()


@router.get("/measure")
def measure(u: int = Query(..., ge=0), v: int = Query(..., ge=0),
            radius: int = Query(6, ge=0, le=50)) -> dict[str, Any]:
    """Back-project a pixel to 3D (camera frame, and base frame if calibrated)."""
    sw = Stopwatch()
    try:
        with sw.step("nuwa_measure"):
            res = _deps.nuwa().measure(u=u, v=v, radius=radius)
    except Exception as e:
        raise HTTPException(status_code=502, detail=f"Nuwa measure failed: {e}")
    out: dict[str, Any] = {"camera": res}
    T_base_cam = _load_calibration()
    pt = _extract_xyz(res)
    if T_base_cam is not None and pt is not None:
        pt_cv = nuwa_to_cv_frame(np.asarray([pt]))[0]
        base_xyz = (T_base_cam[:3, :3] @ pt_cv + T_base_cam[:3, 3]).tolist()
        out["base"] = {"x": base_xyz[0], "y": base_xyz[1], "z": base_xyz[2]}
    out["timings_ms"] = sw.as_dict()
    return out


# ---------- depth-bias correction ----------
# Sensor depth has a non-trivial offset (see /api/perception/measure_rgb vs
# PnP on a known marker). We model it as: depth_raw = scale * depth_true +
# offset_mm and let the user calibrate `scale, offset` from samples.
# State + samples persisted in data/depth_correction.json.

class DepthCorrectionParams(BaseModel):
    enabled: Optional[bool] = None
    scale: Optional[float] = None
    offset_mm: Optional[float] = None


class DepthCorrectionSampleReq(BaseModel):
    """Add a (raw, true) sample. Either supply both directly, or set
    `from_marker=True` to acquire one automatically from the current
    marker (PnP Z taken as ground truth, raw read from depth sensor at
    the marker centre)."""
    from_marker: bool = False
    marker_id: Optional[int] = None
    size_mm: Optional[float] = None
    depth_raw_mm: Optional[int] = None
    depth_true_mm: Optional[int] = None
    note: str = ""


def _load_dc():
    from ..perception.depth_correction import load
    return load(_deps.depth_correction_path())


def _save_dc(state):
    from ..perception.depth_correction import save
    save(_deps.depth_correction_path(), state)


@router.get("/depth_correction")
def get_depth_correction() -> dict[str, Any]:
    return _load_dc()


@router.put("/depth_correction")
def set_depth_correction(req: DepthCorrectionParams) -> dict[str, Any]:
    state = _load_dc()
    if req.enabled is not None:
        state["enabled"] = bool(req.enabled)
    if req.scale is not None:
        state["scale"] = float(req.scale)
    if req.offset_mm is not None:
        state["offset_mm"] = float(req.offset_mm)
    _save_dc(state)
    return state


@router.delete("/depth_correction/samples")
def clear_dc_samples() -> dict[str, Any]:
    from ..perception.depth_correction import clear_samples
    state = clear_samples(_load_dc())
    _save_dc(state)
    return state


@router.post("/depth_correction/sample")
def add_dc_sample(req: DepthCorrectionSampleReq) -> dict[str, Any]:
    """Add a sample to the depth-correction calibration set.

    `from_marker=True`: detect the configured marker, take its PnP-derived
    Z (in the camera frame) as ground truth, and read the raw sensor depth
    at the marker centre via the same registration pipeline. The marker
    must be in view.
    """
    from ..perception.depth_correction import add_sample
    cfg = get_config()
    state = _load_dc()

    if req.from_marker:
        from ..perception.depth_to_rgb import (
            CameraRegistration, back_project_rgb_pixel, fetch_depth_mm,
        )
        marker_id = int(req.marker_id or cfg["aruco"]["wrist_marker"]["id"])
        size_mm = float(req.size_mm or cfg["aruco"]["wrist_marker"]["size_mm"])
        try:
            mres = _deps.marker().detect_url(
                markers=[{"id": marker_id, "size_mm": size_mm}],
                image_url=_deps.nuwa().rgb_url(),
            )
        except Exception as e:
            raise HTTPException(status_code=502, detail=f"marker-finder failed: {e}")
        entry = next((m for m in mres.get("found", []) if int(m["id"]) == marker_id), None)
        if entry is None:
            raise HTTPException(status_code=422, detail=f"marker id={marker_id} non in vista")
        pnp_z_m = float(np.asarray(entry["position_m"], dtype=np.float64)[2])
        # Sample the raw depth at the centre pixel via the registration pipeline
        # WITHOUT correction (we want the uncorrected sensor reading).
        center_uv = np.asarray(entry["corners"], dtype=np.float64).reshape(-1, 2).mean(axis=0)
        nuwa = _deps.nuwa()
        reg = CameraRegistration.from_nuwa(nuwa.intrinsics())
        depth = fetch_depth_mm(nuwa)
        r = back_project_rgb_pixel(center_uv[0], center_uv[1], depth, reg,
                                   sample_radius=4, depth_correction=None)
        if not r.get("ok"):
            raise HTTPException(status_code=422,
                                detail=f"depth non disponibile al marker: {r.get('error')}")
        depth_raw = int(r["depth_mm_raw"])
        depth_true = int(round(pnp_z_m * 1000.0))
        note = req.note or f"from_marker id={marker_id} size={size_mm}mm pnp_z={depth_true}mm"
    else:
        if req.depth_raw_mm is None or req.depth_true_mm is None:
            raise HTTPException(status_code=422,
                                detail="senza from_marker servono depth_raw_mm e depth_true_mm")
        depth_raw = int(req.depth_raw_mm)
        depth_true = int(req.depth_true_mm)
        note = req.note

    state = add_sample(state, depth_raw_mm=depth_raw, depth_true_mm=depth_true, note=note)
    _save_dc(state)
    return {"ok": True, "n_samples": len(state["samples"]),
            "added": {"depth_raw_mm": depth_raw, "depth_true_mm": depth_true,
                      "delta_mm": depth_raw - depth_true, "note": note}}


@router.post("/depth_correction/fit")
def fit_dc() -> dict[str, Any]:
    """Fit the affine model from collected samples and (optionally) enable it.
    Updates `scale` and `offset_mm` in the persisted state. Does not change
    `enabled` — call PUT /depth_correction with `enabled=true` to activate."""
    from ..perception.depth_correction import fit_affine
    state = _load_dc()
    samples = state.get("samples") or []
    try:
        fit = fit_affine(samples)
    except ValueError as e:
        raise HTTPException(status_code=422, detail=str(e))
    state["scale"] = fit["scale"]
    state["offset_mm"] = fit["offset_mm"]
    state["fit_meta"] = fit
    _save_dc(state)
    return state


@router.get("/measure_rgb")
def measure_rgb(u: int = Query(..., ge=0, description="pixel column in the RGB image"),
                v: int = Query(..., ge=0, description="pixel row in the RGB image"),
                radius: int = Query(4, ge=0, le=50)) -> dict[str, Any]:
    """Back-project an **RGB** pixel to 3D, properly handling IR↔RGB
    registration.

    `Nuwa /api/measure` interprets `(u,v)` as IR-image pixel coords and
    back-projects with IR intrinsics; for an RGB pixel that yields a wrong
    point. This endpoint:
      1. samples depth at the corresponding IR pixel (computed iteratively
         from the RGB ray + extrinsics);
      2. back-projects with IR intrinsics;
      3. transforms to the RGB camera frame.

    Output is in the **OpenCV-convention RGB camera frame** (X-right, Y-down,
    Z-forward), the same frame the hand-eye matrix `T_base_camera` expects.
    """
    sw = Stopwatch()
    nuwa = _deps.nuwa()
    try:
        from ..perception.depth_to_rgb import (
            CameraRegistration, back_project_rgb_pixel, fetch_depth_mm,
        )
        with sw.step("intrinsics"):
            reg = CameraRegistration.from_nuwa(nuwa.intrinsics())
        with sw.step("depth_fetch"):
            depth = fetch_depth_mm(nuwa)
        with sw.step("back_project"):
            r = back_project_rgb_pixel(u, v, depth, reg, sample_radius=radius,
                                       depth_correction=_load_dc())
    except Exception as e:
        raise HTTPException(status_code=502, detail=f"measure_rgb failed: {e}")
    if not r["ok"]:
        return {"found": False, "rgb_pixel": [u, v], "error": r.get("error"),
                "timings_ms": sw.as_dict()}

    pt_rgb_m = r["point_rgb_m"]
    out: dict[str, Any] = {
        "found": True,
        "rgb_pixel": r["rgb_pixel"],
        "ir_pixel": r["ir_pixel"],
        "depth_mm": r["depth_mm"],
        "point_rgb_m": {"x": pt_rgb_m[0], "y": pt_rgb_m[1], "z": pt_rgb_m[2]},
        "point_ir_m":  {"x": r["point_ir_m"][0], "y": r["point_ir_m"][1], "z": r["point_ir_m"][2]},
        "distance_m": float(np.linalg.norm(pt_rgb_m)),
    }
    T_base_cam = _load_calibration()
    if T_base_cam is not None:
        pt_arr = np.asarray(pt_rgb_m)
        base_xyz = (T_base_cam[:3, :3] @ pt_arr + T_base_cam[:3, 3]).tolist()
        out["point_base_m"] = {"x": base_xyz[0], "y": base_xyz[1], "z": base_xyz[2]}
    out["timings_ms"] = sw.as_dict()
    return out


def _bbox_to_roi(bbox: dict[str, Any]) -> tuple[int, int, int, int]:
    """Convert a YOLO bbox {x1,y1,x2,y2} to integer ROI (u_min,v_min,u_max,v_max)."""
    u_min = max(0, int(np.floor(bbox["x1"])))
    v_min = max(0, int(np.floor(bbox["y1"])))
    u_max = int(np.ceil(bbox["x2"]))
    v_max = int(np.ceil(bbox["y2"]))
    if u_max <= u_min:
        u_max = u_min + 1
    if v_max <= v_min:
        v_max = v_min + 1
    return u_min, v_min, u_max, v_max


def _extract_xyz(res: Any) -> list[float] | None:
    if not isinstance(res, dict):
        return None
    for key in ("point", "xyz", "camera_xyz"):
        v = res.get(key)
        if isinstance(v, (list, tuple)) and len(v) >= 3:
            return [float(v[0]), float(v[1]), float(v[2])]
    if all(k in res for k in ("x", "y", "z")):
        return [float(res["x"]), float(res["y"]), float(res["z"])]
    return None