Hugging Face's logo

myned-ai
/
facepoints

Keypoint Detection
PyTorch
Safetensors
face-landmarks
dense-landmarks
face-alignment
3d-face
3dmm
gnm
avatar
hrnet
Eval Results (legacy)
Model card Files
xet
 Community
facepoints / train /data /augment.py
antonios-makro's picture
antonios-makro
Facepoints 1.0.0: code, model, prep data and card
88607d8 verified
Raw History Blame Contribute Delete
9.72 kB
"""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