"""文本场景比对管线：OCR提取 + 语义嵌入 + SSIM结构比对 + 印章/手写检测""" import time import cv2 import numpy as np import torch from skimage.metrics import structural_similarity as ssim from module.config import ( TEXT_SIM_WEIGHT, SSIM_WEIGHT, SEAL_WEIGHT, TEXT_SCENE_THRESHOLD, TEXT_SCENE_BGE_THRESHOLD, ) def _cv_imread(path, flags=cv2.IMREAD_COLOR): """支持中文路径的cv2.imread""" return cv2.imdecode(np.fromfile(path, dtype=np.uint8), flags) # ==================== OCR文字提取 ==================== def extract_text(image_path, ocr_engine): """使用PaddleOCR提取图片中的全部文字 Args: image_path: 图片路径 ocr_engine: PaddleOCR引擎实例 Returns: full_text: 拼接后的完整文本 text_lines: 每行文字的列表（含置信度） """ result = ocr_engine.ocr(image_path, cls=True) text_lines = [] if result and result[0]: for line in result[0]: text = line[1][0] conf = line[1][1] text_lines.append({'text': text, 'confidence': float(conf)}) full_text = ' '.join([t['text'] for t in text_lines]) return full_text, text_lines # ==================== 语义嵌入比对 ==================== def compute_text_similarity(text1, text2, bge_tokenizer, bge_model): """使用BGE-small-zh计算两段文本的语义相似度 Args: text1, text2: 待比较的两段文本 bge_tokenizer: BGE分词器实例 bge_model: BGE模型实例 Returns: similarity: 余弦相似度，范围[0,1] """ if not text1.strip() or not text2.strip(): return 0.0 def _encode(text): """编码单段文本为归一化向量""" encoded = bge_tokenizer(text, padding=True, truncation=True, return_tensors='pt', max_length=512) if torch.cuda.is_available(): encoded = {k: v.cuda() for k, v in encoded.items()} with torch.no_grad(): outputs = bge_model(**encoded) cls_embedding = outputs.last_hidden_state[:, 0] cls_embedding = torch.nn.functional.normalize(cls_embedding, p=2, dim=1) return cls_embedding.cpu().numpy() emb1 = _encode(text1) emb2 = _encode(text2) sim = float(np.dot(emb1[0], emb2[0])) return sim # ==================== SSIM结构比对 ==================== def compute_ssim(image_path1, image_path2): """计算两张图片的SSIM结构相似度 将图片缩放到相同尺寸后转为灰度计算SSIM Returns: ssim_score: SSIM值，范围[0,1] """ img1 = _cv_imread(image_path1) img2 = _cv_imread(image_path2) if img1 is None or img2 is None: return 0.0 h = min(img1.shape[0], img2.shape[0]) w = min(img1.shape[1], img2.shape[1]) img1 = cv2.resize(img1, (w, h)) img2 = cv2.resize(img2, (w, h)) gray1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY) gray2 = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY) score, _ = ssim(gray1, gray2, full=True) return float(score) # ==================== 印章/手写检测 ==================== def detect_seals_and_handwriting(image_path, min_area_ratio=0.001): """检测图片中的红色印章和手写笔迹区域 印章检测原理：HSV颜色空间中提取红色区域 Args: image_path: 图片路径 min_area_ratio: 最小区域面积占比，过滤噪声 Returns: seals: 印章区域列表，每项含 {'bbox', 'area', 'area_ratio'} seal_total_ratio: 印章总面积占比 """ img = _cv_imread(image_path) if img is None: return [], 0.0 h, w = img.shape[:2] total_area = h * w # --- 印章检测：HSV红色区域 --- hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV) # 红色在HSV中有两段：0-10 和 160-180 mask1 = cv2.inRange(hsv, np.array([0, 70, 50]), np.array([10, 255, 255])) mask2 = cv2.inRange(hsv, np.array([160, 70, 50]), np.array([180, 255, 255])) red_mask = cv2.bitwise_or(mask1, mask2) # 形态学操作去噪 kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5)) red_mask = cv2.morphologyEx(red_mask, cv2.MORPH_CLOSE, kernel, iterations=2) red_mask = cv2.morphologyEx(red_mask, cv2.MORPH_OPEN, kernel, iterations=1) # 查找轮廓 contours, _ = cv2.findContours(red_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) seals = [] min_area = total_area * min_area_ratio for cnt in contours: area = cv2.contourArea(cnt) if area >= min_area: x, y, bw, bh = cv2.boundingRect(cnt) seals.append({ 'bbox': (x, y, bw, bh), 'area': int(area), 'area_ratio': round(area / total_area, 4), }) seal_total_ratio = sum(s['area_ratio'] for s in seals) return seals, round(seal_total_ratio, 4) def compare_seals(seals1, seals2): """比对两张图的印章区域 比较逻辑： - 两张图都有印章 → 比较数量和面积占比差异 - 一张有一张无 → 印章不匹配 - 两张都无 → 印章项跳过 Returns: match: bool或None，印章是否匹配（None表示两张都无印章） detail: 详细信息字典 """ has_seal1 = len(seals1) > 0 has_seal2 = len(seals2) > 0 ratio1 = sum(s['area_ratio'] for s in seals1) ratio2 = sum(s['area_ratio'] for s in seals2) detail = { 'img1_seal_count': len(seals1), 'img2_seal_count': len(seals2), 'img1_seal_ratio': round(ratio1, 4), 'img2_seal_ratio': round(ratio2, 4), } if not has_seal1 and not has_seal2: return None, detail if has_seal1 != has_seal2: detail['reason'] = '一张有印章一张无' return False, detail # 两张都有印章：比较数量和面积 count_match = len(seals1) == len(seals2) ratio_diff = abs(ratio1 - ratio2) area_match = ratio_diff < 0.02 match = count_match and area_match detail['count_match'] = count_match detail['area_diff'] = round(ratio_diff, 4) detail['area_match'] = area_match return match, detail def compare_text_scene_cached(image_path, cached_h, ocr_engine, bge_tokenizer, bge_model, precomputed_text=None, precomputed_bge=None): """文本场景缓存比对: 仅实时 OCR/BGE/Seal 查询图，历史图从缓存读取 Args: image_path: 查询图片路径 cached_h: 预存储的历史图缓存 dict (ocr_text, bge_vector, seal_count, seal_ratio) ocr_engine: PaddleOCR引擎 bge_tokenizer: BGE分词器 bge_model: BGE模型 Returns: result dict (与 compare_text_scene() 兼容的格式) """ start_time = time.time() # 1. OCR 仅查询图 (优先使用预计算结果) if precomputed_text is not None: text1 = precomputed_text lines1 = [] ocr_time = 0.0 else: text1, lines1 = extract_text(image_path, ocr_engine) ocr_time = time.time() - start_time text2 = cached_h.get('ocr_text', '') # 2. 语义嵌入 (优先使用预计算向量) if precomputed_bge is not None: emb_q = precomputed_bge emb_h = cached_h.get('bge_vector', np.zeros(512, dtype=np.float32)) text_sim = float(np.dot(emb_q, emb_h)) if text1.strip() else 0.0 embed_time = 0.0 elif not text1.strip(): text_sim = 0.0 embed_time = 0.0 else: t_embed = time.time() def _encode_q(text): encoded = bge_tokenizer(text, padding=True, truncation=True, return_tensors='pt', max_length=512) if torch.cuda.is_available(): encoded = {k: v.cuda() for k, v in encoded.items()} with torch.no_grad(): outputs = bge_model(**encoded) cls_embedding = outputs.last_hidden_state[:, 0] cls_embedding = torch.nn.functional.normalize(cls_embedding, p=2, dim=1) return cls_embedding.cpu().numpy()[0] emb_q = _encode_q(text1) emb_h = cached_h.get('bge_vector', np.zeros(512, dtype=np.float32)) text_sim = float(np.dot(emb_q, emb_h)) embed_time = time.time() - t_embed # 3. SSIM 仍需两张图 t_ssim = time.time() ssim_score = compute_ssim(image_path, cached_h['path']) ssim_time = time.time() - t_ssim # 4. 印章：仅查询图（历史图用缓存） t_seal = time.time() seals_q, ratio_q = detect_seals_and_handwriting(image_path) cnt_h = cached_h.get('seal_count', 0) ratio_h = cached_h.get('seal_ratio', 0.0) has_q, has_h = len(seals_q) > 0, cnt_h > 0 if not has_q and not has_h: seal_bonus = 0.5 elif has_q != has_h: seal_bonus = 0.0 else: count_match = len(seals_q) == cnt_h area_match = abs(ratio_q - ratio_h) < 0.02 seal_bonus = 1.0 if (count_match and area_match) else 0.0 seal_time = time.time() - t_seal # 5. 综合判定 final_score = text_sim * TEXT_SIM_WEIGHT + ssim_score * SSIM_WEIGHT + seal_bonus * SEAL_WEIGHT is_same = final_score >= TEXT_SCENE_THRESHOLD total_time = time.time() - start_time result = { 'is_same_scene': bool(is_same), 'scene_type': 'text', 'final_score': round(final_score, 4), 'text_similarity': round(text_sim, 4), 'ssim_score': round(ssim_score, 4), 'seal_bonus': seal_bonus, 'seal_match': None if (not has_q and not has_h) else (has_q and has_h), 'seal_detail': { 'img1_seal_count': len(seals_q), 'img2_seal_count': cnt_h, 'img1_seal_ratio': ratio_q, 'img2_seal_ratio': ratio_h, }, 'img1_text_length': len(text1), 'img2_text_length': len(text2), 'img1_seal_count': len(seals_q), 'img2_seal_count': cnt_h, 'img1_seal_ratio': ratio_q, 'img2_seal_ratio': ratio_h, 'similarity_score': round(final_score, 4), 'match_count': len(seals_q) + cnt_h, 'inlier_ratio': round(text_sim, 4), 'avg_confidence': round(ssim_score, 4), 'time_breakdown': { 'ocr': round(ocr_time, 2), 'embedding': round(embed_time, 2), 'ssim': round(ssim_time, 2), 'seal': round(seal_time, 2), 'total': round(total_time, 2), }, } return result def compare_text_scene_pure_bge(image_path, cached_h, ocr_engine, bge_tokenizer, bge_model, precomputed_text=None, precomputed_bge=None): """Pure BGE: only OCR->BGE cosine similarity, no SSIM/seal""" start_time = time.time() print(f' [TIME-PURE-BGE] enter, precomputed_text={precomputed_text is not None}, precomputed_bge={precomputed_bge is not None}') if precomputed_text is not None: text1 = precomputed_text ocr_time = 0.0 else: text1, _ = extract_text(image_path, ocr_engine) ocr_time = time.time() - start_time if precomputed_bge is not None: emb_q = precomputed_bge emb_h = cached_h.get('_vector', cached_h.get('bge_vector', np.zeros(512, dtype=np.float32))) text_sim = float(np.dot(emb_q, emb_h)) if text1.strip() else 0.0 embed_time = 0.0 elif not text1.strip(): text_sim = 0.0 embed_time = 0.0 else: t_embed = time.time() def _encode_q(text): encoded = bge_tokenizer(text, padding=True, truncation=True, return_tensors='pt', max_length=512) if torch.cuda.is_available(): encoded = {k: v.cuda() for k, v in encoded.items()} with torch.no_grad(): outputs = bge_model(**encoded) cls_embedding = outputs.last_hidden_state[:, 0] cls_embedding = torch.nn.functional.normalize(cls_embedding, p=2, dim=1) return cls_embedding.cpu().numpy()[0] emb_q = _encode_q(text1) emb_h = cached_h.get('_vector', cached_h.get('bge_vector', np.zeros(512, dtype=np.float32))) text_sim = float(np.dot(emb_q, emb_h)) embed_time = time.time() - t_embed is_same = text_sim >= TEXT_SCENE_BGE_THRESHOLD total_time = time.time() - start_time print(f' [TIME-PURE-BGE] text_sim={text_sim:.4f} is_same={is_same} total={total_time:.3f}s') result = { 'is_same_scene': bool(is_same), 'scene_type': 'text', 'final_score': round(text_sim, 4), 'text_similarity': round(text_sim, 4), 'ssim_score': 0.0, 'seal_bonus': 0.0, 'similarity_score': round(text_sim, 4), 'match_count': 0, 'inlier_ratio': round(text_sim, 4), 'avg_confidence': 0.0, 'time_breakdown': { 'ocr': round(ocr_time, 2), 'embedding': round(embed_time, 2), 'ssim': 0.0, 'seal': 0.0, 'total': round(total_time, 2), }, } return result # ==================== 综合判定 ==================== def compare_text_scene(image_path1, image_path2, ocr_engine, bge_tokenizer, bge_model): """文本场景比对主函数 Args: image_path1: 第一张图片路径 image_path2: 第二张图片路径 ocr_engine: PaddleOCR引擎实例 bge_tokenizer: BGE分词器实例 bge_model: BGE模型实例 Returns: result: dict，含各项分数和最终判定 """ start_time = time.time() # 1. OCR文字提取 text1, lines1 = extract_text(image_path1, ocr_engine) text2, lines2 = extract_text(image_path2, ocr_engine) ocr_time = time.time() - start_time # 2. 语义嵌入比对 text_sim = compute_text_similarity(text1, text2, bge_tokenizer, bge_model) embed_time = time.time() - start_time - ocr_time # 3. SSIM结构比对 ssim_score = compute_ssim(image_path1, image_path2) ssim_time = time.time() - start_time - ocr_time - embed_time # 4. 印章检测与比对 seals1, seal_ratio1 = detect_seals_and_handwriting(image_path1) seals2, seal_ratio2 = detect_seals_and_handwriting(image_path2) seal_match, seal_detail = compare_seals(seals1, seals2) seal_time = time.time() - start_time - ocr_time - embed_time - ssim_time # 5. 综合判定 if seal_match is True: seal_bonus = 1.0 elif seal_match is False: seal_bonus = 0.0 else: seal_bonus = 0.5 # 两张图都没有印章，中性 final_score = text_sim * TEXT_SIM_WEIGHT + ssim_score * SSIM_WEIGHT + seal_bonus * SEAL_WEIGHT is_same = final_score >= TEXT_SCENE_THRESHOLD total_time = time.time() - start_time result = { 'is_same_scene': bool(is_same), 'scene_type': 'text', 'final_score': round(final_score, 4), 'text_similarity': round(text_sim, 4), 'ssim_score': round(ssim_score, 4), 'seal_bonus': seal_bonus, 'seal_match': seal_match, 'seal_detail': seal_detail, 'img1_text_length': len(text1), 'img2_text_length': len(text2), 'img1_seal_count': len(seals1), 'img2_seal_count': len(seals2), 'img1_seal_ratio': seal_ratio1, 'img2_seal_ratio': seal_ratio2, 'similarity_score': round(final_score, 4), 'match_count': len(seals1) + len(seals2), # 文本场景用印章数代替 'inlier_ratio': round(text_sim, 4), 'avg_confidence': round(ssim_score, 4), 'time_breakdown': { 'ocr': round(ocr_time, 2), 'embedding': round(embed_time, 2), 'ssim': round(ssim_time, 2), 'seal': round(seal_time, 2), 'total': round(total_time, 2), }, } return result