multiple methods added

main.py

@@ -1,30 +1,36 @@
+import copy
 import sys
 import os
 
+# Keep project-specific submodule path if needed
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'submodules', 'gesture_detection'))
 
+import numpy as np
 import cv2
+import csv
 from config import Config
 from skeleton.mediapipe_detector import MediaPipeDetector
 from skeleton.oak_pose_detector import OakPoseDetector
 from gesture_control.special_gestures import SpecialGestureDetector
-from gesture_control.arm_control import ArmController
+from gesture_control.arm_control import (
+    ArmControllerMethod1, ArmControllerMethod2, ArmControllerMethod3, ArmControllerMethod4
+)
 from ml_gestures_dynamic.predict import DynamicGesturePredictor
 
-#from geom_controll.arm_control import ArmController
 from robot.simulated_robot import DummySimRobot
 from robot.debug_robot import DebugRobot
 
 from camera.factory import create_camera
 
+
 def main():
-    # Все параметры
+    light = False
     cfg = Config()
-    
-    # Детектор скелета
+
+    # Skeleton detector
     if cfg.POSE_DETECTOR == 'oak':
         if cfg.CAMERA_TYPE == 'web':
-            print(f"При вычислении скелета на oak, можно пользоваться только этой же камерой для захвата видео!")
+            print("При вычислении скелета на oak, можно пользоваться только этой же камерой для захвата видео!")
             return
         detector = OakPoseDetector(
             model_type=cfg.OAK_MODEL_TYPE,
@@ -34,17 +40,16 @@ def main():
         camera = None
     else:
         detector = MediaPipeDetector()
-        # Захват кадра с камеры
         camera = create_camera(camera_type=cfg.CAMERA_TYPE, camera_id=cfg.CAMERA_ID, mirror=cfg.MIRROR_CAMERA)
-    
-    # Детектор статичных жестов
+
+    # Static gestures detector
     special_detector = SpecialGestureDetector(
         mode=cfg.SPECIAL_GESTURE_MODE,
         model_path=cfg.ML_GESTURE_MODEL,
         class_names=cfg.ML_GESTURE_CLASSES
     )
-    
-    # Динамические жесты
+
+    # Dynamic gestures
     dynamic_predictor = None
     if cfg.DYNAMIC_GESTURE['enabled']:
         dynamic_predictor = DynamicGesturePredictor(
@@ -53,36 +58,43 @@ def main():
             cfg.DYNAMIC_GESTURE['window_size'],
             cfg.DYNAMIC_GESTURE['threshold']
         )
-        
-    # Геометрический контроллер скоростей
-    arm_control = ArmController(cfg.ARM_CONTROL, mirror=cfg.MIRROR_CAMERA)
 
-    # Создаем папку для статистики
+    # Controllers for all 4 methods
+    controller_method1 = ArmControllerMethod1(cfg.ARM_CONTROL, mirror=cfg.MIRROR_CAMERA)
+    controller_method2 = ArmControllerMethod2(cfg.ARM_CONTROL, mirror=cfg.MIRROR_CAMERA)
+    controller_method3 = ArmControllerMethod3(cfg.ARM_CONTROL, mirror=cfg.MIRROR_CAMERA)
+    controller_method4 = ArmControllerMethod4(cfg.ARM_CONTROL, mirror=cfg.MIRROR_CAMERA)
+
+    active_method = 1  # default method on start: 1
+
+    # Stats folder
     stats_dir = './stats'
     if not os.path.exists(stats_dir):
         os.makedirs(stats_dir)
         print(f"Создана дирректория для статистики: {stats_dir}")
-    
-    # Тип робота: симуляция - с игрой, дебаг - просто печать команды в консоль
+
+    # Robot
     if cfg.ROBOT_MODE == 'simulator':
         robot = DummySimRobot(cfg, stats_path=stats_dir)
     else:
         robot = DebugRobot()
-        
+
     enabled = False
-    # включение управления жестами
     cross_cooldown = 0
+    light_cooldown = 0
+    LIGHT_COOLDOWN_FRAMES = 30
     CROSS_COOLDOWN_FRAMES = 30
     last_dynamic = None
     dynamic_counter = 0
-    dome_processed = False
-    
+    fire_latched = False  # to fire once per gesture
+
     while True:
-        # Получение кадра и скелета в зависимости от типа скелетного детектора
+        # Frame + skeleton
         if cfg.POSE_DETECTOR == 'oak':
             frame, landmarks = detector.get_frame_and_pose()
             if frame is None:
                 continue
+            h, w, _ = frame.shape
         else:
             frame = camera.get_frame()
             if frame is None:
@@ -92,35 +104,65 @@ def main():
                 landmarks = result['landmarks']
             else:
                 landmarks = None
-                
+            h, w, _ = frame.shape
+
+        linear, angular = 0.0, 0.0
+
         if landmarks is not None:
-            
-            # Статический жест
+            # Special static gesture
             special = special_detector.predict(landmarks)
-            
-            # Вкл/выкл управления жестами
-            if special == 'dome' and cross_cooldown == 0:
+
+            # Gesture roles per method:
+            # Methods 1-3: 'cross' toggles Move, 'dome' fires
+            # Method 4: 'dome' toggles Move, 'cross' fires
+            if active_method == 4:
+                toggle_gesture = 'dome'
+                fire_gesture = 'cross'
+            else:
+                toggle_gesture = 'cross'
+                fire_gesture = 'dome'
+
+            # Toggle Move ON/OFF via gestures
+            if special == toggle_gesture and cross_cooldown == 0:
                 enabled = not enabled
                 cross_cooldown = CROSS_COOLDOWN_FRAMES
+                print(f"Move toggled to: {'ON' if enabled else 'OFF'}")
             if cross_cooldown > 0:
                 cross_cooldown -= 1
-                
-            # Тушение
-            if enabled:
-                if special == 'cross' and not dome_processed:
-                    robot.send_fire_ext_burst_cmd()
-                    dome_processed = True
-                elif special != 'cross':
-                    dome_processed = False
-                
-            # Скорости из рук
-            linear, angular = arm_control.compute_speeds(landmarks)
+
+            # Fire burst (latched once per gesture instance)
+            if special == fire_gesture and not fire_latched:
+                robot.send_fire_ext_burst_cmd()
+                fire_latched = True
+                print("Fire burst!")
+            elif special != fire_gesture:
+                fire_latched = False
+
+            # Lights toggle on 'light' for any method
+            if special == 'light' and light_cooldown == 0:
+                light = not light
+                light_cooldown = LIGHT_COOLDOWN_FRAMES
+                print(f"Lights: {'On' if light else 'Off'}")
+            if light_cooldown > 0:
+                light_cooldown -= 1
+
+            # Speeds from active method
+            if active_method == 1:
+                linear, angular = controller_method1.compute_speeds(landmarks, frame_shape=frame.shape)
+            elif active_method == 2:
+                linear, angular = controller_method2.compute_speeds(landmarks, frame_shape=frame.shape)
+            elif active_method == 3:
+                linear, angular = controller_method3.compute_speeds(landmarks, frame_shape=frame.shape)
+            elif active_method == 4:
+                linear, angular = controller_method4.compute_speeds(landmarks, frame_shape=frame.shape)
+
+            # Send speeds
             if enabled:
                 robot.send_speed_cmd(linear, angular)
             else:
                 robot.send_speed_cmd(0.0, 0.0)
-            
-            # Динамический жест
+
+            # Dynamic gesture
             dynamic_gesture = None
             if dynamic_predictor:
                 dynamic_predictor.add_frame(landmarks)
@@ -128,50 +170,103 @@ def main():
                 if dynamic_gesture:
                     last_dynamic = dynamic_gesture
                     dynamic_counter = 30
-            
-            '''
-            if enabled:
-                if dynamic_gesture == 'wave_left':
-                    robot.reset_position()
-                elif dynamic_gesture == 'wave_right':
-                    robot.restart()
-            '''
-                
-            # Отрисовка
+
+            # Drawing
             if cfg.POSE_DETECTOR == 'oak':
                 vis = detector.draw_landmarks(frame, landmarks) if landmarks is not None else frame
             else:
                 vis = detector.draw_landmarks(frame, result['pose_landmarks']) if landmarks is not None else frame
-                    
-            cv2.putText(vis, f"Move: {'ON' if enabled else 'OFF'}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0,255,0) if enabled else (0,0,255), 2)
-            cv2.putText(vis, f"L:{linear:.2f} A:{angular:.2f}", (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255,255,0), 2)
+
+            # Method-specific overlay
+            if active_method == 1:
+                vis = controller_method1.draw_overlay(vis, landmarks)
+            elif active_method == 2:
+                vis = controller_method2.draw_overlay(vis, landmarks)
+            elif active_method == 3:
+                vis = controller_method3.draw_overlay(vis, landmarks)
+            elif active_method == 4:
+                vis = controller_method4.draw_overlay(vis, landmarks)
+
+            # HUD
+            cv2.putText(vis, f"Move: {'ON' if enabled else 'OFF'}", (10, 30),
+                        cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0) if enabled else (0, 0, 255), 2)
+            cv2.putText(vis, f"L:{linear:.2f} A:{angular:.2f}", (10, 60),
+                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 0), 2)
             if special != 'none':
