python视频抽帧

某些情况需要根据视频生成图片，使用视频抽帧。
fps 每秒抽取的帧数。

def select_frame(video, fps=1, save_dir='/media/s4/ae13bf8c-93ea-4157-bb06-1ef3acc87a19/rgb/rgb_test'):
    # select frames of video contained traffic lights, save frames selected into save_dir
    cap = cv2.VideoCapture(video)
    video_name = os.path.basename(video)

    video_FPS = cap.get(cv2.CAP_PROP_FPS)
    # print(video_FPS)
    gap_frame = int(video_FPS / fps)

    frame_cnt = 0
    num = 0

    while (True):
        success, frame = cap.read()
        if not success:
            break

        if frame_cnt % gap_frame == 0:
            # save fram into save_dir
            path = os.path.join(save_dir, video_name.split('.')[0] + '_' + str(num) + '.jpg')
            cv2.imwrite(path, frame)
            num += 1
            pass

        frame_cnt += 1
    cap.release()
    # print('select video %s finished!' % video_name)
    pass


videodir = '/media/s4/新加卷1/downloadvideos'
for videofile in tqdm(os.listdir(videodir)):
    videopath = os.path.join(videodir, videofile)
    if not os.path.exists('/media/s4/新加卷1/downloadimg'):
        os.makedirs('/media/s4/新加卷1/downloadimg')
    select_frame(videopath, '/media/s4/新加卷1/downloadimg')