提交demo.py

tool/demo.py

@@ -0,0 +1,89 @@
+from PP_TSMv2_infer import *
+from mediapipe_detection import mediapipe_detect
+import mediapipe as mp
+import cv2
+mp_holistic = mp.solutions.holistic
+
+
+def main(input_path,output_path, face_b, left_hand_b, right_hand_b):
+
+    cap = cv2.VideoCapture(input_path)
+    config = 'D:/download/PaddleVideo1/output/output/pptsm_lcnet_k400_16frames_uniform.yaml'
+    model_file = 'D:/download/PaddleVideo1/output/output/ppTSMv2.pdmodel'  # 推理模型存放地址
+    params_file = 'D:/download/PaddleVideo1/output/output/ppTSMv2.pdiparams'  # 推理模型参数存放地址
+    batch_size = 1  # 输出推理模型
+    infer,predictor = PP_TSMv2_predict().create_inference_model(config,model_file,params_file)
+    res = PP_TSMv2_predict().predict(config, input_path, batch_size, predictor,infer)
+    label = res["topk_class"]
+    if label == 0:
+        label = "Nodding!"
+    elif label == 1:
+        label = "not playing phone!"
+    elif label == 2:
+        label = "not sleep!"
+    elif label == 3:
+        label = "playing phone!"
+    elif label == 4:
+        label = "sleep!"
+    else:
+        pass
+    fps_video = cap.get(cv2.CAP_PROP_FPS)
+    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    codec = cv2.VideoWriter_fourcc(*'XVID')
+    video_name = os.path.basename(input_path)
+    out = cv2.VideoWriter(output_path + "/" + video_name, codec, fps_video, (frame_width, frame_height))
+    with mp_holistic.Holistic(model_complexity=2,min_detection_confidence=0.5, min_tracking_confidence=0.5) as holistic:
+        while True:
+            ret, frame = cap.read()
+            if not ret:
+                break
+            image, results = mediapipe_detect().mediapipe_detection(frame, holistic)
+            cv2.namedWindow("mediapipe_detections", cv2.WINDOW_AUTOSIZE)
+            if label == "Nodding!":
+                image, res = mediapipe_detect().get_bbox(image, results, face_b, left_hand_b, right_hand_b,label)
+                cv2.putText(image, "the person's head is " + label, (0, 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0),
+                            1)
+            elif label == "sleep!":
+                image, res = mediapipe_detect().get_bbox(image, results, face_b, left_hand_b, right_hand_b,label)
+                cv2.putText(image, "the person is " + label, (0, 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0),
+                            1)
+            elif label == "not sleep!":
+                image, res = mediapipe_detect().get_bbox(image, results, face_b, left_hand_b, right_hand_b,label)
+                cv2.putText(image, "the person is " + label, (0, 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0),
+                            1)
+            elif label == "playing phone!":
+                image, res = mediapipe_detect().get_bbox(image, results, face_b, left_hand_b, right_hand_b,label)
+                cv2.putText(image, "the person'hand is " + label, (0, 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0),
+                            1)
+            elif label == "not playing phone!":
+                image, res = mediapipe_detect().get_bbox(image, results, face_b, left_hand_b, right_hand_b,label)
+                cv2.putText(image, "the person'hand is " + label, (0, 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0),
+                            1)
+
+            cv2.imshow("mediapipe_detections", image)
+            out.write(image)
+            if cv2.waitKey(10) & 0xFF == ord('q'):
+                break
+            out.write(image)
+            # print(res)
+    cap.release()
+    out.release()
+    cv2.destroyAllWindows()
+
+
+if __name__=="__main__":
+    # input = 'D:/download/PaddleVideo1/output/output/after_1/0711-1_0_1.avi'
+    # input = 'D:/download/PaddleVideo1/output/output/after_1/0711-1_398_1.avi'
+    # input = 'D:/download/PaddleVideo1/output/output/after_1/0711-1_597_0.avi'  #正例
+    # input = 'D:/download/PaddleVideo1/output/output/after_1/0711-1_597_1.avi'
+    input = 'D:/download/PaddleVideo1/output/output/after_1/0711-1_796_0.avi'    #正例，推理成功
+    # input = 'D:/download/PaddleVideo1/output/output/after_1/0711-1_796_1.avi'
+    # input = 'D:/download/PaddleVideo1/output/output/after_1/0711-3_0_0.avi'
+    # input = 'D:/download/PaddleVideo1/output/output/after_1/0711-3_1400_0.avi'
+
+    output = "D:/download/PaddleVideo1/output/output1"
+    face_b = 50
+    left_hand_b = 7
+    right_hand_b = 7
+    main(input,output,face_b,left_hand_b,right_hand_b)

