folder_data_provider.py

import numpy as np
import os
import cv2
import random
from util import get_image_center
from data_provider import DataProvider

class FolderDataProvider(DataProvider):

  def __init__(self,
               folder,
               read_limit=-1,
               main_size=80,
               crop_size=64,
               augmentation_factor=4,
               *args,
               **kwargs):
    files = os.listdir(folder)
    files = sorted(files)

    if read_limit != -1:
      files = files[:read_limit]
    data = []
    files.sort()
    for f in files:
      image = (cv2.imread(os.path.join(folder, f))[:, :, ::-1] /
               255.0).astype(np.float32)
      image = get_image_center(image)
      # image = cv2.resize(image, (64, 64), interpolation=cv2.INTER_AREA)
      # data.append(image)
      image = cv2.resize(
          image, (main_size, main_size), interpolation=cv2.INTER_AREA)
      for i in range(augmentation_factor):
        new_image = image
        if random.random() < 0.5:
          new_image = new_image[:, ::-1, :]
        sx, sy = random.randrange(main_size - crop_size + 1), random.randrange(
            main_size - crop_size + 1)
        data.append(new_image[sx:sx + crop_size, sy:sy + crop_size])
    data = np.stack(data, axis=0)
    print("# image after augmentation =", len(data))
    super(FolderDataProvider, self).__init__(data, *args, bnw=False,
                                             augmentation=1.0,
                                             output_size=crop_size,
                                             **kwargs)

def test():
  dp = FolderDataProvider('data/sintel/outputs')
  while True:
    d = dp.get_next_batch(64)
    cv2.imshow('img', d[0][0, :, :, ::-1])
    cv2.waitKey(0)


if __name__ == '__main__':
  test()
  # preprocess()