Spaces:
Runtime error
Runtime error
File size: 881 Bytes
d7dbcdd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 |
import torch.utils.data as data
from PIL import Image
import os
import os.path
import numpy as np
import glob
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):
l0_train = []
l1_train = []
flow_train = []
for flow_map in sorted(glob.glob(os.path.join(filepath,'*_flow.flo'))):
root_filename = flow_map[:-9]
img1 = root_filename+'_img1.ppm'
img2 = root_filename+'_img2.ppm'
if not (os.path.isfile(os.path.join(filepath,img1)) and os.path.isfile(os.path.join(filepath,img2))):
continue
l0_train.append(img1)
l1_train.append(img2)
flow_train.append(flow_map)
return l0_train, l1_train, flow_train
|