tool/mediapipe_detection.py

@@ -7,7 +7,6 @@ import numpy as np
 from mediapipe.framework.formats import landmark_pb2
 import os
 
-
 mp_holistic = mp.solutions.holistic
 
 _PRESENCE_THRESHOLD = 0.5
@@ -163,7 +162,6 @@ class mediapipe_detect:
         result = np.array(result)
         b = bias
         if result.any():
-
             rect = cv2.boundingRect(result)  # 返回值, 左上角的坐标[x,y, w,h]
 
             bbox = [[rect[0] - b, rect[1] - b], [rect[0] + rect[2] + b, rect[1] - b],
@@ -171,7 +169,7 @@ class mediapipe_detect:
 
             return bbox
 
-    def get_bbox(self,image,results,face_b,left_hand_b,right_hand_b):
+    def get_bbox(self, image, results, face_b, left_hand_b, right_hand_b, label):
 
         '''
             主要是根据关键点坐标，绘制矩形框
@@ -214,6 +212,12 @@ class mediapipe_detect:
         lh_bbox = self.Drawing_bbox(left_hand_location, left_hand_b)
         rh_bbox = self.Drawing_bbox(right_hand_location, right_hand_b)
 
+        if label == "Nodding" or label == "sleep!" or label == "not playing phone!":
+            lh_bbox = None
+            rh_bbox = None
+        elif label == "not sleep!" or label == "playing phone!":
+            fl_bbox = None
+
         """调整头部检测框的大小"""
         if fl_bbox is not None:
             """对头部动作检测框微调"""
@@ -260,17 +264,12 @@ class mediapipe_detect:
                         pass
             cv2.rectangle(image, rh_bbox[0], rh_bbox[3], DrawingSpec.color, DrawingSpec.thickness)
 
-
         res = {'face_bbox': fl_bbox, 'hand_bbox': [lh_bbox, rh_bbox]}
 
-
-        # print(res)
-
         return image, res
 
 
 def main(input_path, output_path, face_b, left_hand_b, right_hand_b):
-
     cap = cv2.VideoCapture(input_path)
     frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
     frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
@@ -278,13 +277,14 @@ def main(input_path,output_path,face_b,left_hand_b,right_hand_b):
     codec = cv2.VideoWriter_fourcc(*'XVID')
     video_name = os.path.basename(input_path)
     out = cv2.VideoWriter(output_path + "/" + video_name, codec, fps, (frame_width, frame_height))
+    label = ""
     with mp_holistic.Holistic(min_detection_confidence=0.5, min_tracking_confidence=0.5) as holistic:
         while True:
             ret, frame = cap.read()
             if not ret:
                 break
             image, results = mediapipe_detect().mediapipe_detection(frame, holistic)
-             image,res = mediapipe_detect().get_bbox(image,results,face_b,left_hand_b,right_hand_b)
+            image, res = mediapipe_detect().get_bbox(image, results, face_b, left_hand_b, right_hand_b, label)
             out.write(image)
             cv2.namedWindow("mediapipe_detections", cv2.WINDOW_AUTOSIZE)
             cv2.imshow("mediapipe_detections", image)
@@ -296,18 +296,17 @@ def main(input_path,output_path,face_b,left_hand_b,right_hand_b):
     out.release()
     cv2.destroyAllWindows()
 
+
 if __name__ == "__main__":
-    input = 'D:/download/PaddleVideo1/output/output/after_1/test02_0.avi'
+    # input = 'D:/download/PaddleVideo1/output/output/after_1/test02_0.avi'
     # input = 'D:/download/PaddleVideo1/output/output/after_1/0711-1_0_1.avi'
     # input = 'D:/download/PaddleVideo1/output/output/after_1/0711-3_1400_0.avi'
     # input = "C:/Users/Administrator/Pictures/video_seg_re_hand/test01_3.avi"
     # input = 'C:/Users/Administrator/Pictures/video3.0/sleep/0711-3_7_01_5.avi'
-    # input = " D:/download/PaddleVideo1/output/output/after_1/0711-1_199_0.avi"
+    input = " D:/download/PaddleVideo1/output/output/after_1/0711-1_199_0.avi"
     # input = 'D:/download/PaddleVideo1/output/output/after_1/test05_10750_1.avi'
     output = 'D:/download/PaddleVideo1/output/output/output'
     face_b = 50  # 头部标注框修正值
     left_hand_b = 7  # 左手部分标注框修正值
     right_hand_b = 7  # 右手部分标注框修正值
     main(input, output, face_b, left_hand_b, right_hand_b)
-
-