ProcessPoolExecutor+ThreadPoolExecutor 高效率抽帧 同时用YOLOV5 进行裁人！

import sys
sys.path.append("/media/daniel/disk3/yrq/hub/video_recongnition")
from genericpath import exists
import glob
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2"
import cv2
import torch
import math
import threading
import time
from concurrent.futures import ThreadPoolExecutor,ProcessPoolExecutor
import concurrent.futures



def do_file(frame ,video,new_label,i,zuobiao,model):
    # print('do_file  ***',)
    results = model(frame, size=640)
    pd  = results.pandas().xyxy[0]
    pd = pd[pd["class"] == 0]
    arr2d = pd.iloc[:,[0,1,2,3,5]].values
    for k in range(len(arr2d)) : 
        x0 = math.ceil(arr2d[k][0])
        y0 = math.ceil(arr2d[k][1])
        x1 = math.floor(arr2d[k][2])
        y1 = math.floor(arr2d[k][3])
        single = new_label+'/'+str(video).split(".")[0]+' '+'img_{:05d}.jpg'.format(i+1)+' '+str(x0)+' '+str(y0)+' '+str(x1)+' '+str(y1)
        zuobiao.append(single)
        h = y1 - y0 
        w = x1 - x0
        diff = abs(h-w)//2
        if h>w:
            x0 = x0 - diff
            x1 = x1 + diff
        elif h<w:
            y0 = y0 - diff
            y1 = y1 + diff  

        # 还需要进行一些边缘处理
        if x0<0:
            x0 += diff
            x1 -= diff
            x1 += (h-w)
            cropped = frame[y0:y1, x0:x0+h] 
        elif x1>1920:
            x0 += diff
            x1 -= diff
            x0 -= (h-w)
            cropped = frame[y0:y1, x0:x0+h] 

        elif y1>1080:
            y0 += diff
            y1 -= diff
            y0 -= (w-h)
            cropped = frame[y0:y0+w, x0:x1] 
        else:
            cropped = frame[y0:y1, x0:x1] #向下取整 裁剪坐标为[y0:y1, x0:x1]
        
        
        try:
            cv2.imwrite(filename=os.path.join(path,file,new_label,str(video).split(".")[0],'img_{:05d}.jpg'.format(i+1)), img=cropped)
        except Exception as e:
            print(e)
        cv2.imwrite(filename=os.path.join(path,raw_file,new_label,str(video).split(".")[0],'img_{:05d}.jpg'.format(i+1)), img =frame)
        i+=1
    return i

def do_video(video,label,model):
    print('do_video  ***',video,label)
    # 信号量：信号量是产生的一堆进程/线程，即产生了多个任务都去抢那一把锁
    # thread_max_num.acquire()  #锁定线程的最大数量
    zuobiao = []
    new_label = label +'_'+ str(video).split(".")[0].split("_")[-1]

    if not os.path.exists(os.path.join(path,file,new_label)):

        os.mkdir(os.path.join(path,file,new_label))
    if not os.path.exists(os.path.join(path,file,new_label,str(video).split(".")[0])):

        os.mkdir(os.path.join(path,file,new_label,str(video).split(".")[0]))

    if not os.path.exists(os.path.join(path,raw_file,new_label)):

        os.mkdir(os.path.join(path,raw_file,new_label))
    if not os.path.exists(os.path.join(path,raw_file,new_label,str(video).split(".")[0])):

        os.mkdir(os.path.join(path,raw_file,new_label,str(video).split(".")[0]))
        
    video_name = os.path.join("/media/daniel/disk3/yrq/data/aist203",label , video)

    capture = cv2.VideoCapture(video_name)

    frame_count = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))

    EXTRACT_FREQUENCY = 4
    if frame_count // EXTRACT_FREQUENCY <= 160:
        EXTRACT_FREQUENCY -= 1
        if frame_count // EXTRACT_FREQUENCY <= 160:
            EXTRACT_FREQUENCY -= 1
            if frame_count // EXTRACT_FREQUENCY <= 160:
                EXTRACT_FREQUENCY -= 1
    count = 0
    i = 0
    retaining = True

    while (count < frame_count and retaining):
        retaining, frame = capture.read()
        if frame is None:
            continue
        if count % EXTRACT_FREQUENCY == 0:
            # 一个图是最快的！！！
            i = do_file(frame,video,new_label,i,zuobiao,model)
        count += 1
    capture.release()

    print(label , video)
    #  处理完一个视频添加一次
    if len(zuobiao) >= 0:
        with open('./zuobiao2.txt','a') as f:
            for i in range(len(zuobiao)):
                f.write(zuobiao[i]+'\n')
        zuobiao = []
    # thread_max_num.release()

def do_label(label):
    model = torch.hub.load('ultralytics/yolov5', 'custom', path='./yolov5x.pt' ) 
    print('do_label  ***',label)
    videos = sorted(os.listdir(label))
    label = label.split("/")[-1]
    # for video in videos:
    with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
        executor.map(do_video, videos,[label] * len(videos),[model]* len(videos))
           
# IO密集型 使用 多线程
# cpu 密集型 使用 多进程
if __name__ == '__main__':

    path = "/media/daniel/disk3/yrq/hub/video_recongnition/data/square/"
    #  抽帧
    videos_dir = sorted(glob.glob(r'/media/daniel/disk3/yrq/data/aist203/g*'))
    print(videos_dir)
    videos_dir[:3]
    file = 'frame_speed2'

    raw_file = 'rawframe_speed2'

    
    
    if not os.path.exists(os.path.join(path,file)):
        os.mkdir(os.path.join(path,file))
    if not os.path.exists(os.path.join(path,raw_file)):
        os.mkdir(os.path.join(path,raw_file))
    
    with concurrent.futures.ProcessPoolExecutor(max_workers=3) as executor:
        executor.map(do_label, videos_dir)
    # 利用with 会自动调用 executor.shutdown()
    
           
        
            
           
            
           

同时开启多进程 和 多线程！！！
速度提升非常快！！

这个只针对我这个数据集，具体的数据集可以自己进行更改！！！
同时指定进程数和线程数可以自己进行更改！达到最佳的性能

人物裁出来的效果：