| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667 |
- """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✅ 测试完成!")
|