更新检测人后检测头部和手部代码

tool/person-head-coord.py

@@ -0,0 +1,61 @@
+import cv2
+from yolov8_det import analysis_yolov8
+from ultralytics import YOLO
+
+class VideoAnalysis():
+    @staticmethod
+    def detect_head(image, person_coord, confidence, yolo_head_model, yolo_hand_model):
+        """
+        截取后图像保存视频
+        """
+        x_min, y_min, x_max, y_max = person_coord['person']
+        # 提取出人的区域
+        person_frame = image[int(y_min)-100:int(y_max)+100, int(x_min)-100:int(x_max)+100, :]
+
+        # 使用 YOLO 模型分析人体中头部和手部区域
+        head_result = analysis_yolov8(person_frame, YOLO(yolo_head_model), confidence)
+        hand_result = analysis_yolov8(person_frame, YOLO(yolo_hand_model), confidence)
+
+        # 如果检测到头部，返回头部坐标
+        if not head_result:
+            return head_result[1]
+
+        # 如果检测到手部，返回手部坐标
+        elif not head_result:
+            return hand_result[3]
+
+
+    @staticmethod
+    def frame_analysis(test_video_path, yolo_model, confidence, yolo_head_model, yolo_hand_model):
+        # 读取测试视频
+        video_capture = cv2.VideoCapture(test_video_path)
+        coords = []
+        while True:
+            result, frame_img = video_capture.read()
+            if result:
+                # 使用 YOLO 模型分析人的区域
+                person_result = analysis_yolov8(frame_img, YOLO(yolo_model), confidence)
+                # 检测头部，返回头部和手部坐标
+                coord = VideoAnalysis.detect_head(frame_img, person_result[0], confidence, yolo_head_model, yolo_hand_model)
+
+                if coord is not None:
+                    coords.append(coord)
+            else:
+                break
+        video_capture.release()
+        print(coords)
+
+        # 返回坐标信息列表
+        return coords
+
+def main():
+    test_video_path = r"E:\git-Xznsh\XZNSH-Code-AI\xznsh_flow2\video\test01_0.avi"
+    yolo_model = r"E:\git-Xznsh\XZNSH-Code-AI\xznsh_flow2\all_labels.pt"
+    yolo_head_model = r"E:\git-Xznsh\XZNSH-Code-AI\xznsh_flow2\head.pt"
+    yolo_hand_model = r"E:\git-Xznsh\XZNSH-Code-AI\7.27\best.pt"
+    confidence = {"person": 0.5, "head": 0.5, "phone": 0.5, "hand": 0.5}
+    VideoAnalysis.frame_analysis(test_video_path, yolo_model, confidence, yolo_head_model, yolo_hand_model)
+
+
+if __name__ == '__main__':
+    main()