import torch.utils.data as data from PIL import Image import os import os.path import numpy as np IMG_EXTENSIONS = [ '.jpg', '.JPG', '.jpeg', '.JPEG', '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP', ] def is_image_file(filename): return any(filename.endswith(extension) for extension in IMG_EXTENSIONS) def dataloader(filepath): left_fold = 'image_2/' flow_noc = 'flow_occ/' train = [img for img in os.listdir(filepath+left_fold) if img.find('_10') > -1] # train = [i for i in train if int(i.split('_')[0])%5!=0] with open('/data/gengshay/kitti_scene/devkit/mapping/train_mapping.txt','r') as f: flags = [True if len(i)>1 else False for i in f.readlines()] train = [fn for (it,fn) in enumerate(sorted(train)) if flags[it] ][100:] l0_train = [filepath+left_fold+img for img in train] l1_train = [filepath+left_fold+img.replace('_10','_11') for img in train] flow_train = [filepath+flow_noc+img for img in train] return sorted(l0_train), sorted(l1_train), sorted(flow_train)