import numpy as np
from PIL import Image

class ImagePathDataset(object):
    def __init__(self, root, transform=None):
        self.transform = transform
        self.inputs = []
        self.outputs = []

        with open(root) as f:
            lines = f.readlines()

        for line in lines:
            input, output = line.split(' ')
            self.inputs.append(input)
            self.outputs.append(int(output))

    def __getitem__(self, index):
        input = Image.open(self.inputs[index])
        if self.transform is not None:
            input = self.transform(input)

        return input, self.outputs[index]

    def __len__(self):
        return len(self.inputs)

dataset = ImagePathDataset('train.txt')