dummy_simulation/main.py

import sys
import os

sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'submodules', 'gesture_detection'))

import cv2
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 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():
    # Все параметры
    cfg = Config()
    
    # Детектор скелета
    if cfg.POSE_DETECTOR == 'oak':
        if cfg.CAMERA_TYPE == 'web':
            print(f"При вычислении скелета на 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)
    
    # Детектор статичных жестов
    special_detector = SpecialGestureDetector(
        mode=cfg.SPECIAL_GESTURE_MODE,
        model_path=cfg.ML_GESTURE_MODEL,
        class_names=cfg.ML_GESTURE_CLASSES
    )
    
    # Динамические жесты
    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']
        )
        
    # Геометрический контроллер скоростей
    arm_control = ArmController(cfg.ARM_CONTROL, mirror=cfg.MIRROR_CAMERA)
    
    # Тип робота: симуляция - с игрой, дебаг - просто печать команды в консоль
    if cfg.ROBOT_MODE == 'simulator':
        robot = DummySimRobot(cfg)
    else:
        robot = DebugRobot()
        
    enabled = False
    # включение управления жестами
    cross_cooldown = 0
    CROSS_COOLDOWN_FRAMES = 30
    last_dynamic = None
    dynamic_counter = 0
    dome_processed = False
    
    while True:
        # Получение кадра и скелета в зависимости от типа скелетного детектора
        if cfg.POSE_DETECTOR == 'oak':
            frame, landmarks = detector.get_frame_and_pose()
            if frame is None:
                continue
        else:
            frame = camera.get_frame()
            if frame is None:
                continue
            result = detector.detect(frame)
            if result['success']:
                landmarks = result['landmarks']
            else:
                landmarks = None
                
        if landmarks is not None:
            
            # Статический жест
            special = special_detector.predict(landmarks)
            
            # Вкл/выкл управления жестами
            if special == 'cross' and cross_cooldown == 0:
                enabled = not enabled
                cross_cooldown = CROSS_COOLDOWN_FRAMES
            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)
            if enabled:
                robot.send_speed_cmd(linear, angular)
            else:
                robot.send_speed_cmd(0.0, 0.0)
            
            # Динамический жест
            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
            
            '''
            if enabled:
                if dynamic_gesture == 'wave_left':
                    robot.reset_position()
                elif dynamic_gesture == 'wave_right':
                    robot.restart()
            '''
                
            # Отрисовка
            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)
            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
            
        else:
            vis = frame
            cv2.putText(vis, "No pose detected", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0,0,255), 2)
        
        cv2.imshow('Camera', vis)
        
        if not robot.step():
            break
        
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    
    if cfg.POSE_DETECTOR == 'oak':
        detector.release()
    else:
        camera.release()
        detector.pose.close()
        
    cv2.destroyAllWindows()
    robot.quit()
    
if __name__ == '__main__':
    main()