|
|
@@ -10,6 +10,7 @@ from typing import Dict, Any, Union
|
|
|
from pathlib import Path
|
|
|
import numpy as np
|
|
|
from PIL import Image
|
|
|
+import cv2
|
|
|
from loguru import logger
|
|
|
|
|
|
from .base import BaseAdapter
|
|
|
@@ -73,27 +74,30 @@ class PaddleTableClassifier(BaseAdapter):
|
|
|
|
|
|
def classify(
|
|
|
self,
|
|
|
- image: Union[np.ndarray, Image.Image]
|
|
|
+ image: Union[np.ndarray, Image.Image],
|
|
|
+ use_line_detection: bool = True
|
|
|
) -> Dict[str, Any]:
|
|
|
"""
|
|
|
分类单个表格图像
|
|
|
|
|
|
Args:
|
|
|
image: 表格图像(numpy数组或PIL图像)
|
|
|
+ use_line_detection: 是否使用线条检测辅助判断(默认True)
|
|
|
|
|
|
Returns:
|
|
|
分类结果字典:
|
|
|
{
|
|
|
'table_type': 'wired' | 'wireless',
|
|
|
'confidence': float,
|
|
|
- 'raw_label': str # AtomicModel.WiredTable 或 AtomicModel.WirelessTable
|
|
|
+ 'raw_label': str, # AtomicModel.WiredTable 或 AtomicModel.WirelessTable
|
|
|
+ 'line_detection': dict # 线条检测结果(如果启用)
|
|
|
}
|
|
|
"""
|
|
|
if self.model is None:
|
|
|
raise RuntimeError("Model not initialized. Call initialize() first.")
|
|
|
|
|
|
try:
|
|
|
- # 调用 MinerU 的预测接口
|
|
|
+ # Step 1: 调用 MinerU 的预测接口
|
|
|
label, confidence = self.model.predict(image)
|
|
|
|
|
|
# 转换标签为简化形式
|
|
|
@@ -111,7 +115,37 @@ class PaddleTableClassifier(BaseAdapter):
|
|
|
'raw_label': str(label)
|
|
|
}
|
|
|
|
|
|
- logger.debug(f"Table classified as '{table_type}' (confidence: {confidence:.3f})")
|
|
|
+ # Step 2: 使用线条检测辅助判断(覆盖低置信度结果)
|
|
|
+ if use_line_detection:
|
|
|
+ line_info = self._detect_table_lines(image)
|
|
|
+ result['line_detection'] = line_info
|
|
|
+
|
|
|
+ # 🔑 关键逻辑:只有横线没有竖线 → 强制无线表格
|
|
|
+ has_horizontal = line_info['horizontal_lines'] > 0
|
|
|
+ has_vertical = line_info['vertical_lines'] > 0
|
|
|
+
|
|
|
+ if (has_horizontal and not has_vertical) or (not has_horizontal and has_vertical):
|
|
|
+ # 只有横线或者竖线,强制判断为无线表格
|
|
|
+ if table_type != 'wireless':
|
|
|
+ logger.info(
|
|
|
+ f"📊 Line detection override: {table_type}→wireless "
|
|
|
+ f"(H={line_info['horizontal_lines']}, V={line_info['vertical_lines']})"
|
|
|
+ )
|
|
|
+ result['table_type'] = 'wireless'
|
|
|
+ result['override_reason'] = 'only_horizontal_lines' if has_horizontal else 'only_vertical_lines'
|
|
|
+ elif not has_horizontal and not has_vertical:
|
|
|
+ # 没有线条,可能是纯文本表格,判断为无线
|
|
|
+ if table_type != 'wireless':
|
|
|
+ logger.info(f"📊 Line detection override: {table_type}→wireless (no lines detected)")
|
|
|
+ result['table_type'] = 'wireless'
|
|
|
+ result['override_reason'] = 'no_lines_detected'
|
|
|
+
|
|
|
+ logger.debug(
|
|
|
+ f"Table classified as '{result['table_type']}' "
|
|
|
+ f"(confidence: {confidence:.3f}, "
|
|
|
+ f"H={result.get('line_detection', {}).get('horizontal_lines', 'N/A')}, "
|
|
|
+ f"V={result.get('line_detection', {}).get('vertical_lines', 'N/A')})"
|
|
|
+ )
|
|
|
return result
|
|
|
|
|
|
except Exception as e:
|
|
|
@@ -124,6 +158,48 @@ class PaddleTableClassifier(BaseAdapter):
|
|
|
'error': str(e)
|
|
|
}
|
|
|
|
|
|
+ def _detect_table_lines(self, image: Union[np.ndarray, Image.Image]) -> Dict[str, int]:
|
|
|
+ """
|
|
|
+ 检测表格图像中的横线和竖线数量
|
|
|
+
|
|
|
+ Args:
|
|
|
+ image: 表格图像
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ {'horizontal_lines': int, 'vertical_lines': int}
|
|
|
+ """
|
|
|
+ # 转换为numpy数组
|
|
|
+ if isinstance(image, Image.Image):
|
|
|
+ img_array = np.array(image)
|
|
|
+ else:
|
|
|
+ img_array = image.copy()
|
|
|
+
|
|
|
+ # 转换为灰度图
|
|
|
+ if len(img_array.shape) == 3:
|
|
|
+ gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
|
|
|
+ else:
|
|
|
+ gray = img_array
|
|
|
+
|
|
|
+ # 二值化
|
|
|
+ _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
|
|
|
+
|
|
|
+ h, w = binary.shape
|
|
|
+
|
|
|
+ # 检测横线
|
|
|
+ horizontal_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (max(20, w//30), 1))
|
|
|
+ horizontal_mask = cv2.morphologyEx(binary, cv2.MORPH_OPEN, horizontal_kernel)
|
|
|
+ horizontal_lines = cv2.findContours(horizontal_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0]
|
|
|
+
|
|
|
+ # 检测竖线
|
|
|
+ vertical_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, max(20, h//30)))
|
|
|
+ vertical_mask = cv2.morphologyEx(binary, cv2.MORPH_OPEN, vertical_kernel)
|
|
|
+ vertical_lines = cv2.findContours(vertical_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0]
|
|
|
+
|
|
|
+ return {
|
|
|
+ 'horizontal_lines': len(horizontal_lines),
|
|
|
+ 'vertical_lines': len(vertical_lines)
|
|
|
+ }
|
|
|
+
|
|
|
def batch_classify(
|
|
|
self,
|
|
|
images: list[Union[np.ndarray, Image.Image]]
|