File size: 9,723 Bytes
88607d8 | 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 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 | """Augmentation of a face crop and its labels: affine, low resolution, photometric, occluders."""
# Training only. One random affine that the labels follow, with a constant random border;
# photometric jitter (brightness, contrast, per-channel gain, gamma, hue/saturation, grey, gaussian
# and directional motion blur, noise, JPEG re-encode); zero to two occluders. The in-the-wild extras
# (background swap behind the person mask, relighting with cast shadows, sunglasses, real-texture
# occluders) are applied through a hook right after the affine (data.wild.Wild).
import functools
from typing import NamedTuple
import cv2
import numpy as np
class Affine(NamedTuple):
"""The random affine: scale range, rotation range in degrees and its probability."""
scale_min: float
scale_max: float
rotation_deg: float
rotation_p: float
class LowRes(NamedTuple):
"""Low-resolution branch: smallest face size (px), probability, JPEG probability and quality."""
min_px: int
p: float
jpeg_p: float
jpeg_q: int
def photometric(im, rng, grey_p=0.05):
"""Return the image with random photometric changes: brightness to JPEG re-encoding."""
x = im.astype(np.float32) / 255.0
if rng.random() < 0.9:
x = x * rng.uniform(0.6, 1.4)
m = x.mean()
x = (x - m) * rng.uniform(0.6, 1.4) + m
x = x * rng.uniform(0.85, 1.15, 3).astype(np.float32)
x = np.clip(x, 0, 1) ** rng.uniform(0.7, 1.4)
if rng.random() < 0.5: # hue / saturation
hsv = cv2.cvtColor((np.clip(x, 0, 1) * 255).astype(np.uint8), cv2.COLOR_BGR2HSV)
hsv = hsv.astype(np.float32)
hsv[:, :, 0] = (hsv[:, :, 0] + rng.uniform(-8, 8)) % 180
hsv[:, :, 1] = np.clip(hsv[:, :, 1] * rng.uniform(0.7, 1.3), 0, 255)
x = cv2.cvtColor(hsv.astype(np.uint8), cv2.COLOR_HSV2BGR).astype(np.float32) / 255.0
if rng.random() < grey_p:
x = np.repeat(x.mean(2, keepdims=True), 3, axis=2)
if rng.random() < 0.3:
x = cv2.GaussianBlur(x, (0, 0), rng.uniform(0.5, 2.0))
if rng.random() < 0.2: # directional motion blur, length proportional to the side
n = int(rng.integers(3, max(4, x.shape[0] // 20)))
k = np.zeros((n, n), np.float32)
cv2.line(k, (0, n // 2), (n - 1, n // 2), 1.0, 1)
k = cv2.warpAffine(
k, cv2.getRotationMatrix2D((n / 2, n / 2), rng.uniform(0, 180), 1), (n, n)
)
x = cv2.filter2D(x, -1, k / max(k.sum(), 1e-6))
if rng.random() < 0.3:
x = x + rng.normal(0, rng.uniform(0.01, 0.05), x.shape).astype(np.float32)
x = (np.clip(x, 0, 1) * 255).astype(np.uint8)
if rng.random() < 0.3:
x = jpeg(x, int(rng.integers(25, 90)))
return x
def jpeg(im, quality):
"""Return the image after a JPEG encode-decode round trip at the given quality."""
return cv2.imdecode(
cv2.imencode(".jpg", im, [cv2.IMWRITE_JPEG_QUALITY, quality])[1], cv2.IMREAD_COLOR
)
@functools.lru_cache(maxsize=4)
def grid(R):
"""Return the (yy, xx) pixel coordinates of an R x R image as float32 (cached per size)."""
return np.mgrid[0:R, 0:R].astype(np.float32)
def swap_bg(im, mask, bg, rng):
"""Return the person (mask 255) over a background crop, with a feathered edge."""
a = cv2.GaussianBlur(mask.astype(np.float32) / 255.0, (0, 0), rng.uniform(0.3, 1.0))[:, :, None]
if rng.random() < 0.5:
bg = bg[:, ::-1]
return (im * a + bg * (1 - a)).astype(np.uint8)
def relight(im, rng, mask=None):
"""Return the image relit by a directional light with random level, colour cast and shadow."""
# gain ramp e^(k d.p) across the image, level 0.25-1.1 of the original, with p 0.3 a soft-edged
# shadow polygon; applied to the person only when a mask is given
R = im.shape[0]
yy, xx = grid(R) / (R - 1) * 2 - 1
t = rng.uniform(0, 2 * np.pi)
g = np.exp(rng.uniform(0.3, 1.6) * (np.cos(t) * xx + np.sin(t) * yy))
g = g / g.mean() * rng.uniform(0.25, 1.1)
if rng.random() < 0.3:
m = np.zeros((R, R), np.uint8)
pts = (rng.uniform(-0.2, 1.2, (int(rng.integers(3, 6)), 2)) * R).astype(np.int32)
cv2.fillPoly(m, [cv2.convexHull(pts)], 1)
m = cv2.GaussianBlur(m.astype(np.float32), (0, 0), rng.uniform(1, 6))
g = g * (1 - m * rng.uniform(0.3, 0.7))
g = g[:, :, None] * rng.uniform(0.8, 1.2, 3).astype(np.float32)
lit = np.clip(im.astype(np.float32) * g, 0, 255)
if mask is not None:
a = (mask.astype(np.float32) / 255.0)[:, :, None]
lit = lit * a + im * (1 - a)
return lit.astype(np.uint8)
def glasses(im, rng, e1, e2):
"""Return the image with dark sunglasses over the two eye centres, one highlight on a lens."""
d = float(np.linalg.norm(e2 - e1))
if not (np.isfinite(d) and d > 4):
return im
ang = float(np.degrees(np.arctan2(e2[1] - e1[1], e2[0] - e1[0])))
ax = (int(d * rng.uniform(0.3, 0.45)), int(d * rng.uniform(0.2, 0.35)))
m = np.zeros(im.shape[:2], np.uint8)
for e in (e1, e2):
cv2.ellipse(m, (int(e[0]), int(e[1])), ax, ang, 0, 360, 1, -1)
cv2.line(m, tuple(int(v) for v in e1), tuple(int(v) for v in e2), 1, max(1, int(d * 0.06)))
lens = np.full(im.shape, rng.integers(0, 70, 3), np.float32)
yy, xx = grid(im.shape[0])
hl = np.exp(-(((xx - e1[0]) * 0.7 + (yy - e1[1])) ** 2) / (2 * (0.15 * d) ** 2))
lens = lens + hl[:, :, None] * rng.uniform(0, 120)
a = m[:, :, None].astype(np.float32) * rng.uniform(0.8, 1.0)
return np.clip(im * (1 - a) + lens * a, 0, 255).astype(np.uint8)
def occlude_real(im, rng, src, centre, size):
"""Return the image with a real-texture occluder (ellipse, polygon or band) near the face."""
# the occluder spans 20-50 % of the face size and sits within half a face size of the centre
R = im.shape[0]
m = np.zeros((R, R), np.uint8)
c = centre + rng.uniform(-0.5, 0.5, 2) * size
w, h = rng.uniform(0.1, 0.25, 2) * size
k = rng.random()
if k < 0.4:
cv2.ellipse(
m, (int(c[0]), int(c[1])), (int(w), int(h)), float(rng.uniform(0, 180)), 0, 360, 1, -1
)
elif k < 0.8:
pts = c + rng.uniform(-1, 1, (int(rng.integers(4, 8)), 2)) * [w, h]
cv2.fillPoly(m, [cv2.convexHull(pts.astype(np.int32))], 1)
else: # a band from one side, its edge 10-45 % of the face size off-centre (scarf, collar)
t = rng.uniform(0, 2 * np.pi)
n = np.array([np.cos(t), np.sin(t)])
e = centre + n * rng.uniform(0.1, 0.45) * size
yy, xx = grid(R)
m = (((xx - e[0]) * n[0] + (yy - e[1]) * n[1]) > 0).astype(np.uint8)
a = cv2.GaussianBlur(m.astype(np.float32), (0, 0), rng.uniform(0.3, 1.0))[:, :, None]
return (im * (1 - a) + src * a).astype(np.uint8)
def augment(im, uv, rng, res, train, aff, lowres, hook=None, shift=0.1, occ_p=1.0, grey_p=0.05):
"""Return (res x res image, labels) after affine, hook, low-res, photometric and occluders."""
# aff: Affine; lowres: LowRes; hook(im, uv, M) runs after the affine (data.wild.Wild.apply).
# shift, occ_p and grey_p at their defaults reproduce the studio recipe's random stream exactly.
S = im.shape[0]
if train:
s = rng.uniform(aff.scale_min, aff.scale_max)
a = (
rng.uniform(-aff.rotation_deg, aff.rotation_deg)
if rng.random() < aff.rotation_p
else 0.0
)
t = rng.uniform(-shift, shift, 2) * res
else:
s, a, t = 1.0, 0.0, np.zeros(2)
M = cv2.getRotationMatrix2D((S / 2, S / 2), a, s * res / S)
M[:, 2] += (res - S) / 2 + t
border = [int(v) for v in rng.integers(0, 256, 3)] if train else [0, 0, 0]
im = cv2.warpAffine(im, M, (res, res), flags=cv2.INTER_LINEAR, borderValue=border)
uv = uv @ M[:, :2].T + M[:, 2]
if hook is not None:
im = hook(im, uv, M)
if train:
if lowres.min_px < res and rng.random() < lowres.p:
# a small face in phone video: down to s px and back up (labels unmoved); with
# probability jpeg_p, JPEG-compress AT the small size (video is encoded at its own
# resolution, so its blocks are large after upscaling)
s = int(rng.integers(lowres.min_px, res))
im = cv2.resize(im, (s, s), interpolation=cv2.INTER_AREA)
if lowres.jpeg_p > 0 and rng.random() < lowres.jpeg_p:
im = jpeg(im, int(rng.integers(lowres.jpeg_q, 91)))
im = cv2.resize(im, (res, res))
im = photometric(im, rng, grey_p)
if occ_p >= 1.0 or rng.random() < occ_p:
im = occlude(im, rng)
return im, uv
def occlude(im, rng):
"""Return the image with zero to two random occluders (solid or noise; rectangle or ellipse)."""
if rng.random() > 0.35:
return im
R = im.shape[0]
for _ in range(int(rng.integers(1, 3))):
w, h = (rng.uniform(0.1, 0.35, 2) * R).astype(int)
cx, cy = (rng.uniform(0.1, 0.9, 2) * R).astype(int)
if rng.random() < 0.6:
patch = np.full((h, w, 3), rng.integers(0, 256, 3), np.uint8)
else:
patch = rng.integers(0, 256, (h, w, 3)).astype(np.uint8)
mask = np.zeros((R, R), np.uint8)
if rng.random() < 0.5:
cv2.rectangle(mask, (cx - w // 2, cy - h // 2), (cx + w // 2, cy + h // 2), 1, -1)
else:
cv2.ellipse(mask, (cx, cy), (w // 2, h // 2), float(rng.uniform(0, 180)), 0, 360, 1, -1)
full = np.zeros_like(im)
y0, x0 = max(cy - h // 2, 0), max(cx - w // 2, 0)
ph, pw = min(h, R - y0), min(w, R - x0)
full[y0 : y0 + ph, x0 : x0 + pw] = patch[:ph, :pw]
im = np.where(mask[:, :, None] > 0, full, im)
return im
|