-                cv2.putText(vis, f"Special: {special}", (10, 90), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255,255,0), 2)
+                cv2.putText(vis, f"Special: {special}", (10, 90),
+                            cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 0), 2)
             if dynamic_counter > 0 and last_dynamic:
-                cv2.putText(vis, f"Dynamic: {last_dynamic}", (10, 120), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0,255,255), 2)
+                cv2.putText(vis, f"Dynamic: {last_dynamic}", (10, 120),
+                            cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2)
                 dynamic_counter -= 1
-            
+            cv2.putText(vis, f"Method: {active_method}", (10, 150),
+                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
+            cv2.putText(vis, f"Lights: {'On' if light else 'Off'}", (10, 180),
+                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
+
         else:
             vis = frame
-            cv2.putText(vis, "No pose detected", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0,0,255), 2)
-        
-        cv2.imshow('Camera', vis)
-        
+            cv2.putText(vis, "No pose detected", (10, 30),
+                        cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
+
+        big = cv2.resize(vis, (1200, 900))
+        cv2.imshow('Camera', big)
+
         if not robot.step():
             break
-        
-        if cv2.waitKey(1) & 0xFF == ord('q'):
+
+        curr_key = cv2.waitKey(1)
+
+        if curr_key & 0xFF == ord('q'):
             break
-    
+
+        # Movement keyboard toggle on 'S' or 's'
+        if curr_key & 0xFF in (ord('s'), ord('S')):
+            enabled = not enabled
+            print(f"Move toggled to: {'ON' if enabled else 'OFF'}")
+
+        # Method switching keys (1, 2, 3, 4)
+        if curr_key & 0xFF == ord('1'):
+            active_method = 1
+            enabled = False
+            print("Switched to Method 1")
+        elif curr_key & 0xFF == ord('2'):
+            active_method = 2
+            enabled = False
+            print("Switched to Method 2")
+        elif curr_key & 0xFF == ord('3'):
+            active_method = 3
+            enabled = False
+            print("Switched to Method 3")
+        elif curr_key & 0xFF == ord('4'):
+            active_method = 4
+            enabled = False
+            print("Switched to Method 4 (grid, lm19)")
+
+        # Pose saving (exclude method-switch keys 1-4)
+        pose_codes = [ord('0') + i for i in range(10) if i not in (1, 2, 3, 4)]
+
+        if curr_key & 0xFF in pose_codes and 'landmarks' in locals() and landmarks is not None:
+            print('Button was pressed: ', curr_key)
+            landmarks2 = copy.deepcopy(landmarks)
+            correct_rows = landmarks2[:, 3] > 0.5
+            for i in range(landmarks2.shape[0]):
+                if not correct_rows[i]:
+                    landmarks2[i, :] = 0.
+            data = list(np.ravel(landmarks2[:, 0:3]))
+            with open("poses.csv", 'a') as f:
+                writer = csv.writer(f, delimiter=';')
+                writer.writerow(data + [curr_key - ord('0')])
+
     if cfg.POSE_DETECTOR == 'oak':
         detector.release()
     else:
         camera.release()
         detector.pose.close()
-        
+
     cv2.destroyAllWindows()
     robot.quit()
-    
+
+
 if __name__ == '__main__':
     main()
-