dummy_simulation/main.py

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()