XZNSH-Code-AI/Bank_second_part/detect_process/video_process.py

import cv2
import os
import time
import mediapipe as mp

from ultralytics import YOLO 
import queue

import threading
from config import Q_SZ

from personDet import analysis_yolov8
import tools_function
from holisticDet import MediapipeProcess
import mediapipe_detection_image
from PP_TSMv2_infer import PP_TSMv2_predict
import shutil
import json



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.predictor = pptsmv2_model[1]
        self.infer = pptsmv2_model[0]
        self.batch_size = 1

        # 队列
        self.videoQueue = queue.Queue(maxsize=Q_SZ)
        self.videoQueue2 = queue.Queue(maxsize=Q_SZ)
        self.cutbboxQueue = queue.Queue(maxsize=0)
        self.videodetQueue = queue.Queue(maxsize=0)
        self.videoQueue3 = queue.Queue(maxsize=0)
        self.videoreturnQueue = 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.write_videoThread = threading.Thread(target=self.write_video)
        self.head_hands_detThread = threading.Thread(target=self.head_hands_det)
        self.video_select_dectThread = threading.Thread(target=self.video_select_dect)
        self.select_video_pathThread = threading.Thread(target=self.select_video_path)
        self.analysis_return_meassageThread = threading.Thread(target=self.analysis_return_meassage)
        
        

    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():

                time.sleep(1)
            
            else:
                
                video_path = self.videoQueue.get()  
                
                # video_basename = os.path.basename(video_path).split('.')[0]

                # print('video_path:',video_path)

                cap = cv2.VideoCapture(video_path)
                video_fps = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
                
                # frame_list = []
                count_fps = 0
                frame_result_contact = []
                count_fps_del = 0

                while cap.isOpened():
                    success, frame = cap.read()
                    if not success:
                        print(video_path,"Ignoring empty camera frame.")
                        # print('video_fps:',video_fps,'count_fps:',count_fps)
                        break
                    
                    # print('count_fps_read_video=',count_fps)
                    imgsize = frame.shape

                    person_det = analysis_yolov8(frame=frame,
                                                     model_coco=self.person_model,
                                                     confidence_set=0.5)                    

                    person_list = tools_function.get_dict_values(person_det)

                    if frame_result_contact:
                            start_fps = frame_result_contact[0]['fps']
                        
                    else:
                            start_fps = count_fps

                    if count_fps == (video_fps - 1):

                        video_end = True

                    else:

                        video_end = False

                    if person_list:

                        count_fps_del_re,updata_result_contact = self.analysis_by_bbox(imgsize=imgsize,
                                                                                            detect_result=person_list,
                                                                                            dertpara=4,
                                                                                            start_fps=start_fps,
                                                                                            now_fps=count_fps,
                                                                                            label_name='person',
                                                                                            video_path=video_path,
                                                                                            frame_result_contact=frame_result_contact,
                                                                                            parameter_fps=150,
                                                                                            count_fps_del=count_fps_del,
                                                                                            video_end=video_end
                                                                                            )
                        count_fps_del = count_fps_del_re
                        frame_result_contact = updata_result_contact

                    count_fps  += 1

    def head_hands_det(self):

        # print('head_hands_detaohgaogh')

        while True:

            if self.videoQueue3.empty():

                time.sleep(1)
            else:

                t0 = time.time()
                video_path = self.videoQueue3.get()  


                # print('video_path_head_hands_det:',video_path)

                cap = cv2.VideoCapture(video_path)
                video_fps = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

                # frame_list = []
                count_fps = 0

                head_result_contact = []
                hands_result_contact = []
                count_fps_del_head = 0
                count_fps_del_hand = 0

                while cap.isOpened():
                    success, frame = cap.read()
                    if not success:
                        print(video_path,"Ignoring empty camera frame.")
                        # print('count_fps:',count_fps,'video_fps:',video_fps)
                        break

                    # print('count_fps_read_video=',count_fps)
                    imgsize = frame.shape

                    # 模型推理
                    hh_result = MediapipeProcess.mediapipe_det(image=frame,
                                                                        holistic=self.mediapipe_model)
                    hh_result_dict = MediapipeProcess.get_analysis_result(image=frame,results=hh_result)

                    # # 获得当前坐标列表
                    head_result = hh_result_dict['face_bbox']
                    head_result_1 = tools_function.select_list(head_result)
                    hands_result = hh_result_dict['hand_bbox']
                    hands_result_1 = tools_function.select_list(hands_result)

                    
                    if count_fps == (video_fps - 1):

                        # print('count_fps:',count_fps,'video_fps:',video_fps)

                        video_end = True
                    else:

                        video_end = False

                    # 统一修正坐标,分别对头和手进行分析
                    if head_result_1:

                        if head_result_contact:
                            start_fps = head_result_contact[0]['fps']
                        else:
                            start_fps = count_fps


                        
                        count_fps_del_re,updata_result_contact = self.analysis_by_bbox(imgsize=imgsize,
                                                                                        detect_result=head_result_1,
                                                                                        dertpara=1,
                                                                                        start_fps=start_fps,
                                                                                        now_fps=count_fps,
                                                                                        label_name='head',
                                                                                        video_path=video_path,
                                                                                        frame_result_contact=head_result_contact,
                                                                                        parameter_fps=80,
                                                                                        count_fps_del=count_fps_del_head,
                                                                                        video_end=video_end
                                                                                        )
                        count_fps_del_head = count_fps_del_re
                        head_result_contact = updata_result_contact

                    if hands_result_1:

                        if hands_result_contact:
                            start_fps = hands_result_contact[0]['fps']
                        else:
                            start_fps = count_fps

                        count_fps_del_re,updata_result_contact = self.analysis_by_bbox(imgsize=imgsize,
                                                                                        detect_result=hands_result_1,
                                                                                        dertpara=2,
                                                                                        start_fps=start_fps,
                                                                                        now_fps=count_fps,
                                                                                        label_name='hands',
                                                                                        video_path=video_path,
                                                                                        frame_result_contact=hands_result_contact,
                                                                                        parameter_fps=80,
                                                                                        count_fps_del=count_fps_del_hand,
                                                                                        video_end=video_end
                                                                                        )
                        count_fps_del_hand = count_fps_del_re
                        hands_result_contact = updata_result_contact
                        # print(count_fps,'----------------hands_result_1-------------------------------',hands_result_1)
                        # print(count_fps,"-------------------------updata_result_contact----------------------:",updata_result_contact)

                    count_fps  += 1


    def video_select_dect(self):

        while True:

            if self.videodetQueue.empty():
                time.sleep(5)
            else:

                video_path = self.videodetQueue.get()

                try:

                    result_list = PP_TSMv2_predict().predict(input_f=video_path,
                                                            batch_size=self.batch_size,
                                                            predictor=self.predictor,
                                                            InferenceHelper=self.infer)
                    
                    if result_list['topk_scores'] > 0.9:
                                
                        video_base_name = os.path.basename(video_path)
                        video_save_select_path = self.video_save_file + '/' + 'video_select_dect/'+ str(result_list['topk_class']) 
                        os.makedirs(video_save_select_path, exist_ok=True)
                        video_save = os.path.join(video_save_select_path, video_base_name)
                            
                        os.rename(video_path, video_save)

                        self.videoreturnQueue.put(video_save)

                        print("result_list_video_select_dect:",result_list)
                
                except Exception as e:
                    print(e)

    def analysis_by_bbox(self,imgsize,detect_result,dertpara,start_fps,now_fps,label_name,video_path,frame_result_contact,parameter_fps,count_fps_del,video_end):
        '''
        imgsize:图片的尺寸,
        detect_result:检测到的图像的结果,bboxlist
        dertpara:scale_factor：缩放因子，大于 1 表示放大，小于 1 表示缩小
        start_fps: 对比列表中的起始帧
        now_fpsl:当前帧率
        label_name:用于分析的检测类别
        video_path:视频路径
        frame_result_contact:对比列表
        parameter_fps:统计截止时间
        count_fps_del:统计前后帧未出现次数
        
        '''

        bbox_list = tools_function.para_list_correction(images_size=imgsize,bbox_list=detect_result,dertpara=dertpara)

        count_fps_del_re,update_frame_result_contact = self.get_cut_message(fps1=now_fps,
                                                                        label_name = label_name,
                                                                        re_list=bbox_list,
                                                                        video_path=video_path,
                                                                        frame_result_contact=frame_result_contact,
                                                                        parameter_fps=parameter_fps,
                                                                        count_fps_del=count_fps_del,
                                                                        video_end=video_end)
        
        # count_fps_del_re,updata_result_contact = self.get_continue_keys(count_fps_del=count_fps_del,
        #                                                              continue_para=continue_para,
        #                                                              start_fps=start_fps,
        #                                                              now_fps=now_fps,
        #                                                              frame_result_contact=frame_result_contact,
        #                                                              update_frame_result_contact=update_frame_result_contact)\
        
        return count_fps_del_re,update_frame_result_contact
    

    def get_cut_message(self,fps1,label_name,re_list,video_path,frame_result_contact,parameter_fps,count_fps_del,video_end):

        # continue_para = False

        if not frame_result_contact:

            bbox_list_all = tools_function.change_list_dict(fps1=fps1,re_list=re_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 = tools_function.analysis_re01_list(example_list=example_dict_list,
                                                                                result_list=re_list)
            # 有目标减少情况
            if example_lst:
  
                # 截图保存视频
                # continue_para = True

                # cut_dict = {'video_path':video_path,'label_name':label_name,"stop_fps":fps1,'bbox_list':example_lst}
                
                start_fps = example_lst[0]['fps']

                if count_fps_del <= 5:

                    frame_result_contact = frame_result_contact
                    count_fps_del = count_fps_del + 1

                # else:

                    if (fps1 - start_fps) < 15:

                        frame_result_contact = frame_result_contact

                else:   
                    cut_dict = {'video_path':video_path,'label_name':label_name,"stop_fps":fps1,'bbox_list':example_lst}             
                    frame_result_contact = [item for item in frame_result_contact if item not in example_lst]
                    self.cutbboxQueue.put(cut_dict)
                                
            # 有新添加目标情况
            if re_dict_lst:

                # 对比示例列表更新
                update_list = tools_function.change_list_dict(fps1=fps1,re_list=re_dict_lst)

                frame_result_contact = frame_result_contact + update_list

            # 统计截止时间
            time_out_list = tools_function.statistics_fps(fps_now=fps1,re_list=frame_result_contact,parameter=parameter_fps)
                            
            if time_out_list:

                # 裁剪保存视频
                # bbox_list = Process_tools.change_dict_list(time_out_list)
 
                cut_dict = {'video_path':video_path,'label_name':label_name,"stop_fps":fps1,'bbox_list':time_out_list}

                # 添加到新的队列
                self.cutbboxQueue.put(cut_dict)

                # 对比示例列表更新
                frame_result_contact = [item for item in frame_result_contact if item not in time_out_list]
            
            if video_end:

                cut_dict = {'video_path':video_path,'label_name':label_name,"stop_fps":fps1,'bbox_list':frame_result_contact}

                self.cutbboxQueue.put(cut_dict)

                frame_result_contact.clear()

            # print('frame_result_contact:',frame_result_contact)
 
        return count_fps_del,frame_result_contact


    def analysis_return_meassage(self):

        # big_add_list = []
        # big_list = []

        while True:

            if self.videoreturnQueue.empty():

                time.sleep(5)
            else:
                video_message_path = self.videoreturnQueue.get()

                directory = os.path.dirname(video_message_path)
                labels_pptsm = directory.split('/')[-1]

                video_basename = os.path.basename(video_message_path).split('.')[0]

                small_anno_infor = video_basename.split('__')[-1]
                big_anno_infor = video_basename.split('__')[-2]
                video_base_name =  video_basename.split('__')[0]

                #保存的json文件格式
                file_path = self.video_save_file + '/' + video_base_name + '.json'

                # 对小图上的坐标和帧率进行分析
                small_startfps,small_stopfps,small_fps = small_anno_infor.split('_')[0].split('-')
                small_bbox_0,small_bbox_1,small_bbox_2,small_bbox_3 = small_anno_infor.split('_')[1].split('-')

                big_startfps,big_stopfps,big_fps = big_anno_infor.split('_')[0].split('-')
                big_bbox_0,big_bbox_1,big_bbox_2,big_bbox_3 = big_anno_infor.split('_')[1].split('-')

                big_add_startfps = int(big_startfps) + int(small_startfps)
                big_add_stopfps = int(big_startfps) + int(small_stopfps)
                big_add_bbox_0 = int(big_bbox_0) + int(small_bbox_0)
                big_add_bbox_1 = int(big_bbox_1) + int(small_bbox_1)
                big_add_bbox_2 = int(big_bbox_0) + int(small_bbox_2)
                big_add_bbox_3 = int(big_bbox_1) + int(small_bbox_3)

                big_add_dict = {'labels':labels_pptsm,'startfps':big_add_startfps,'stopfps':big_add_stopfps,'bbox':[big_add_bbox_0,big_add_bbox_1,big_add_bbox_2,big_add_bbox_3]}
                big_person_dict = {'labels':'person','startfps':big_startfps,'stopfps':big_stopfps,'bbox':[big_bbox_0,big_bbox_1,big_bbox_2,big_bbox_3]}



                if os.path.isfile(file_path):
                    # 如果文件已存在，读取其中的字典数据
                    with open(file_path, "r") as json_file:
                        data = json.load(json_file)
                        data['big_dict'].append(big_add_dict)

                        if tools_function.compare_dicts(data['big_dict'], big_person_dict):
                            data['big_dict'].append(big_person_dict)
                        
                    with open(file_path, "w") as json_file:

                        json.dump(data, json_file)
                    
                    # # 访问和处理字典数据
                    # print(data)
                else:
                    # 如果文件不存在，创建一个新的字典并保存到文件中
                    bbox_dict = {'big_dict':[big_add_dict,big_person_dict]}
                    with open(file_path, "w") as json_file:
                        json.dump(bbox_dict, json_file)              

    

    def write_video(self): 
        # print('write_videoafagragr')

        '''  
        保存成视频
        '''
        while True:
            if self.cutbboxQueue.empty():
                time.sleep(2)
            else:
                video_frame_dict = self.cutbboxQueue.get()
                # 视频路径
                video_path = video_frame_dict['video_path']
                video_basename = os.path.basename(video_path).split('.')[0]
                file_name = video_frame_dict['label_name']
                # 原视频帧率和尺寸
                cap = cv2.VideoCapture(video_path)
                fps = cap.get(cv2.CAP_PROP_FPS)
                video_fps = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

                # print(video_path,'fps:',fps,'video_fps:',video_fps)
                # 获得起始
                stop_fps = video_frame_dict['stop_fps']
                # 裁剪信息
                result_list = video_frame_dict['bbox_list']
                if  cap.isOpened():

                    for i,bbox_dict in enumerate(result_list):
                        start_fps = bbox_dict['fps']

                        if start_fps >= stop_fps:

                            # print('start_fps:',start_fps,'stop_fps:',stop_fps)
                            break

                        else:
                            bbox_list = bbox_dict['result']
                            bbox_int_list = [int(bbox_list[0]),int(bbox_list[1]),int(bbox_list[2]),int(bbox_list[3])]
                            w = bbox_int_list[2] - bbox_int_list[0]
                            h = bbox_int_list[3] - bbox_int_list[1]
                            size = [w,h]
                            if tools_function.determine_zero(size):
                                size = (w,h)

                                # 根据标签保存不同视频分类
                                # bbox_name = '{}-{}-{}_{}'.format(int(bbox_list[0]), int(bbox_list[1]), int(bbox_list[2]), int(bbox_list[3]))
                                video_name_save = '{}__{}-{}-{}_{}-{}-{}-{}.avi'.format(video_basename, start_fps, stop_fps, video_fps,int(bbox_list[0]), int(bbox_list[1]), int(bbox_list[2]), int(bbox_list[3]))
                                video_save_file = self.video_save_file + '/' + file_name
                                os.makedirs(video_save_file, exist_ok=True)
                                video_save_path = os.path.join(video_save_file, video_name_save)
                                videoWriter =cv2.VideoWriter(video_save_path,cv2.VideoWriter_fourcc('X','V','I','D'),fps,size)

                                tools_function.save_seg_video(video_name=video_path,
                                                                frameToStart=start_fps,
                                                                frametoStop=stop_fps,
                                                                videoWriter=videoWriter,
                                                                bbox=bbox_int_list,
                                                                size=size)
                                videoWriter.release()
                                self.videoQueue2.put(video_save_path)
                            
                            else:
                                print('-----------------agrag-----------------',size,'-----------------agag-----------------')

                    cap.release()
                
                else:
                    print(video_path)
                    break


    def select_video_path(self):

        while True:

            if self.videoQueue2.empty():
                time.sleep(5)
            else:
                video_path = self.videoQueue2.get()
                directory = os.path.dirname(video_path)
                labels = directory.split('/')[-1]

                # print('video_pathagfg:',video_path)

                # print(labels)

                if labels == 'person':

                    self.videoQueue3.put(video_path)

                # if labels == 'head' or labels == 'hands':
                if labels == 'hands':
                    self.videodetQueue.put(video_path)

                else:
                    pass         

    def run(self):

        self.get_video_listThread.start()
        self.get_video_frameThread.start()
        self.write_videoThread.start()
        self.head_hands_detThread.start()
        self.video_select_dectThread.start()
        self.select_video_pathThread.start()
        self.analysis_return_meassageThread.start()



if __name__ == '__main__':  
    t1 = time.time()
    video = "E:/Bank_files/Bank_02/dataset/video_kf/02.mp4"
    video_save = 'test_video'

    # 初始化目标检测
    person_model = YOLO("model_file/yolov8x.pt")

    # 初始化pptsmv2
    config = 'model_file/inference/pptsm_lcnet_k400_16frames_uniform.yaml'  # 配置文件地址
    model_file = 'model_file/inference_hands_2/ppTSMv2.pdmodel'  # 推理模型存放地址
    params_file = 'model_file/inference_hands_2/ppTSMv2.pdiparams'
    # batch_size= 1
    infer,predictor = PP_TSMv2_predict().create_inference_model(config,model_file,params_file)
    # PP_TSMv2_predict().predict(config,input_file,batch_size,predictor,infer)

    # 初始化mediapipe
    mp_holistic = mp.solutions.holistic
    holistic = mp_holistic.Holistic(
    min_detection_confidence=0.5,
    min_tracking_confidence=0.5)

    # 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=holistic,pptsmv2_model=[infer,predictor])
    deal.run()

    t2 = time.time()

    # print('总时间：',t2-t1)