195 lines
6.5 KiB
Python
195 lines
6.5 KiB
Python
#!/usr/bin/env python3
|
||
"""KOCR v4 — 图像预处理 + 自动图像分析调参"""
|
||
|
||
import cv2
|
||
import numpy as np
|
||
from dataclasses import dataclass, field
|
||
from typing import Optional
|
||
|
||
|
||
@dataclass
|
||
class ImageQuality:
|
||
"""图像质量分析结果"""
|
||
brightness: float = 0.0 # 平均亮度 [0-255]
|
||
contrast: float = 0.0 # 标准偏差
|
||
blur_score: float = 0.0 # 拉普拉斯方差（越小越模糊）
|
||
shadow_ratio: float = 0.0 # 阴影面积占比
|
||
photo_type: str = 'normal' # flash / dark / blurry / normal
|
||
needs_shadow_removal: bool = True
|
||
needs_sharpen: bool = False
|
||
needs_deskew: bool = True
|
||
|
||
|
||
def analyze_image(img: np.ndarray) -> ImageQuality:
|
||
"""
|
||
自动分析图像质量，返回各项指标和推荐的照片类型
|
||
|
||
Args:
|
||
img: 输入图像（BGR格式）
|
||
|
||
Returns:
|
||
ImageQuality: 包含所有分析结果
|
||
"""
|
||
h, w = img.shape[:2]
|
||
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
||
|
||
# 1. 亮度分析（取直方图均值）
|
||
brightness = float(np.mean(gray))
|
||
|
||
# 2. 对比度（灰度标准偏差）
|
||
contrast = float(np.std(gray))
|
||
|
||
# 3. 模糊检测（拉普拉斯方差）
|
||
laplacian = cv2.Laplacian(gray, cv2.CV_64F)
|
||
blur_score = float(laplacian.var())
|
||
|
||
# 4. 阴影检测
|
||
# 用大核膨胀分离背景和前景
|
||
dilated = cv2.dilate(gray, np.ones((15, 15), np.uint8))
|
||
bg = cv2.medianBlur(dilated, 31)
|
||
diff = cv2.absdiff(gray.astype(np.float32), bg.astype(np.float32))
|
||
shadow_mask = diff < 20 # 与背景差异小的区域视为阴影
|
||
shadow_ratio = float(np.sum(shadow_mask) / (h * w))
|
||
|
||
# 5. 判定照片类型
|
||
if brightness >= 200:
|
||
photo_type = 'flash'
|
||
elif brightness <= 60:
|
||
photo_type = 'dark'
|
||
elif blur_score < 80:
|
||
photo_type = 'blurry'
|
||
else:
|
||
photo_type = 'normal'
|
||
|
||
# 6. 预处理策略建议
|
||
needs_shadow_removal = shadow_ratio > 0.3 and photo_type != 'flash'
|
||
needs_sharpen = blur_score < 150
|
||
needs_deskew = photo_type not in ('flash',) # 闪光灯照片也可能歪
|
||
|
||
return ImageQuality(
|
||
brightness=round(brightness, 1),
|
||
contrast=round(contrast, 1),
|
||
blur_score=round(blur_score, 1),
|
||
shadow_ratio=round(shadow_ratio, 3),
|
||
photo_type=photo_type,
|
||
needs_shadow_removal=needs_shadow_removal,
|
||
needs_sharpen=needs_sharpen,
|
||
needs_deskew=needs_deskew,
|
||
)
|
||
|
||
|
||
class VoucherPreprocessor:
|
||
"""记账凭证图像预处理器，支持自动调参"""
|
||
|
||
def __init__(self, max_side=1600, denoise=True, deskew=True,
|
||
remove_shadow=True, sharpen=True):
|
||
self.max_side = max_side
|
||
self.denoise = denoise
|
||
self.deskew = deskew
|
||
self.remove_shadow = remove_shadow
|
||
self.sharpen = sharpen
|
||
|
||
@classmethod
|
||
def from_config(cls, config: dict, photo_type: str = 'normal'):
|
||
"""根据配置和照片类型创建预处理器实例
|
||
|
||
Args:
|
||
config: config.yaml 中 preprocessing 章节
|
||
photo_type: 照片类型（用于覆盖默认参数）
|
||
"""
|
||
params = {
|
||
'max_side': config.get('max_side', 1600),
|
||
'denoise': config.get('denoise', True),
|
||
'deskew': config.get('deskew', True),
|
||
'remove_shadow': config.get('remove_shadow', True),
|
||
'sharpen': config.get('sharpen', True),
|
||
}
|
||
# 应用照片类型覆盖
|
||
overrides = config.get('photo_type_overrides', {}).get(photo_type, {})
|
||
params.update(overrides)
|
||
return cls(**params)
|
||
|
||
def _resize(self, img):
|
||
"""统一缩放，长边不超过max_side，保持比例"""
|
||
h, w = img.shape[:2]
|
||
if max(h, w) > self.max_side:
|
||
scale = self.max_side / max(h, w)
|
||
new_w, new_h = int(w * scale), int(h * scale)
|
||
img = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)
|
||
return img
|
||
|
||
def _remove_shadow(self, img):
|
||
"""去除拍照阴影、纸张背景不均"""
|
||
rgb_planes = cv2.split(img)
|
||
result_planes = []
|
||
for plane in rgb_planes:
|
||
dilated_img = cv2.dilate(plane, np.ones((7, 7), np.uint8))
|
||
bg_img = cv2.medianBlur(dilated_img, 21)
|
||
diff_img = 255 - cv2.absdiff(plane, bg_img)
|
||
norm_img = cv2.normalize(diff_img, None, alpha=0, beta=255,
|
||
norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8UC1)
|
||
result_planes.append(norm_img)
|
||
return cv2.merge(result_planes)
|
||
|
||
def _deskew(self, img):
|
||
"""基于霍夫直线的倾斜矫正"""
|
||
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
||
edges = cv2.Canny(gray, 50, 150, apertureSize=3)
|
||
lines = cv2.HoughLines(edges, 1, np.pi / 180, 200)
|
||
if lines is None:
|
||
return img
|
||
angles = []
|
||
for line in lines:
|
||
_, theta = line[0]
|
||
angle = (theta * 180 / np.pi) - 90
|
||
if abs(angle) < 15:
|
||
angles.append(angle)
|
||
if not angles:
|
||
return img
|
||
median_angle = np.median(angles)
|
||
h, w = img.shape[:2]
|
||
center = (w // 2, h // 2)
|
||
M = cv2.getRotationMatrix2D(center, median_angle, 1.0)
|
||
rotated = cv2.warpAffine(img, M, (w, h), flags=cv2.INTER_CUBIC,
|
||
borderMode=cv2.BORDER_REPLICATE)
|
||
return rotated
|
||
|
||
def _denoise(self, img):
|
||
"""高斯去噪"""
|
||
return cv2.GaussianBlur(img, (3, 3), 0)
|
||
|
||
def _sharpen_op(self, img):
|
||
"""USM 锐化"""
|
||
blur = cv2.GaussianBlur(img, (0, 0), 3)
|
||
sharp = cv2.addWeighted(img, 1.5, blur, -0.5, 0)
|
||
return sharp
|
||
|
||
def process(self, img_or_path):
|
||
"""完整预处理入口
|
||
|
||
Args:
|
||
img_or_path: 图像路径或 numpy 数组
|
||
|
||
Returns:
|
||
预处理后的图像
|
||
"""
|
||
if isinstance(img_or_path, str):
|
||
img = cv2.imread(img_or_path)
|
||
if img is None:
|
||
raise ValueError(f"无法读取图片：{img_or_path}")
|
||
else:
|
||
img = img_or_path.copy()
|
||
|
||
img = self._resize(img)
|
||
|
||
if self.remove_shadow:
|
||
img = self._remove_shadow(img)
|
||
if self.deskew:
|
||
img = self._deskew(img)
|
||
if self.denoise:
|
||
img = self._denoise(img)
|
||
if self.sharpen:
|
||
img = self._sharpen_op(img)
|
||
|
||
return img
|