sasha mod

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.dead_zone = config.get('dead_zone', 0.15)
-        self.linear_arm = config.get('linear_arm', 'right')
+        self.debug = config.get('debug', False)
-        self.angular_arm = config.get('angular_arm', 'left')
 
-        self.wrist_idx = {'left': 15, 'right': 16}
+        # UI speeds in 0..100; used to map from wrist offset to final robot speeds
-        self.shoulder_idx = {'left': 11, 'right': 12}
+        self.max_speed = config.get('max_speed', 100)
-        self.hip_idx = {'left': 23, 'right': 24}
 
-        self.dead_zone = config.get('dead_zone', 0.2)
+        # Robot kinematics limits
-        self.debug = config.get('debug', False)
+        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
+    def update_frame_size(self, w, h):
-        self.max_speed_linear = config.get('max_speed_linear', 1.0)
+        self.screen_width = int(w)
-        self.max_speed_angular = config.get('max_speed_angular', 1.0)
+        self.screen_height = int(h)
 
-    def _get_side_indices(self, side):
+    def _get_wrist_index(self, side):
-        """Returns the wrist landmark index for a logical side, accounting for mirror."""
+        """
+        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 15 if side == 'left' else 16
-        return w_idx
 
-    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.
+        Return (x, y) of the requested wrist, or None if not visible.
-        Uses shoulders as horizontal reference and shoulder/hip as vertical reference.
+        Only wrist landmarks are used; all other landmarks are ignored.
-        Returns None if the wrist is not visible.
-        landmarks: array [N, 4] -> [x, y, z, visibility], x/y are normalized 0..1 (MediaPipe).
         """
-        w_idx = self._get_side_indices(side)
+        idx = self._get_wrist_index(side)
-
+        if idx >= len(landmarks):
-        if landmarks[w_idx][3] < 0.5:
             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)
 
-        wx = float(landmarks[w_idx][0])
+    def _both_wrists_in_top_center(self, landmarks):
-        wy = float(landmarks[w_idx][1])
+        """
+        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])
 
-        # Reference points (body center)
+    def _scale_speed(self, norm_offset):
-        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]
+        norm_offset is in [0..1].
-        lh_y = landmarks[self.hip_idx['left']][1]
+        Convert from [dead_zone..1] to [0..1], then scale to 0..max_speed.
-        rh_y = landmarks[self.hip_idx['right']][1]
+        """
+        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)
 
-        center_x = (ls_x + rs_x) / 2.0
+    def _compute_speeds(self, x, y):
-        shoulder_y = (ls_y + rs_y) / 2.0
+        """
-        hip_y = (lh_y + rh_y) / 2.0
+        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
 
-        shoulder_width = abs(ls_x - rs_x)
+        # Normalize offset: -1..1
-        torso_height = abs(hip_y - shoulder_y)
+        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))
 
-        # avoid division by zero
+        # Horizontal
-        shoulder_width = max(shoulder_width, 1e-3)
+        if abs(dx) < self.dead_zone:
-        torso_height = max(torso_height, 1e-3)
+            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))
 
-        # normalized offset from body center, scaled by body size
+        return h_dir, h_speed, v_dir, v_speed
-        nx = (wx - center_x) / shoulder_width
-        # vertical: up = positive; measured from shoulder line
-        ny = (shoulder_y - wy) / torso_height
 
-        nx = float(np.clip(nx, -1.0, 1.0))
+    def get_command(self, landmarks):
-        ny = float(np.clip(ny, -1.0, 1.0))
+        """
+        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'}
 
-        return nx, ny
+        # 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,
+            }
 
-    def _apply_dead_zone(self, value):
+        return {'command': 'stop'}
-        """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 compute_speeds(self, landmarks):
+    def to_linear_angular(self, cmd):
         """
-        Main API for the ROS node.
+        Convert a 'move' command (0..100 UI speeds) to robot linear/angular velocities.
-        Returns (linear, angular), both in the range -1..1.
 
-        - linear   : controlled by the vertical position of the linear_arm wrist (up = forward)
+        Mapping:
-        - angular  : controlled by the horizontal position of the angular_arm wrist
+          - 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) ----
+        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
-        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
 
-        # ---- Angular speed from angular_arm (horizontal axis) ----
+        # 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
-        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
 
-        linear = float(np.clip(linear, -1.0, 1.0))
+        return float(linear), float(angular)
-        angular = float(np.clip(angular, -1.0, 1.0))
+
+    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"[ArmController] linear={linear:.2f} angular={angular:.2f}")
+            print(f"V: {v_dir} {v_speed} | H: {h_dir} {h_speed}")
 
-        return linear, angular
+        return frame