File size: 3,064 Bytes
c793f45 | 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 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | """Expectation-over-transformation ops applied inside the PGD loop so the
perturbation survives what a scraper/trainer does: JPEG, resize, crop, blur.
JPEG is non-differentiable, so we use BPDA / straight-through: forward pass uses
real PIL JPEG, backward pass is identity. That makes the optimizer account for
recompression in expectation instead of finding a fragile solution.
"""
from __future__ import annotations
import io
import math
import random
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
def _to_pil_batch(x: torch.Tensor) -> list[Image.Image]:
a = (x.clamp(0, 1).permute(0, 2, 3, 1).detach().cpu().float().numpy() * 255)
a = a.round().astype(np.uint8)
return [Image.fromarray(im) for im in a]
def jpeg_ste(x: torch.Tensor, quality: int) -> torch.Tensor:
imgs = _to_pil_batch(x)
out = []
for im in imgs:
buf = io.BytesIO()
im.save(buf, format="JPEG", quality=int(quality))
buf.seek(0)
out.append(np.asarray(Image.open(buf).convert("RGB"), dtype=np.float32) / 255.0)
xj = torch.from_numpy(np.stack(out)).permute(0, 3, 1, 2).to(x.device, x.dtype)
return x + (xj - x).detach() # straight-through
def _gauss_kernel(sigma: float, device, dtype) -> torch.Tensor:
r = max(1, int(math.ceil(3 * sigma)))
xs = torch.arange(-r, r + 1, device=device, dtype=torch.float32)
k = torch.exp(-(xs ** 2) / (2 * sigma * sigma))
k = (k / k.sum()).to(dtype)
return k
def gaussian_blur(x: torch.Tensor, sigma: float) -> torch.Tensor:
k = _gauss_kernel(sigma, x.device, x.dtype)
ck = k.view(1, 1, 1, -1).repeat(x.shape[1], 1, 1, 1)
pad = (k.numel() // 2)
x = F.conv2d(F.pad(x, (pad, pad, 0, 0), mode="reflect"), ck, groups=x.shape[1])
ck = k.view(1, 1, -1, 1).repeat(x.shape[1], 1, 1, 1)
x = F.conv2d(F.pad(x, (0, 0, pad, pad), mode="reflect"), ck, groups=x.shape[1])
return x
def resize_roundtrip(x: torch.Tensor, scale: float) -> torch.Tensor:
_, _, h, w = x.shape
nh, nw = max(8, int(h * scale)), max(8, int(w * scale))
down = F.interpolate(x, size=(nh, nw), mode="bilinear", align_corners=False, antialias=True)
return F.interpolate(down, size=(h, w), mode="bilinear", align_corners=False)
def random_crop_pad(x: torch.Tensor, keep: float) -> torch.Tensor:
_, _, h, w = x.shape
ch, cw = int(h * keep), int(w * keep)
top = random.randint(0, h - ch)
left = random.randint(0, w - cw)
return x[:, :, top:top + ch, left:left + cw]
def apply_eot(x: torch.Tensor, rng: random.Random, strong: bool = True) -> torch.Tensor:
"""Sample a random subset of transforms. Always returns a valid [0,1] image."""
if rng.random() < 0.85:
x = jpeg_ste(x, rng.randint(55, 92))
if rng.random() < 0.6:
x = resize_roundtrip(x, rng.uniform(0.55, 1.0))
if strong and rng.random() < 0.4:
x = gaussian_blur(x, rng.uniform(0.4, 1.4))
if strong and rng.random() < 0.35:
x = random_crop_pad(x, rng.uniform(0.85, 0.98))
return x.clamp(0, 1)
|