# --- scripts/02_create_labels.py --- # # **重要**: 此脚本用于训练 "阶段一" 的区域检测器. # 它只生成 1 个类别 (class_id = 0), 即 "payment_slip" 的 *整个* 区域. # 它不使用您的高级 classify_text 逻辑, 而是使用启发式规则 (关键词, 线条) 来找到大区域. # import os import pandas as pd import numpy as np import cv2 import re import shutil import json from sklearn.cluster import DBSCAN # --- 锚点词典 (技术二) --- # 这些是用来 *定位* 凭证区域的词, 不是用来提取的 KEYWORDS = [ "Bankgirot", "PlusGirot", "OCR-nummer", "Att betala", "Mottagare", "Betalningsmottagare", "Tillhanda senast", "Förfallodag", "Belopp", "BG-nr", "PG-nr", "Meddelande", "OCR-kod", "Inbetalningskort", "Betalningsavi" ] KEYWORDS_REGEX = '|'.join(KEYWORDS) REGEX_PATTERNS = { 'bg_pg': r'(\b\d{2,4}[- ]\d{4}\b)|(\b\d{2,7}[- ]\d\b)', # Bankgiro/PlusGiro 'long_num': r'\b\d{10,}\b', # 可能是 OCR 'machine_code': r'#[0-9\s>#]+#' # 机读码 } FULL_REGEX = f"({KEYWORDS_REGEX})|{REGEX_PATTERNS['bg_pg']}|{REGEX_PATTERNS['long_num']}|{REGEX_PATTERNS['machine_code']}" # --- 路径设置 --- BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) OCR_DIR = os.path.join(BASE_DIR, "data", "ocr_results") IMG_DIR = os.path.join(BASE_DIR, "data", "processed_images") # Use temp directories for initial label generation (before train/val split) LABEL_DIR = os.path.join(BASE_DIR, "data", "yolo_dataset", "temp_all_labels") IMAGE_DIR = os.path.join(BASE_DIR, "data", "yolo_dataset", "temp_all_images") os.makedirs(LABEL_DIR, exist_ok=True) os.makedirs(IMAGE_DIR, exist_ok=True) def find_text_anchors(text_boxes, img_height): """技术一 (位置) + 技术二 (词典)""" anchors = [] if not text_boxes: return [] # 技术一: 只在页面下半部分 (40% 处开始) 查找 page_midpoint = img_height * 0.4 for box in text_boxes: # 检查位置 if box["bbox"]["y_min"] > page_midpoint: # 检查文本内容 if re.search(FULL_REGEX, box["text"], re.IGNORECASE): bbox = box["bbox"] anchors.append(pd.Series({ 'left': bbox["x_min"], 'top': bbox["y_min"], 'width': bbox["x_max"] - bbox["x_min"], 'height': bbox["y_max"] - bbox["y_min"] })) return anchors def find_visual_anchors(image, img_height, img_width): """技术三 (视觉锚点 - 找线)""" anchors = [] try: # 1. 只看下半页 crop_y_start = img_height // 2 crop = image[crop_y_start:, :] gray = cv2.cvtColor(crop, cv2.COLOR_BGR2GRAY) thresh = cv2.threshold(gray, 150, 255, cv2.THRESH_BINARY_INV)[1] # 2. 查找长的水平线 min_line_length = img_width // 4 horizontal_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (min_line_length, 1)) detected_lines = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, horizontal_kernel, iterations=2) contours, _ = cv2.findContours(detected_lines, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) for c in contours: x, y, w, h = cv2.boundingRect(c) if w > min_line_length: # 确保线足够长 original_y = y + crop_y_start # 转换回原图坐标 anchors.append(pd.Series({ 'left': x, 'top': original_y, 'width': w, 'height': h })) except Exception as e: print(f"查找视觉锚点时出错: {e}") return anchors def cluster_anchors(all_anchors, img_height): """技术四 (聚类)""" if len(all_anchors) < 2: return all_anchors # 锚点太少，无法聚类 # 1. 获取中心点 points = [] for anchor in all_anchors: x_center = anchor.left + anchor.width / 2 y_center = anchor.top + anchor.height / 2 points.append((x_center, y_center)) points = np.array(points) # 2. 运行 DBSCAN eps_dist = img_height * 0.2 clustering = DBSCAN(eps=eps_dist, min_samples=2).fit(points) labels = clustering.labels_ if len(labels) == 0 or np.all(labels == -1): return all_anchors # 聚类失败 # 3. 找到最大的那个簇 unique_labels, counts = np.unique(labels[labels != -1], return_counts=True) if len(counts) == 0: return all_anchors # 只有噪声 largest_cluster_label = unique_labels[np.argmax(counts)] # 4. 只返回属于最大簇的锚点 main_cluster_anchors = [ anchor for i, anchor in enumerate(all_anchors) if labels[i] == largest_cluster_label ] return main_cluster_anchors def create_labels(): print("开始生成弱标签 (用于区域检测器)...") processed_count = 0 skipped_count = 0 for ocr_filename in os.listdir(OCR_DIR): if not ocr_filename.endswith(".json"): continue base_name = os.path.splitext(ocr_filename)[0] json_path = os.path.join(OCR_DIR, ocr_filename) img_path = os.path.join(IMG_DIR, f"{base_name}.png") if not os.path.exists(img_path): continue try: # 1. 加载数据 image = cv2.imread(img_path) img_h, img_w = image.shape[:2] with open(json_path, 'r', encoding='utf-8') as f: ocr_data = json.load(f) text_boxes = ocr_data.get("text_boxes", []) # 2. 查找所有锚点 text_anchors = find_text_anchors(text_boxes, img_h) visual_anchors = find_visual_anchors(image, img_h, img_w) all_anchors = text_anchors + visual_anchors if not all_anchors: print(f"SKIPPING: {base_name} (未找到任何锚点)") skipped_count += 1 continue # 3. 聚类锚点 final_anchors = cluster_anchors(all_anchors, img_h) if not final_anchors: print(f"SKIPPING: {base_name} (未找到有效聚类)") skipped_count += 1 continue # 4. 聚合坐标 min_x = min(a.left for a in final_anchors) min_y = min(a.top for a in final_anchors) max_x = max(a.left + a.width for a in final_anchors) max_y = max(a.top + a.height for a in final_anchors) # 5. 添加边距 (Padding) padding = 10 min_x = max(0, min_x - padding) min_y = max(0, min_y - padding) max_x = min(img_w, max_x + padding) max_y = min(img_h, max_y + padding) # 6. 转换为 YOLO 格式 box_w = max_x - min_x box_h = max_y - min_y x_center = (min_x + box_w / 2) / img_w y_center = (min_y + box_h / 2) / img_h norm_w = box_w / img_w norm_h = box_h / img_h # ** 关键: 类别 ID 永远是 0 ** class_id = 0 # 唯一的类别: 'payment_slip' yolo_label = f"{class_id} {x_center} {y_center} {norm_w} {norm_h}\n" # 7. 保存标签和图片 label_path = os.path.join(LABEL_DIR, f"{base_name}.txt") with open(label_path, 'w', encoding='utf-8') as f: f.write(yolo_label) shutil.copy( img_path, os.path.join(IMAGE_DIR, f"{base_name}.png") ) processed_count += 1 if processed_count % 20 == 0: print(f"已生成 {processed_count} 个标签...") except Exception as e: print(f"处理 {base_name} 时出错: {e}") skipped_count += 1 print("--- 弱标签生成完成 ---") print(f"成功生成: {processed_count} 个") print(f"跳过: {skipped_count} 个") if __name__ == "__main__": create_labels()