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

import numpy as np
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
from 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.cutbboxQueue = queue.Queue(maxsize=0)
        self.videoframeQueue = queue.Queue(maxsize=0)
        self.videohandsQueue = queue.Queue(maxsize=0)
        self.videoheadQueue = queue.Queue(maxsize=0)
        self.videopersonQueue = 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)
        self.select_video_frameThread = threading.Thread(target=self.select_video_frame)
        self.head_hands_detThread = threading.Thread(target=self.head_hands_det)
        

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

    #     pass

    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
                        # print('count_fps_read_video=',count_fps)

                        frame_dict = {'fps':count_fps,'frame':frame}
                        frame_list.append(frame_dict)
                        

                    video_dict = {'video_path':video_path,'frame_list':frame_list}
                
                    self.frameQueue.put(video_dict)
                    # time.sleep(30)


                except Exception as e:
                    print(e)


    def person_det(self):

        '''
        从队列中获取视频帧frame，进行第一步人员的检测
        '''

        while True: 

            if  ~self.videoframeQueue.empty():

                video_frame_dict = self.videoframeQueue.get()

                frame_list = video_frame_dict['frame_list']
                video_path = video_frame_dict['video_path']

                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 = tools.get_dict_values(person_det)

                        label_name = list(person_det[0].keys())[0]

                        update_frame_result_contact = self.get_cut_message(fps1=frame_list[i]["fps"],
                                                                            label_name = label_name,
                                                                            re_list=person_list,
                                                                            video_path=video_path,
                                                                            frame_result_contact=frame_result_contact)
                        
                        frame_result_contact = update_frame_result_contact
                        # print('frame_result_contact:',frame_result_contact)


    def head_hands_det(self):

        while True:

            if ~self.videopersonQueue.empty():

                person_frame_dict = self.videopersonQueue.get()


                person_frame_list = person_frame_dict['frame_list']
                video_path = person_frame_dict['video_path']

                head_result_contact = []
                hands_result_contact = []

                for i in range(len(person_frame_list)):

                    if person_frame_list[i]["fps"] == i + 1: 

                        image = person_frame_list[i]["frame"]

                        imgsize = image.shape

                        # print(type(image))

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


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

                        print('head_result_1:',head_result_1)
                        print('head_result_1:',hands_result_1)

                        # 统一修正坐标,分别对头和手进行分析
                        if head_result_1:
                            head_bbox_list = tools.para_list_correction(images_size=imgsize,
                                                                            bbox_list=head_result_1,
                                                                            dertpara=[])


                            update_head_result_contact = self.get_cut_message(fps1=person_frame_list[i]["fps"],
                                                                                label_name = 'head',
                                                                                re_list=head_bbox_list,
                                                                                video_path=video_path,
                                                                                frame_result_contact=head_result_contact)
                            head_result_contact = update_head_result_contact


                        if hands_result_1:

                            hands_bbox_list = tools.para_list_correction(images_size=imgsize,
                                                                            bbox_list=hands_result_1,
                                                                            dertpara=[]) 
                           
                            update_hands_result_contact = self.get_cut_message(fps1=person_frame_list[i]["fps"],
                                                                                label_name = 'hands',
                                                                                re_list=hands_bbox_list,
                                                                                video_path=video_path,
                                                                                frame_result_contact=hands_result_contact)
                        
                            hands_result_contact = update_hands_result_contact

                        # print("head_result_contact:",head_result_contact)
                        # print("hands_result_contact:",hands_result_contact)

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


        if not frame_result_contact:

            bbox_list_all = tools.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
            print('example_dict_list:',example_dict_list)
            print('re_list:',re_list)

            cut_list,example_lst,re_dict_lst = tools.analysis_re01_list(example_list=example_dict_list,
                                                                                result_list=re_list)
                            
            # print('cut_list:',cut_list)
            # print('example_sorted_lst:',example_lst)
            # print('re_dict_sorted_lst:',re_dict_lst)

                        
            # 有目标减少情况
            if example_lst:

                # 截图保存视频

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

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

                frame_result_contact = [item for item in frame_result_contact if item not in example_lst]
                                
                # 有新添加目标情况
            if re_dict_lst:

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

                frame_result_contact = frame_result_contact + update_list

            # 统计截止时间
            time_out_list = tools.statistics_fps(fps_now=fps1,re_list=frame_result_contact,parameter=20)
                            

            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]

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


    def write_video(self): 

        '''  
        保存成视频
        '''

        while True:

             if ~self.cutbboxQueue.empty():

                video_frame_dict = self.cutbboxQueue.get()

                # print('video_frame_dict:',video_frame_dict)

                # 视频路径
                video_path = video_frame_dict['video_path']
                video_basename = os.path.basename(video_path).split('.')[0]
                file_name = video_frame_dict['label_name']
                # video_name_save = os.path.join(self.video_save_file, video_basename)

                # 原视频帧率和尺寸
                cap = cv2.VideoCapture(video_path)
                fps = cap.get(cv2.CAP_PROP_FPS)

                # 获得起始帧
                stop_fps = video_frame_dict['stop_fps']

                # 裁剪信息
                result_list = video_frame_dict['bbox_list']

                for i,bbox_dict in enumerate(result_list):

                    start_fps = bbox_dict['fps']
                    bbox_list = bbox_dict['result']

                    w = int(bbox_list[2]) - int(bbox_list[0])
                    h = int(bbox_list[3]) - int(bbox_list[1])

                    size = (w,h)

                    # 根据标签保存不同视频分类
                    video_name_save = video_basename + '_' + str(start_fps)  + '_' + str(stop_fps) + '_'  + str(i) + '.avi'
                    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.save_seg_video(video_name=video_path,
                                                    frameToStart=start_fps,
                                                    frametoStop=stop_fps,
                                                    videoWriter=videoWriter,
                                                    bbox=bbox_list)
                        
                    videoWriter.release()

                    self.videoQueue.put(video_save_path)

                cap.release()


    def select_video_frame(self):

        while True:

            if  ~self.frameQueue.empty():

                video_dict = self.frameQueue.get()
                video_path = video_dict["video_path"]
                directory = os.path.dirname(video_path)
                labels = directory.split('/')[-1]

                print('labels:',labels)

                if labels == 'person':

                    self.videopersonQueue.put(video_dict)

                if labels == 'head':

                    # print('youshou')

                    self.videoheadQueue.put(video_dict)
                
                if labels == 'hands':

                    # print('youshou')

                    self.videohandsQueue.put(video_dict)

                else:

                    self.videoframeQueue.put(video_dict)


    def run(self):

        self.get_video_listThread.start()
        self.get_video_frameThread.start()
        self.person_detThread.start()
        self.write_videoThread.start()
        self.select_video_frameThread.start()
        self.head_hands_detThread.start()



if __name__ == '__main__':  
  


    video = "E:/Bank_files/Bank_02/dataset/video_test/test03_3.avi"
    video_save = 'test_video'

    person_model = YOLO("model_file/yolov8x_person.pt")

    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='model_file/yolov8x_person.pt')
    deal.run()