From 5dc31ce4fc509c162c28ae95380ed08dd00dad36 Mon Sep 17 00:00:00 2001
From: fanpt <320622572@qq.com>
Date: Thu, 3 Aug 2023 18:03:33 +0800
Subject: [PATCH] =?UTF-8?q?=E6=9B=B4=E6=96=B0=E6=A3=80=E6=B5=8B=E4=BA=BA?=
 =?UTF-8?q?=E5=90=8E=E6=A3=80=E6=B5=8B=E5=A4=B4=E9=83=A8=E5=92=8C=E6=89=8B?=
 =?UTF-8?q?=E9=83=A8=E4=BB=A3=E7=A0=81?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 tool/person-head-coord.py | 61 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 61 insertions(+)
 create mode 100644 tool/person-head-coord.py

diff --git a/tool/person-head-coord.py b/tool/person-head-coord.py
new file mode 100644
index 0000000..7d264b7
--- /dev/null
+++ b/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()