PythonDetection/carpet_projection.py at main · Rukki13/PythonDetection · GitHub

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
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
"""Project carpet grab geometry (meters, tilt-corrected frame) onto the color
image, using the Vzense color intrinsics + depth->color extrinsics.

Pure numpy: no camera, no I/O. See
docs/superpowers/specs/2026-07-17-carpet-grabline-png-overlay-design.md.
"""
from __future__ import annotations
import numpy as np


def _tilt_matrix(x_deg: float, y_deg: float, z_deg: float) -> np.ndarray:
    """Rebuild the SAME rotation carpet.pipeline.rotate_xyz composes (Rz@Ry@Rx)."""
    rx, ry, rz = np.radians([x_deg, y_deg, z_deg])
    cx, sx = np.cos(rx), np.sin(rx)
    cy, sy = np.cos(ry), np.sin(ry)
    cz, sz = np.cos(rz), np.sin(rz)
    Rx = np.array([[1, 0, 0], [0, cx, -sx], [0, sx, cx]])
    Ry = np.array([[cy, 0, sy], [0, 1, 0], [-sy, 0, cy]])
    Rz = np.array([[cz, -sz, 0], [sz, cz, 0], [0, 0, 1]])
    return Rz @ Ry @ Rx


def _undo_tilt(pts_tilted: np.ndarray, tilt: dict) -> np.ndarray:
    """rotate_xyz does pts @ R.T, so the inverse is pts @ R."""
    R = _tilt_matrix(tilt.get("x_deg", 0.0), tilt.get("y_deg", 0.0), tilt.get("z_deg", 0.0))
    return pts_tilted @ R


def project_points(points_m, calib: dict) -> dict:
    ci = calib["color_intrinsics"]
    ext = calib["depth_to_color"]
    tilt = calib.get("tilt", {})
    R = np.asarray(ext["R"], float).reshape(3, 3)
    t = np.asarray(ext["t"], float)                      # mm
    pts = np.asarray(points_m, float).reshape(-1, 3)

    depth_mm = _undo_tilt(pts, tilt) * 1000.0            # m -> mm, raw depth frame
    Pc = depth_mm @ R.T + t                              # depth -> color frame
    Z = Pc[:, 2]
    valid = Z > 1e-6
    Zs = np.where(valid, Z, 1.0)
    xn = Pc[:, 0] / Zs
    yn = Pc[:, 1] / Zs
    r2 = xn * xn + yn * yn
    r4 = r2 * r2
    r6 = r4 * r2
    k1, k2, p1, p2, k3, k4, k5, k6 = ci["dist"]
    radial = (1 + k1 * r2 + k2 * r4 + k3 * r6) / (1 + k4 * r2 + k5 * r4 + k6 * r6)
    xd = xn * radial + 2 * p1 * xn * yn + p2 * (r2 + 2 * xn * xn)
    yd = yn * radial + p1 * (r2 + 2 * yn * yn) + 2 * p2 * xn * yn
    u = ci["fx"] * xd + ci["cx"]
    v = ci["fy"] * yd + ci["cy"]
    w, h = ci["width"], ci["height"]
    inframe = valid & (u >= 0) & (u < w) & (v >= 0) & (v < h)
    return {"px": np.stack([u, v], axis=1), "valid": inframe}


def _r(v):
    return [round(float(v[0]), 1), round(float(v[1]), 1)]


def project_grab(grab, calib: dict, arrow_len: float = 0.08):
    if not grab or "midpoint" not in grab:
        return None
    mid = np.asarray(grab["midpoint"], float)
    dirv = np.asarray(grab["direction"], float)
    if grab.get("polyline"):
        line_pts = np.asarray(grab["polyline"], float)
    else:
        half = 0.5 * float(grab.get("length", 0.0) or 0.0)
        line_pts = np.asarray([mid - half * dirv, mid + half * dirv])
    arrow_end = mid + arrow_len * dirv
    n_line = len(line_pts)
    allpts = np.vstack([mid[None, :], line_pts, arrow_end[None, :]])
    out = project_points(allpts, calib)
    px, valid = out["px"], out["valid"]
    ci = calib["color_intrinsics"]
    return {
        "grab_px": _r(px[0]),
        "line_px": [_r(px[1 + i]) for i in range(n_line)],
        "dir_px": _r(px[1 + n_line]),
        "valid": bool(valid[0]),
        "img_w": int(ci["width"]),
        "img_h": int(ci["height"]),
    }


def nearest_point_to_pixel(points_m, calib, px, max_px=14.0):
    """Photo pixel -> nearest in-frame cloud point. Returns (index, [x,y,z]) in
    the same (tilt-corrected) frame as points_m, or (None, None) if the closest
    projected point is farther than max_px or nothing projects in-frame."""
    pts = np.asarray(points_m, float).reshape(-1, 3)
    if len(pts) == 0:
        return None, None
    proj = project_points(pts, calib)
    valid = proj["valid"]
    if not bool(valid.any()):
        return None, None
    pix = proj["px"]
    du = pix[:, 0] - float(px[0])
    dv = pix[:, 1] - float(px[1])
    d2 = np.where(valid, du * du + dv * dv, np.inf)
    idx = int(np.argmin(d2))
    if not np.isfinite(d2[idx]) or d2[idx] > float(max_px) * float(max_px):
        return None, None
    return idx, [round(float(v), 4) for v in pts[idx]]


def manual_grab_from_polyline(polyline_m, grab_index):
    """Ordered tilt-corrected 3D points + chosen index -> grab pose mirroring the
    detector's grab dict. Reuses the detector's _look_dir so the approach axis
    convention is identical to an auto grab."""
    from carpet.furrow import _look_dir  # read-only reuse of detector convention
    poly = np.asarray(polyline_m, float).reshape(-1, 3)
    n = len(poly)
    if n == 0:
        raise ValueError("empty polyline")
    i = max(0, min(int(grab_index), n - 1))
    mid = poly[i]
    if n == 1:
        direction = np.array([1.0, 0.0, 0.0])
    elif i == 0:
        direction = poly[1] - poly[0]
    elif i == n - 1:
        direction = poly[-1] - poly[-2]
    else:
        direction = poly[i + 1] - poly[i - 1]
    dn = float(np.linalg.norm(direction))
    direction = np.array([1.0, 0.0, 0.0]) if dn < 1e-9 else direction / dn
    look = _look_dir(direction)
    yaw = float(np.degrees(np.arctan2(direction[2], direction[0])) % 180.0)
    length = float(np.linalg.norm(poly[-1] - poly[0])) if n >= 2 else 0.0
    return {
        "midpoint": [round(float(v), 4) for v in mid],
        "direction": [round(float(v), 4) for v in direction],
        "look": [round(float(v), 4) for v in look],
        "yaw_deg": round(yaw, 2),
        "length": round(length, 4),
        "depth": round(float(mid[2]), 4),
        "polyline": [[round(float(v), 4) for v in p] for p in poly],
    }