| from __future__ import annotations | |
| import io | |
| import random | |
| import numpy as np | |
| from PIL import Image, ImageEnhance, ImageFilter, ImageOps | |
| import torch | |
| from torchvision import transforms as T | |
| IMAGENET_MEAN = (0.485, 0.456, 0.406) | |
| IMAGENET_STD = (0.229, 0.224, 0.225) | |
| def jpeg_recompress(img: Image.Image, quality: int) -> Image.Image: | |
| b = io.BytesIO() | |
| img.convert("RGB").save(b, format="JPEG", quality=int(quality), optimize=False) | |
| b.seek(0) | |
| return Image.open(b).convert("RGB") | |
| def sensor_noise(img: Image.Image, sigma: float) -> Image.Image: | |
| a = np.asarray(img.convert("RGB"), dtype=np.float32) / 255.0 | |
| read = np.random.normal(0, sigma, a.shape).astype(np.float32) | |
| shot = np.random.normal(0, np.sqrt(np.clip(a, 0, 1) + 1e-4) * sigma * 0.35).astype(np.float32) | |
| a = np.clip(a + read + shot, 0, 1) | |
| return Image.fromarray((a * 255 + .5).astype(np.uint8), mode="RGB") | |
| class SymmetricAugment: | |
| def __call__(self, image: Image.Image) -> Image.Image: | |
| im = ImageOps.exif_transpose(image).convert("RGB") | |
| if random.random() < .70: | |
| im = jpeg_recompress(im, random.randint(35, 100)) | |
| if random.random() < .15: | |
| im = jpeg_recompress(im, random.randint(45, 92)) | |
| if random.random() < .52: | |
| w, h = im.size; s = random.uniform(.45, 1.25) | |
| im = im.resize((max(64, int(w*s)), max(64, int(h*s))), random.choice([ | |
| Image.Resampling.BICUBIC, Image.Resampling.BILINEAR, Image.Resampling.LANCZOS])) | |
| if random.random() < .20: | |
| im = im.filter(ImageFilter.GaussianBlur(random.uniform(.1, 1.25))) | |
| if random.random() < .12: | |
| im = im.filter(ImageFilter.UnsharpMask(radius=1, percent=random.randint(60, 130), threshold=2)) | |
| if random.random() < .25: | |
| im = ImageEnhance.Brightness(im).enhance(random.uniform(.68, 1.30)) | |
| if random.random() < .25: | |
| im = ImageEnhance.Contrast(im).enhance(random.uniform(.72, 1.28)) | |
| if random.random() < .15: | |
| im = ImageEnhance.Color(im).enhance(random.uniform(.72, 1.28)) | |
| if random.random() < .18: | |
| im = sensor_noise(im, random.uniform(.002, .016)) | |
| if random.random() < .18: | |
| w, h = im.size; s = random.uniform(.60, .92) | |
| tmp = im.resize((max(64, int(w*s)), max(64, int(h*s))), Image.Resampling.BICUBIC) | |
| im = tmp.resize((w, h), Image.Resampling.BILINEAR) | |
| if random.random() < .7: | |
| im = jpeg_recompress(im, random.randint(55, 92)) | |
| return im | |
| def foundation_transform(size: int = 224): | |
| return T.Compose([ | |
| T.Resize(size + 32, interpolation=T.InterpolationMode.BICUBIC), | |
| T.CenterCrop(size), | |
| T.ToTensor(), | |
| T.Normalize(IMAGENET_MEAN, IMAGENET_STD), | |
| ]) | |