import numpy as np import cv2 import os import time from ultralytics import YOLO import queue import threading from config import Q_SZ from ModelDet.personDet import analysis_yolov8 from tools import Process_tools # from ModelDet.holisticDet import MediapipeProcess class DealVideo(): def __init__(self,video_file,video_save_file,person_model,mediapipe_model,pptsmv2_model): ''' 加载数据 ''' self.video_file = video_file self.video_save_file = video_save_file # 初始化模型 self.person_model = person_model self.mediapipe_model = mediapipe_model self.pptsmv2_model = pptsmv2_model # 图片检测后队列 self.videoQueue = queue.Queue(maxsize=Q_SZ) self.frameQueue = queue.Queue(maxsize=0) #线程 self.get_video_listThread = threading.Thread(target=self.get_video_list) self.get_video_frameThread = threading.Thread(target=self.get_video_frame) self.person_detThread = threading.Thread(target=self.person_det) self.write_videoThread = threading.Thread(target=self.write_video) def get_video_list(self): ''' 获取数据文件 ''' if os.path.isdir(self.video_file): video_ext = [".mp4", ".avi",".MP4"] for maindir, subdir, file_name_list in os.walk(self.video_file): for filename in file_name_list: apath = os.path.join(maindir, filename) ext = os.path.splitext(apath)[1] if ext in video_ext: self.videoQueue.put(apath) else: self.videoQueue.put(self.video_file) def get_video_frame(self): ''' 对视频进行分帧、每一帧都保存队列 ''' while True: if ~self.videoQueue.empty(): try: video_path = self.videoQueue.get() video_basename = os.path.basename(video_path).split('.')[0] cap = cv2.VideoCapture(video_path) frame_list = [] count_fps = 0 while cap.isOpened(): success, frame = cap.read() if not success: print(video_path,"Ignoring empty camera frame.") break count_fps += 1 frame_dict = {'fps':count_fps,'frame':frame} frame_list.append(frame_dict) video_dict = {'video_path':video_path,'frame_list':frame_list,'cap':cap} self.frameQueue.put(video_dict) except Exception as e: print(e) def person_det(self): ''' 从队列中获取视频帧frame，进行第一步人员的检测 ''' while True: if ~self.frameQueue.empty(): video_frame_dict = self.frameQueue.get() frame_list = video_frame_dict['frame_list'] frame_result_contact = [] for i in range(len(frame_list)): if frame_list[i]["fps"] == i + 1: person_det = analysis_yolov8(frame=frame_list[i]['frame'], model_coco=self.person_model, confidence_set=0.5) # 当前帧检测的结果列表，只包含bboxlist person_list = Process_tools.get_dict_values(person_det) # 保存第一帧结果为对比坐标 if not frame_result_contact: bbox_list_all = Process_tools.change_list_dict(fps1=frame_list[i]["fps"],re_list=person_list) frame_result_contact = bbox_list_all print("frame_result_contact:",frame_result_contact) else: example_dict_list = frame_result_contact cut_list,example_lst,re_dict_lst = Process_tools.analysis_re01_list(example_list=example_dict_list, result_list=person_list) print('cut_list:',cut_list) print('example_sorted_lst:',example_lst) print('re_dict_sorted_lst:',re_dict_lst) # 统计截止时间 time_out_list = Process_tools.statistics_fps(fps_now=frame_list[i]["fps"],re_list=frame_result_contact) if time_out_list: # bbox_list = Process_tools.change_dict_list(time_out_list) # 裁剪保存视频 # cut_dict = {"start_fps":time_out_list[0]['fps'],"stop_fps":frame_list[i]["fps"],'bbox_list':bbox_list} frame_result_contact = [item for item in frame_result_contact if item not in time_out_list] # 有目标减少情况 if example_lst: # cut_dict = {"start_fps":frame_result_contact[0]['fps'],"stop_fps":frame_list[i]["fps"],'bbox_list':example_lst} frame_result_contact = [item for item in frame_result_contact if item not in example_lst] # 有新添加目标情况 if re_dict_lst: update_list = Process_tools.change_list_dict(fps1=frame_list[i]["fps"],re_list=re_dict_lst) frame_result_contact = frame_result_contact + update_list print('frame_result_contact:',frame_result_contact) def write_video(self): ''' 保存成视频 ''' while True: if ~self.frameQueue.empty(): video_frame_dict = self.frameQueue.get() video_basename = os.path.basename(video_frame_dict['video_path']) video_name_save = os.path.join(self.video_save_file, video_basename) # 原视频帧率和尺寸 cap = video_frame_dict['cap'] fps = cap.get(cv2.CAP_PROP_FPS) size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))) # 保存视频 videoWriter =cv2.VideoWriter(video_name_save,cv2.VideoWriter_fourcc('X','V','I','D'),fps,size) frame_list = video_frame_dict['frame_list'] for i in range(len(frame_list)): if frame_list[i]["fps"] == i + 1: videoWriter.write(frame_list[i]["frame"]) else: break def run(self): self.get_video_listThread.start() self.get_video_frameThread.start() self.person_detThread.start() # self.write_videoThread.start() if __name__ == '__main__': # 每个视频的时长（单位秒） dertTime = 5 video = "E:/Bank_files/Bank_02/dataset/video_test" video_save = 'videos_codes_2' person_model = YOLO("model_file/yolov8x_person.pt") # get_seg_video(video_file=video,video_save_path=video_save,dertTime=dertTime) deal = DealVideo(video_file=video,video_save_file=video_save,person_model=person_model,mediapipe_model='model_file/yolov8x_person.pt',pptsmv2_model='model_file/yolov8x_person.pt') deal.run()