单模型多标签

xznsh_flow/cal_utils.py

@@ -19,3 +19,8 @@ def same_position_judgment(be_coord, af_coord):
     distance_l = math.sqrt((af_coord[0] - be_coord[0]) ** 2 + (af_coord[1] - be_coord[1]) ** 2)
     distance_r = math.sqrt((af_coord[2] - be_coord[2]) ** 2 + (af_coord[3] - be_coord[3]) ** 2)
     return (distance_l + distance_r) / 2
+
+class LabelTreat():
+    def __init__(self,label):
+        self.
+

xznsh_flow/xznsh_main.py

@@ -23,10 +23,14 @@ class FrameToVideo(Thread):
         super(FrameToVideo, 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'])}
+        # self.yolo_model = {'person': YOLO(cfg_dict['model_path']['person']),
+        #                    'head': YOLO(cfg_dict['model_path']['head'])}
+
+        self.yolo_model = YOLO(cfg_dict['model_path']['all'])
+
         self.person_target_list = []
         self.head_target_list = []
+        self.phone_target_list = []
 
     @staticmethod
     def save_video(frames, fps, fourcc, video_path, w_h_size):
@@ -85,8 +89,10 @@ class FrameToVideo(Thread):
                 # label that self define maybe different from model
                 if label == 'head':
                     coord = line['head']
-                else:
+                elif label == 'person':
                     coord = line['person']
+                else:
+                    coord = line['phone']
                 split_x, split_y = self.boundary_treat(frame_x, frame_y, coord)
 
                 # 裁剪大一圈,固定裁剪范围
@@ -138,11 +144,16 @@ class FrameToVideo(Thread):
 
             new_person_target_list = []
             new_head_target_list = []
+            new_phone_target_list = []
             # 调用模型，逐帧检测
-            person_coord_list = analysis_yolov8(frame=frame_img, model_coco=self.yolo_model['person'],
-                                                confidence=cfg_dict['confidence']['person'])
-            head_coord_list = analysis_yolov8(frame=frame_img, model_coco=self.yolo_model['head'],
-                                              confidence=cfg_dict['confidence']['head'])
+            # person_coord_list = analysis_yolov8(frame=frame_img, model_coco=self.yolo_model['person'],
+            #                                     confidence=cfg_dict['confidence']['person'])
+            # head_coord_list = analysis_yolov8(frame=frame_img, model_coco=self.yolo_model['head'],
+            #                                   confidence=cfg_dict['confidence']['head'])
+
+            person_coord_list, head_coord_list, phone_coord_list = analysis_yolov8(frame=frame_img,
+                                                                                   model_coco=self.yolo_model,
+                                                                                   confidence=cfg_dict['confidence'])
 
             frame_y, frame_x, _ = frame_img.shape
             logger.debug(f'帧尺寸，y:{frame_y},x:{frame_x}')
@@ -150,6 +161,8 @@ class FrameToVideo(Thread):
                                  frame_x, frame_y, frame_img, 'person')
             self.target_analysis(self.head_target_list, new_head_target_list, head_coord_list, frame_x,
                                  frame_y, frame_img, 'head')
+            self.target_analysis(self.phone_target_list, new_phone_target_list, phone_coord_list, frame_x,
+                                 frame_y, frame_img, 'phone')
 
     def run(self):
         self.frame_analysis()

xznsh_flow/yolov8_det.py

@@ -4,12 +4,23 @@ def analysis_yolov8(frame, model_coco, confidence):
     if not results_coco:
         return []
     boxes = results_coco[0].boxes
-    result_list = []
+    person_result_list = []
+    head_result_list = []
+    phone_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
+        if labels_name == 'person':
+            if float(box.conf) >= confidence['person']:
+                label_xyxy_dict = {labels_name: box.xyxy[0].tolist()}
+                person_result_list.append(label_xyxy_dict)
+
+        elif labels_name == 'head':
+            if float(box.conf) >= confidence['head']:
+                label_xyxy_dict = {labels_name: box.xyxy[0].tolist()}
+                head_result_list.append(label_xyxy_dict)
+        else:
+            if float(box.conf) >= confidence['phone']:
+                label_xyxy_dict = {labels_name: box.xyxy[0].tolist()}
+                phone_result_list.append(label_xyxy_dict)
+
+    return person_result_list,head_result_list,phone_result_list