expansion/dataloader/hd1klist.py

import torch.utils.data as data

from PIL import Image
import os
import os.path
import numpy as np
import pdb

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/'
  train = [img for img in os.listdir(filepath+left_fold) if img.find('HD1K2018') > -1]
  train = sorted(train)

  l0_train  = [filepath+left_fold+img for img in train]
  l0_train  = [img for img in l0_train if '%s_%s.png'%(img.rsplit('_',1)[0],'%04d'%(1+int(img.split('.')[0].split('_')[-1])) ) in l0_train ]
  l1_train = ['%s_%s.png'%(img.rsplit('_',1)[0],'%04d'%(1+int(img.split('.')[0].split('_')[-1])) ) for img in l0_train]
  flow_train = [img.replace('image_2','flow_occ') for img in l0_train]

  return l0_train, l1_train, flow_train