File size: 6,136 Bytes
4dce182 | 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 | """Image degradations that mimic scans, phone photos, faxes and faded thermal receipts.
Geometric transforms return the homography so boxes can be moved with the pixels.
"""
import cv2
import numpy as np
from ..geometry import transform_quads
def _smooth_field(rng, h, w, cells=4):
f = rng.random((cells, cells)).astype(np.float32)
return cv2.resize(f, (w, h), interpolation=cv2.INTER_CUBIC)
def photometric(img, rng, strength=1.0, thermal=False):
"""img: uint8 HxW or HxWx3. Returns a degraded copy of the same shape."""
out = img.astype(np.float32)
h, w = out.shape[:2]
if thermal or rng.random() < 0.15 * strength: # faded print: lift ink toward paper, unevenly
fade = 0.25 + 0.6 * _smooth_field(rng, h, w, int(rng.integers(2, 6)))
if out.ndim == 3:
fade = fade[..., None]
paper = np.percentile(out, 90)
out = paper - (paper - out) * (1 - fade * rng.uniform(0.3, 0.9))
if rng.random() < 0.5 * strength: # contrast / brightness
out = (out - 128) * rng.uniform(0.6, 1.3) + 128 + rng.uniform(-40, 30)
if rng.random() < 0.3 * strength: # lighting gradient / shadow
g = 0.55 + 0.45 * _smooth_field(rng, h, w, int(rng.integers(2, 4)))
out *= g[..., None] if out.ndim == 3 else g
if out.ndim == 3 and rng.random() < 0.3 * strength: # colour cast
out *= rng.uniform(0.85, 1.1, 3)
out = np.clip(out, 0, 255).astype(np.uint8)
if rng.random() < 0.3 * strength: # resolution loss
f = rng.uniform(0.35, 0.8)
small = cv2.resize(out, (max(1, int(w * f)), max(1, int(h * f))), interpolation=cv2.INTER_AREA)
out = cv2.resize(small, (w, h), interpolation=rng.choice([cv2.INTER_LINEAR, cv2.INTER_NEAREST, cv2.INTER_CUBIC]))
if rng.random() < 0.4 * strength: # blur
if rng.random() < 0.8:
out = cv2.GaussianBlur(out, (0, 0), rng.uniform(0.3, 1.3))
else:
k = int(rng.integers(3, 7))
kern = np.zeros((k, k), np.float32)
kern[k // 2, :] = 1.0 / k
m = cv2.getRotationMatrix2D((k / 2 - 0.5, k / 2 - 0.5), rng.uniform(0, 180), 1)
out = cv2.filter2D(out, -1, cv2.warpAffine(kern, m, (k, k)))
if rng.random() < 0.5 * strength: # sensor / scan noise
noise = rng.normal(0, rng.uniform(2, 14), out.shape)
out = np.clip(out.astype(np.float32) + noise, 0, 255).astype(np.uint8)
if rng.random() < 0.08 * strength: # salt & pepper speckle
mask = rng.random(out.shape[:2])
out[mask < 0.002] = 0
out[mask > 0.998] = 255
if rng.random() < 0.05 * strength: # fax-like binarisation
g = out if out.ndim == 2 else cv2.cvtColor(out, cv2.COLOR_RGB2GRAY)
_, b = cv2.threshold(g, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
out = b if out.ndim == 2 else cv2.cvtColor(b, cv2.COLOR_GRAY2RGB)
if rng.random() < 0.6 * strength: # JPEG
ok, enc = cv2.imencode(".jpg", out, [cv2.IMWRITE_JPEG_QUALITY, int(rng.integers(25, 95))])
out = cv2.imdecode(enc, cv2.IMREAD_UNCHANGED)
return out
def _background(rng, h, w):
r = rng.random()
if r < 0.4:
base = rng.integers(20, 200, 3)
bg = np.ones((h, w, 3), np.float32) * base
bg += (rng.random((h // 8 + 1, w // 8 + 1, 3)) * 40 - 20).repeat(8, 0).repeat(8, 1)[:h, :w]
elif r < 0.7: # wood-ish stripes
base = rng.integers(60, 180, 3).astype(np.float32)
stripes = np.sin(np.linspace(0, rng.uniform(20, 80), w))[None, :, None] * rng.uniform(5, 25)
bg = np.ones((h, w, 3), np.float32) * base + stripes
else:
bg = np.ones((h, w, 3), np.float32) * rng.integers(0, 255, 3)
bg *= (0.6 + 0.4 * _smooth_field(rng, h, w, 3))[..., None]
return np.clip(bg, 0, 255).astype(np.uint8)
def geometric(img, rng, mode=None):
"""Returns (image, 3x3 homography from input to output coordinates).
modes: none | rotate (scanner skew) | photo (page on a background with perspective)
"""
h, w = img.shape[:2]
mode = mode or rng.choice(["none", "rotate", "photo"], p=[0.35, 0.4, 0.25])
if mode == "none":
return img, np.eye(3)
if mode == "rotate":
ang = rng.normal(0, 2.0) if rng.random() < 0.8 else rng.uniform(-8, 8)
m = cv2.getRotationMatrix2D((w / 2, h / 2), ang, 1.0)
cos, sin = abs(m[0, 0]), abs(m[0, 1])
nw, nh = int(h * sin + w * cos), int(h * cos + w * sin)
m[0, 2] += nw / 2 - w / 2
m[1, 2] += nh / 2 - h / 2
border = tuple(int(v) for v in np.median(img.reshape(-1, img.shape[-1]) if img.ndim == 3 else img.reshape(-1, 1), 0))
out = cv2.warpAffine(img, m, (nw, nh), flags=cv2.INTER_LINEAR, borderValue=border)
return out, np.vstack([m, [0, 0, 1]])
# photo: page occupies 55-95% of a larger frame, corners jittered
s = rng.uniform(1.05, 1.6)
W, H = int(w * s), int(h * s)
ox, oy = rng.uniform(0, W - w), rng.uniform(0, H - h)
j = 0.04 * min(w, h)
src = np.float32([[0, 0], [w, 0], [w, h], [0, h]])
dst = np.float32([[ox, oy], [ox + w, oy], [ox + w, oy + h], [ox, oy + h]]) + rng.uniform(-j, j, (4, 2)).astype(np.float32)
m = cv2.getPerspectiveTransform(src, dst)
bg = _background(rng, H, W) if img.ndim == 3 else cv2.cvtColor(_background(rng, H, W), cv2.COLOR_RGB2GRAY)
warped = cv2.warpPerspective(img, m, (W, H), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=0)
mask = cv2.warpPerspective(np.full((h, w), 255, np.uint8), m, (W, H), flags=cv2.INTER_NEAREST)
out = np.where((mask > 0)[..., None] if img.ndim == 3 else mask > 0, warped, bg)
return out, m
def augment_page(img, quads, ignore_quads, rng, strength=1.0, thermal=False, geo_mode=None):
img, m = geometric(img, rng, geo_mode)
quads = transform_quads(quads, m)
ignore_quads = transform_quads(ignore_quads, m)
img = photometric(img, rng, strength, thermal)
return img, quads, ignore_quads
|