2.extract_dms_depth_at_same_time.py

import pytesseract
import cv2
import glob
import os
# 加速处理开关，打开后可以加速处理，但缺点是容错率低
OPEN_ACCELERATE = True
# 处理结果保存路径
DMS_IMG_SAVE_DIR = os.path.join(os.getcwd(), 'result/dms_img')
DEPTH_IMG_SAVE_DIR = os.path.join(os.getcwd(), 'result/depth_img')

if __name__ == '__main__':
    # -----------拿到所有视频的字典------------------------------------------
    dms_video_dir = os.path.join(os.getcwd(), 'dms_video')
    depth_video_dir = os.path.join(os.getcwd(), 'depth_video')
    dms_video_list = glob.glob(dms_video_dir+'/*.mp4', recursive=True)
    depth_video_list = glob.glob(depth_video_dir+'/*.mp4', recursive=True)
    dms_video_dict = {}
    depth_video_dict = {}
    for item in dms_video_list:
        base = os.path.basename(item)
        dms_video_dict[base] = item
    for item in depth_video_list:
        base = os.path.basename(item)
        depth_video_dict[base] = item
    # -----------对每个视频进行处理，提取图片对------------------------------------------
    for dms_video_item in dms_video_list:
        # get relevant depth video
        base_name = os.path.basename(dms_video_item)  # dms_id_point.mp4
        dms_filename, extname = os.path.splitext(base_name)
        id_point = dms_filename.split('_')
        id_point[0] = "depth"
        depth_filename = '_'.join(id_point)
        depth_video_item = depth_filename + extname
        if depth_video_item not in depth_video_dict:
            continue

        dms_capture = cv2.VideoCapture(dms_video_item)
        depth_capture = cv2.VideoCapture(depth_video_dict[depth_video_item])

        frame_num = 0
        while True:
            dms_success, dms_frame = dms_capture.read()
            if not dms_success:
                break
            dms_digit = pytesseract.image_to_string(dms_frame)
            while True:
                depth_success, depth_frame = depth_capture.read()
                if not depth_success:
                    depth_capture.set(cv2.CAP_PROP_POS_FRAMES, 0)
                    break
                depth_digit = pytesseract.image_to_string(depth_frame)
                if depth_digit == dms_digit:
                    # ----保存dms图片和depth图片--------------
                    dms_save_dir = os.path.join(DMS_IMG_SAVE_DIR, id_point[1], id_point[2],
                                                id_point[1]+id_point[2]+str(frame_num)+'.jpg')
                    cv2.imwrite(dms_save_dir, dms_frame)
                    depth_save_dir = os.path.join(DEPTH_IMG_SAVE_DIR, id_point[1], id_point[2],
                                                  id_point[1]+id_point[2]+str(frame_num)+'.jpg')
                    cv2.imwrite(depth_save_dir, depth_frame)
                    frame_num += 1
                    # ----设置视频下次搜索起始点---------------
                    if OPEN_ACCELERATE:
                        depth_start_frame = depth_capture.get(cv2.CAP_PROP_POS_FRAMES) + 1
                    else:
                        depth_start_frame = 0
                    depth_capture.set(cv2.CAP_PROP_POS_FRAMES, depth_start_frame)
                    break
                elif dms_digit < depth_digit and OPEN_ACCELERATE:
                    depth_capture.set(cv2.CAP_PROP_POS_FRAMES, 0)
                    break
        dms_capture.release()
        depth_capture.release()