Spaces:
Build error
Build error
| from types import SimpleNamespace | |
| import torch | |
| from torch.utils.data import DataLoader | |
| from torchvision.datasets import ImageFolder | |
| from torchvision.transforms import v2 | |
| DATASET_MEAN = 0.5077385902404785 | |
| DATASET_STD = 0.255077600479126 | |
| class PreprocessedImageFolder(ImageFolder): | |
| def __init__(self, root, transform=None): | |
| super().__init__(root, transform=transform) | |
| self.preprocess = v2.Compose( | |
| [ | |
| v2.Grayscale(), | |
| v2.PILToTensor(), | |
| v2.ToDtype(torch.float32, scale=True), | |
| v2.Normalize(mean=(DATASET_MEAN,), std=(DATASET_STD,)), | |
| ] | |
| ) | |
| processed_samples = [] | |
| for path, target in self.samples: | |
| sample = self.loader(path) | |
| processed_sample = self.preprocess(sample) | |
| processed_samples.append((processed_sample, target)) | |
| self.samples = processed_samples | |
| def __getitem__(self, index): | |
| sample, target = self.samples[index] | |
| if self.transform is not None: | |
| sample = self.transform(sample) | |
| return sample, target | |
| augmentations = v2.Compose( | |
| [ | |
| v2.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2), | |
| v2.RandomHorizontalFlip(), | |
| v2.RandomResizedCrop(size=48, scale=(0.9, 1.1), antialias=True), | |
| v2.RandomAffine(degrees=(-10, 10), translate=(0.1, 0.1), scale=(0.9, 1.1)), | |
| v2.GaussianBlur(kernel_size=(3, 3), sigma=(0.1, 1.0)), | |
| ] | |
| ) | |
| def make_dls(train_ds, valid_ds, batch_size=64, num_workers=2): | |
| train_dl = DataLoader( | |
| train_ds, | |
| batch_size=batch_size, | |
| shuffle=True, | |
| num_workers=num_workers, | |
| pin_memory=True, | |
| persistent_workers=True, | |
| ) | |
| valid_dl = DataLoader( | |
| valid_ds, | |
| batch_size=batch_size, | |
| shuffle=True, | |
| num_workers=num_workers, | |
| pin_memory=True, | |
| persistent_workers=True, | |
| ) | |
| dls = SimpleNamespace(**{"train": train_dl, "valid": valid_dl}) | |
| return dls | |