File size: 2,879 Bytes
30d60f0 | 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 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | 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),
])
|