You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
325 lines
11 KiB
Python
325 lines
11 KiB
Python
import numpy as np
|
|
import cv2
|
|
|
|
|
|
def _clip_unit(v):
|
|
return float(np.clip(v, -1.0, 1.0))
|
|
|
|
|
|
def _apply_dead_zone(v, dz):
|
|
return 0.0 if abs(v) < dz else v
|
|
|
|
|
|
def _robust_metrics(landmarks, min_conf=0.5):
|
|
"""
|
|
Compute shoulder_center, shoulder_width, torso_height robustly.
|
|
Uses hips if available; otherwise falls back to nose/shoulder geometry.
|
|
|
|
Returns:
|
|
shoulder_center (np.array shape (2,))
|
|
shoulder_width (float)
|
|
torso_height (float)
|
|
ok (bool)
|
|
"""
|
|
def pt(i):
|
|
return np.array(landmarks[i][:2], dtype=float), float(landmarks[i][3])
|
|
|
|
l_sh, c_lsh = pt(11)
|
|
r_sh, c_rsh = pt(12)
|
|
if c_lsh < min_conf or c_rsh < min_conf:
|
|
return None, 0.0, 0.0, False
|
|
|
|
shoulder_center = (l_sh + r_sh) / 2.0
|
|
shoulder_width = float(np.linalg.norm(r_sh - l_sh))
|
|
if shoulder_width < 1e-3:
|
|
return shoulder_center, 0.0, 0.0, False
|
|
|
|
# Try hips
|
|
l_hip, c_lhip = pt(23)
|
|
r_hip, c_rhip = pt(24)
|
|
|
|
if c_lhip >= min_conf and c_rhip >= min_conf:
|
|
hip_center = (l_hip + r_hip) / 2.0
|
|
torso_height = float(np.linalg.norm(hip_center - shoulder_center))
|
|
if torso_height >= 1e-3:
|
|
return shoulder_center, shoulder_width, torso_height, True
|
|
|
|
# Fallbacks (upper-body only)
|
|
nose, c_nose = pt(0)
|
|
if c_nose >= min_conf:
|
|
nose_to_shoulder = abs(nose[1] - shoulder_center[1])
|
|
torso_height = max(1.6 * nose_to_shoulder, 0.9 * shoulder_width)
|
|
else:
|
|
torso_height = max(1.2 * shoulder_width, 1.0)
|
|
|
|
return shoulder_center, shoulder_width, float(torso_height), True
|
|
|
|
|
|
class ArmControllerMethod1:
|
|
"""
|
|
Method 1: Single-hand driving with right wrist.
|
|
- Linear: vertical offset of right wrist from shoulder center (normalized by torso height)
|
|
- Angular: horizontal offset of right wrist from shoulder center (normalized by shoulder width)
|
|
"""
|
|
def __init__(self, config, mirror=False):
|
|
self.config = config
|
|
self.mirror = mirror
|
|
self.dead_zone = config.get('dead_zone', 0.1)
|
|
self.debug = config.get('debug', False)
|
|
self.min_conf = config.get('min_conf', 0.5)
|
|
|
|
def compute_speeds(self, landmarks, frame_shape=None):
|
|
if landmarks is None:
|
|
return 0.0, 0.0
|
|
|
|
# Require: shoulders + right wrist
|
|
need = [11, 12, 16]
|
|
if any(landmarks[i][3] < self.min_conf for i in need):
|
|
return 0.0, 0.0
|
|
|
|
shoulder_center, shoulder_width, torso_height, ok = _robust_metrics(landmarks, self.min_conf)
|
|
if not ok or shoulder_width < 1e-3 or torso_height < 1e-3:
|
|
return 0.0, 0.0
|
|
|
|
r_wr = landmarks[16][:2]
|
|
|
|
# Positive linear when wrist above shoulder center (forward)
|
|
linear = (shoulder_center[1] - r_wr[1]) / torso_height
|
|
# Positive angular when wrist to the right of shoulder center
|
|
angular = (r_wr[0] - shoulder_center[0]) / shoulder_width
|
|
|
|
if self.mirror:
|
|
angular = -angular
|
|
|
|
linear = _clip_unit(_apply_dead_zone(linear, self.dead_zone))
|
|
angular = _clip_unit(_apply_dead_zone(angular, self.dead_zone))
|
|
|
|
if self.debug:
|
|
print(f"[M1] L:{linear:.2f} A:{angular:.2f}")
|
|
return linear, angular
|
|
|
|
def draw_overlay(self, frame, landmarks=None):
|
|
if frame is None:
|
|
return frame
|
|
h, w = frame.shape[:2]
|
|
# Draw center cross (screen center approximation)
|
|
cv2.line(frame, (w // 2, 0), (w // 2, h), (0, 0, 0), 1)
|
|
cv2.line(frame, (0, h // 2), (w, h // 2), (0, 0, 0), 1)
|
|
# Draw right wrist
|
|
if landmarks is not None and landmarks[16][3] > 0.5:
|
|
x, y = int(landmarks[16][0]), int(landmarks[16][1])
|
|
cv2.circle(frame, (x, y), 8, (0, 255, 255), -1)
|
|
return frame
|
|
|
|
|
|
class ArmControllerMethod2:
|
|
"""
|
|
Method 2: Two-hand blended control.
|
|
- Linear: average vertical offset of both wrists from shoulder center (normalized by torso height)
|
|
- Angular: horizontal balance of wrists around shoulder center (normalized by shoulder width)
|
|
"""
|
|
def __init__(self, config, mirror=False):
|
|
self.config = config
|
|
self.mirror = mirror
|
|
self.dead_zone = config.get('dead_zone', 0.1)
|
|
self.debug = config.get('debug', False)
|
|
self.min_conf = config.get('min_conf', 0.5)
|
|
|
|
def compute_speeds(self, landmarks, frame_shape=None):
|
|
if landmarks is None:
|
|
return 0.0, 0.0
|
|
|
|
# Require: shoulders + both wrists
|
|
need = [11, 12, 15, 16]
|
|
if any(landmarks[i][3] < self.min_conf for i in need):
|
|
return 0.0, 0.0
|
|
|
|
shoulder_center, shoulder_width, torso_height, ok = _robust_metrics(landmarks, self.min_conf)
|
|
if not ok or shoulder_width < 1e-3 or torso_height < 1e-3:
|
|
return 0.0, 0.0
|
|
|
|
l_wr = landmarks[15][:2]
|
|
r_wr = landmarks[16][:2]
|
|
|
|
# Linear: average elevation of both wrists
|
|
lin_l = (shoulder_center[1] - l_wr[1]) / torso_height
|
|
lin_r = (shoulder_center[1] - r_wr[1]) / torso_height
|
|
linear = 0.5 * (lin_l + lin_r)
|
|
|
|
# Angular: horizontal balance
|
|
ang = ((r_wr[0] - shoulder_center[0]) - (shoulder_center[0] - l_wr[0])) / shoulder_width
|
|
angular = ang
|
|
|
|
if self.mirror:
|
|
angular = -angular
|
|
|
|
linear = _clip_unit(_apply_dead_zone(linear, self.dead_zone))
|
|
angular = _clip_unit(_apply_dead_zone(angular, self.dead_zone))
|
|
|
|
if self.debug:
|
|
print(f"[M2] L:{linear:.2f} A:{angular:.2f}")
|
|
return linear, angular
|
|
|
|
def draw_overlay(self, frame, landmarks=None):
|
|
if frame is None:
|
|
return frame
|
|
if landmarks is not None:
|
|
for idx, color in [(15, (255, 0, 255)), (16, (0, 255, 255))]:
|
|
if landmarks[idx][3] > 0.5:
|
|
x, y = int(landmarks[idx][0]), int(landmarks[idx][1])
|
|
cv2.circle(frame, (x, y), 8, color, -1)
|
|
return frame
|
|
|
|
|
|
class ArmControllerMethod3:
|
|
"""
|
|
Method 3: Elbow-augmented control.
|
|
- Linear: average vertical offset of elbows (normalized by torso height)
|
|
- Angular: wrist horizontal balance (normalized by shoulder width)
|
|
"""
|
|
def __init__(self, config, mirror=False):
|
|
self.config = config
|
|
self.mirror = mirror
|
|
self.dead_zone = config.get('dead_zone', 0.1)
|
|
self.debug = config.get('debug', False)
|
|
self.min_conf = config.get('min_conf', 0.5)
|
|
|
|
def compute_speeds(self, landmarks, frame_shape=None):
|
|
if landmarks is None:
|
|
return 0.0, 0.0
|
|
|
|
# Require shoulders; prefer elbows for linear; wrists for angular.
|
|
need_base = [11, 12]
|
|
if any(landmarks[i][3] < self.min_conf for i in need_base):
|
|
return 0.0, 0.0
|
|
|
|
elbows_ok = (landmarks[13][3] >= self.min_conf and landmarks[14][3] >= self.min_conf)
|
|
wrists_ok = (landmarks[15][3] >= self.min_conf and landmarks[16][3] >= self.min_conf)
|
|
|
|
if not elbows_ok and not wrists_ok:
|
|
return 0.0, 0.0
|
|
|
|
shoulder_center, shoulder_width, torso_height, ok = _robust_metrics(landmarks, self.min_conf)
|
|
if not ok or shoulder_width < 1e-3 or torso_height < 1e-3:
|
|
return 0.0, 0.0
|
|
|
|
# Linear: prefer elbows, fallback to wrists average if elbows missing
|
|
if elbows_ok:
|
|
l_el = landmarks[13][:2]
|
|
r_el = landmarks[14][:2]
|
|
lin_l = (shoulder_center[1] - l_el[1]) / torso_height
|
|
lin_r = (shoulder_center[1] - r_el[1]) / torso_height
|
|
linear = 0.5 * (lin_l + lin_r)
|
|
else:
|
|
l_wr = landmarks[15][:2]
|
|
r_wr = landmarks[16][:2]
|
|
lin_l = (shoulder_center[1] - l_wr[1]) / torso_height
|
|
lin_r = (shoulder_center[1] - r_wr[1]) / torso_height
|
|
linear = 0.5 * (lin_l + lin_r)
|
|
|
|
# Angular: use wrists if available, else 0
|
|
if wrists_ok:
|
|
l_wr = landmarks[15][:2]
|
|
r_wr = landmarks[16][:2]
|
|
angular = ((r_wr[0] + l_wr[0]) - 2 * shoulder_center[0]) / shoulder_width
|
|
else:
|
|
angular = 0.0
|
|
|
|
if self.mirror:
|
|
angular = -angular
|
|
|
|
linear = _clip_unit(_apply_dead_zone(linear, self.dead_zone))
|
|
angular = _clip_unit(_apply_dead_zone(angular, self.dead_zone))
|
|
|
|
if self.debug:
|
|
print(f"[M3] L:{linear:.2f} A:{angular:.2f}")
|
|
return linear, angular
|
|
|
|
def draw_overlay(self, frame, landmarks=None):
|
|
if frame is None:
|
|
return frame
|
|
if landmarks is not None:
|
|
for idx, color in [(13, (0, 200, 0)), (14, (0, 200, 0)), (15, (0, 255, 255)), (16, (255, 0, 255))]:
|
|
if landmarks[idx][3] > 0.5:
|
|
x, y = int(landmarks[idx][0]), int(landmarks[idx][1])
|
|
cv2.circle(frame, (x, y), 6, color, -1)
|
|
return frame
|
|
|
|
|
|
class ArmControllerMethod4:
|
|
"""
|
|
Method 4: 3x3 grid based on landmark 19 (right index finger tip).
|
|
Screen split at 2/5 and 3/5 (both axes). Center band = 0.
|
|
Proportional speed away from the center bands.
|
|
"""
|
|
def __init__(self, config, mirror=False):
|
|
self.config = config
|
|
self.mirror = mirror
|
|
self.finger_idx = 19 # right index finger tip
|
|
self.debug = config.get('debug', False)
|
|
|
|
def compute_speeds(self, landmarks, frame_shape=None):
|
|
linear = 0.0
|
|
angular = 0.0
|
|
|
|
if frame_shape is None or landmarks is None:
|
|
return 0.0, 0.0
|
|
if landmarks[self.finger_idx][3] < 0.5:
|
|
return 0.0, 0.0
|
|
|
|
h, w = int(frame_shape[0]), int(frame_shape[1])
|
|
x = landmarks[self.finger_idx][0]
|
|
y = landmarks[self.finger_idx][1]
|
|
|
|
# Angular (horizontal): center band 2/5..3/5 = 0
|
|
if 2 * w / 5 <= x <= 3 * w / 5:
|
|
angular = 0.0
|
|
elif x > 3 * w / 5:
|
|
angular = (x * 5) / (2 * w) - 1
|
|
else:
|
|
angular = (x - 3 * w / 5) / (2 * w / 5)
|
|
|
|
# Linear (vertical): center band 2/5..3/5 = 0
|
|
if 2 * h / 5 <= y <= 3 * h / 5:
|
|
linear = 0.0
|
|
elif y > 3 * h / 5:
|
|
linear = -((y * 5) / (2 * h) - 1)
|
|
else:
|
|
linear = -(y - 3 * h / 5) / (2 * h / 5)
|
|
|
|
if self.mirror:
|
|
angular = -angular
|
|
|
|
if self.debug:
|
|
print(f"[M4] L:{linear:.2f} A:{angular:.2f}")
|
|
return _clip_unit(linear), _clip_unit(angular)
|
|
|
|
def draw_overlay(self, frame, landmarks=None):
|
|
if frame is None:
|
|
return frame
|
|
|
|
h, w = frame.shape[:2]
|
|
x1, x2 = int(w * 2 / 5), int(w * 3 / 5)
|
|
y1, y2 = int(h * 2 / 5), int(h * 3 / 5)
|
|
|
|
# Grid lines
|
|
cv2.line(frame, (x1, 0), (x1, h), (0, 0, 0), 2)
|
|
cv2.line(frame, (x2, 0), (x2, h), (0, 0, 0), 2)
|
|
cv2.line(frame, (0, y1), (w, y1), (0, 0, 0), 2)
|
|
cv2.line(frame, (0, y2), (w, y2), (0, 0, 0), 2)
|
|
|
|
# Highlight active cell + finger
|
|
if landmarks is not None and landmarks[self.finger_idx][3] > 0.5:
|
|
fx, fy = int(landmarks[self.finger_idx][0]), int(landmarks[self.finger_idx][1])
|
|
cx0, cx1 = (0, x1) if fx < x1 else ((x2, w) if fx > x2 else (x1, x2))
|
|
cy0, cy1 = (0, y1) if fy < y1 else ((y2, h) if fy > y2 else (y1, y2))
|
|
|
|
overlay = frame.copy()
|
|
cv2.rectangle(overlay, (cx0, cy0), (cx1, cy1), (0, 255, 255), -1)
|
|
frame = cv2.addWeighted(overlay, 0.2, frame, 0.8, 0)
|
|
|
|
cv2.circle(frame, (fx, fy), 8, (0, 255, 255), -1)
|
|
cv2.circle(frame, (fx, fy), 12, (0, 120, 120), 2)
|
|
|
|
return frame
|