流程代码开发

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"}}
+{"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}}

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()

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