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 ( ArmControllerMethod1, ArmControllerMethod2, ArmControllerMethod3, ArmControllerMethod4 ) from ml_gestures_dynamic.predict import DynamicGesturePredictor 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("При вычислении скелета на oak, можно пользоваться только этой же камерой для захвата видео!") return detector = OakPoseDetector( model_type=cfg.OAK_MODEL_TYPE, detection_threshold=cfg.OAK_DETECTION_THRESHOLD, shaves=cfg.OAK_SHAVES ) 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( cfg.DYNAMIC_GESTURE['model_path'], cfg.DYNAMIC_GESTURE['classes_path'], cfg.DYNAMIC_GESTURE['window_size'], cfg.DYNAMIC_GESTURE['threshold'] ) # 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 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: continue result = detector.detect(frame) if result['success']: 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) # 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 # 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) dynamic_gesture = dynamic_predictor.predict() if dynamic_gesture: last_dynamic = dynamic_gesture dynamic_counter = 30 # 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 # 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) 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) 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) big = cv2.resize(vis, (1200, 900)) cv2.imshow('Camera', big) if not robot.step(): break 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()