"""Docling Layout Predictor 适配器 (符合 BaseLayoutDetector 规范) 基于 HuggingFace transformers 直接加载 Docling 布局模型,不依赖 docling-ibm-models 包。 使用 MinerU 环境中的 transformers 库。 支持的 Hugging Face 模型仓库: - ds4sd/docling-layout-old - ds4sd/docling-layout-heron - ds4sd/docling-layout-heron-101 - ds4sd/docling-layout-egret-medium - ds4sd/docling-layout-egret-large - ds4sd/docling-layout-egret-xlarge """ import cv2 import numpy as np import threading from pathlib import Path from typing import Dict, List, Union, Any, Optional from PIL import Image try: from .base import BaseLayoutDetector except ImportError: from base import BaseLayoutDetector # 全局锁,防止模型初始化时的线程问题 _model_init_lock = threading.Lock() class DoclingLayoutDetector(BaseLayoutDetector): """Docling Layout Predictor 适配器 直接使用 transformers 库加载 Docling 布局模型,无需安装 docling-ibm-models。 """ # Docling 原始类别定义(来自 docling-ibm-models/layoutmodel/labels.py) # 使用 shifted canonical(带 Background)的版本,因为 RT-DETR 模型使用这个 DOCLING_LABELS = { 0: 'Background', 1: 'Caption', 2: 'Footnote', 3: 'Formula', 4: 'List-item', 5: 'Page-footer', 6: 'Page-header', 7: 'Picture', 8: 'Section-header', 9: 'Table', 10: 'Text', 11: 'Title', 12: 'Document Index', 13: 'Code', 14: 'Checkbox-Selected', 15: 'Checkbox-Unselected', 16: 'Form', 17: 'Key-Value Region', } # 类别映射:Docling LayoutLabels → MinerU/EnhancedDocPipeline 类别体系 # 参考: # - MinerU: mineru/utils/enum_class.py (BlockType, CategoryId) # - Pipeline: universal_doc_parser/core/pipeline_manager_v2.py (EnhancedDocPipeline 类别定义) CATEGORY_MAP = { 'Caption': 'image_caption', # Caption -> image_caption (IMAGE_TEXT_CATEGORIES) 'Footnote': 'page_footnote', # Footnote -> page_footnote (TEXT_CATEGORIES) 'Formula': 'interline_equation', # Formula -> interline_equation (EQUATION_CATEGORIES) 'List-item': 'text', # List-item -> text (TEXT_CATEGORIES) 'Page-footer': 'footer', # Page-footer -> footer (TEXT_CATEGORIES) 'Page-header': 'header', # Page-header -> header (TEXT_CATEGORIES) 'Picture': 'image_body', # Picture -> image_body (IMAGE_BODY_CATEGORIES) 'Section-header': 'title', # Section-header -> title (TEXT_CATEGORIES) 'Table': 'table_body', # Table -> table_body (TABLE_BODY_CATEGORIES) 'Text': 'text', # Text -> text (TEXT_CATEGORIES) 'Title': 'title', # Title -> title (TEXT_CATEGORIES) 'Document Index': 'text', # Document Index -> text (TEXT_CATEGORIES) 'Code': 'code', # Code -> code (CODE_CATEGORIES) 'Checkbox-Selected': 'abandon', # Checkbox -> abandon (DISCARD_CATEGORIES) 'Checkbox-Unselected': 'abandon', # Checkbox -> abandon (DISCARD_CATEGORIES) 'Form': 'abandon', # Form -> abandon (DISCARD_CATEGORIES) 'Key-Value Region': 'text', # Key-Value Region -> text (TEXT_CATEGORIES) 'Background': 'abandon', # Background -> abandon (DISCARD_CATEGORIES) } def __init__(self, config: Dict[str, Any]): """ 初始化 Docling Layout 检测器 Args: config: 配置字典,支持以下参数: - model_dir: 模型目录路径或 HuggingFace 仓库 ID - device: 运行设备 ('cpu', 'cuda', 'mps') - conf: 置信度阈值 (默认 0.3) - num_threads: CPU 线程数 (默认 4) """ super().__init__(config) self.model = None self.image_processor = None self._device = None self._threshold = 0.3 self._num_threads = 4 self._model_path = None # RT-DETR 使用 shifted labels,label_offset = 1 self._label_offset = 1 def initialize(self): """初始化模型""" try: import torch from transformers import AutoModelForObjectDetection, RTDetrImageProcessor from huggingface_hub import snapshot_download model_dir = self.config.get('model_dir', 'ds4sd/docling-layout-old') device = self.config.get('device', 'cpu') self._threshold = self.config.get('conf', 0.3) self._num_threads = self.config.get('num_threads', 4) # 设置设备 self._device = torch.device(device) if device == 'cpu': torch.set_num_threads(self._num_threads) # 判断是本地路径还是 HuggingFace 仓库 model_path = Path(model_dir) if model_path.exists() and model_path.is_dir(): # 本地路径 self._model_path = str(model_path) print(f"📂 Loading model from local path: {self._model_path}") else: # 从 HuggingFace 下载 print(f"📥 Downloading model from HuggingFace: {model_dir}") self._model_path = snapshot_download(repo_id=model_dir) # 检查必要文件 processor_config = Path(self._model_path) / "preprocessor_config.json" model_config = Path(self._model_path) / "config.json" safetensors_file = Path(self._model_path) / "model.safetensors" if not processor_config.exists(): raise FileNotFoundError(f"Missing preprocessor_config.json in {self._model_path}") if not model_config.exists(): raise FileNotFoundError(f"Missing config.json in {self._model_path}") if not safetensors_file.exists(): raise FileNotFoundError(f"Missing model.safetensors in {self._model_path}") # 加载图像处理器 self.image_processor = RTDetrImageProcessor.from_json_file(str(processor_config)) # 加载模型(使用锁防止线程问题) with _model_init_lock: self.model = AutoModelForObjectDetection.from_pretrained( self._model_path, config=str(model_config), device_map=self._device ) self.model.eval() # 检测模型类型 model_name = type(self.model).__name__ print(f"✅ Docling Layout Detector initialized") print(f" - Model: {model_name}") print(f" - Device: {self._device}") print(f" - Threshold: {self._threshold}") print(f" - Image size: {self.image_processor.size}") except ImportError as e: print(f"❌ Failed to import required libraries: {e}") print(" Please ensure transformers and torch are installed") raise except Exception as e: print(f"❌ Failed to initialize Docling Layout Detector: {e}") raise def cleanup(self): """清理资源""" self.model = None self.image_processor = None self._model_path = None def detect(self, image: Union[np.ndarray, Image.Image]) -> List[Dict[str, Any]]: """ 检测布局 Args: image: 输入图像 (numpy数组或PIL图像) Returns: 检测结果列表,每个元素包含: - category: MinerU类别名称 - bbox: [x1, y1, x2, y2] - confidence: 置信度 - raw: 原始检测结果 """ import torch if self.model is None: raise RuntimeError("Model not initialized. Call initialize() first.") # 转换为 PIL Image if isinstance(image, np.ndarray): # OpenCV BGR -> RGB if len(image.shape) == 3 and image.shape[2] == 3: image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) else: image_rgb = image pil_image = Image.fromarray(image_rgb).convert("RGB") orig_h, orig_w = image.shape[:2] else: pil_image = image.convert("RGB") orig_w, orig_h = image.size # 推理 with torch.inference_mode(): target_sizes = torch.tensor([pil_image.size[::-1]]) # (h, w) inputs = self.image_processor(images=[pil_image], return_tensors="pt").to(self._device) outputs = self.model(**inputs) results = self.image_processor.post_process_object_detection( outputs, target_sizes=target_sizes, threshold=self._threshold, ) # 解析结果 w, h = pil_image.size result = results[0] formatted_results = [] for score, label_id, box in zip(result["scores"], result["labels"], result["boxes"]): score = float(score.item()) # 获取原始标签(考虑 offset) label_id_int = int(label_id.item()) + self._label_offset original_label = self.DOCLING_LABELS.get(label_id_int, f'unknown_{label_id_int}') # 映射到 MinerU 类别 mineru_category = self.CATEGORY_MAP.get(original_label, 'text') # 跳过 Background if original_label == 'Background': continue # 提取边界框 bbox_float = [float(b.item()) for b in box] x1 = min(w, max(0, bbox_float[0])) y1 = min(h, max(0, bbox_float[1])) x2 = min(w, max(0, bbox_float[2])) y2 = min(h, max(0, bbox_float[3])) bbox = [int(x1), int(y1), int(x2), int(y2)] # 计算宽高 width = bbox[2] - bbox[0] height = bbox[3] - bbox[1] # 过滤太小的框 if width < 10 or height < 10: continue # 过滤面积异常大的框 area = width * height img_area = orig_w * orig_h if area > img_area * 0.95: continue # 生成多边形坐标 poly = [ bbox[0], bbox[1], # 左上 bbox[2], bbox[1], # 右上 bbox[2], bbox[3], # 右下 bbox[0], bbox[3], # 左下 ] formatted_results.append({ 'category': mineru_category, 'bbox': bbox, 'confidence': score, 'raw': { 'original_label': original_label, 'original_label_id': label_id_int, 'poly': poly, 'width': width, 'height': height } }) return formatted_results def detect_batch( self, images: List[Union[np.ndarray, Image.Image]] ) -> List[List[Dict[str, Any]]]: """ 批量检测布局(更高效) Args: images: 输入图像列表 Returns: 每个图像的检测结果列表 """ import torch if self.model is None: raise RuntimeError("Model not initialized. Call initialize() first.") if not images: return [] # 转换为 PIL Image 列表 pil_images = [] orig_sizes = [] for image in images: if isinstance(image, np.ndarray): if len(image.shape) == 3 and image.shape[2] == 3: image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) else: image_rgb = image pil_images.append(Image.fromarray(image_rgb).convert("RGB")) orig_sizes.append((image.shape[1], image.shape[0])) # (w, h) else: pil_images.append(image.convert("RGB")) orig_sizes.append(image.size) # (w, h) # 批量推理 with torch.inference_mode(): target_sizes = torch.tensor([img.size[::-1] for img in pil_images]) inputs = self.image_processor(images=pil_images, return_tensors="pt").to(self._device) outputs = self.model(**inputs) results_list = self.image_processor.post_process_object_detection( outputs, target_sizes=target_sizes, threshold=self._threshold, ) # 转换结果 all_formatted_results = [] for pil_img, results, (orig_w, orig_h) in zip(pil_images, results_list, orig_sizes): w, h = pil_img.size formatted_results = [] for score, label_id, box in zip(results["scores"], results["labels"], results["boxes"]): score = float(score.item()) label_id_int = int(label_id.item()) + self._label_offset original_label = self.DOCLING_LABELS.get(label_id_int, f'unknown_{label_id_int}') mineru_category = self.CATEGORY_MAP.get(original_label, 'text') if original_label == 'Background': continue bbox_float = [float(b.item()) for b in box] x1 = min(w, max(0, bbox_float[0])) y1 = min(h, max(0, bbox_float[1])) x2 = min(w, max(0, bbox_float[2])) y2 = min(h, max(0, bbox_float[3])) bbox = [int(x1), int(y1), int(x2), int(y2)] width = bbox[2] - bbox[0] height = bbox[3] - bbox[1] if width < 10 or height < 10: continue area = width * height img_area = orig_w * orig_h if area > img_area * 0.95: continue poly = [ bbox[0], bbox[1], bbox[2], bbox[1], bbox[2], bbox[3], bbox[0], bbox[3], ] formatted_results.append({ 'category': mineru_category, 'bbox': bbox, 'confidence': score, 'raw': { 'original_label': original_label, 'original_label_id': label_id_int, 'poly': poly, 'width': width, 'height': height } }) all_formatted_results.append(formatted_results) return all_formatted_results def visualize( self, img: np.ndarray, results: List[Dict], output_path: str = None, show_confidence: bool = True, min_confidence: float = 0.0 ) -> np.ndarray: """ 可视化检测结果 Args: img: 输入图像 (BGR 格式) results: 检测结果 (MinerU 格式) output_path: 输出路径(可选) show_confidence: 是否显示置信度 min_confidence: 最小置信度阈值 Returns: 标注后的图像 """ import random vis_img = img.copy() # 预定义类别颜色(与 EnhancedDocPipeline 保持一致) predefined_colors = { # 文本类 'text': (153, 0, 76), 'title': (102, 102, 255), 'header': (128, 128, 128), 'footer': (128, 128, 128), 'page_footnote': (200, 200, 200), # 表格类 'table_body': (204, 204, 0), 'table_caption': (255, 255, 102), # 图片类 'image_body': (153, 255, 51), 'image_caption': (102, 178, 255), # 公式类 'interline_equation': (0, 255, 0), # 代码类 'code': (102, 0, 204), # 丢弃类 'abandon': (100, 100, 100), } # 过滤低置信度结果 filtered_results = [ res for res in results if res['confidence'] >= min_confidence ] if not filtered_results: print(f"⚠️ No results to visualize (min_confidence={min_confidence})") return vis_img # 为每个出现的类别分配颜色 category_colors = {} for res in filtered_results: cat = res['category'] if cat not in category_colors: if cat in predefined_colors: category_colors[cat] = predefined_colors[cat] else: category_colors[cat] = ( random.randint(50, 255), random.randint(50, 255), random.randint(50, 255) ) # 绘制检测框 for res in filtered_results: bbox = res['bbox'] x1, y1, x2, y2 = bbox cat = res['category'] confidence = res['confidence'] color = category_colors[cat] # 获取原始标签 original_label = res.get('raw', {}).get('original_label', cat) # 绘制矩形边框 cv2.rectangle(vis_img, (x1, y1), (x2, y2), color, 2) # 构造标签文本 if show_confidence: label = f"{original_label}->{cat} {confidence:.2f}" else: label = f"{original_label}->{cat}" # 计算标签尺寸 label_size, baseline = cv2.getTextSize( label, cv2.FONT_HERSHEY_SIMPLEX, 0.4, 1 ) label_w, label_h = label_size # 绘制标签背景 cv2.rectangle( vis_img, (x1, y1 - label_h - 4), (x1 + label_w, y1), color, -1 ) # 绘制标签文字 cv2.putText( vis_img, label, (x1, y1 - 2), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1, cv2.LINE_AA ) # 添加图例 if category_colors: self._draw_legend(vis_img, category_colors, len(filtered_results)) # 保存可视化结果 if output_path: output_path = Path(output_path) output_path.parent.mkdir(parents=True, exist_ok=True) cv2.imwrite(str(output_path), vis_img) print(f"💾 Visualization saved to: {output_path}") return vis_img def _draw_legend( self, img: np.ndarray, category_colors: Dict[str, tuple], total_count: int ): """在图像上绘制图例""" legend_x = img.shape[1] - 200 legend_y = 20 line_height = 25 # 绘制半透明背景 overlay = img.copy() cv2.rectangle( overlay, (legend_x - 10, legend_y - 10), (img.shape[1] - 10, legend_y + len(category_colors) * line_height + 30), (255, 255, 255), -1 ) cv2.addWeighted(overlay, 0.7, img, 0.3, 0, img) # 绘制标题 cv2.putText( img, f"Legend ({total_count} total)", (legend_x, legend_y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA ) # 绘制每个类别 y_offset = legend_y + line_height for cat, color in sorted(category_colors.items()): cv2.rectangle( img, (legend_x, y_offset - 10), (legend_x + 15, y_offset), color, -1 ) cv2.rectangle( img, (legend_x, y_offset - 10), (legend_x + 15, y_offset), (0, 0, 0), 1 ) cv2.putText( img, cat, (legend_x + 20, y_offset - 2), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 0, 0), 1, cv2.LINE_AA ) y_offset += line_height # 测试代码 if __name__ == "__main__": import sys # 测试配置 - 使用 HuggingFace 下载模型 config = { 'model_dir': 'ds4sd/docling-layout-old', # HuggingFace 仓库 ID 'device': 'cpu', 'conf': 0.3, 'num_threads': 4 } # 初始化检测器 print("🔧 Initializing Docling Layout Detector...") detector = DoclingLayoutDetector(config) detector.initialize() # 读取测试图像 img_path = "/Users/zhch158/workspace/data/流水分析/康强_北京农村商业银行/ppstructurev3_client_results/康强_北京农村商业银行/康强_北京农村商业银行_page_001.png" print(f"\n📖 Loading image: {img_path}") img = cv2.imread(img_path) if img is None: print(f"❌ Failed to load image: {img_path}") sys.exit(1) print(f" Image shape: {img.shape}") # 执行检测 print("\n🔍 Detecting layout...") results = detector.detect(img) print(f"\n✅ 检测到 {len(results)} 个区域:") for i, res in enumerate(results, 1): print(f" [{i}] {res['category']}: " f"score={res['confidence']:.3f}, " f"bbox={res['bbox']}, " f"original={res['raw']['original_label']}") # 统计各类别 category_counts = {} for res in results: cat = res['category'] category_counts[cat] = category_counts.get(cat, 0) + 1 print(f"\n📊 类别统计 (MinerU格式):") for cat, count in sorted(category_counts.items()): print(f" - {cat}: {count}") # 可视化 if len(results) > 0: print("\n🎨 Generating visualization...") output_dir = Path(__file__).parent.parent.parent / "tests" / "output" output_dir.mkdir(parents=True, exist_ok=True) output_path = output_dir / f"{Path(img_path).stem}_docling_layout_vis.jpg" vis_img = detector.visualize( img, results, output_path=str(output_path), show_confidence=True, min_confidence=0.0 ) print(f"💾 Visualization saved to: {output_path}") # 清理 detector.cleanup() print("\n✅ 测试完成!")