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.
79 lines
2.5 KiB
Python
79 lines
2.5 KiB
Python
import cv2
|
|
import json
|
|
import os
|
|
|
|
def draw_video(video_path,labels_file,video_save):
|
|
|
|
cap = cv2.VideoCapture(video_path)
|
|
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))
|
|
size = (width,height)
|
|
video_basename = os.path.basename(video_path)
|
|
|
|
labels_pptsm_list = ["Nodding",'Not_Playing_Mobile','Not_sleet','Playing_Mobile','slppe']
|
|
|
|
with open(labels_file, "r") as json_file:
|
|
data_dict = json.load(json_file)
|
|
labels_list = data_dict['big_dict']
|
|
|
|
video_save_file = video_save
|
|
os.makedirs(video_save_file, exist_ok=True)
|
|
video_save_path = os.path.join(video_save_file, video_basename)
|
|
videoWriter =cv2.VideoWriter(video_save_path,cv2.VideoWriter_fourcc('X','V','I','D'),fps,size)
|
|
|
|
count_fps = 0
|
|
while cap.isOpened():
|
|
success, frame = cap.read()
|
|
print('count_fps:',count_fps)
|
|
if not success:
|
|
print(video_path,"Ignoring empty camera frame.")
|
|
# print('video_fps:',video_fps,'count_fps:',count_fps)
|
|
break
|
|
|
|
re_anno_list = get_bbox_list(count_fps=count_fps,bbox_dict_list=labels_list)
|
|
|
|
for re_dic in re_anno_list:
|
|
|
|
re_txt = re_dic[0]
|
|
re_bbox = re_dic[1]
|
|
|
|
cv2.rectangle(frame, (int(re_bbox[0]), int(re_bbox[1])),(int(re_bbox[2]), int(re_bbox[3])), (0, 255, 255), 2)
|
|
cv2.putText(frame, re_txt, (int(re_bbox[0]) - 10, int(re_bbox[1]) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 0, 255), 2)
|
|
|
|
if re_txt != 'person':
|
|
re_txt_1 = labels_pptsm_list[int(re_txt)]
|
|
cv2.putText(frame, re_txt_1, (int(50), (int(50) + int(re_txt)*30)), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (255, 0, 255), 3)
|
|
|
|
videoWriter.write(frame)
|
|
|
|
count_fps += 1
|
|
|
|
|
|
|
|
videoWriter.release()
|
|
cap.release()
|
|
|
|
def get_bbox_list(count_fps,bbox_dict_list):
|
|
|
|
re_list = []
|
|
|
|
for i,bbox_dict in enumerate(bbox_dict_list):
|
|
|
|
startfps = int(bbox_dict['startfps'])
|
|
stopfps = int(bbox_dict['stopfps'])
|
|
|
|
if count_fps in range(startfps,stopfps):
|
|
|
|
labels = bbox_dict['labels']
|
|
bbox_list = bbox_dict['bbox']
|
|
|
|
re_list.append([labels,bbox_list])
|
|
|
|
return re_list
|
|
|
|
if __name__ == '__main__':
|
|
|
|
draw_video(video_path='E:/Bank_files/Bank_02/dataset/video_kf/02.mp4',
|
|
labels_file='E:/Bank_files/Bank_02/process_file/test_video/02.json',
|
|
video_save='E:/Bank_files/Bank_02/process_file/test_video') |