File size: 1,225 Bytes
803ef9e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
import random
from torchvision.datasets import ImageFolder
import torchvision.transforms as T
from PIL import ImageFilter
from .transforms import MultiSample, aug_transform
from .base import BaseDataset
class RandomBlur:
def __init__(self, r0, r1):
self.r0, self.r1 = r0, r1
def __call__(self, image):
r = random.uniform(self.r0, self.r1)
return image.filter(ImageFilter.GaussianBlur(radius=r))
def base_transform():
return T.Compose(
[T.ToTensor(), T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))]
)
class ImageNet(BaseDataset):
def ds_train(self):
aug_with_blur = aug_transform(
224,
base_transform,
self.aug_cfg,
extra_t=[T.RandomApply([RandomBlur(0.1, 2.0)], p=0.5)],
)
t = MultiSample(aug_with_blur, n=self.aug_cfg.num_samples)
return ImageFolder(root=self.aug_cfg.imagenet_path + "train", transform=t)
def ds_clf(self):
t = base_transform()
return ImageFolder(root=self.aug_cfg.imagenet_path + "clf", transform=t)
def ds_test(self):
t = base_transform()
return ImageFolder(root=self.aug_cfg.imagenet_path + "test", transform=t)
|