583 lines
19 KiB
Python
583 lines
19 KiB
Python
#!/usr/bin/env python
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
CMA Code Extraction using Template Matching (Primary Method)
|
||
|
||
This module uses template matching to locate the CMA logo, then extracts
|
||
the CMA code from the region around the logo using OCR.
|
||
|
||
This is the PRIMARY method for CMA extraction, with fallback to full-page OCR.
|
||
|
||
Author: Claude Code
|
||
Date: 2025-02-16
|
||
"""
|
||
|
||
import os
|
||
import re
|
||
import cv2
|
||
import numpy as np
|
||
import logging
|
||
from pathlib import Path
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# CMA code patterns
|
||
PATTERN_PRIMARY = r'2[0-9]{10}' # 11 digits starting with 2
|
||
PATTERN_FALLBACK = r'[0-9]{11}' # any 11 digits
|
||
|
||
|
||
def imread_unicode(path, flags=cv2.IMREAD_COLOR):
|
||
"""
|
||
cv2.imread replacement that supports paths with non-ASCII characters.
|
||
|
||
Args:
|
||
path: Image file path (may contain Chinese characters)
|
||
flags: cv2.IMREAD_* flags
|
||
|
||
Returns:
|
||
Image as numpy array or None if failed
|
||
"""
|
||
try:
|
||
data = np.fromfile(str(path), dtype=np.uint8)
|
||
img = cv2.imdecode(data, flags)
|
||
return img
|
||
except Exception as e:
|
||
logger.error(f"Failed to read image {path}: {e}")
|
||
return None
|
||
|
||
|
||
def load_cma_template(template_path='template/CMA_Logo.png'):
|
||
"""
|
||
加载 CMA logo 模板图像
|
||
|
||
Args:
|
||
template_path: 模板图像路径
|
||
|
||
Returns:
|
||
template: 模板图像（灰度）
|
||
template_rgb: 模板图像（RGB，用于可视化）
|
||
"""
|
||
if not os.path.exists(template_path):
|
||
logger.error(f"模板文件不存在: {template_path}")
|
||
return None, None
|
||
|
||
# 读取模板图像（灰度）
|
||
template = cv2.imread(template_path, cv2.IMREAD_GRAYSCALE)
|
||
if template is None:
|
||
logger.error(f"无法读取模板文件: {template_path}")
|
||
return None, None
|
||
|
||
logger.debug(f"加载模板: {template_path}, 尺寸: {template.shape}")
|
||
|
||
return template, template
|
||
|
||
|
||
def match_template(page_img, template, method=cv2.TM_CCOEFF_NORMED):
|
||
"""
|
||
使用 cv2.matchTemplate 进行模板匹配
|
||
|
||
Args:
|
||
page_img: 页面图像（灰度或彩色）
|
||
template: CMA logo 模板（灰度）
|
||
method: 匹配方法（默认 TM_CCOEFF_NORMED）
|
||
|
||
Returns:
|
||
result: 匹配结果字典，包含匹配区域、最大值、位置
|
||
"""
|
||
# 转换为灰度（如果是彩色图像）
|
||
if len(page_img.shape) == 3:
|
||
page_gray = cv2.cvtColor(page_img, cv2.COLOR_BGR2GRAY)
|
||
else:
|
||
page_gray = page_img
|
||
|
||
# 执行模板匹配
|
||
result = cv2.matchTemplate(page_gray, template, method=method)
|
||
|
||
if result is None:
|
||
logger.warning("模板匹配失败")
|
||
return None
|
||
|
||
# 获取匹配结果
|
||
min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
|
||
|
||
# 对于 TM_SQDIFF 方法，最小值是最佳匹配
|
||
if method in [cv2.TM_SQDIFF, cv2.TM_SQDIFF_NORMED]:
|
||
top_left = min_loc
|
||
match_value = 1 - min_val # 转换为相似度
|
||
else:
|
||
top_left = max_loc
|
||
match_value = max_val
|
||
|
||
# 计算匹配区域的中心
|
||
template_h, template_w = template.shape[:2]
|
||
center_x = top_left[0] + template_w // 2
|
||
center_y = top_left[1] + template_h // 2
|
||
|
||
logger.info(f"[TM] Match confidence: {match_value:.3f} (threshold: 0.4)")
|
||
logger.info(f"[TM] Logo detected at center ({center_x}, {center_y}) in image {page_gray.shape[1]}x{page_gray.shape[0]}")
|
||
|
||
return {
|
||
'max_val': float(match_value),
|
||
'top_left': top_left,
|
||
'center': (center_x, center_y),
|
||
'template_size': (template_w, template_h)
|
||
}
|
||
|
||
|
||
def extract_cma_from_roi(roi_img, ocr_engine, output_dir=None, debug_prefix=""):
|
||
"""
|
||
在指定的 ROI 区域内进行 OCR 提取 CMA 码
|
||
|
||
Args:
|
||
roi_img: ROI 区域图像
|
||
ocr_engine: OCR 引擎
|
||
output_dir: 输出目录
|
||
debug_prefix: 调试信息前缀
|
||
|
||
Returns:
|
||
result: 提取结果字典
|
||
"""
|
||
result = {
|
||
'code': None,
|
||
'confidence': 0.0,
|
||
'raw_text': '',
|
||
'position': (0, 0),
|
||
'box': None,
|
||
'success': False
|
||
}
|
||
|
||
if roi_img is None or roi_img.size == 0:
|
||
logger.error(f"{debug_prefix}Invalid ROI image")
|
||
return result
|
||
|
||
h, w = roi_img.shape[:2]
|
||
logger.info(f"{debug_prefix}ROI: (0, 0) -> ({w}, {h})")
|
||
logger.info(f"{debug_prefix}ROI size: {w}x{h}")
|
||
|
||
# 运行 OCR
|
||
try:
|
||
# 检查是否为 PaddleOCRVL
|
||
if hasattr(ocr_engine, 'predict'):
|
||
raw_result = ocr_engine.predict(roi_img)
|
||
else:
|
||
raw_result = ocr_engine.ocr(roi_img)
|
||
|
||
if raw_result is None or len(raw_result) == 0:
|
||
logger.error(f"{debug_prefix}OCR returned empty result")
|
||
return result
|
||
|
||
except Exception as e:
|
||
logger.error(f"{debug_prefix}OCR failed: {e}")
|
||
return result
|
||
|
||
# 处理 OCR 结果
|
||
rec_texts = []
|
||
rec_scores = []
|
||
rec_boxes = []
|
||
|
||
# 检查结果格式
|
||
if isinstance(raw_result[0], dict):
|
||
# 新 API: raw_result[0] 是 OCRResult 对象
|
||
ocr_data = raw_result[0]
|
||
rec_texts = list(ocr_data.get('rec_texts', []))
|
||
rec_scores = list(ocr_data.get('rec_scores', []))
|
||
rec_boxes = list(ocr_data.get('rec_boxes', []))
|
||
logger.info(f"{debug_prefix}Using predict() API format, found {len(rec_texts)} lines")
|
||
elif isinstance(raw_result[0], list):
|
||
# 旧 API: raw_result[0] 是 [ [box, (text, score)], ... ]
|
||
for item in raw_result[0]:
|
||
if item and len(item) >= 2:
|
||
box = item[0]
|
||
text_info = item[1]
|
||
if text_info and len(text_info) >= 2:
|
||
text = text_info[0]
|
||
score = text_info[1]
|
||
|
||
# 计算边界框 (从4个角点)
|
||
if isinstance(box, list) and len(box) >= 4:
|
||
x_coords = [p[0] for p in box]
|
||
y_coords = [p[1] for p in box]
|
||
x1, y1, x2, y2 = min(x_coords), min(y_coords), max(x_coords), max(y_coords)
|
||
rec_boxes.append([x1, y1, x2, y2])
|
||
else:
|
||
rec_boxes.append(box)
|
||
|
||
rec_texts.append(text)
|
||
rec_scores.append(score)
|
||
logger.info(f"{debug_prefix}Using legacy ocr() API format, found {len(rec_texts)} lines")
|
||
else:
|
||
logger.warning(f"{debug_prefix}Unknown OCR result format: {type(raw_result[0])}")
|
||
return result
|
||
|
||
if not rec_texts:
|
||
logger.warning(f"{debug_prefix}No text recognized in ROI")
|
||
return result
|
||
|
||
logger.info(f"{debug_prefix}OCR found {len(rec_texts)} text lines")
|
||
|
||
# 打印所有识别的文本（调试）
|
||
for i, (text, score) in enumerate(zip(rec_texts, rec_scores)):
|
||
logger.info(f"{debug_prefix}Line {i}: '{text}' (score: {score:.2f})")
|
||
|
||
# 提取 CMA 码候选
|
||
cma_candidates = []
|
||
|
||
for i, text in enumerate(rec_texts):
|
||
if not text:
|
||
continue
|
||
|
||
# 提取所有数字序列（优先匹配12位，其次是11位）
|
||
numbers = re.findall(r'\d{12}', str(text))
|
||
if not numbers:
|
||
numbers = re.findall(r'\d{11}', str(text))
|
||
|
||
# Debug: print what we found
|
||
if numbers and any('210020349' in n for n in numbers):
|
||
logger.debug(f"[DEBUG] Found numbers in '{text}': {numbers}")
|
||
|
||
for num in numbers:
|
||
# 获取对应的边界框和分数
|
||
box = rec_boxes[i] if i < len(rec_boxes) else None
|
||
score = rec_scores[i] if i < len(rec_scores) else 0.5
|
||
|
||
# 计算位置 (边界框中心)
|
||
if box is not None and len(box) >= 4:
|
||
position = ((box[0] + box[2]) / 2, (box[1] + box[3]) / 2)
|
||
else:
|
||
position = (0, 0)
|
||
|
||
cma_candidates.append({
|
||
'code': num,
|
||
'confidence': score,
|
||
'text': str(text),
|
||
'position': position,
|
||
'box': box,
|
||
})
|
||
|
||
# 选择最佳候选
|
||
if cma_candidates:
|
||
# 按分数排序（考虑位置和长度）
|
||
cma_candidates.sort(key=lambda x: (
|
||
x['confidence'] * 100
|
||
+ (30 if x['position'][0] > w / 3 and x['position'][1] < h / 3 else 0) # 右上角加分
|
||
+ (10 if len(x['code']) == 11 else 0)
|
||
- (20 if x['code'].startswith('2') else 0)
|
||
), reverse=True)
|
||
|
||
best = cma_candidates[0]
|
||
result['code'] = best['code']
|
||
result['confidence'] = best['confidence']
|
||
result['raw_text'] = best['text']
|
||
result['position'] = best['position']
|
||
result['box'] = best['box']
|
||
result['success'] = True
|
||
|
||
logger.info(f"{debug_prefix}Best CMA candidate: {best['code']} (conf: {best['confidence']:.2f})")
|
||
else:
|
||
logger.warning(f"{debug_prefix}No CMA code candidates found in ROI text")
|
||
|
||
# 保存可视化结果
|
||
box = result.get('box')
|
||
if output_dir and result['success'] and box is not None:
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
vis_roi = roi_img.copy()
|
||
if box is not None and len(box) >= 4:
|
||
# box is [x1, y1, x2, y2] format
|
||
cv2.rectangle(vis_roi, (int(box[0]), int(box[1])),
|
||
(int(box[2]), int(box[3])), (0, 255, 0), 2)
|
||
# 在边界框上方显示文本
|
||
text_pos = (int(box[0]), max(10, int(box[1]) - 10))
|
||
cv2.putText(vis_roi, f"CMA: {result['code']}", text_pos,
|
||
cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 0), 2)
|
||
cv2.imwrite(os.path.join(output_dir, f"{debug_prefix.strip()}cma_roi_extraction.png"), vis_roi)
|
||
logger.info(f"{debug_prefix}Saved ROI extraction visualization")
|
||
|
||
return result
|
||
|
||
|
||
def extract_cma_code_fullpage(page_img, ocr_engine, template_path='template/CMA_Logo.png',
|
||
output_dir=None, use_template_matching=True):
|
||
"""
|
||
使用模板匹配提取 CMA 码的完整流程
|
||
|
||
Args:
|
||
page_img: 页面图像
|
||
ocr_engine: OCR 引擎
|
||
template_path: CMA logo 模板路径
|
||
output_dir: 输出目录
|
||
use_template_matching: 是否使用模板匹配（False则直接全页OCR）
|
||
|
||
Returns:
|
||
result: CMA 提取结果
|
||
"""
|
||
result = {
|
||
'code': None,
|
||
'confidence': 0.0,
|
||
'raw_text': '',
|
||
'position': (0, 0),
|
||
'box': None,
|
||
'success': False,
|
||
'method': 'none'
|
||
}
|
||
|
||
# 加载图像
|
||
if isinstance(page_img, str):
|
||
image = imread_unicode(page_img, cv2.IMREAD_COLOR)
|
||
elif isinstance(page_img, np.ndarray):
|
||
image = page_img
|
||
else:
|
||
logger.error(f"Invalid image type: {type(page_img)}")
|
||
return result
|
||
|
||
if image is None or image.size == 0:
|
||
logger.error("Failed to load image or empty image")
|
||
return result
|
||
|
||
h, w = image.shape[:2]
|
||
|
||
# 加载模板
|
||
if use_template_matching:
|
||
template, _ = load_cma_template(template_path)
|
||
if template is None:
|
||
logger.warning("Cannot load template, falling back to full-page OCR")
|
||
use_template_matching = False
|
||
|
||
# 方法1: 模板匹配 + ROI OCR
|
||
template_match_success = False
|
||
if use_template_matching:
|
||
logger.info("[TM] Starting template matching extraction...")
|
||
match_result = match_template(image, template)
|
||
|
||
if match_result is None:
|
||
logger.warning("[TM] Template matching failed")
|
||
else:
|
||
match_value = match_result['max_val']
|
||
|
||
# 检查匹配置信度
|
||
if match_value < 0.4:
|
||
logger.warning(f"[TM] Match confidence too low: {match_value:.3f}")
|
||
else:
|
||
# 模板匹配成功，尝试ROI提取
|
||
template_match_success = True
|
||
|
||
# 确定 ROI（关键：ROI 应该在 logo 的右侧，而不是以 logo 为中心）
|
||
center_x, center_y = match_result['center']
|
||
template_w, template_h = match_result['template_size']
|
||
|
||
# 修正：ROI应该在logo的右侧，因为CMA编号通常在logo右边
|
||
# 而不是以logo为中心
|
||
roi_x1 = max(0, center_x) # 从logo中心开始向右
|
||
roi_y1 = max(0, center_y - template_h // 2) # 上下与logo对齐
|
||
roi_x2 = min(w, center_x + min(600, w - center_x)) # 向右扩展最多600px
|
||
roi_y2 = min(h, center_y + template_h // 2 + template_h) # 向下扩展一些
|
||
|
||
# 确保ROI在图像范围内
|
||
roi_x1 = max(roi_x1, 0)
|
||
roi_y1 = max(roi_y1, 0)
|
||
roi_x2 = min(w, roi_x2)
|
||
roi_y2 = min(h, roi_y2)
|
||
|
||
logger.info(f"[TM] ROI: ({roi_x1}, {roi_y1}) -> ({roi_x2}, {roi_y2})")
|
||
|
||
roi_img = image[roi_y1:roi_y2, roi_x1:roi_x2]
|
||
|
||
# 在 ROI 内提取 CMA 码
|
||
result = extract_cma_from_roi(roi_img, ocr_engine, output_dir, debug_prefix="[TM] ")
|
||
|
||
if result['success']:
|
||
result['method'] = 'template_matching'
|
||
logger.info(f"[TM] Template matching SUCCESS: {result['code']} (conf: {result['confidence']:.2f})")
|
||
return result
|
||
else:
|
||
logger.warning("[TM] Template matching found logo, but OCR failed to extract CMA code")
|
||
|
||
# 模板匹配失败，尝试全页OCR作为fallback
|
||
logger.info("[FALLBACK] Template matching failed, trying full-page OCR...")
|
||
result = extract_cma_fullpage_fallback(image, ocr_engine, output_dir)
|
||
result['method'] = 'fullpage_fallback'
|
||
return result
|
||
|
||
|
||
def extract_cma_fullpage_fallback(page_img, ocr_engine, output_dir=None):
|
||
"""
|
||
全页OCR fallback方法 - 当模板匹配失败时使用
|
||
|
||
Args:
|
||
page_img: 页面图像
|
||
ocr_engine: OCR 引擎
|
||
output_dir: 输出目录
|
||
|
||
Returns:
|
||
result: CMA 提取结果
|
||
"""
|
||
result = {
|
||
'code': None,
|
||
'confidence': 0.0,
|
||
'raw_text': '',
|
||
'position': (0, 0),
|
||
'box': None,
|
||
'success': False
|
||
}
|
||
|
||
if isinstance(page_img, str):
|
||
image = imread_unicode(page_img, cv2.IMREAD_COLOR)
|
||
elif isinstance(page_img, np.ndarray):
|
||
image = page_img
|
||
else:
|
||
logger.error(f"Invalid image type: {type(page_img)}")
|
||
return result
|
||
|
||
if image is None or image.size == 0:
|
||
logger.error("Failed to load image or empty image")
|
||
return result
|
||
|
||
h, w = image.shape[:2]
|
||
|
||
# 运行全页OCR
|
||
logger.info("[FALLBACK] Running full-page OCR...")
|
||
try:
|
||
raw_result = ocr_engine.ocr(image)
|
||
except Exception as e:
|
||
logger.error(f"[FALLBACK] OCR failed: {e}")
|
||
return result
|
||
|
||
# 处理OCR结果
|
||
rec_texts = []
|
||
rec_scores = []
|
||
rec_boxes = []
|
||
|
||
if raw_result and len(raw_result) > 0:
|
||
first = raw_result[0]
|
||
if isinstance(first, dict):
|
||
rec_texts = list(first.get('rec_texts', []))
|
||
rec_scores = list(first.get('rec_scores', []))
|
||
rec_boxes = list(first.get('rec_boxes', []))
|
||
elif isinstance(first, list):
|
||
for item in first:
|
||
if item and len(item) >= 2:
|
||
box = item[0]
|
||
text_info = item[1]
|
||
if text_info and len(text_info) >= 2:
|
||
text = text_info[0]
|
||
score = text_info[1]
|
||
|
||
if isinstance(box, list) and len(box) >= 4:
|
||
x_coords = [p[0] for p in box]
|
||
y_coords = [p[1] for p in box]
|
||
x1, y1, x2, y2 = min(x_coords), min(y_coords), max(x_coords), max(y_coords)
|
||
rec_boxes.append([x1, y1, x2, y2])
|
||
else:
|
||
rec_boxes.append(box)
|
||
|
||
rec_texts.append(text)
|
||
rec_scores.append(score)
|
||
|
||
logger.info(f"[FALLBACK] Found {len(rec_texts)} text lines")
|
||
|
||
# 提取CMA码候选
|
||
cma_candidates = []
|
||
|
||
for i, text in enumerate(rec_texts):
|
||
if not text:
|
||
continue
|
||
|
||
# 提取所有数字序列（优先匹配12位，其次是11位）
|
||
numbers = re.findall(r'\d{12}', str(text))
|
||
if not numbers:
|
||
numbers = re.findall(r'\d{11}', str(text))
|
||
|
||
for num in numbers:
|
||
box = rec_boxes[i] if i < len(rec_boxes) else None
|
||
score = rec_scores[i] if i < len(rec_scores) else 0.5
|
||
|
||
if box is not None and len(box) >= 4:
|
||
position = ((box[0] + box[2]) / 2, (box[1] + box[3]) / 2)
|
||
else:
|
||
position = (0, 0)
|
||
|
||
cma_candidates.append({
|
||
'code': num,
|
||
'confidence': score,
|
||
'text': str(text),
|
||
'position': position,
|
||
'box': box,
|
||
})
|
||
|
||
if not cma_candidates:
|
||
logger.warning("[FALLBACK] No CMA code candidates found")
|
||
return result
|
||
|
||
# 评分和排序（优先右上角，优先以2开头的）
|
||
cma_candidates.sort(key=lambda x: (
|
||
x['confidence'] * 100
|
||
+ (50 if x['code'].startswith('2') else 0) # 以2开头的优先
|
||
+ (30 if x['position'][0] > w / 2 and x['position'][1] < h / 3 else 0) # 右上角加分
|
||
+ (10 if len(x['code']) == 11 else 0)
|
||
), reverse=True)
|
||
|
||
best = cma_candidates[0]
|
||
result['code'] = best['code']
|
||
result['confidence'] = best['confidence']
|
||
result['raw_text'] = best['text']
|
||
result['position'] = best['position']
|
||
result['box'] = best['box']
|
||
result['success'] = True
|
||
|
||
logger.info(f"[FALLBACK] CMA extracted: {best['code']} (conf: {best['confidence']:.2f})")
|
||
|
||
return result
|
||
|
||
|
||
if __name__ == "__main__":
|
||
import argparse
|
||
|
||
parser = argparse.ArgumentParser(description='CMA Logo 模板匹配提取')
|
||
parser.add_argument('--pdf', help='PDF 文件路径')
|
||
parser.add_argument('--template', default='template/CMA_Logo.png', help='CMA logo 模板路径')
|
||
parser.add_argument('--output', default='template_match_debug', help='输出目录')
|
||
|
||
args = parser.parse_args()
|
||
|
||
# 检查文件
|
||
if not os.path.exists(args.pdf):
|
||
print(f"错误: PDF 文件不存在: {args.pdf}")
|
||
sys.exit(1)
|
||
|
||
if not os.path.exists(args.template):
|
||
print(f"错误: 模板文件不存在: {args.template}")
|
||
sys.exit(1)
|
||
|
||
# 加载 OCR 引擎
|
||
os.environ["DISABLE_MODEL_SOURCE_CHECK"] = "True"
|
||
os.environ["PADDLE_PDX_DISABLE_MODEL_SOURCE_CHECK"] = "True"
|
||
|
||
from paddleocr import PaddleOCR
|
||
ocr_engine = PaddleOCR(use_angle_cls=True, lang='ch', use_gpu=False)
|
||
|
||
# 处理 PDF 的第一页
|
||
import fitz
|
||
doc = fitz.open(args.pdf)
|
||
page = doc[0]
|
||
pix = page.get_pixmap(matrix=fitz.Matrix(300/72, 300/72))
|
||
img = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.height, pix.width, 3)
|
||
img_rgb = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
|
||
|
||
print(f"PDF 尺寸: {pix.width}x{pix.height}")
|
||
print(f"图像尺寸: {img_rgb.shape}")
|
||
|
||
# 执行模板匹配提取
|
||
result = extract_cma_code_fullpage(img_rgb, ocr_engine, args.template, args.output)
|
||
|
||
# 输出结果
|
||
print()
|
||
print("="*80)
|
||
print("CMA 提取结果:")
|
||
print("-"*80)
|
||
print(f" 方法: {result.get('method', 'unknown')}")
|
||
print(f" CMA码: {result.get('code', 'N/A')}")
|
||
print(f" 置信度: {result.get('confidence', 0.0):.2f}")
|
||
print(f" 位置: {result.get('position', 'N/A')}")
|
||
print("-"*80)
|
||
print(f" 提取成功: {result.get('success', False)}")
|
||
print("="*80)
|