diff --git a/src/presa_mano/api/grasp_routes.py b/src/presa_mano/api/grasp_routes.py index 2cb541c..ab9e1d8 100644 --- a/src/presa_mano/api/grasp_routes.py +++ b/src/presa_mano/api/grasp_routes.py @@ -2,20 +2,31 @@ from __future__ import annotations +import base64 from typing import Any, Optional import numpy as np from fastapi import APIRouter, HTTPException -from pydantic import BaseModel +from pydantic import BaseModel, Field from ..calibration import CalibrationStore -from ..clients import G1Unavailable +from ..clients import G1Unavailable, Open3DModelError, Open3DModelUnavailable from ..config import get_config from ..geometry import transform_points from ..grasp import GraspPlanner, compute_axes from ..perception import ObjectDetector, filter_outliers, nuwa_to_cv_frame -from ..perception.sam21_bridge import Sam21Fallback, points_via_sam21 -from ..perception.work_plane import apply_below_plane_filter, get_work_plane +from ..perception.projection import K_from_intrinsics, project_camera_to_pixels +from ..perception.sam21_bridge import ( + Sam21Fallback, + points_via_sam21, + segment_with_sam21, +) +from ..perception.stl_sampler import sample_surface +from ..perception.work_plane import ( + apply_below_plane_filter, + get_work_plane, + partition_by_plane, +) from .perception_routes import _bbox_to_roi from ..timing import Stopwatch from . import _deps @@ -172,3 +183,154 @@ def execute(req: ExecuteReq) -> dict[str, Any]: raise HTTPException(status_code=502, detail=f"G1 unavailable: {e}") return {"plan": g.to_dict(), "steps": steps, "timings_ms": sw.as_dict(), **meta} + +class FindModelReq(BaseModel): + recipe_id: str = Field(..., min_length=1) + conf: Optional[float] = None + class_name: Optional[str] = None + + +_FIT_DOMAIN_CODES = { + "REGISTRATION_FAILED", "EMPTY_CLOUD", "MODEL_MISSING", + "INVALID_MESH", "NO_CORRESPONDENCES", +} + + +@router.post("/find_model") +def find_model(req: FindModelReq) -> dict[str, Any]: + sw = Stopwatch() + cfg = get_config() + fm_cfg = cfg.get("grasp", {}).get("find_model", {}) + min_fitness = float(fm_cfg.get("min_fitness", 0.5)) + max_rmse_mm = float(fm_cfg.get("max_rmse_mm", 5.0)) + sample_n = int(fm_cfg.get("sample_points", 2000)) + + side = cfg.get("robot", {}).get("wrist_side", "right") + T_base_cam = CalibrationStore(_deps.calibration_file_path(side)).T_base_camera() + if T_base_cam is None: + raise HTTPException(status_code=412, detail=f"calibration missing for side='{side}'") + + nuwa = _deps.nuwa() + yolo = _deps.yolo() + sam21 = _deps.sam21() + o3d = _deps.open3dmodel() + + # 1. YOLO + detector = ObjectDetector(yolo, nuwa) + try: + with sw.step("yolo"): + dets = detector.detect(conf=req.conf) + except HTTPException: + raise + 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 = _bbox_to_roi(top["bbox"]) + + # 2. SAM21 (no plane filter) + try: + seg = segment_with_sam21(nuwa, sam21, bbox, sw) + except Sam21Fallback as e: + raise HTTPException(status_code=422, detail=f"SAM21 fail: {e}") + + # 3. Partition by plane + with sw.step("partition_plane"): + accepted_m, excluded_m, plane_info = partition_by_plane( + seg.points_cv_m, get_work_plane() + ) + if accepted_m.shape[0] < 10: + raise HTTPException( + status_code=422, + detail=f"too few accepted points after plane filter: {accepted_m.shape[0]}" + f" (before plane: {plane_info['n_points_before_plane']})" + ) + + # 4. Open3DModel pose + fit_failed = False + fit_reason: str | None = None + pose_data: dict[str, Any] = {} + try: + with sw.step("open3d_pose"): + pose_data = o3d.estimate_pose(req.recipe_id, accepted_m, seg.intrinsics_rgb) + except Open3DModelUnavailable as e: + raise HTTPException(status_code=502, detail=f"Open3DModel unreachable: {e}") + except Open3DModelError as e: + if e.code == "RECIPE_NOT_FOUND": + raise HTTPException(status_code=404, detail=f"recipe_id not found: {req.recipe_id}") + if e.code in {"HTTP_401", "HTTP_403"}: + raise HTTPException(status_code=502, detail=f"Open3DModel auth: {e.message}") + if e.code in _FIT_DOMAIN_CODES: + fit_failed = True + fit_reason = f"{e.code}: {e.message}" + else: + raise HTTPException(status_code=502, detail=f"Open3DModel error: {e.code}: {e.message}") + + fitness = float(pose_data.get("fitness", 0.0)) + rmse_mm = float(pose_data.get("inlier_rmse_mm", 1e9)) + transform_4x4 = pose_data.get("transform_4x4") + + if not fit_failed: + if fitness < min_fitness: + fit_failed = True + fit_reason = f"low fitness {fitness:.2f} < {min_fitness}" + elif rmse_mm > max_rmse_mm: + fit_failed = True + fit_reason = f"rmse {rmse_mm:.2f}mm > {max_rmse_mm}" + + # 5. Mesh + sample + project + W = int(seg.intrinsics_rgb["width"]) + H = int(seg.intrinsics_rgb["height"]) + K = K_from_intrinsics(seg.intrinsics_rgb) + + model_px = np.empty((0, 2), dtype=np.int32) + if not fit_failed and transform_4x4 is not None: + try: + with sw.step("recipe_mesh"): + mesh = _deps.recipe_cache().get(req.recipe_id, client=o3d) + with sw.step("sample"): + model_pts_mm = sample_surface(mesh.vertices, mesh.triangles, sample_n) + with sw.step("transform"): + T = np.asarray(transform_4x4, dtype=np.float64) + # transform_points expects (4,4) acting on column vectors -> use directly + model_pts_cam_mm = transform_points(T, model_pts_mm) + with sw.step("project_model"): + model_px = project_camera_to_pixels(model_pts_cam_mm, K, (W, H)) + except (Open3DModelError, Open3DModelUnavailable) as e: + raise HTTPException(status_code=502, detail=f"STL fetch/parse failed: {e}") + + with sw.step("project_points"): + accepted_px = project_camera_to_pixels(accepted_m, K, (W, H)) + excluded_px = project_camera_to_pixels(excluded_m, K, (W, H)) + + with sw.step("encode_rgb"): + rgb_b64 = base64.b64encode(seg.rgb_jpeg).decode("ascii") + + fit_ok = not fit_failed + + return { + "ok": True, + "rgb_jpeg_b64": rgb_b64, + "rgb_size": [W, H], + "detection": top, + "recipe": {"id": req.recipe_id, + "name": pose_data.get("recipe_name"), + "symmetry_type": pose_data.get("symmetry_type")}, + "model_points_px": model_px.tolist(), + "accepted_points_px": accepted_px.tolist(), + "excluded_points_px": excluded_px.tolist(), + "n_accepted": int(accepted_m.shape[0]), + "n_excluded": int(excluded_m.shape[0]), + "plane_applied": bool(plane_info["plane_applied"]), + "fit_ok": fit_ok, + "fit_failed": fit_failed, + "fit_reason": fit_reason, + "fitness": fitness if not fit_failed or pose_data else None, + "inlier_rmse_mm": rmse_mm if not fit_failed or pose_data else None, + "transform_camera_4x4": transform_4x4, + "timings_ms": sw.as_dict(), + }