scripts/pdf_preprocess.py
13.0 KB| 1 | #!/usr/bin/env python3 |
| 2 | """ |
| 3 | pdf_preprocess.py — 文档预处理工具（扫描件 OCR 前增强） |
| 4 | |
| 5 | 功能： |
| 6 | 1. 倾斜校正 (deskew) — 自动检测文本角度, 旋转修正 |
| 7 | 2. 文档展平 (unwarp) — 透视变换矫正弯曲/变形的文档照片 |
| 8 | 3. 方向分类 (orientation) — 检测并自动旋转倒置/侧向的页面 |
| 9 | |
| 10 | 用法： |
| 11 | # 安装依赖 |
| 12 | pip install opencv-python-headless numpy Pillow pdf2image |
| 13 | |
| 14 | # 单独使用 |
| 15 | python pdf_preprocess.py input.jpg -o ./output # 单图片 |
| 16 | python pdf_preprocess.py input.pdf -o ./output --all # PDF 全页 |
| 17 | |
| 18 | # 选择预处理步骤 |
| 19 | python pdf_preprocess.py input.jpg -o ./output --deskew # 仅倾斜校正 |
| 20 | python pdf_preprocess.py input.jpg -o ./output --unwarp # 仅展平 |
| 21 | python pdf_preprocess.py input.jpg -o ./output --orient # 仅方向校正 |
| 22 | python pdf_preprocess.py input.pdf -o ./output --all --dpi 400 # 全部 + 高DPI |
| 23 | """ |
| 24 | |
| 25 | import argparse |
| 26 | import os |
| 27 | import sys |
| 28 | from pathlib import Path |
| 29 | |
| 30 | try: |
| 31 | import cv2 |
| 32 | import numpy as np |
| 33 | from PIL import Image |
| 34 | except ImportError: |
| 35 | print("请先安装依赖: pip install opencv-python-headless numpy Pillow pdf2image") |
| 36 | sys.exit(1) |
| 37 | |
| 38 | |
| 39 | # ── 倾斜校正 (Deskew) ────────────────────────────────────────── |
| 40 | |
| 41 | def deskew_image(img: np.ndarray) -> np.ndarray: |
| 42 | """ |
| 43 | 检测图像中文本的倾斜角度并旋转校正。 |
| 44 | 适用于扫描件稍微偏斜的场景（< 15度）。 |
| 45 | """ |
| 46 | gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) |
| 47 | gray = cv2.bitwise_not(gray) |
| 48 | thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)[1] |
| 49 | |
| 50 | # 用 Hough 变换检测直线, 计算角度 |
| 51 | edges = cv2.Canny(thresh, 50, 150, apertureSize=3) |
| 52 | lines = cv2.HoughLines(edges, 1, np.pi / 180, 300) |
| 53 | |
| 54 | if lines is None: |
| 55 | return img # 无法检测直线, 原图返回 |
| 56 | |
| 57 | angles = [] |
| 58 | for rho, theta in lines[:, 0]: |
| 59 | angle = np.degrees(theta) - 90 # 转为水平偏移角 |
| 60 | if -45 <= angle <= 45: |
| 61 | angles.append(angle) |
| 62 | |
| 63 | if not angles: |
| 64 | return img |
| 65 | |
| 66 | median_angle = np.median(angles) |
| 67 | |
| 68 | # 只有角度超过 0.5 度才校正 |
| 69 | if abs(median_angle) < 0.5: |
| 70 | return img |
| 71 | |
| 72 | h, w = img.shape[:2] |
| 73 | center = (w // 2, h // 2) |
| 74 | rot = cv2.getRotationMatrix2D(center, median_angle, 1.0) |
| 75 | |
| 76 | # 计算旋转后的画布尺寸, 避免裁剪 |
| 77 | cos = abs(rot[0, 0]) |
| 78 | sin = abs(rot[0, 1]) |
| 79 | nw = int(h * sin + w * cos) |
| 80 | nh = int(h * cos + w * sin) |
| 81 | rot[0, 2] += nw / 2 - center[0] |
| 82 | rot[1, 2] += nh / 2 - center[1] |
| 83 | |
| 84 | corrected = cv2.warpAffine(img, rot, (nw, nh), flags=cv2.INTER_CUBIC, |
| 85 | borderMode=cv2.BORDER_REPLICATE) |
| 86 | return corrected |
| 87 | |
| 88 | |
| 89 | # ── 方向校正 (Orientation) ───────────────────────────────────── |
| 90 | |
| 91 | def correct_orientation(img: np.ndarray) -> np.ndarray: |
| 92 | """ |
| 93 | 检测图像方向并自动旋转到正位 (0°/90°/180°/270°)。 |
| 94 | 利用 PaddleOCR 的方向分类模型或基于文本行检测。 |
| 95 | """ |
| 96 | gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) |
| 97 | |
| 98 | # 方法: 检测文本区域的分布方向 |
| 99 | # 对二值图像做投影分析, 判断文本是横向还是纵向 |
| 100 | _, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU) |
| 101 | h, w = thresh.shape |
| 102 | |
| 103 | # 计算水平和垂直投影 |
| 104 | hor_proj = np.sum(thresh, axis=1) # 每行像素和 |
| 105 | ver_proj = np.sum(thresh, axis=0) # 每列像素和 |
| 106 | |
| 107 | # 判断文本排列方向 |
| 108 | hor_var = np.var(hor_proj) |
| 109 | ver_var = np.var(ver_proj) |
| 110 | |
| 111 | # 如果竖排方差更大, 说明文本可能是纵向排列, 需要旋转 |
| 112 | if ver_var > hor_var * 1.5: |
| 113 | # 判断上下方向: 简单方法 - 取上1/4和下1/4区域比较文本密度 |
| 114 | top_quarter = np.sum(thresh[:h//4, :]) |
| 115 | bot_quarter = np.sum(thresh[3*h//4:, :]) |
| 116 | |
| 117 | if top_quarter > bot_quarter: |
| 118 | # 文本在上方 → 逆时针旋转90度 |
| 119 | img = cv2.rotate(img, cv2.ROTATE_90_COUNTERCLOCKWISE) |
| 120 | else: |
| 121 | img = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE) |
| 122 | |
| 123 | # 已正位, 返回 |
| 124 | return img |
| 125 | |
| 126 | |
| 127 | # ── 文档展平 (Unwarp / Perspective Correction) ──────────────── |
| 128 | |
| 129 | def find_document_contour(img: np.ndarray) -> np.ndarray | None: |
| 130 | """ |
| 131 | 检测文档边界（最大四边形轮廓）。 |
| 132 | 返回 4 个角点坐标（左上、右上、右下、左下）。 |
| 133 | """ |
| 134 | gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) |
| 135 | # 自适应阈值, 适应不同光照 |
| 136 | thresh = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, |
| 137 | cv2.THRESH_BINARY_INV, 11, 2) |
| 138 | # 形态学闭运算, 连接断裂边缘 |
| 139 | kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5)) |
| 140 | closed = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel) |
| 141 | |
| 142 | contours, _ = cv2.findContours(closed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) |
| 143 | if not contours: |
| 144 | return None |
| 145 | |
| 146 | # 找最大轮廓 |
| 147 | largest = max(contours, key=cv2.contourArea) |
| 148 | peri = cv2.arcLength(largest, True) |
| 149 | approx = cv2.approxPolyDP(largest, 0.02 * peri, True) |
| 150 | |
| 151 | # 如果近似轮廓是四边形, 返回其4个角点 |
| 152 | if len(approx) == 4: |
| 153 | return _order_corners(approx.reshape(4, 2)) |
| 154 | |
| 155 | # 非四边形, 尝试找最小外接矩形 |
| 156 | rect = cv2.minAreaRect(largest) |
| 157 | corners = cv2.boxPoints(rect) |
| 158 | return _order_corners(corners) |
| 159 | |
| 160 | |
| 161 | def _order_corners(pts: np.ndarray) -> np.ndarray: |
| 162 | """将 4 个点排序为: 左上、右上、右下、左下""" |
| 163 | rect = np.zeros((4, 2), dtype=np.float32) |
| 164 | s = pts.sum(axis=1) |
| 165 | rect[0] = pts[np.argmin(s)] # 左上 (sum 最小) |
| 166 | rect[2] = pts[np.argmax(s)] # 右下 (sum 最大) |
| 167 | diff = np.diff(pts, axis=1) |
| 168 | rect[1] = pts[np.argmin(diff)] # 右上 (diff 最小) |
| 169 | rect[3] = pts[np.argmax(diff)] # 左下 (diff 最大) |
| 170 | return rect |
| 171 | |
| 172 | |
| 173 | def unwarp_image(img: np.ndarray, padding: float = 0.05) -> np.ndarray: |
| 174 | """ |
| 175 | 透视变换校正文档照片。 |
| 176 | 自动检测文档边界并展平为矩形。 |
| 177 | padding: 输出图像边界的留白比例(0~1) |
| 178 | """ |
| 179 | corners = find_document_contour(img) |
| 180 | if corners is None: |
| 181 | return img # 无法检测边界, 原图返回 |
| 182 | |
| 183 | # 计算目标矩形尺寸 |
| 184 | (tl, tr, br, bl) = corners |
| 185 | w_top = np.linalg.norm(tr - tl) |
| 186 | w_bot = np.linalg.norm(br - bl) |
| 187 | max_w = int(max(w_top, w_bot)) |
| 188 | h_left = np.linalg.norm(bl - tl) |
| 189 | h_right = np.linalg.norm(br - tr) |
| 190 | max_h = int(max(h_left, h_right)) |
| 191 | |
| 192 | # 加留白 |
| 193 | pad_x = int(max_w * padding) |
| 194 | pad_y = int(max_h * padding) |
| 195 | |
| 196 | dst = np.array([ |
| 197 | [pad_x, pad_y], |
| 198 | [max_w + pad_x, pad_y], |
| 199 | [max_w + pad_x, max_h + pad_y], |
| 200 | [pad_x, max_h + pad_y] |
| 201 | ], dtype=np.float32) |
| 202 | |
| 203 | matrix = cv2.getPerspectiveTransform(corners, dst) |
| 204 | result = cv2.warpPerspective(img, matrix, |
| 205 | (max_w + 2 * pad_x, max_h + 2 * pad_y), |
| 206 | flags=cv2.INTER_CUBIC, |
| 207 | borderMode=cv2.BORDER_REPLICATE) |
| 208 | return result |
| 209 | |
| 210 | |
| 211 | # ── 完整管线 ──────────────────────────────────────────────────── |
| 212 | |
| 213 | def preprocess_image(img: np.ndarray, deskew: bool = True, |
| 214 | orient: bool = True, unwarp: bool = False) -> np.ndarray: |
| 215 | """ |
| 216 | 完整的预处理流水线: |
| 217 | 1. 方向校正 (先做, 保证方向正确) |
| 218 | 2. 文档展平 (再做, 修正透视) |
| 219 | 3. 倾斜校正 (最后, 微调旋转) |
| 220 | """ |
| 221 | if orient: |
| 222 | img = correct_orientation(img) |
| 223 | if unwarp: |
| 224 | img = unwarp_image(img) |
| 225 | if deskew: |
| 226 | img = deskew_image(img) |
| 227 | return img |
| 228 | |
| 229 | |
| 230 | def save_result(img: np.ndarray, output_path: str): |
| 231 | """保存预处理后的图片""" |
| 232 | os.makedirs(os.path.dirname(output_path) or '.', exist_ok=True) |
| 233 | cv2.imwrite(output_path, img, |
| 234 | [cv2.IMWRITE_JPEG_QUALITY, 95]) |
| 235 | |
| 236 | |
| 237 | # ── PDF 处理 ──────────────────────────────────────────────────── |
| 238 | |
| 239 | def pdf_to_images(pdf_path: str, dpi: int = 300) -> list[np.ndarray]: |
| 240 | """将 PDF 每页转为 OpenCV 图像列表""" |
| 241 | try: |
| 242 | from pdf2image import convert_from_path |
| 243 | except ImportError: |
| 244 | print("需要安装 pdf2image: pip install pdf2image") |
| 245 | print("Windows 还需要安装 poppler: https://github.com/oschwartz10612/poppler-windows/releases/") |
| 246 | print("或在 PATH 中添加 poppler bin 目录") |
| 247 | sys.exit(1) |
| 248 | |
| 249 | images = convert_from_path(pdf_path, dpi=dpi) |
| 250 | return [cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR) for img in images] |
| 251 | |
| 252 | |
| 253 | # ── CLI ───────────────────────────────────────────────────────── |
| 254 | |
| 255 | def main(): |
| 256 | parser = argparse.ArgumentParser( |
| 257 | description="文档预处理工具 — 倾斜校正 / 方向校正 / 文档展平", |
| 258 | formatter_class=argparse.RawDescriptionHelpFormatter, |
| 259 | epilog=""" |
| 260 | 示例: |
| 261 | %(prog)s input.jpg -o ./output |
| 262 | %(prog)s input.jpg -o ./output --deskew # 仅倾斜校正 |
| 263 | %(prog)s input.jpg -o ./output --unwarp # 仅文档展平 |
| 264 | %(prog)s input.pdf -o ./output --all # PDF 全页全部预处理 |
| 265 | %(prog)s input.pdf -o ./output --all --dpi 400 # 高DPI扫描 |
| 266 | %(prog)s input.pdf -o ./output --preview # 预览效果(不保存) |
| 267 | """ |
| 268 | ) |
| 269 | parser.add_argument('input', help='输入文件 (PDF / 图片)') |
| 270 | parser.add_argument('-o', '--output', default='./preprocessed', |
| 271 | help='输出目录 (默认: ./preprocessed)') |
| 272 | |
| 273 | # 预处理选项 |
| 274 | parser.add_argument('--all', action='store_true', |
| 275 | help='启用全部预处理: 方向校正 + 文档展平 + 倾斜校正') |
| 276 | parser.add_argument('--deskew', action='store_true', help='倾斜校正') |
| 277 | parser.add_argument('--orient', action='store_true', help='方向校正') |
| 278 | parser.add_argument('--unwarp', action='store_true', help='文档展平 (透视变换)') |
| 279 | |
| 280 | # 高级选项 |
| 281 | parser.add_argument('--dpi', type=int, default=300, |
| 282 | help='PDF 转图片 DPI (默认: 300)') |
| 283 | parser.add_argument('--preview', action='store_true', |
| 284 | help='预览模式: 显示处理后图片（不保存, GUI 模式）') |
| 285 | |
| 286 | args = parser.parse_args() |
| 287 | |
| 288 | input_path = args.input |
| 289 | if not os.path.exists(input_path): |
| 290 | print(f"❌ 文件不存在: {input_path}") |
| 291 | sys.exit(1) |
| 292 | |
| 293 | # 确定预处理步骤 |
| 294 | if args.all: |
| 295 | do_deskew = do_orient = do_unwarp = True |
| 296 | else: |
| 297 | do_deskew = args.deskew or (not args.orient and not args.unwarp and not args.all) |
| 298 | do_orient = args.orient |
| 299 | do_unwarp = args.unwarp |
| 300 | |
| 301 | # 如果没有任何选项被指定, 默认只做 deskew |
| 302 | if not (args.deskew or args.orient or args.unwarp or args.all): |
| 303 | do_deskew = True |
| 304 | |
| 305 | print(f"📄 输入: {input_path}") |
| 306 | print(f" 预处理: {'倾斜校正 ' if do_deskew else ''}" |
| 307 | f"{'方向校正 ' if do_orient else ''}" |
| 308 | f"{'文档展平 ' if do_unwarp else ''}") |
| 309 | |
| 310 | ext = input_path.lower() |
| 311 | if ext.endswith('.pdf'): |
| 312 | print(f" 正在转换 PDF → 图片 (DPI={args.dpi})...") |
| 313 | images = pdf_to_images(input_path, dpi=args.dpi) |
| 314 | print(f" 共 {len(images)} 页") |
| 315 | else: |
| 316 | img = cv2.imread(input_path) |
| 317 | if img is None: |
| 318 | print(f"❌ 无法读取图片: {input_path}") |
| 319 | sys.exit(1) |
| 320 | images = [img] |
| 321 | |
| 322 | output_dir = args.output |
| 323 | basename = os.path.splitext(os.path.basename(input_path))[0] |
| 324 | |
| 325 | for i, img in enumerate(images): |
| 326 | # 预处理 |
| 327 | processed = preprocess_image(img, deskew=do_deskew, |
| 328 | orient=do_orient, unwarp=do_unwarp) |
| 329 | |
| 330 | if args.preview: |
| 331 | # 显示原图 vs 处理后的对比 |
| 332 | preview = np.hstack([ |
| 333 | cv2.resize(img, (640, int(640 * img.shape[0] / img.shape[1]) )), |
| 334 | cv2.resize(processed, (640, int(640 * processed.shape[0] / processed.shape[1]) )) |
| 335 | ]) |
| 336 | cv2.imshow(f'预处理对比 — 第{i+1}页 (原图 | 处理后)', preview) |
| 337 | cv2.waitKey(0) |
| 338 | cv2.destroyAllWindows() |
| 339 | else: |
| 340 | if len(images) > 1: |
| 341 | out_name = f"{basename}_p{i+1:03d}.jpg" |
| 342 | else: |
| 343 | out_name = f"{basename}_preprocessed.jpg" |
| 344 | out_path = os.path.join(output_dir, out_name) |
| 345 | save_result(processed, out_path) |
| 346 | print(f" ✅ 已保存: {out_path}") |
| 347 | |
| 348 | if not args.preview: |
| 349 | print(f"\n✅ 预处理完成！结果保存在: {os.path.abspath(output_dir)}") |
| 350 | print(f" 提示: 可以将预处理后的图片输入到 'pdf2md' 管线获得更好效果") |
| 351 | |
| 352 | |
| 353 | if __name__ == '__main__': |
| 354 | main() |
文件内容是平台检查时读出的文本，只供查看。平台不会运行包里的任何脚本或程序。
