ts-model-refining/OneFlow-KAE/kamal/vision/datasets/flowers102.py

import flowvision
import os


class Flowers102(flowvision.datasets.ImageFolder):
    def __init__(self, root, split='train'):
        self.split = split
        self.num_classes = 102
        path = os.path.join(root, 'train') if self.split=='train' else os.path.join(root, 'valid')
        super().__init__(path)
        print('Flowers-102, Split: %s, Size: %d' % (self.split, len(self.imgs)))