diff --git a/gesture_control/arm_control.py b/gesture_control/arm_control.py index c2d6bb8..0100a35 100644 --- a/gesture_control/arm_control.py +++ b/gesture_control/arm_control.py @@ -1,115 +1,262 @@ +# 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 - # arms assignment - self.linear_arm = config.get('linear_arm', 'right') - self.angular_arm = config.get('angular_arm', 'left') + self.dead_zone = config.get('dead_zone', 0.15) + self.debug = config.get('debug', False) - self.wrist_idx = {'left': 15, 'right': 16} - self.shoulder_idx = {'left': 11, 'right': 12} - self.hip_idx = {'left': 23, 'right': 24} + # UI speeds in 0..100; used to map from wrist offset to final robot speeds + self.max_speed = config.get('max_speed', 100) - self.dead_zone = config.get('dead_zone', 0.2) - self.debug = config.get('debug', False) + # 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) - # separate max speeds - self.max_speed_linear = config.get('max_speed_linear', 1.0) - self.max_speed_angular = config.get('max_speed_angular', 1.0) + def update_frame_size(self, w, h): + self.screen_width = int(w) + self.screen_height = int(h) - def _get_side_indices(self, side): - """Returns the wrist landmark index for a logical side, accounting for mirror.""" + 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: - w_idx = self.wrist_idx['right'] if side == 'left' else self.wrist_idx['left'] + return 16 if side == 'left' else 15 else: - w_idx = self.wrist_idx[side] - return w_idx + return 15 if side == 'left' else 16 - def _get_wrist_norm(self, landmarks, side): + def _get_wrist_position(self, landmarks, side): """ - Returns normalized wrist position in the range -1..1 relative to the body. - Uses shoulders as horizontal reference and shoulder/hip as vertical reference. - Returns None if the wrist is not visible. - landmarks: array [N, 4] -> [x, y, z, visibility], x/y are normalized 0..1 (MediaPipe). + Return (x, y) of the requested wrist, or None if not visible. + Only wrist landmarks are used; all other landmarks are ignored. """ - w_idx = self._get_side_indices(side) - - if landmarks[w_idx][3] < 0.5: + 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 - wx = float(landmarks[w_idx][0]) - wy = float(landmarks[w_idx][1]) + 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) - # Reference points (body center) - ls_x, ls_y = landmarks[self.shoulder_idx['left']][0], landmarks[self.shoulder_idx['left']][1] - rs_x, rs_y = landmarks[self.shoulder_idx['right']][0], landmarks[self.shoulder_idx['right']][1] - lh_y = landmarks[self.hip_idx['left']][1] - rh_y = landmarks[self.hip_idx['right']][1] + 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]) - center_x = (ls_x + rs_x) / 2.0 - shoulder_y = (ls_y + rs_y) / 2.0 - hip_y = (lh_y + rh_y) / 2.0 + 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) - shoulder_width = abs(ls_x - rs_x) - torso_height = abs(hip_y - shoulder_y) + 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 - # avoid division by zero - shoulder_width = max(shoulder_width, 1e-3) - torso_height = max(torso_height, 1e-3) + # 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)) - # normalized offset from body center, scaled by body size - nx = (wx - center_x) / shoulder_width - # vertical: up = positive; measured from shoulder line - ny = (shoulder_y - wy) / torso_height + # 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)) - nx = float(np.clip(nx, -1.0, 1.0)) - ny = float(np.clip(ny, -1.0, 1.0)) + # 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 nx, ny + return h_dir, h_speed, v_dir, v_speed - def _apply_dead_zone(self, value): - """Removes dead zone and rescales -1..1 -> -1..1.""" - if abs(value) < self.dead_zone: - return 0.0 - sign = 1.0 if value > 0 else -1.0 - scaled = (abs(value) - self.dead_zone) / (1.0 - self.dead_zone) - return sign * float(np.clip(scaled, 0.0, 1.0)) + 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 compute_speeds(self, landmarks): + def to_linear_angular(self, cmd): """ - Main API for the ROS node. - Returns (linear, angular), both in the range -1..1. + Convert a 'move' command (0..100 UI speeds) to robot linear/angular velocities. - - linear : controlled by the vertical position of the linear_arm wrist (up = forward) - - angular : controlled by the horizontal position of the angular_arm wrist + Mapping: + - Vertical axis controls linear velocity: up -> +linear (forward), down -> -linear (backward) + - Horizontal axis controls angular velocity: left -> +angular (CCW), right -> -angular (CW) """ - if landmarks is None: + if cmd.get('command') != 'move': return 0.0, 0.0 - # ---- Linear speed from linear_arm (vertical axis) ---- - linear = 0.0 - lin_pos = self._get_wrist_norm(landmarks, self.linear_arm) - if lin_pos is not None: - _, ny = lin_pos - linear = self._apply_dead_zone(ny) * self.max_speed_linear + 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) - # ---- Angular speed from angular_arm (horizontal axis) ---- - angular = 0.0 - ang_pos = self._get_wrist_norm(landmarks, self.angular_arm) - if ang_pos is not None: - nx, _ = ang_pos - # invert so that hand-to-the-left = turn left (adjust sign if needed) - angular = -self._apply_dead_zone(nx) * self.max_speed_angular + v_dir, v_speed = cmd['v_dir'], cmd['v_speed'] + h_dir, h_speed = cmd['h_dir'], cmd['h_speed'] - linear = float(np.clip(linear, -1.0, 1.0)) - angular = float(np.clip(angular, -1.0, 1.0)) + 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"[ArmController] linear={linear:.2f} angular={angular:.2f}") + print(f"V: {v_dir} {v_speed} | H: {h_dir} {h_speed}") - return linear, angular + return frame