| """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: |
| 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: |
| out = (out - 128) * rng.uniform(0.6, 1.3) + 128 + rng.uniform(-40, 30) |
| if rng.random() < 0.3 * strength: |
| 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: |
| out *= rng.uniform(0.85, 1.1, 3) |
| out = np.clip(out, 0, 255).astype(np.uint8) |
| if rng.random() < 0.3 * strength: |
| 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: |
| 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: |
| 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: |
| mask = rng.random(out.shape[:2]) |
| out[mask < 0.002] = 0 |
| out[mask > 0.998] = 255 |
| if rng.random() < 0.05 * strength: |
| 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: |
| 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: |
| 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]]) |
| |
| 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 |
|
|
