mpisintel.py 文件源码

def make_dataset(dir, split, dataset_type='clean'):
    training_dir = os.path.join(dir,'training')
    flow_dir = 'flow'
    assert(os.path.isdir(os.path.join(training_dir,flow_dir)))
    img_dir = dataset_type
    assert(os.path.isdir(os.path.join(training_dir,img_dir)))

    images = []
    for flow_map in glob.iglob(os.path.join(dir,flow_dir,'*','*.flo')):
        flow_map = os.path.relpath(flow_map,os.path.join(dir,flow_dir))
        root_filename = flow_map[:-8]
        frame_nb = int(flow_map[-8:-4])
        img1 = os.path.join(img_dir,root_filename+str(frame_nb).zfill(4)+'.png')
        img2 = os.path.join(img_dir,root_filename+str(frame_nb+1).zfill(4)+'.png')
        flow_map = os.path.join(flow_dir,flow_map)
        if not (os.path.isfile(os.path.join(dir,img1)) or os.path.isfile(os.path.join(dir,img2))):
            continue
        images.append([[img1,img2],flow_map])

    return split2list(images, split, default_split=0.87)