# 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