vid_ocr.py 文件源码

def vidocr_to_csv(video,vcoords,tbcoords,f1app=True):
    # inputs: 
    # video = video file as a string
    # vcoords = array of pixel coordinates [top left x, top left y, bottom right x, bottom right y] for velocity
    # tbcoords = array of pixel coordinates [top left x, top left y, bottom right x, bottom right y] for throttle/brake
    # f1app = boolean, default = True, use True for video from the F1 app, False for onboard video.
    # outputs .csv file with each line as a row of each extracted parameter. 

    # capture video via opencv
    vid = cv2.VideoCapture(video)
    s,frm = vid.read()

    v_all = []
    t_all = []
    thr_all = []
    brk_all = []

    step = 1
    total_frames = vid.get(cv2.CAP_PROP_FRAME_COUNT);
    print(total_frames)
    i = int(total_frames*0);
    vid.set(0,i)

    # go through each frame and extract data
    while s:
        if i >= int(total_frames):
            break
        s,frm = vid.read()
        if i%step == 0 or i == total_frames-1:
            v_temp = velocity_ocr(frm,vcoords,f1app)
            t_temp = vid.get(cv2.CAP_PROP_POS_MSEC)/1e3
            v_all += [v_temp]
            # thr_temp = throttle_ocr(frm,tbcoords)
            # brk_temp = brake_ocr(frm,tbcoords)
            # thr_all += [thr_temp]
            # brk_all += [brk_temp]
        if i%200 == 0:
            print(v_temp,t_temp,i)
        i += 1

    t_all = get_timestamps(video)
    # save data to .csv with same filename as videofile
    with open(video[0:-4]+'.csv', 'w', newline='') as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow(v_all)
        writer.writerow(t_all)
        # writer.writerow(thr_all)
        # writer.writerow(brk_all)
        writer.writerow([])
        writer.writerow([])

    print(video,"completed.")