更新mediapipe标注框修正问题

tool/mediapipe_detection.py

@@ -7,6 +7,7 @@ import numpy as np
 from mediapipe.framework.formats import landmark_pb2
 import os
 
+
 mp_holistic = mp.solutions.holistic
 
 _PRESENCE_THRESHOLD = 0.5
@@ -167,7 +168,9 @@ class mediapipe_detect:
 
             bbox = [[rect[0] - b, rect[1] - b], [rect[0] + rect[2] + b, rect[1] - b],
                     [rect[0] - b, rect[1] + rect[3] + b], [rect[0] + rect[2] + b, rect[1] + rect[3] + b]]  #四个角的坐标
+            
             return bbox
+        
     def get_bbox(self,image,results,face_b,left_hand_b,right_hand_b):
 
         '''
@@ -178,6 +181,8 @@ class mediapipe_detect:
 
         image.flags.writeable = True
         image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+        h,w,g = image.shape
+        # print("h:",h,"w:",w,"g:",g)
 
         """获取头部、手部关键点"""
         face_location = draw_landmarks(
@@ -216,16 +221,51 @@ class mediapipe_detect:
             fl_bbox[3][1] = fl_bbox[3][1] - 30
             fl_bbox[0][0] = fl_bbox[0][0] + 30
             fl_bbox[0][1] = fl_bbox[0][1] + 5
+            # print(fl_bbox)
+            for i in range(0,4):
+                for j in range(0,2):
+                    if fl_bbox[i][j] < 0:
+                        fl_bbox[i][j] = 0
+                    elif fl_bbox[i][0] > w:
+                        fl_bbox[i][0] = w
+                    elif fl_bbox[i][1] > h :
+                        fl_bbox[i][1] = h
+                    else:
+                        pass
             cv2.rectangle(image, fl_bbox[0], fl_bbox[3],DrawingSpec.color, DrawingSpec.thickness)
+
         if lh_bbox is not None:
+            for i in range(0,4):
+                for j in range(0,2):
+                    if lh_bbox[i][j] < 0:
+                        lh_bbox[i][j] = 0
+                    elif lh_bbox[i][0] > w:
+                        lh_bbox[i][0] = w
+                    elif lh_bbox[i][1] > h :
+                        lh_bbox[i][1] = h
+                    else:
+                        pass
             cv2.rectangle(image, lh_bbox[0], lh_bbox[3],DrawingSpec.color, DrawingSpec.thickness)
+
         if rh_bbox is not None:
+            for i in range(0,4):
+                for j in range(0,2):
+                    if rh_bbox[i][j] < 0:
+                        rh_bbox[i][j] = 0
+                    elif rh_bbox[i][0] > w:
+                        rh_bbox[i][0] = w
+                    elif rh_bbox[i][1] > h:
+                        rh_bbox[i][1] = h
+                    else:
+                        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]}
-        cv2.namedWindow("mediapipe_detections", cv2.WINDOW_AUTOSIZE)
-        cv2.imshow("mediapipe_detections",image)
-        # print(result_dict)
+
+
+        # print(res)
+
         return image,res
 
 
@@ -236,8 +276,8 @@ def main(input_path,output_path,face_b,left_hand_b,right_hand_b):
      frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
      fps = int(cap.get(cv2.CAP_PROP_FPS))
      codec = cv2.VideoWriter_fourcc(*'XVID')
-     video_name = os.path.basename(input_path).split('.')[0]
-     out = cv2.VideoWriter(output_path +"/"+ video_name+".avi", codec, fps, (frame_width, frame_height))
+     video_name = os.path.basename(input_path)
+     out = cv2.VideoWriter(output_path +"/"+ video_name, codec, fps, (frame_width, frame_height))
      with mp_holistic.Holistic(min_detection_confidence=0.5, min_tracking_confidence=0.5) as holistic:
          while True:
              ret, frame = cap.read()
@@ -246,20 +286,28 @@ def main(input_path,output_path,face_b,left_hand_b,right_hand_b):
              image, results = mediapipe_detect().mediapipe_detection(frame,holistic)
              image,res = mediapipe_detect().get_bbox(image,results,face_b,left_hand_b,right_hand_b)
              out.write(image)
-
+             cv2.namedWindow("mediapipe_detections", cv2.WINDOW_AUTOSIZE)
+             cv2.imshow("mediapipe_detections", image)
+             # print(res)
              if cv2.waitKey(10) & 0xFF == ord('q'):
                  break
 
-
      cap.release()
      out.release()
      cv2.destroyAllWindows()
 
 if __name__ == "__main__":
-    input = 'D:/inference/mediapipe/mediapipe/python/video/test/test02_2.avi'
-    output = 'D:/inference/mediapipe/mediapipe/python/video/output_video'
+    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/test05_10750_1.avi'
+    output = 'D:/download/PaddleVideo1/output/output/output'
     face_b = 50          #头部标注框修正值
-    left_hand_b = 7      #头部标注框修正值
-    right_hand_b = 7     #头部标注框修正值
+    left_hand_b = 7      #左手部分标注框修正值
+    right_hand_b = 7     #右手部分标注框修正值
     main(input,output,face_b,left_hand_b,right_hand_b)
 
+