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.
263 lines
9.7 KiB
Python
263 lines
9.7 KiB
Python
# gesture_control/arm_control.py
|
|
import numpy as np
|
|
import cv2
|
|
|
|
class ArmController:
|
|
def __init__(self, config, mirror=False):
|
|
"""
|
|
Wrist-only controller using pose landmarks (indices 15 and 16).
|
|
Behavior:
|
|
- RIGHT wrist controls motion.
|
|
- EXTINGUISHING only when BOTH wrists are visible AND BOTH are inside the TOP-CENTER rectangle:
|
|
y < H/3 AND w/3 <= x <= 2w/3
|
|
- Otherwise:
|
|
* If RIGHT wrist visible -> 'move'
|
|
* If no RIGHT wrist -> 'stop'
|
|
"""
|
|
self.config = config
|
|
self.mirror = mirror
|
|
|
|
self.dead_zone = config.get('dead_zone', 0.15)
|
|
self.debug = config.get('debug', False)
|
|
|
|
# UI speeds in 0..100; used to map from wrist offset to final robot speeds
|
|
self.max_speed = config.get('max_speed', 100)
|
|
|
|
# Robot kinematics limits
|
|
self.max_speed_linear = config.get('max_speed_linear', 0.8)
|
|
self.max_speed_angular = config.get('max_speed_angular', 1.5)
|
|
|
|
self.screen_width = config.get('screen_width', 640)
|
|
self.screen_height = config.get('screen_height', 480)
|
|
|
|
def update_frame_size(self, w, h):
|
|
self.screen_width = int(w)
|
|
self.screen_height = int(h)
|
|
|
|
def _get_wrist_index(self, side):
|
|
"""
|
|
Wrist indices (MediaPipe Pose):
|
|
- left wrist: 15
|
|
- right wrist: 16
|
|
With mirror=True, interpret 'left'/'right' as in mirrored preview.
|
|
"""
|
|
if self.mirror:
|
|
return 16 if side == 'left' else 15
|
|
else:
|
|
return 15 if side == 'left' else 16
|
|
|
|
def _get_wrist_position(self, landmarks, side):
|
|
"""
|
|
Return (x, y) of the requested wrist, or None if not visible.
|
|
Only wrist landmarks are used; all other landmarks are ignored.
|
|
"""
|
|
idx = self._get_wrist_index(side)
|
|
if idx >= len(landmarks):
|
|
return None
|
|
# If visibility exists (4th value), require >= 0.5
|
|
if len(landmarks[idx]) >= 4 and landmarks[idx][3] < 0.5:
|
|
return None
|
|
x = float(landmarks[idx][0])
|
|
y = float(landmarks[idx][1])
|
|
return x, y
|
|
|
|
def _in_top_center_section(self, x, y):
|
|
"""
|
|
Top-center rectangle:
|
|
- y in [0, H/3)
|
|
- x in [W/3, 2W/3)
|
|
"""
|
|
w_third = self.screen_width / 3.0
|
|
h_third = self.screen_height / 3.0
|
|
return (y < h_third) and (x >= w_third) and (x < 2.0 * w_third)
|
|
|
|
def _both_wrists_in_top_center(self, landmarks):
|
|
"""
|
|
Returns True only if BOTH wrists are visible AND both are inside the top-center rectangle.
|
|
"""
|
|
lw = self._get_wrist_position(landmarks, 'left')
|
|
rw = self._get_wrist_position(landmarks, 'right')
|
|
if lw is None or rw is None:
|
|
return False
|
|
return self._in_top_center_section(lw[0], lw[1]) and self._in_top_center_section(rw[0], rw[1])
|
|
|
|
def _scale_speed(self, norm_offset):
|
|
"""
|
|
norm_offset is in [0..1].
|
|
Convert from [dead_zone..1] to [0..1], then scale to 0..max_speed.
|
|
"""
|
|
eff = (norm_offset - self.dead_zone) / (1.0 - self.dead_zone)
|
|
eff = np.clip(eff, 0.0, 1.0)
|
|
return round(eff * self.max_speed)
|
|
|
|
def _compute_speeds(self, x, y):
|
|
"""
|
|
Compute directions and 0..100 speeds from a wrist position relative to screen center.
|
|
Returns (h_dir, h_speed, v_dir, v_speed)
|
|
- h_dir: 'left' | 'right' | 'center'
|
|
- v_dir: 'up' | 'down' | 'center'
|
|
"""
|
|
cx = self.screen_width / 2.0
|
|
cy = self.screen_height / 2.0
|
|
|
|
# Normalize offset: -1..1
|
|
dx = (x - cx) / (self.screen_width / 2.0)
|
|
dy = (y - cy) / (self.screen_height / 2.0)
|
|
dx = float(np.clip(dx, -1.0, 1.0))
|
|
dy = float(np.clip(dy, -1.0, 1.0))
|
|
|
|
# Horizontal
|
|
if abs(dx) < self.dead_zone:
|
|
h_dir, h_speed = 'center', 0
|
|
elif dx > 0:
|
|
h_dir = 'right'
|
|
h_speed = self._scale_speed(abs(dx))
|
|
else:
|
|
h_dir = 'left'
|
|
h_speed = self._scale_speed(abs(dx))
|
|
|
|
# Vertical (image y grows down; up is dy < 0)
|
|
if abs(dy) < self.dead_zone:
|
|
v_dir, v_speed = 'center', 0
|
|
elif dy < 0:
|
|
v_dir = 'up'
|
|
v_speed = self._scale_speed(abs(dy))
|
|
else:
|
|
v_dir = 'down'
|
|
v_speed = self._scale_speed(abs(dy))
|
|
|
|
return h_dir, h_speed, v_dir, v_speed
|
|
|
|
def get_command(self, landmarks):
|
|
"""
|
|
Logic:
|
|
- If BOTH wrists are visible AND BOTH are inside the top-center rectangle (the 'UP' section)
|
|
-> {'command': 'extinguishing fire'}
|
|
- Else if RIGHT wrist is visible -> {'command': 'move', ...} using RIGHT wrist position
|
|
- Else -> {'command': 'stop'}
|
|
"""
|
|
if self._both_wrists_in_top_center(landmarks):
|
|
if self.debug:
|
|
print("Two wrists in TOP-CENTER ('UP' section) -> EXTINGUISHING FIRE")
|
|
return {'command': 'extinguishing fire'}
|
|
|
|
# Otherwise, control with RIGHT wrist only
|
|
right_wrist = self._get_wrist_position(landmarks, 'right')
|
|
if right_wrist is not None:
|
|
x, y = right_wrist
|
|
h_dir, h_speed, v_dir, v_speed = self._compute_speeds(x, y)
|
|
return {
|
|
'command': 'move',
|
|
'x': x, 'y': y,
|
|
'h_dir': h_dir, 'h_speed': h_speed,
|
|
'v_dir': v_dir, 'v_speed': v_speed,
|
|
}
|
|
|
|
return {'command': 'stop'}
|
|
|
|
def to_linear_angular(self, cmd):
|
|
"""
|
|
Convert a 'move' command (0..100 UI speeds) to robot linear/angular velocities.
|
|
|
|
Mapping:
|
|
- Vertical axis controls linear velocity: up -> +linear (forward), down -> -linear (backward)
|
|
- Horizontal axis controls angular velocity: left -> +angular (CCW), right -> -angular (CW)
|
|
"""
|
|
if cmd.get('command') != 'move':
|
|
return 0.0, 0.0
|
|
|
|
v_dir, v_speed = cmd['v_dir'], float(cmd['v_speed'])
|
|
h_dir, h_speed = cmd['h_dir'], float(cmd['h_speed'])
|
|
|
|
# Normalize speeds to 0..1
|
|
v_norm = v_speed / float(self.max_speed) if self.max_speed > 0 else 0.0
|
|
h_norm = h_speed / float(self.max_speed) if self.max_speed > 0 else 0.0
|
|
|
|
# Linear: forward/backward
|
|
if v_dir == 'up':
|
|
linear = +v_norm * self.max_speed_linear
|
|
elif v_dir == 'down':
|
|
linear = -v_norm * self.max_speed_linear
|
|
else:
|
|
linear = 0.0
|
|
|
|
# Angular: left/right
|
|
# Convention: left turn -> +angular, right turn -> -angular
|
|
if h_dir == 'left':
|
|
angular = +h_norm * self.max_speed_angular
|
|
elif h_dir == 'right':
|
|
angular = -h_norm * self.max_speed_angular
|
|
else:
|
|
angular = 0.0
|
|
|
|
return float(linear), float(angular)
|
|
|
|
def draw_hands(self, frame, landmarks):
|
|
"""
|
|
Overlay control UI on the frame.
|
|
- Shows 3x3 grid.
|
|
- Extinguish only when both wrists are in the TOP-CENTER rectangle (the 'UP' section).
|
|
- Otherwise, RIGHT wrist controls motion.
|
|
"""
|
|
cmd = self.get_command(landmarks)
|
|
w, h = self.screen_width, self.screen_height
|
|
|
|
# 3x3 grid
|
|
grid_color = (200, 200, 200)
|
|
cv2.line(frame, (w // 3, 0), (w // 3, h), grid_color, 1)
|
|
cv2.line(frame, (2 * w // 3, 0), (2 * w // 3, h), grid_color, 1)
|
|
cv2.line(frame, (0, h // 3), (w, h // 3), grid_color, 1)
|
|
cv2.line(frame, (0, 2 * h // 3), (w, 2 * h // 3), grid_color, 1)
|
|
|
|
# Highlight the top-center rectangle (UP section)
|
|
x1, x2 = int(w / 3.0), int(2 * w / 3.0)
|
|
y2 = int(h / 3.0)
|
|
overlay = frame.copy()
|
|
cv2.rectangle(overlay, (x1, 0), (x2, y2), (0, 165, 255), -1) # orange overlay
|
|
cv2.addWeighted(overlay, 0.10, frame, 0.90, 0, frame)
|
|
|
|
# Center dot
|
|
cv2.circle(frame, (w // 2, h // 2), 6, (255, 255, 255), -1)
|
|
|
|
# Draw detected wrist points for feedback (if available)
|
|
lw = self._get_wrist_position(landmarks, 'left')
|
|
rw = self._get_wrist_position(landmarks, 'right')
|
|
if lw is not None:
|
|
cv2.circle(frame, (int(lw[0]), int(lw[1])), 10, (255, 0, 0), -1) # left: blue
|
|
if rw is not None:
|
|
cv2.circle(frame, (int(rw[0]), int(rw[1])), 10, (0, 255, 0), -1) # right: green
|
|
|
|
if cmd['command'] == 'stop':
|
|
cv2.putText(frame, "COMMAND: STOP", (20, 40),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 0, 255), 3)
|
|
if self.debug:
|
|
print("COMMAND: stop")
|
|
return frame
|
|
|
|
if cmd['command'] == 'extinguishing fire':
|
|
cv2.putText(frame, "EXTINGUISHING FIRE (2 WRISTS IN TOP-CENTER 'UP' SECTION)", (20, 40),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 165, 255), 2)
|
|
if self.debug:
|
|
print("COMMAND: extinguishing fire (two wrists in top-center section)")
|
|
return frame
|
|
|
|
# command == move (right wrist visible)
|
|
x, y = int(cmd['x']), int(cmd['y'])
|
|
cv2.circle(frame, (x, y), 12, (0, 255, 0), -1)
|
|
|
|
# Vector from center to wrist
|
|
cv2.line(frame, (w // 2, h // 2), (x, y), (0, 255, 0), 2)
|
|
|
|
v_dir, v_speed = cmd['v_dir'], cmd['v_speed']
|
|
h_dir, h_speed = cmd['h_dir'], cmd['h_speed']
|
|
|
|
cv2.putText(frame, f"V: {v_dir.upper()} {v_speed}", (20, 40),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 255), 2)
|
|
cv2.putText(frame, f"H: {h_dir.upper()} {h_speed}", (20, 75),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 255), 2)
|
|
|
|
if self.debug:
|
|
print(f"V: {v_dir} {v_speed} | H: {h_dir} {h_speed}")
|
|
|
|
return frame
|