增加了一些逻辑判断

xznsh_paddle_scrfd/d_face.py

@@ -103,14 +103,11 @@ class SCRFD():
         kpss[:, :, 1] = (kpss[:, :, 1] - padh) * ratioh
         indices = cv2.dnn.NMSBoxes(bboxes.tolist(), scores.tolist(), self.confThreshold, self.nmsThreshold)
         # 根据阈值拦截后的人脸框元组
-        # print("indices", indices)
         if indices:
             face_flag["face"] += 1
             face_flag["frame"].append(count)
-            # print("frame", count)
         for i in indices:
             i = i[0]
-            # 这个就是xmin, ymin, xamx, ymax （人脸框 ）
             xmin, ymin, xamx, ymax = int(bboxes[i, 0]), int(bboxes[i, 1]), int(bboxes[i, 0] + bboxes[i, 2]), int(bboxes[i, 1] + bboxes[i, 3])
             cv2.rectangle(srcimg, (xmin, ymin), (xamx, ymax), (0, 0, 255), thickness=2)
             for j in range(5):
@@ -120,11 +117,7 @@ class SCRFD():
 
 def face_detection(onnxmodel,image,crop_image,x1,y1,x2,y2,face_flag,count):
 
-    # onnxmodel = r'E:\pythonProject\xznsh\scrfd-opencv\weights\scrfd_10g_kps.onnx'
     mynet = SCRFD(onnxmodel, confThreshold=0.8, nmsThreshold=0.5)
     outimg= mynet.detect(crop_image,face_flag,count)
     image[y1:y2,x1:x2] = outimg
-    # cv2.imshow("img",image)
-    # cv2.waitKey(0)
-    # cv2.imwrite(os.path.join(outpath,file), outimg)
     return image

xznsh_paddle_scrfd/d_people.py

@@ -4,26 +4,16 @@ from tqdm import tqdm
 from ultralytics import YOLO
 from ultralytics.yolo.utils.plotting import Annotator
 from d_face import face_detection
-# from dec_duration import detect_duration
-
-# model_police = YOLO('models/police0508.pt')
-# model_electromobile = YOLO('models/electromobile0509.pt')
-# model_coco = YOLO(r"E:\pythonProject\xznsh\scrfd-opencv\weights\best20230606.pt")
-# action_model = YOLO(r'E:\pythonProject\xznsh\scrfd-opencv\weights\action_recognition.pt')
 
 def analysis_video(source_path,output_path,people_modle_path,face_modle_path,action_modle_path):
-# def analysis_video(source_path, output_path):
     model_coco = YOLO(people_modle_path)
     action_model = YOLO(action_modle_path)
-    # start_time = time.time()
     cap = cv2.VideoCapture(source_path)
 
     # 直接从视频的第 frameToStart 帧开始
-    frameToStart = 1430
+    frameToStart = 0
     cap.set(cv2.CAP_PROP_POS_FRAMES, frameToStart)
 
-    # total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-    # tqdm_desc = 'Processing video frames'
     fps = int(cap.get(cv2.CAP_PROP_FPS))
     width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
     height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
@@ -50,23 +40,18 @@ def analysis_video(source_path,output_path,people_modle_path,face_modle_path,act
         "head":0,
         "tail":0
         }
-    # with tqdm(total=total_frames, desc=tqdm_desc) as pbar:
     while cap.isOpened():
         # Read a frame from the video
         success, frame = cap.read()
-        # cv2.imshow("f",frame)
-        # cv2.waitKey(0)
 
         count += 1
         if success:
             # 第一步：用COCO数据集推理
             results_coco = model_coco(frame)
             action_result = action_model(frame)
-            # print("*"*100, results_coco.masks)
             
             for r in results_coco:
                 annotator = Annotator(frame, line_width=1)   
-                # print("*"*100, r.boxes)
                 boxes = r.boxes
                 for box in boxes:                             
                     b = box.xyxy[0]  # get box coordinates in (x1,y1,x2,y2) format  #tensor([ 677.5757,  147.2737, 1182.3381,  707.2565])
@@ -82,35 +67,31 @@ def analysis_video(source_path,output_path,people_modle_path,face_modle_path,act
                             XJ_dict['head'] = count
                         else:
                             XJ_dict['tail'] = count
-                        # print ("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!巡检人员")
                         crop_img = frame[b_i[1]:b_i[3],b_i[0]:b_i[2]]
                         frame = face_detection(face_modle_path,frame,crop_img,b_i[0],b_i[1],b_i[2],b_i[3],face_flag,count)
-                        # annotator = Annotator(frame, line_width=1)
                     annotator.box_label(b, model_coco.names[int(c)]+str(confidence),(0,0,255))
 
             for r_a in action_result:
                 annotator_a = Annotator(frame, line_width=1)    
-                # print("*"*100, r.boxes)
                 boxes_a = r_a.boxes
-                if boxes_a:
-                    action_flag["action"] += 1
-                    action_flag["action_frame"].append(count)
                 for box_a in boxes_a:                             
                     b_a = box_a.xyxy[0]  # get box coordinates in (x1,y1,x2,y2) format  #tensor([ 677.5757,  147.2737, 1182.3381,  707.2565])
-                    # b_i_a = b_a.int() + 1
                     c_a = box_a.cls  # tensor([0.])
+                    
                     confidence_a = float(box_a.conf)
                     confidence_a = round(confidence_a, 2)
                     # 过滤置信度0.5以下目标
                     if confidence_a < 0.5:
                         continue
-                    # annotator = Annotator(frame, line_width=1)
+                    # 当类别为check
+                    if c_a.int() == 1:
+                        action_flag["action"] += 1
+                        action_flag["action_frame"].append(count)
+
                     annotator_a.box_label(b_a, action_model.names[int(c_a)]+str(confidence_a),(255,0,0))
                     
-            annotated_frame_coco = annotator.result()
             annotated_a_frame_coco = annotator_a.result()
             output_movie.write(annotated_a_frame_coco)
-            # pbar.update(1)
         else:
             # Break the loop if the end of the video is reached
             break
@@ -124,7 +105,4 @@ def analysis_video(source_path,output_path,people_modle_path,face_modle_path,act
     action_frame = [round(_ /fps,2) for _ in action_flag["action_frame"]] 
     s_action = ', '.join(map(str, action_frame))
 
-    return diff,face_flag,s,action_flag,s_action
-
-# if __name__ == '__main__':
-#     analysis_video(r"E:\pythonProject\xznsh\scrfd-opencv\data\xj_video.mp4", r"E:\pythonProject\xznsh\scrfd-opencv\result\video.mp4")
+    return diff,face_flag,s,action_flag,s_action