collections dataset test

collections_dataset_test.py

class CollectionsDatasetTest(Dataset):

    def __init__(self, 
                 csv_file, 
                 root_dir, 
                 transform=None):
        self.data = pd.read_csv(csv_file)
        self.root_dir = root_dir
        self.transform = transform

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

    def __getitem__(self, idx):
        img_name = os.path.join(self.root_dir, 
                                self.data.loc[idx, 'id'] + '.png')
        image = Image.open(img_name)
        if self.transform:
            image = self.transform(image)

        return {'image': image}