local files

reachy_mini_red_light_green_light/main.py

@@ -15,51 +15,10 @@ import cv2
 import mediapipe as mp
 from importlib.resources import files
 import reachy_mini_red_light_green_light
+from reachy_mini_red_light_green_light.people_detector import PeopleDetector
 
 UPPER_BODY_LANDMARKS = [11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22]
 
-
-class PeopleDetector:
-    def __init__(self, model_path="pose_landmarker_full.task", max_poses=5):
-        # Initialisation de MediaPipe Tasks
-        BaseOptions = mp.tasks.BaseOptions
-        PoseLandmarker = mp.tasks.vision.PoseLandmarker
-        PoseLandmarkerOptions = mp.tasks.vision.PoseLandmarkerOptions
-        VisionRunningMode = mp.tasks.vision.RunningMode
-
-        self.pose_landmarker = PoseLandmarker.create_from_options(
-            PoseLandmarkerOptions(
-                base_options=BaseOptions(model_asset_path=model_path),
-                running_mode=VisionRunningMode.IMAGE,
-                num_poses=max_poses,
-            )
-        )
-
-    def detect(self, frame, draw=True):
-        mp_image = mp.Image(
-            image_format=mp.ImageFormat.SRGB,
-            data=cv2.cvtColor(frame, cv2.COLOR_BGR2RGB),
-        )
-        result = self.pose_landmarker.detect(mp_image)
-
-        h, w, _ = frame.shape
-        detected_people = []
-
-        if result.pose_landmarks:
-            for person_landmarks in result.pose_landmarks:
-                person = {"landmarks": []}
-                for idx, lm in enumerate(person_landmarks):
-                    if idx not in UPPER_BODY_LANDMARKS:
-                        continue
-                    cx, cy = int(lm.x * w), int(lm.y * h)
-                    person["landmarks"].append((cx, cy))
-                    if draw:
-                        cv2.circle(frame, (cx, cy), 3, (0, 255, 0), -1)
-                detected_people.append(person)
-
-        return frame, detected_people
-
-
 class RedLightGreenLightApp(ReachyMiniApp):
     def run(self, reachy_mini: ReachyMini, stop_event: threading.Event):
         # cap = find_camera()