diff --git a/main.py b/main.py index 6e1a987..89922ee 100644 --- a/main.py +++ b/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 == 'cross' 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 == 'dome' and not dome_processed: - robot.send_fire_ext_burst_cmd() - dome_processed = True - elif special != 'dome': - 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() - diff --git a/submodules/gesture_detection b/submodules/gesture_detection index e74f1b0..08c5b66 160000 --- a/submodules/gesture_detection +++ b/submodules/gesture_detection @@ -1 +1 @@ -Subproject commit e74f1b09d71973da345aeff691daeda954e5c080 +Subproject commit 08c5b66acec9ead59b6c40487065d1782b6b6ab6