docling_layout_adapter.py 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667
  1. """Docling Layout Predictor 适配器 (符合 BaseLayoutDetector 规范)
  2. 基于 HuggingFace transformers 直接加载 Docling 布局模型,不依赖 docling-ibm-models 包。
  3. 使用 MinerU 环境中的 transformers 库。
  4. 支持的 Hugging Face 模型仓库:
  5. - ds4sd/docling-layout-old
  6. - ds4sd/docling-layout-heron
  7. - ds4sd/docling-layout-heron-101
  8. - ds4sd/docling-layout-egret-medium
  9. - ds4sd/docling-layout-egret-large
  10. - ds4sd/docling-layout-egret-xlarge
  11. """
  12. import cv2
  13. import numpy as np
  14. import threading
  15. from pathlib import Path
  16. from typing import Dict, List, Union, Any, Optional
  17. from PIL import Image
  18. try:
  19. from .base import BaseLayoutDetector
  20. except ImportError:
  21. from base import BaseLayoutDetector
  22. # 全局锁,防止模型初始化时的线程问题
  23. _model_init_lock = threading.Lock()
  24. class DoclingLayoutDetector(BaseLayoutDetector):
  25. """Docling Layout Predictor 适配器
  26. 直接使用 transformers 库加载 Docling 布局模型,无需安装 docling-ibm-models。
  27. """
  28. # Docling 原始类别定义(来自 docling-ibm-models/layoutmodel/labels.py)
  29. # 使用 shifted canonical(带 Background)的版本,因为 RT-DETR 模型使用这个
  30. DOCLING_LABELS = {
  31. 0: 'Background',
  32. 1: 'Caption',
  33. 2: 'Footnote',
  34. 3: 'Formula',
  35. 4: 'List-item',
  36. 5: 'Page-footer',
  37. 6: 'Page-header',
  38. 7: 'Picture',
  39. 8: 'Section-header',
  40. 9: 'Table',
  41. 10: 'Text',
  42. 11: 'Title',
  43. 12: 'Document Index',
  44. 13: 'Code',
  45. 14: 'Checkbox-Selected',
  46. 15: 'Checkbox-Unselected',
  47. 16: 'Form',
  48. 17: 'Key-Value Region',
  49. }
  50. # 类别映射:Docling LayoutLabels → MinerU/EnhancedDocPipeline 类别体系
  51. # 参考:
  52. # - MinerU: mineru/utils/enum_class.py (BlockType, CategoryId)
  53. # - Pipeline: universal_doc_parser/core/pipeline_manager_v2.py (EnhancedDocPipeline 类别定义)
  54. CATEGORY_MAP = {
  55. 'Caption': 'image_caption', # Caption -> image_caption (IMAGE_TEXT_CATEGORIES)
  56. 'Footnote': 'page_footnote', # Footnote -> page_footnote (TEXT_CATEGORIES)
  57. 'Formula': 'interline_equation', # Formula -> interline_equation (EQUATION_CATEGORIES)
  58. 'List-item': 'text', # List-item -> text (TEXT_CATEGORIES)
  59. 'Page-footer': 'footer', # Page-footer -> footer (TEXT_CATEGORIES)
  60. 'Page-header': 'header', # Page-header -> header (TEXT_CATEGORIES)
  61. 'Picture': 'image_body', # Picture -> image_body (IMAGE_BODY_CATEGORIES)
  62. 'Section-header': 'title', # Section-header -> title (TEXT_CATEGORIES)
  63. 'Table': 'table_body', # Table -> table_body (TABLE_BODY_CATEGORIES)
  64. 'Text': 'text', # Text -> text (TEXT_CATEGORIES)
  65. 'Title': 'title', # Title -> title (TEXT_CATEGORIES)
  66. 'Document Index': 'text', # Document Index -> text (TEXT_CATEGORIES)
  67. 'Code': 'code', # Code -> code (CODE_CATEGORIES)
  68. 'Checkbox-Selected': 'abandon', # Checkbox -> abandon (DISCARD_CATEGORIES)
  69. 'Checkbox-Unselected': 'abandon', # Checkbox -> abandon (DISCARD_CATEGORIES)
  70. 'Form': 'abandon', # Form -> abandon (DISCARD_CATEGORIES)
  71. 'Key-Value Region': 'text', # Key-Value Region -> text (TEXT_CATEGORIES)
  72. 'Background': 'abandon', # Background -> abandon (DISCARD_CATEGORIES)
  73. }
  74. def __init__(self, config: Dict[str, Any]):
  75. """
  76. 初始化 Docling Layout 检测器
  77. Args:
  78. config: 配置字典,支持以下参数:
  79. - model_dir: 模型目录路径或 HuggingFace 仓库 ID
  80. - device: 运行设备 ('cpu', 'cuda', 'mps')
  81. - conf: 置信度阈值 (默认 0.3)
  82. - num_threads: CPU 线程数 (默认 4)
  83. """
  84. super().__init__(config)
  85. self.model = None
  86. self.image_processor = None
  87. self._device = None
  88. self._threshold = 0.3
  89. self._num_threads = 4
  90. self._model_path = None
  91. # RT-DETR 使用 shifted labels,label_offset = 1
  92. self._label_offset = 1
  93. def initialize(self):
  94. """初始化模型"""
  95. try:
  96. import torch
  97. from transformers import AutoModelForObjectDetection, RTDetrImageProcessor
  98. from huggingface_hub import snapshot_download
  99. model_dir = self.config.get('model_dir', 'ds4sd/docling-layout-old')
  100. device = self.config.get('device', 'cpu')
  101. self._threshold = self.config.get('conf', 0.3)
  102. self._num_threads = self.config.get('num_threads', 4)
  103. # 设置设备
  104. self._device = torch.device(device)
  105. if device == 'cpu':
  106. torch.set_num_threads(self._num_threads)
  107. # 判断是本地路径还是 HuggingFace 仓库
  108. model_path = Path(model_dir)
  109. if model_path.exists() and model_path.is_dir():
  110. # 本地路径
  111. self._model_path = str(model_path)
  112. print(f"📂 Loading model from local path: {self._model_path}")
  113. else:
  114. # 从 HuggingFace 下载
  115. print(f"📥 Downloading model from HuggingFace: {model_dir}")
  116. self._model_path = snapshot_download(repo_id=model_dir)
  117. # 检查必要文件
  118. processor_config = Path(self._model_path) / "preprocessor_config.json"
  119. model_config = Path(self._model_path) / "config.json"
  120. safetensors_file = Path(self._model_path) / "model.safetensors"
  121. if not processor_config.exists():
  122. raise FileNotFoundError(f"Missing preprocessor_config.json in {self._model_path}")
  123. if not model_config.exists():
  124. raise FileNotFoundError(f"Missing config.json in {self._model_path}")
  125. if not safetensors_file.exists():
  126. raise FileNotFoundError(f"Missing model.safetensors in {self._model_path}")
  127. # 加载图像处理器
  128. self.image_processor = RTDetrImageProcessor.from_json_file(str(processor_config))
  129. # 加载模型(使用锁防止线程问题)
  130. with _model_init_lock:
  131. self.model = AutoModelForObjectDetection.from_pretrained(
  132. self._model_path,
  133. config=str(model_config),
  134. device_map=self._device
  135. )
  136. self.model.eval()
  137. # 检测模型类型
  138. model_name = type(self.model).__name__
  139. print(f"✅ Docling Layout Detector initialized")
  140. print(f" - Model: {model_name}")
  141. print(f" - Device: {self._device}")
  142. print(f" - Threshold: {self._threshold}")
  143. print(f" - Image size: {self.image_processor.size}")
  144. except ImportError as e:
  145. print(f"❌ Failed to import required libraries: {e}")
  146. print(" Please ensure transformers and torch are installed")
  147. raise
  148. except Exception as e:
  149. print(f"❌ Failed to initialize Docling Layout Detector: {e}")
  150. raise
  151. def cleanup(self):
  152. """清理资源"""
  153. self.model = None
  154. self.image_processor = None
  155. self._model_path = None
  156. def detect(self, image: Union[np.ndarray, Image.Image]) -> List[Dict[str, Any]]:
  157. """
  158. 检测布局
  159. Args:
  160. image: 输入图像 (numpy数组或PIL图像)
  161. Returns:
  162. 检测结果列表,每个元素包含:
  163. - category: MinerU类别名称
  164. - bbox: [x1, y1, x2, y2]
  165. - confidence: 置信度
  166. - raw: 原始检测结果
  167. """
  168. import torch
  169. if self.model is None:
  170. raise RuntimeError("Model not initialized. Call initialize() first.")
  171. # 转换为 PIL Image
  172. if isinstance(image, np.ndarray):
  173. # OpenCV BGR -> RGB
  174. if len(image.shape) == 3 and image.shape[2] == 3:
  175. image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
  176. else:
  177. image_rgb = image
  178. pil_image = Image.fromarray(image_rgb).convert("RGB")
  179. orig_h, orig_w = image.shape[:2]
  180. else:
  181. pil_image = image.convert("RGB")
  182. orig_w, orig_h = image.size
  183. # 推理
  184. with torch.inference_mode():
  185. target_sizes = torch.tensor([pil_image.size[::-1]]) # (h, w)
  186. inputs = self.image_processor(images=[pil_image], return_tensors="pt").to(self._device)
  187. outputs = self.model(**inputs)
  188. results = self.image_processor.post_process_object_detection(
  189. outputs,
  190. target_sizes=target_sizes,
  191. threshold=self._threshold,
  192. )
  193. # 解析结果
  194. w, h = pil_image.size
  195. result = results[0]
  196. formatted_results = []
  197. for score, label_id, box in zip(result["scores"], result["labels"], result["boxes"]):
  198. score = float(score.item())
  199. # 获取原始标签(考虑 offset)
  200. label_id_int = int(label_id.item()) + self._label_offset
  201. original_label = self.DOCLING_LABELS.get(label_id_int, f'unknown_{label_id_int}')
  202. # 映射到 MinerU 类别
  203. mineru_category = self.CATEGORY_MAP.get(original_label, 'text')
  204. # 跳过 Background
  205. if original_label == 'Background':
  206. continue
  207. # 提取边界框
  208. bbox_float = [float(b.item()) for b in box]
  209. x1 = min(w, max(0, bbox_float[0]))
  210. y1 = min(h, max(0, bbox_float[1]))
  211. x2 = min(w, max(0, bbox_float[2]))
  212. y2 = min(h, max(0, bbox_float[3]))
  213. bbox = [int(x1), int(y1), int(x2), int(y2)]
  214. # 计算宽高
  215. width = bbox[2] - bbox[0]
  216. height = bbox[3] - bbox[1]
  217. # 过滤太小的框
  218. if width < 10 or height < 10:
  219. continue
  220. # 过滤面积异常大的框
  221. area = width * height
  222. img_area = orig_w * orig_h
  223. if area > img_area * 0.95:
  224. continue
  225. # 生成多边形坐标
  226. poly = [
  227. bbox[0], bbox[1], # 左上
  228. bbox[2], bbox[1], # 右上
  229. bbox[2], bbox[3], # 右下
  230. bbox[0], bbox[3], # 左下
  231. ]
  232. formatted_results.append({
  233. 'category': mineru_category,
  234. 'bbox': bbox,
  235. 'confidence': score,
  236. 'raw': {
  237. 'original_label': original_label,
  238. 'original_label_id': label_id_int,
  239. 'poly': poly,
  240. 'width': width,
  241. 'height': height
  242. }
  243. })
  244. return formatted_results
  245. def detect_batch(
  246. self,
  247. images: List[Union[np.ndarray, Image.Image]]
  248. ) -> List[List[Dict[str, Any]]]:
  249. """
  250. 批量检测布局(更高效)
  251. Args:
  252. images: 输入图像列表
  253. Returns:
  254. 每个图像的检测结果列表
  255. """
  256. import torch
  257. if self.model is None:
  258. raise RuntimeError("Model not initialized. Call initialize() first.")
  259. if not images:
  260. return []
  261. # 转换为 PIL Image 列表
  262. pil_images = []
  263. orig_sizes = []
  264. for image in images:
  265. if isinstance(image, np.ndarray):
  266. if len(image.shape) == 3 and image.shape[2] == 3:
  267. image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
  268. else:
  269. image_rgb = image
  270. pil_images.append(Image.fromarray(image_rgb).convert("RGB"))
  271. orig_sizes.append((image.shape[1], image.shape[0])) # (w, h)
  272. else:
  273. pil_images.append(image.convert("RGB"))
  274. orig_sizes.append(image.size) # (w, h)
  275. # 批量推理
  276. with torch.inference_mode():
  277. target_sizes = torch.tensor([img.size[::-1] for img in pil_images])
  278. inputs = self.image_processor(images=pil_images, return_tensors="pt").to(self._device)
  279. outputs = self.model(**inputs)
  280. results_list = self.image_processor.post_process_object_detection(
  281. outputs,
  282. target_sizes=target_sizes,
  283. threshold=self._threshold,
  284. )
  285. # 转换结果
  286. all_formatted_results = []
  287. for pil_img, results, (orig_w, orig_h) in zip(pil_images, results_list, orig_sizes):
  288. w, h = pil_img.size
  289. formatted_results = []
  290. for score, label_id, box in zip(results["scores"], results["labels"], results["boxes"]):
  291. score = float(score.item())
  292. label_id_int = int(label_id.item()) + self._label_offset
  293. original_label = self.DOCLING_LABELS.get(label_id_int, f'unknown_{label_id_int}')
  294. mineru_category = self.CATEGORY_MAP.get(original_label, 'text')
  295. if original_label == 'Background':
  296. continue
  297. bbox_float = [float(b.item()) for b in box]
  298. x1 = min(w, max(0, bbox_float[0]))
  299. y1 = min(h, max(0, bbox_float[1]))
  300. x2 = min(w, max(0, bbox_float[2]))
  301. y2 = min(h, max(0, bbox_float[3]))
  302. bbox = [int(x1), int(y1), int(x2), int(y2)]
  303. width = bbox[2] - bbox[0]
  304. height = bbox[3] - bbox[1]
  305. if width < 10 or height < 10:
  306. continue
  307. area = width * height
  308. img_area = orig_w * orig_h
  309. if area > img_area * 0.95:
  310. continue
  311. poly = [
  312. bbox[0], bbox[1],
  313. bbox[2], bbox[1],
  314. bbox[2], bbox[3],
  315. bbox[0], bbox[3],
  316. ]
  317. formatted_results.append({
  318. 'category': mineru_category,
  319. 'bbox': bbox,
  320. 'confidence': score,
  321. 'raw': {
  322. 'original_label': original_label,
  323. 'original_label_id': label_id_int,
  324. 'poly': poly,
  325. 'width': width,
  326. 'height': height
  327. }
  328. })
  329. all_formatted_results.append(formatted_results)
  330. return all_formatted_results
  331. def visualize(
  332. self,
  333. img: np.ndarray,
  334. results: List[Dict],
  335. output_path: str = None,
  336. show_confidence: bool = True,
  337. min_confidence: float = 0.0
  338. ) -> np.ndarray:
  339. """
  340. 可视化检测结果
  341. Args:
  342. img: 输入图像 (BGR 格式)
  343. results: 检测结果 (MinerU 格式)
  344. output_path: 输出路径(可选)
  345. show_confidence: 是否显示置信度
  346. min_confidence: 最小置信度阈值
  347. Returns:
  348. 标注后的图像
  349. """
  350. import random
  351. vis_img = img.copy()
  352. # 预定义类别颜色(与 EnhancedDocPipeline 保持一致)
  353. predefined_colors = {
  354. # 文本类
  355. 'text': (153, 0, 76),
  356. 'title': (102, 102, 255),
  357. 'header': (128, 128, 128),
  358. 'footer': (128, 128, 128),
  359. 'page_footnote': (200, 200, 200),
  360. # 表格类
  361. 'table_body': (204, 204, 0),
  362. 'table_caption': (255, 255, 102),
  363. # 图片类
  364. 'image_body': (153, 255, 51),
  365. 'image_caption': (102, 178, 255),
  366. # 公式类
  367. 'interline_equation': (0, 255, 0),
  368. # 代码类
  369. 'code': (102, 0, 204),
  370. # 丢弃类
  371. 'abandon': (100, 100, 100),
  372. }
  373. # 过滤低置信度结果
  374. filtered_results = [
  375. res for res in results
  376. if res['confidence'] >= min_confidence
  377. ]
  378. if not filtered_results:
  379. print(f"⚠️ No results to visualize (min_confidence={min_confidence})")
  380. return vis_img
  381. # 为每个出现的类别分配颜色
  382. category_colors = {}
  383. for res in filtered_results:
  384. cat = res['category']
  385. if cat not in category_colors:
  386. if cat in predefined_colors:
  387. category_colors[cat] = predefined_colors[cat]
  388. else:
  389. category_colors[cat] = (
  390. random.randint(50, 255),
  391. random.randint(50, 255),
  392. random.randint(50, 255)
  393. )
  394. # 绘制检测框
  395. for res in filtered_results:
  396. bbox = res['bbox']
  397. x1, y1, x2, y2 = bbox
  398. cat = res['category']
  399. confidence = res['confidence']
  400. color = category_colors[cat]
  401. # 获取原始标签
  402. original_label = res.get('raw', {}).get('original_label', cat)
  403. # 绘制矩形边框
  404. cv2.rectangle(vis_img, (x1, y1), (x2, y2), color, 2)
  405. # 构造标签文本
  406. if show_confidence:
  407. label = f"{original_label}->{cat} {confidence:.2f}"
  408. else:
  409. label = f"{original_label}->{cat}"
  410. # 计算标签尺寸
  411. label_size, baseline = cv2.getTextSize(
  412. label,
  413. cv2.FONT_HERSHEY_SIMPLEX,
  414. 0.4,
  415. 1
  416. )
  417. label_w, label_h = label_size
  418. # 绘制标签背景
  419. cv2.rectangle(
  420. vis_img,
  421. (x1, y1 - label_h - 4),
  422. (x1 + label_w, y1),
  423. color,
  424. -1
  425. )
  426. # 绘制标签文字
  427. cv2.putText(
  428. vis_img,
  429. label,
  430. (x1, y1 - 2),
  431. cv2.FONT_HERSHEY_SIMPLEX,
  432. 0.4,
  433. (255, 255, 255),
  434. 1,
  435. cv2.LINE_AA
  436. )
  437. # 添加图例
  438. if category_colors:
  439. self._draw_legend(vis_img, category_colors, len(filtered_results))
  440. # 保存可视化结果
  441. if output_path:
  442. output_path = Path(output_path)
  443. output_path.parent.mkdir(parents=True, exist_ok=True)
  444. cv2.imwrite(str(output_path), vis_img)
  445. print(f"💾 Visualization saved to: {output_path}")
  446. return vis_img
  447. def _draw_legend(
  448. self,
  449. img: np.ndarray,
  450. category_colors: Dict[str, tuple],
  451. total_count: int
  452. ):
  453. """在图像上绘制图例"""
  454. legend_x = img.shape[1] - 200
  455. legend_y = 20
  456. line_height = 25
  457. # 绘制半透明背景
  458. overlay = img.copy()
  459. cv2.rectangle(
  460. overlay,
  461. (legend_x - 10, legend_y - 10),
  462. (img.shape[1] - 10, legend_y + len(category_colors) * line_height + 30),
  463. (255, 255, 255),
  464. -1
  465. )
  466. cv2.addWeighted(overlay, 0.7, img, 0.3, 0, img)
  467. # 绘制标题
  468. cv2.putText(
  469. img,
  470. f"Legend ({total_count} total)",
  471. (legend_x, legend_y),
  472. cv2.FONT_HERSHEY_SIMPLEX,
  473. 0.5,
  474. (0, 0, 0),
  475. 1,
  476. cv2.LINE_AA
  477. )
  478. # 绘制每个类别
  479. y_offset = legend_y + line_height
  480. for cat, color in sorted(category_colors.items()):
  481. cv2.rectangle(
  482. img,
  483. (legend_x, y_offset - 10),
  484. (legend_x + 15, y_offset),
  485. color,
  486. -1
  487. )
  488. cv2.rectangle(
  489. img,
  490. (legend_x, y_offset - 10),
  491. (legend_x + 15, y_offset),
  492. (0, 0, 0),
  493. 1
  494. )
  495. cv2.putText(
  496. img,
  497. cat,
  498. (legend_x + 20, y_offset - 2),
  499. cv2.FONT_HERSHEY_SIMPLEX,
  500. 0.4,
  501. (0, 0, 0),
  502. 1,
  503. cv2.LINE_AA
  504. )
  505. y_offset += line_height
  506. # 测试代码
  507. if __name__ == "__main__":
  508. import sys
  509. # 测试配置 - 使用 HuggingFace 下载模型
  510. config = {
  511. 'model_dir': 'ds4sd/docling-layout-old', # HuggingFace 仓库 ID
  512. 'device': 'cpu',
  513. 'conf': 0.3,
  514. 'num_threads': 4
  515. }
  516. # 初始化检测器
  517. print("🔧 Initializing Docling Layout Detector...")
  518. detector = DoclingLayoutDetector(config)
  519. detector.initialize()
  520. # 读取测试图像
  521. img_path = "/Users/zhch158/workspace/data/流水分析/康强_北京农村商业银行/ppstructurev3_client_results/康强_北京农村商业银行/康强_北京农村商业银行_page_001.png"
  522. print(f"\n📖 Loading image: {img_path}")
  523. img = cv2.imread(img_path)
  524. if img is None:
  525. print(f"❌ Failed to load image: {img_path}")
  526. sys.exit(1)
  527. print(f" Image shape: {img.shape}")
  528. # 执行检测
  529. print("\n🔍 Detecting layout...")
  530. results = detector.detect(img)
  531. print(f"\n✅ 检测到 {len(results)} 个区域:")
  532. for i, res in enumerate(results, 1):
  533. print(f" [{i}] {res['category']}: "
  534. f"score={res['confidence']:.3f}, "
  535. f"bbox={res['bbox']}, "
  536. f"original={res['raw']['original_label']}")
  537. # 统计各类别
  538. category_counts = {}
  539. for res in results:
  540. cat = res['category']
  541. category_counts[cat] = category_counts.get(cat, 0) + 1
  542. print(f"\n📊 类别统计 (MinerU格式):")
  543. for cat, count in sorted(category_counts.items()):
  544. print(f" - {cat}: {count}")
  545. # 可视化
  546. if len(results) > 0:
  547. print("\n🎨 Generating visualization...")
  548. output_dir = Path(__file__).parent.parent.parent / "tests" / "output"
  549. output_dir.mkdir(parents=True, exist_ok=True)
  550. output_path = output_dir / f"{Path(img_path).stem}_docling_layout_vis.jpg"
  551. vis_img = detector.visualize(
  552. img,
  553. results,
  554. output_path=str(output_path),
  555. show_confidence=True,
  556. min_confidence=0.0
  557. )
  558. print(f"💾 Visualization saved to: {output_path}")
  559. # 清理
  560. detector.cleanup()
  561. print("\n✅ 测试完成!")