Spaces:
Running
Running
Download processor/app/processing/modules/rgb_world_tracking.py from Stouching-Robotics/DataFlowing: direct link, hf CLI and curl.
- Browser
- Download file 18.9 kB
-
https://huggingface.co/spaces/Stouching-Robotics/DataFlowing/resolve/main/processor/app/processing/modules/rgb_world_tracking.py
- Command line
-
hf download hf://spaces/Stouching-Robotics/DataFlowing/processor/app/processing/modules/rgb_world_tracking.py
-
curl -L -o rgb_world_tracking.py https://huggingface.co/spaces/Stouching-Robotics/DataFlowing/resolve/main/processor/app/processing/modules/rgb_world_tracking.py
18.9 kB
| """RGB-only hand-to-camera tracking used by the Hand Skeleton module. | |
| This module is deliberately self-contained. It does not import the legacy | |
| ``Python`` project or any code outside ``processor``. | |
| MediaPipe gives us two useful pieces of information from an RGB frame: | |
| * normalized image landmarks (the 2D observation), and | |
| * a hand-local 3D landmark model (the shape used as the PnP object model). | |
| The PnP solution gives a stable camera-relative placement when the model and | |
| the image points are good enough. A scale-from-image fallback keeps the hand | |
| visible when PnP is temporarily ill-conditioned. The result is intentionally | |
| named ``rgb_estimated_meters``: it is a camera-relative estimate, not a depth | |
| sensor measurement and must not be presented as metric ground truth. | |
| """ | |
| from __future__ import annotations | |
| from dataclasses import dataclass | |
| import cv2 | |
| import numpy as np | |
| # MediaPipe image landmarks expose ``z`` as an optional image-space hint. | |
| # Some RGB/black-glove detections provide a valid 21x3 array but set every | |
| # image-space z to zero. Feeding that planar model to PnP produces a hand | |
| # whose joints have identical depth, so keep a small anatomical depth prior | |
| # as a visual/pose-estimation fallback. This is deliberately relative and | |
| # must not be interpreted as measured depth. | |
| _RGB_FALLBACK_Z_PROFILE = np.array([ | |
| 0.000, # wrist | |
| 0.004, 0.009, 0.014, 0.018, # thumb | |
| 0.000, 0.006, 0.012, 0.017, # index | |
| 0.000, 0.007, 0.014, 0.020, # middle | |
| 0.000, 0.006, 0.012, 0.017, # ring | |
| 0.000, 0.005, 0.010, 0.014, # pinky | |
| ], dtype=np.float64) | |
| def _cohere_thumb_depth_direction(points: np.ndarray) -> np.ndarray: | |
| """Keep the RGB thumb on the same relative depth side as the fingers. | |
| RGB image landmarks do not provide reliable metric Z. Depending on the | |
| detector/model, the thumb's image-space Z hint can have the opposite sign | |
| from the four finger chains. That produces a visually mirrored thumb in | |
| the 3D preview even though the 2D keypoints are correct. Only the thumb | |
| chain's relative Z is corrected; its X/Y image geometry and all measured | |
| depth paths remain untouched. | |
| """ | |
| out = np.asarray(points, dtype=np.float64).copy() | |
| finger_z = out[[8, 12, 16, 20], 2] | |
| thumb_z = out[[1, 2, 3, 4], 2] | |
| finger_median = float(np.median(finger_z)) | |
| thumb_median = float(np.median(thumb_z)) | |
| if (abs(finger_median) > 1e-5 and abs(thumb_median) > 1e-5 | |
| and finger_median * thumb_median < 0.0): | |
| out[1:5, 2] *= -1.0 | |
| return out | |
| def camera_matrix(width: int, height: int, focal_scale: float = 0.9) -> np.ndarray: | |
| """Return a useful pinhole approximation when no camera calibration exists.""" | |
| w = max(1, int(width)) | |
| h = max(1, int(height)) | |
| focal = max(1.0, float(focal_scale) * max(w, h)) | |
| return np.array([[focal, 0.0, w / 2.0], | |
| [0.0, focal, h / 2.0], | |
| [0.0, 0.0, 1.0]], dtype=np.float64) | |
| class _OneEuro: | |
| """Small vector One-Euro filter for camera position.""" | |
| def __init__(self, freq_min: float = 5.0, beta: float = 0.05, | |
| dcutoff: float = 1.0): | |
| self.freq_min = max(0.1, float(freq_min)) | |
| self.beta = max(0.0, float(beta)) | |
| self.dcutoff = max(0.1, float(dcutoff)) | |
| self.prev_x: np.ndarray | None = None | |
| self.prev_dx = np.zeros(3, dtype=np.float64) | |
| self.prev_ts: float | None = None | |
| def _alpha(cutoff: float, dt: float) -> float: | |
| tau = 1.0 / (2.0 * np.pi * max(1e-9, cutoff)) | |
| return float(dt / (dt + tau)) | |
| def __call__(self, value: np.ndarray, timestamp_s: float) -> np.ndarray: | |
| x = np.asarray(value, dtype=np.float64) | |
| if self.prev_x is None or self.prev_ts is None: | |
| self.prev_x = x.copy() | |
| self.prev_ts = float(timestamp_s) | |
| self.prev_dx.fill(0.0) | |
| return x.copy() | |
| dt = max(1e-5, float(timestamp_s) - self.prev_ts) | |
| dx = (x - self.prev_x) / dt | |
| alpha_d = self._alpha(self.dcutoff, dt) | |
| dx_hat = alpha_d * dx + (1.0 - alpha_d) * self.prev_dx | |
| cutoff = self.freq_min + self.beta * np.abs(dx_hat) | |
| alpha = np.array([self._alpha(float(c), dt) for c in cutoff]) | |
| filtered = alpha * x + (1.0 - alpha) * self.prev_x | |
| self.prev_x = filtered | |
| self.prev_dx = dx_hat | |
| self.prev_ts = float(timestamp_s) | |
| return filtered.copy() | |
| class _Track: | |
| position: np.ndarray | None = None | |
| timestamp_s: float | None = None | |
| velocity: np.ndarray = None # type: ignore[assignment] | |
| position_filter: _OneEuro | None = None | |
| def __post_init__(self): | |
| if self.velocity is None: | |
| self.velocity = np.zeros(3, dtype=np.float64) | |
| class RGBWorldTracker: | |
| """Estimate hand camera-relative 3D from one RGB stream.""" | |
| def __init__(self, width: int, height: int, *, camera_matrix_=None, | |
| distortion=None, focal_scale: float = 0.9, | |
| freq_min: float = 5.0, beta: float = 0.05, | |
| max_depth: float = 2.5, min_depth: float = 0.20): | |
| self.width = max(1, int(width)) | |
| self.height = max(1, int(height)) | |
| self.K = (np.asarray(camera_matrix_, dtype=np.float64).reshape(3, 3) | |
| if camera_matrix_ is not None else | |
| camera_matrix(self.width, self.height, focal_scale)) | |
| self.dist = (np.asarray(distortion, dtype=np.float64).reshape(-1, 1) | |
| if distortion is not None else np.zeros((5, 1), np.float64)) | |
| self.freq_min = float(freq_min) | |
| self.beta = float(beta) | |
| self.max_depth = float(max_depth) | |
| self.min_depth = float(min_depth) | |
| self._tracks: dict[str, _Track] = {} | |
| # RGB handedness can flip briefly during fast motion or self-occlusion. | |
| # Keep a small per-track vote state so the preview root anchor does | |
| # not jump from the left hand slot to the right hand slot. | |
| self._hand_labels: dict[str, dict[str, object]] = {} | |
| def stabilize_handedness(self, key: str, label: str | None, | |
| votes_required: int = 3) -> str: | |
| """Return a short-term latched left/right label for one RGB track.""" | |
| value = str(label or "").strip().lower() | |
| if value not in {"left", "right"}: | |
| state = self._hand_labels.get(str(key)) | |
| return str(state["stable"]).title() if state else "" | |
| track_key = str(key) | |
| state = self._hand_labels.setdefault( | |
| track_key, {"stable": value, "candidate": value, "count": 0}) | |
| stable = str(state["stable"]) | |
| if value == stable: | |
| state["candidate"] = value | |
| state["count"] = 0 | |
| else: | |
| if state.get("candidate") == value: | |
| state["count"] = int(state.get("count", 0)) + 1 | |
| else: | |
| state["candidate"] = value | |
| state["count"] = 1 | |
| if int(state["count"]) >= max(1, int(votes_required)): | |
| state["stable"] = value | |
| state["candidate"] = value | |
| state["count"] = 0 | |
| return str(state["stable"]).title() | |
| def normalize_local_world(points) -> np.ndarray | None: | |
| """Center MediaPipe world landmarks on the wrist. | |
| The Y/Z flips put the points in the display convention used by the | |
| project: X right, Y up, Z forward. The returned model remains in the | |
| original hand scale (approximately metres), which is useful for PnP. | |
| """ | |
| arr = np.asarray(points, dtype=np.float64) | |
| if arr.shape != (21, 3) or not np.isfinite(arr).all(): | |
| return None | |
| out = arr - arr[0] | |
| # MediaPipe can return a correctly shaped all-zero world model when | |
| # the optional world-landmark estimate is unavailable. It is not a | |
| # valid PnP object model; returning None lets the caller construct the | |
| # independent RGB/image-scale model instead. | |
| if (float(np.max(np.abs(out))) <= 1e-6 | |
| or float(np.ptp(out[:, 2])) <= 1e-5): | |
| return None | |
| out[:, 1] *= -1.0 | |
| out[:, 2] *= -1.0 | |
| # MediaPipe world landmarks are still an RGB-relative estimate here, | |
| # not a depth-sensor measurement. Keep its thumb depth direction | |
| # consistent with the four finger chains for the shared preview. | |
| return _cohere_thumb_depth_direction(out) | |
| def orient_palm_facing(points, handedness: str | None) -> np.ndarray: | |
| """Orient an RGB-relative model toward its semantic palm side. | |
| RGB landmarks do not contain a reliable absolute palm/back signal. | |
| For the shared estimated-3D convention, keep the four finger chains | |
| on one forward depth side of the wrist. This prevents an RGB/PnP | |
| mirror solution from making the fingers point away from the hand | |
| while preserving every landmark's X/Y image geometry. This is a | |
| display/estimated-3D convention, not a depth-sensor measurement. | |
| """ | |
| out = np.asarray(points, dtype=np.float64).copy() | |
| if out.shape != (21, 3) or not np.isfinite(out).all(): | |
| return out | |
| # Keep the handedness argument in this shared API because callers | |
| # already have it and future calibration can use it. It does not | |
| # change the relative-Z sign: left and right hands use the same | |
| # camera-relative depth convention. | |
| _ = str(handedness or "").strip().lower() | |
| # Select one stable RGB convention: fingertips lie on the positive | |
| # relative-Z side of the wrist/palm. This is the direction used by | |
| # the fallback anatomical prior and removes the mirrored PnP choice. | |
| finger_depth = float(np.median(out[[8, 12, 16, 20], 2]) - out[0, 2]) | |
| if finger_depth < -1e-6: | |
| out[:, 2] *= -1.0 | |
| return _cohere_thumb_depth_direction(out) | |
| def local_from_image(points_2d) -> np.ndarray | None: | |
| """Build a non-planar relative hand model when World Landmarks fail. | |
| The 2D detector's ``p.z`` is not guaranteed to be populated. When | |
| it is all zero/constant, using it as depth creates a flat skeleton. | |
| Preserve a real image-space z signal when available; otherwise use a | |
| small topology-based prior so PnP and the right-side preview retain a | |
| relative third dimension. This remains RGB-estimated 3D, not depth | |
| ground truth. | |
| """ | |
| arr = np.asarray(points_2d, dtype=np.float64) | |
| if arr.shape != (21, 3) or not np.isfinite(arr[:, :2]).all(): | |
| return None | |
| out = np.zeros((21, 3), dtype=np.float64) | |
| out[:, 0] = (arr[:, 0] - arr[0, 0]) * 0.22 | |
| out[:, 1] = -(arr[:, 1] - arr[0, 1]) * 0.22 | |
| image_z = (-(arr[:, 2] - arr[0, 2]) * 0.08 | |
| if np.isfinite(arr[:, 2]).all() else None) | |
| if image_z is not None and float(np.ptp(image_z)) > 1e-5: | |
| out[:, 2] = image_z | |
| else: | |
| image_span = max( | |
| float(np.ptp(arr[:, 0])), float(np.ptp(arr[:, 1])), 1e-3) | |
| # Scale the prior with the observed hand size, while keeping it | |
| # close to the x/y model scale and bounded for distant hands. | |
| prior_scale = float(np.clip(image_span / 0.14, 0.55, 1.45)) | |
| out[:, 2] = _RGB_FALLBACK_Z_PROFILE * prior_scale | |
| return _cohere_thumb_depth_direction(out) | |
| def black_glove_reference_local(points_2d) -> np.ndarray | None: | |
| """Build the same RGB local model used by Black Glove RGB3D. | |
| The bare-hand MediaPipe path also exposes ``p.z`` and world | |
| landmarks, but those values use a different local orientation from | |
| the YOLO/RTMPose black-glove path. Feeding that optional z hint made | |
| the two RGB workflows render different palm directions. The | |
| black-glove reference derives depth from the shared anatomical prior | |
| and observed 2D geometry only. | |
| """ | |
| arr = np.asarray(points_2d, dtype=np.float64).copy() | |
| if arr.shape != (21, 3): | |
| return None | |
| arr[:, 2] = 0.0 | |
| return RGBWorldTracker.local_from_image(arr) | |
| def _pixels(self, points_2d) -> np.ndarray | None: | |
| arr = np.asarray(points_2d, dtype=np.float64) | |
| if arr.shape != (21, 3) or not np.isfinite(arr[:, :2]).all(): | |
| return None | |
| return arr[:, :2] * np.array([self.width, self.height], dtype=np.float64) | |
| def _image_position(self, pixels: np.ndarray, | |
| local: np.ndarray) -> np.ndarray | None: | |
| # Estimate hand depth from the median ratio of corresponding 3D and | |
| # image bone lengths. This remains useful when PnP loses an inlier. | |
| model_pairs = ((0, 5), (0, 9), (0, 13), (0, 17), | |
| (5, 9), (9, 13), (13, 17), | |
| (5, 8), (9, 12), (13, 16), (17, 20)) | |
| ratios = [] | |
| for a, b in model_pairs: | |
| model_len = float(np.linalg.norm(local[a] - local[b])) | |
| image_len = float(np.linalg.norm(pixels[a] - pixels[b])) | |
| if model_len > 1e-5 and image_len > 1.0: | |
| ratios.append(model_len / image_len) | |
| if len(ratios) < 3: | |
| return None | |
| depth = float(np.median(ratios) * self.K[0, 0]) | |
| depth = float(np.clip(depth, self.min_depth, self.max_depth)) | |
| u, v = pixels[0] | |
| x = (u - self.K[0, 2]) * depth / self.K[0, 0] | |
| y_down = (v - self.K[1, 2]) * depth / self.K[1, 1] | |
| return np.array([x, -y_down, depth], dtype=np.float64) | |
| def _pnp(self, pixels: np.ndarray, local: np.ndarray): | |
| valid = np.isfinite(local).all(axis=1) & np.isfinite(pixels).all(axis=1) | |
| if int(valid.sum()) < 6: | |
| return None | |
| object_points = np.ascontiguousarray(local[valid], dtype=np.float64) | |
| image_points = np.ascontiguousarray(pixels[valid], dtype=np.float64) | |
| try: | |
| ok, rvec, tvec, inliers = cv2.solvePnPRansac( | |
| object_points, image_points, self.K, self.dist, | |
| iterationsCount=80, reprojectionError=8.0, | |
| confidence=0.98, | |
| flags=cv2.SOLVEPNP_EPNP) | |
| except Exception: | |
| ok = False | |
| rvec = None | |
| tvec = None | |
| inliers = None | |
| if not ok: | |
| try: | |
| ok, rvec, tvec = cv2.solvePnP( | |
| object_points, image_points, self.K, self.dist, | |
| flags=cv2.SOLVEPNP_SQPNP) | |
| inliers = None | |
| except Exception: | |
| return None | |
| if not ok or rvec is None or tvec is None: | |
| return None | |
| try: | |
| rotation, _ = cv2.Rodrigues(rvec) | |
| projected, _ = cv2.projectPoints(object_points, rvec, tvec, | |
| self.K, self.dist) | |
| error = float(np.mean(np.linalg.norm( | |
| projected.reshape(-1, 2) - image_points, axis=1))) | |
| except Exception: | |
| return None | |
| if not np.isfinite(error): | |
| return None | |
| return rotation, np.asarray(tvec, dtype=np.float64).reshape(3), error, inliers | |
| def update(self, key: str, points_2d, local_points, | |
| timestamp_s: float, *, preserve_model_geometry: bool = True) -> dict | None: | |
| """Estimate camera-relative hand points. | |
| RGB-only 3D is intentionally model-geometry-preserving by default. | |
| PnP is retained for the camera-relative wrist position/depth estimate, | |
| but its frame-by-frame rotation is not applied to the hand joints. | |
| This avoids mirrored or rapidly changing PnP rotations becoming | |
| visible after the frontend fixes the wrist root. | |
| """ | |
| pixels = self._pixels(points_2d) | |
| local = np.asarray(local_points, dtype=np.float64) | |
| if pixels is None or local.shape != (21, 3) or not np.isfinite(local).all(): | |
| return None | |
| track = self._tracks.setdefault( | |
| str(key), _Track(position_filter=_OneEuro(self.freq_min, self.beta))) | |
| pnp = self._pnp(pixels, local) | |
| source = "rgb_image_scale" | |
| reprojection_error = None | |
| rotation = np.eye(3, dtype=np.float64) | |
| position = None | |
| if pnp is not None: | |
| rotation, tvec, reprojection_error, _inliers = pnp | |
| if reprojection_error <= 18.0 and self.min_depth <= tvec[2] <= self.max_depth: | |
| position = np.array([tvec[0], -tvec[1], tvec[2]], dtype=np.float64) | |
| source = "rgb_pose_pnp" | |
| if position is None: | |
| position = self._image_position(pixels, local) | |
| if position is None and track.position is not None: | |
| previous_timestamp = (track.timestamp_s | |
| if track.timestamp_s is not None | |
| else float(timestamp_s)) | |
| dt = max(0.0, float(timestamp_s) - previous_timestamp) | |
| position = track.position + track.velocity * min(dt, 0.10) | |
| source = "rgb_prediction" | |
| if position is None: | |
| return None | |
| if track.position_filter is None: | |
| track.position_filter = _OneEuro(self.freq_min, self.beta) | |
| filtered = track.position_filter(position, float(timestamp_s)) | |
| if track.position is not None and track.timestamp_s is not None: | |
| dt = max(1e-5, float(timestamp_s) - track.timestamp_s) | |
| measured_velocity = (filtered - track.position) / dt | |
| track.velocity = 0.7 * track.velocity + 0.3 * measured_velocity | |
| track.position = filtered.copy() | |
| track.timestamp_s = float(timestamp_s) | |
| # PnP tvec is in the OpenCV camera convention (Y down). Rotate the | |
| # local model into that camera frame, then convert every point to the | |
| # display convention (Y up) so position and landmarks share a frame. | |
| camera_points = (rotation @ local.T).T + np.array( | |
| [filtered[0], -filtered[1], filtered[2]], dtype=np.float64) | |
| landmarks = camera_points.copy() | |
| landmarks[:, 1] *= -1.0 | |
| if source != "rgb_pose_pnp" or preserve_model_geometry: | |
| landmarks = local + filtered | |
| return { | |
| "landmarks_3d": landmarks.astype(np.float32), | |
| "landmarks_3d_local": local.astype(np.float32), | |
| "position": filtered.astype(np.float32), | |
| "tracking_source": source, | |
| "coordinate_frame": "camera_relative", | |
| "unit": "rgb_estimated_meters", | |
| "metric_3d_available": False, | |
| "reprojection_error": (float(reprojection_error) | |
| if reprojection_error is not None else None), | |
| } | |