You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
89 lines
4.4 KiB
Python
89 lines
4.4 KiB
Python
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) |