From ca989dfffd5d9b2a3a2bea698ca121b7d56395c9 Mon Sep 17 00:00:00 2001
From: zhouyang <fjzhouyang@163.com>
Date: Fri, 14 Jul 2023 16:44:56 +0800
Subject: [PATCH] =?UTF-8?q?=E6=B5=81=E7=A8=8B=E4=BB=A3=E7=A0=81=E5=BC=80?=
 =?UTF-8?q?=E5=8F=91?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 xznsh_flow/cfg.json      |  2 +-
 xznsh_flow/xznsh_main.py | 59 ++++++++++++++++++++++++++++++++++++++++
 xznsh_flow/yolov8_det.py | 15 ++++++++++
 3 files changed, 75 insertions(+), 1 deletion(-)
 create mode 100644 xznsh_flow/xznsh_main.py
 create mode 100644 xznsh_flow/yolov8_det.py

diff --git a/xznsh_flow/cfg.json b/xznsh_flow/cfg.json
index afa6a0b..fcd04fd 100644
--- a/xznsh_flow/cfg.json
+++ b/xznsh_flow/cfg.json
@@ -1 +1 @@
-{"log_path": "xznsh.log", "frame": 0.05, "camera": {"camera_01": "rtsp://admin:@192.168.10.18", "camera_02": "rtsp://admin:@192.168.10.12"}}
\ No newline at end of file
+{"log_path": "xznsh.log", "frame": 0.05, "camera": {"camera_01": "rtsp://admin:@192.168.10.18", "camera_02": "rtsp://admin:@192.168.10.12"}, "model_path": {"person": "person.pt", "head": "xxx", "phone": "xxx"}, "confidence": {"person": 0.5, "head": 0.5, "phone": 0.5}}
\ No newline at end of file
diff --git a/xznsh_flow/xznsh_main.py b/xznsh_flow/xznsh_main.py
new file mode 100644
index 0000000..a76f6cf
--- /dev/null
+++ b/xznsh_flow/xznsh_main.py
@@ -0,0 +1,59 @@
+import os
+import time
+
+import json
+from queue import Queue, Empty
+from threading import Thread
+from log import logger
+from ultralytics import YOLO
+from yolov8_det import analysis_yolov8
+from capture_queue import CAMERA_QUEUE, camera_mul_thread
+
+with open('cfg.json', 'r') as f:
+    cfg_dict = json.load(f)
+
+
+class ModelInvoke(Thread):
+    """
+    农商行员工打瞌睡，玩手机分析类
+    """
+
+    def __int__(self, camera_name):
+        super(ModelInvoke, self).__init__()
+        self.camera = camera_name
+        self.queue_img = CAMERA_QUEUE[camera_name]
+        self.yolo_model = {'person': YOLO(cfg_dict['model_path']['person']),
+                           'head': YOLO(cfg_dict['model_path']['head']),
+                           'phone': YOLO(cfg_dict['model_path']['phone'])}
+
+    def frame_analysis(self):
+        while True:
+            try:
+                frame_img = self.queue_img.get_nowait()
+            except Empty:
+                time.sleep(0.01)
+                continue
+
+            # 调用模型，逐帧检测
+            results_img = analysis_yolov8(frame=frame_img, model_coco=self.yolo_model['person'],
+                                          confidence=cfg_dict['confidence']['person'])
+
+        # try:
+        #     self.process_frame(y_frame)
+        # except Exception:
+        #     logger.exception(f"{self.name}出错")
+        #     time.sleep(0.1)
+        #     return
+
+    def run(self):
+        pass
+
+
+def process_run():
+    camera_mul_thread()
+    logger.info('程序启动')
+    # todo 分析流程
+
+
+if __name__ == '__main__':
+    process_run()
diff --git a/xznsh_flow/yolov8_det.py b/xznsh_flow/yolov8_det.py
new file mode 100644
index 0000000..987d986
--- /dev/null
+++ b/xznsh_flow/yolov8_det.py
@@ -0,0 +1,15 @@
+def analysis_yolov8(frame, model_coco, confidence):
+    # 第一步：用COCO数据集推理
+    results_coco = model_coco(frame)
+    if not results_coco:
+        return []
+    boxes = results_coco[0].boxes
+    result_list = []
+    for box in boxes:
+        # 过滤置信度0.5以下目标
+        if float(box.conf) < confidence:
+            continue
+        labels_name = model_coco.names[int(box.cls)]
+        label_xyxy_dict = {labels_name: box.xyxy[0].tolist()}
+        result_list.append(label_xyxy_dict)
+    return result_list