khurramjaved96 / incremental-learning Goto Github PK

def pil_loader(path):
    # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)
    with open(path, 'rb') as f:
        img = Image.open(f)
        img = img.resize((32,32))
        img = img.convert('RGB')
        return img

class GEO(Dataset):
    '''
    :param train_path: /media/gaoya/disk/Applications/pytorch/Incremental Learning/train_dataset
    :param test_path /media/gaoya/disk/Applications/pytorch/Incremental Learning/test_dataset
    '''
    def __init__(self,train_path, test_path):
        super().__init__(classes=5, name="GEO", labels_per_class_train=500, labels_per_class_test=100)

        imgs = numpy.array([])
        self.train_transform = transforms.Compose([
            transforms.Resize(32),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
             ])

        self.test_transform = transforms.Compose([
            transforms.Resize(32),
            transforms.ToTensor(),
            transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
             ])

        self.train_data = datasets.ImageFolder(root=train_path, 
                                                transform=self.train_transform,
                                                loader=pil_loader)

        imgs = numpy.array([self.train_transform(pil_loader(path[0])).numpy().transpose(2,1,0) for path in self.train_data.imgs])
        # img_temp = imgs[0]
        # for k in range(1, len(imgs)):
        #     img_temp = numpy.stack((img_temp, imgs[k]))
        
        self.train_imgs = imgs
        self.train_targets = self.train_data.targets

        self.test_data = datasets.ImageFolder(root=test_path, 
                                                transform=self.test_transform,
                                                loader=pil_loader)

        self.test_imgs = numpy.array([self.test_transform(pil_loader(path[0])).numpy().transpose(2,1,0) for path in self.test_data.imgs])
        self.test_targets = self.test_data.targets