diff --git a/config.py b/config.py index ec1b4f7..54c9a9b 100644 --- a/config.py +++ b/config.py @@ -4,6 +4,12 @@ class Config: CAMERA_TYPE = 'oak' # 'web' или 'oak' MIRROR_CAMERA = True + # ===== Скелет ===== + POSE_DETECTOR = 'mediapipe' # 'oak' или 'mediapipe' + OAK_MODEL_TYPE = 'efficienthrnet2' + OAK_DETECTION_THRESHOLD = 0.1 + OAK_SHAVES = 6 + # ===== Преобразование позы в скорости ===== ARM_CONTROL = { 'linear_arm': 'right', diff --git a/main.py b/main.py index 1fe122c..a17a350 100644 --- a/main.py +++ b/main.py @@ -6,6 +6,7 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'submodules', 'gestur 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 ml_gestures_dynamic.predict import DynamicGesturePredictor @@ -20,7 +21,20 @@ def main(): cfg = Config() # Детектор скелета - detector = MediaPipeDetector() + 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( @@ -47,9 +61,6 @@ def main(): robot = DummySimRobot(cfg) else: robot = DebugRobot() - - # Захват кадра с камеры - camera = create_camera(camera_type=cfg.CAMERA_TYPE, camera_id=cfg.CAMERA_ID, mirror=cfg.MIRROR_CAMERA) enabled = False # включение управления жестами @@ -60,13 +71,22 @@ def main(): dome_processed = False while True: - frame = camera.get_frame() - if frame is None: - continue - - result = detector.detect(frame) - if result['success']: - landmarks = result['landmarks'] + # Получение кадра и скелета в зависимости от типа скелетного детектора + 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) @@ -111,7 +131,11 @@ def main(): ''' # Отрисовка - vis = detector.draw_landmarks(frame, result['pose_landmarks']) + 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': @@ -131,8 +155,13 @@ def main(): if cv2.waitKey(1) & 0xFF == ord('q'): break + + if cfg.POSE_DETECTOR == 'oak': + detector.release() + else: + camera.release() + detector.pose.close() - camera.release() cv2.destroyAllWindows() robot.quit() diff --git a/submodules/gesture_detection b/submodules/gesture_detection index a751f51..e2e24a5 160000 --- a/submodules/gesture_detection +++ b/submodules/gesture_detection @@ -1 +1 @@ -Subproject commit a751f510788aa8cc30a30687eb2270244ec55279 +Subproject commit e2e24a5e7eb09e43cdb163985ec50e4d6ad32863 diff --git a/submodules/robot_controller_base b/submodules/robot_controller_base index 9f65622..b34711e 160000 --- a/submodules/robot_controller_base +++ b/submodules/robot_controller_base @@ -1 +1 @@ -Subproject commit 9f65622b633900e07ee1f9dbe9e33a5add3167ae +Subproject commit b34711ebc690cf9199b15bc3074e08f258303003