guoju0808/ry-ocr/app/core/image_processor.py

"""图像预处理：自动旋转、放大、去噪"""
from __future__ import annotations

from pathlib import Path

import cv2
import numpy as np


def auto_rotate(img_path: str | Path) -> Path:
    """基于方向检测的简易旋转（> 阈值倾斜就转 90°）

    说明：发票多为横向，此处只处理 0/90/180/270 四方向，
    复杂倾斜交给 PaddleOCR 自带的 textline orientation。
    """
    p = Path(img_path)
    img = cv2.imread(str(p))
    if img is None:
        return p

    h, w = img.shape[:2]
    # 横图宽 > 高 * 1.2，认为方向正确；否则旋转
    if h > w * 1.2:
        rotated = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)
        out = p.with_name(f"{p.stem}_rot.png")
        cv2.imwrite(str(out), rotated)
        return out
    return p


def enhance(img_path: str | Path) -> Path:
    """轻度增强：灰度 + 自适应二值化（对手机拍的发票有帮助）"""
    p = Path(img_path)
    img = cv2.imread(str(p))
    if img is None:
        return p

    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    # 弱增强：只对低对比度图做二值化
    if gray.std() < 50:
        binary = cv2.adaptiveThreshold(
            gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 31, 10
        )
        out = p.with_name(f"{p.stem}_enh.png")
        cv2.imwrite(str(out), binary)
        return out
    return p