CustomFolderDataSource.py

from flash.core.data.data_source import DataSource, DefaultDataKeys
class CustomFolderDataSource(DataSource):
  
  def load_data(self, folder, auto_dataset):
    paths, labels = make_dataset(folder)
    return [(p, l) for p, l in zip(paths, labels)]
  
  def load_sample(self, sample, auto_dataset):
    path, label = sample
    return {DefaultDataKeys.INPUT: loader(path), DefaultDataKeys.TARGET: label}
  
  def predict_load_data(self, folder, auto_dataset):
    return predict_make_dataset(folder)
  
  def predict_load_sample(self, path, auto_dataset):
    return {DefaultDataKeys.INPUT: loader(path)}
  def default_transforms(self):
     return None
  def train_default_transforms(self):
     return None