模拟相机拍照——对文档进行数据增强

一. 背景

假如我们有一个标准文件，我们对其进行文字识别、版面分析或者其他下游任务就比较容易。然而，当图片是手机拍照获取的，图片中往往有阴影、摩尔纹、弯曲。
那么，如何通过标准的文档，获得类似相机拍照的图片呢？
这里介绍的就是文档数据增强，用标准文档模拟相机拍照场景。该方法不仅能用于文档各场景的数据增强，用于OCR检测识别等任务；还能合成各种图片训练对，用于文档去阴影、文档去摩尔纹、文档弯曲矫正等各项任务。

二. 效果实现

首先给大家展示的是一个PDF截图和对应的标注（红色为标注框）
在这里插入图片描述
下面给标准图片分别添加阴影、摩尔纹、弯曲，效果如下：

摩尔纹+弯曲，并且把标注点映射到弯曲图片上，如下图所示：

阴影+弯曲，并且把标注点映射到弯曲图片上，如下图所示：

三. 算法原理与代码实现

原理：利用渲染工具（推荐blender），渲染出各种弯曲、阴影、摩尔纹，然后再pdf图片上进行合成。
最后，一定要代码实现(只给初级版本，完整版本比较复杂)：

import os
import cv2
import json
import random
import numpy as np
from scipy.interpolate import LinearNDInterpolator as linterp
from scipy.interpolate import NearestNDInterpolator as nearestclass LinearNDInterpolatorExt(object):def __init__(self, points, values):self.funcinterp = linterp(points, values)self.funcnearest = nearest(points, values)def __call__(self, *args):z = self.funcinterp(*args)chk = np.isnan(z)if chk.any():return np.where(chk, self.funcnearest(*args), z)else:return zdef crop_flow_from_nan(flow):mask = ~np.any(np.isnan(flow), -1)x, y, w, h = cv2.boundingRect(mask.astype(np.uint8))flow = flow[y: y + h, x: x + w]mask = mask[y: y + h, x: x + w]max_nonzero_ratio = 0.9max_crop_size = 20mask_h, mask_w = mask.shape[0], mask.shape[1]y0 = max_crop_sizefor i in range(0, max_crop_size):if np.count_nonzero(mask[i]) / mask_w > max_nonzero_ratio:y0 = ibreaky1 = mask_h - 1 - max_crop_sizefor i in range(mask_h - 1, y1, -1):if np.count_nonzero(mask[i]) / mask_w > max_nonzero_ratio:y1 = ibreakcrop_mask = mask[y0:y1]mask_h, mask_w = crop_mask.shape[0], crop_mask.shape[1]x0 = max_crop_sizefor i in range(0, x0):if np.count_nonzero(mask[:, i]) / mask_h > max_nonzero_ratio:x0 = ibreakx1 = mask_w - 1 - max_crop_sizefor i in range(mask_w - 1, x1, -1):if np.count_nonzero(mask[:, i]) / mask_h > max_nonzero_ratio:x1 = ibreakflow = flow[y0:y1, x0:x1]return flowdef flow_2_points(flow, pts):"""根据flow映射场反向计算点的对应点:param flow: 前向、或后向映射场, range (-1,  1):param pts: 目标图、或原图的坐标点, 点经过归一化 range (0, 1),  shape: (n, 2):return: 原图、或目标图的坐标点, 经过归一化 range (0, 1), shape: (n, 2)"""mask = ~np.any(np.isnan(flow), -1)flow_masked = flow[mask]flow_w, flow_h = flow.shape[1], flow.shape[0]flow_xrange = np.arange(flow_w, dtype=np.float32)flow_yrange = np.arange(flow_h, dtype=np.float32)flow_xgrid, flow_ygrid = np.meshgrid(flow_xrange, flow_yrange)flow_xgrid_masked = flow_xgrid[mask]flow_ygrid_masked = flow_ygrid[mask]src_pts = (pts - 0.5) * 2  # (0-1) to (-1, 1)interpX = LinearNDInterpolatorExt(np.reshape(flow_masked, [-1, 2]), flow_xgrid_masked.reshape(-1))interpY = LinearNDInterpolatorExt(np.reshape(flow_masked, [-1, 2]), flow_ygrid_masked.reshape(-1))fm_x = interpX(src_pts)fm_y = interpY(src_pts)# fm_x, fm_y range is (0, flow_w-1)  and (0, flow_h-1), need convert to (0-1)fm_x = fm_x / (flow_w - 1)fm_y = fm_y / (flow_h - 1)return np.stack((fm_x, fm_y), axis=-1)def warp_img(img, flow, points_list):h, w, _ = img.shapeflow = crop_flow_from_nan(flow)flow = flow.astype(np.float32)flow = cv2.resize(flow, (256, 256))points_list_warp = []for points in points_list:points = points.astype(np.float64)points[:, 0] /= w*1.0points[:, 1] /= h*1.0points_warp = flow_2_points(flow, points)points_warp[:, 0] *= wpoints_warp[:, 1] *= hpoints_list_warp.append(points_warp)bm_flow = flow / 2 + 0.5bm_flow[..., 0] = bm_flow[..., 0] * wbm_flow[..., 1] = bm_flow[..., 1] * hbm_flow = np.nan_to_num(bm_flow, nan=-1)if bm_flow.shape[0] != h or bm_flow.shape[1] != w:bm_flow = cv2.resize(bm_flow, (w, h))warp_img = cv2.remap(img, bm_flow.astype(np.float32), None, cv2.INTER_LINEAR, borderValue=(255, 255, 255))return warp_img, points_list_warpdef json_2_points(json_path):with open(json_path, "r") as f:data = json.load(f)obj_list = []for obj in data[0]['annotations']:obj = obj['coordinates']cx, cy, w, h = obj['x'], obj['y'], obj['width'], obj['height']x1 = cx - 0.5 * wx2 = cx + 0.5 * wy1 = cy - 0.5 * hy2 = cy + 0.5 * hpoints = np.array([[x1,y1], [x2,y1], [x2,y2], [x1,y2]], np.int32)obj_list.append(points)return obj_listdef add_background(img, img_background):height, width, _ = img.shapebackground = cv2.resize(img_background, (width, height))img_res = img * 0.5 + background * 0.5img_res = np.clip(img_res, 0, 255)return img_resif __name__ == "__main__":img = cv2.imread("test.png")shadow = cv2.imread("./background/shadow.jpg")img = add_background(img, shadow)obj_list = json_2_points("test.json")flow = np.load("test.npy")warp_img, points_list_warp = warp_img(img, flow, obj_list)cv2.imwrite("warp_shadow.jpg", warp_img)for points in points_list_warp:cv2.polylines(warp_img, [points.astype(np.int32)], isClosed=True, color=(0, 0, 255), thickness=1)cv2.imwrite("warp_shadow_draw.jpg", warp_img)