feat: add OCR to UnifiedDocument converter for PP-StructureV3 integration
Implements the converter that transforms PP-StructureV3 OCR results into the UnifiedDocument format, enabling consistent output for both OCR and direct extraction tracks. - Create OCRToUnifiedConverter class with full element type mapping - Handle both enhanced (parsing_res_list) and standard markdown results - Support 4-point and simple bbox formats for coordinates - Establish element relationships (captions, lists, headers) - Integrate converter into OCR service dual-track processing - Update tasks.md marking section 3.3 complete 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
egg
@@ -22,10 +22,11 @@ from app.services.office_converter import OfficeConverter, OfficeConverterError
|
||||
try:
|
||||
from app.services.document_type_detector import DocumentTypeDetector, ProcessingTrackRecommendation
|
||||
from app.services.direct_extraction_engine import DirectExtractionEngine
|
||||
from app.services.ocr_to_unified_converter import OCRToUnifiedConverter
|
||||
from app.models.unified_document import (
|
||||
UnifiedDocument, UnifiedDocumentConverter, DocumentMetadata,
|
||||
UnifiedDocument, DocumentMetadata,
|
||||
ProcessingTrack, ElementType, DocumentElement, Page, Dimensions,
|
||||
BoundingBox
|
||||
BoundingBox, ProcessingInfo
|
||||
)
|
||||
DUAL_TRACK_AVAILABLE = True
|
||||
except ImportError as e:
|
||||
@@ -66,11 +67,13 @@ class OCRService:
|
||||
enable_table_detection=True,
|
||||
enable_image_extraction=True
|
||||
)
|
||||
self.ocr_to_unified_converter = OCRToUnifiedConverter()
|
||||
self.dual_track_enabled = True
|
||||
logger.info("Dual-track processing enabled")
|
||||
else:
|
||||
self.document_detector = None
|
||||
self.direct_extraction_engine = None
|
||||
self.ocr_to_unified_converter = None
|
||||
self.dual_track_enabled = False
|
||||
logger.info("Dual-track processing not available, using OCR-only mode")
|
||||
|
||||
@@ -541,6 +544,17 @@ class OCRService:
|
||||
}
|
||||
}
|
||||
|
||||
# If layout data is enhanced, add enhanced results for converter
|
||||
if layout_data and layout_data.get('enhanced'):
|
||||
result['enhanced_results'] = [{
|
||||
'elements': layout_data.get('elements', []),
|
||||
'reading_order': layout_data.get('reading_order', []),
|
||||
'element_types': layout_data.get('element_types', {}),
|
||||
'page': current_page,
|
||||
'width': ocr_width,
|
||||
'height': ocr_height
|
||||
}]
|
||||
|
||||
logger.info(
|
||||
f"OCR completed: {image_path.name} - "
|
||||
f"{len(text_regions)} regions, "
|
||||
@@ -621,7 +635,7 @@ class OCRService:
|
||||
|
||||
def analyze_layout(self, image_path: Path, output_dir: Optional[Path] = None, current_page: int = 0) -> Tuple[Optional[Dict], List[Dict]]:
|
||||
"""
|
||||
Analyze document layout using PP-StructureV3
|
||||
Analyze document layout using PP-StructureV3 with enhanced element extraction
|
||||
|
||||
Args:
|
||||
image_path: Path to image file
|
||||
@@ -634,8 +648,49 @@ class OCRService:
|
||||
try:
|
||||
structure_engine = self.get_structure_engine()
|
||||
|
||||
# Perform structure analysis using predict() method (PaddleOCR 3.x API)
|
||||
logger.info(f"Running layout analysis on {image_path.name}")
|
||||
# Try enhanced processing first
|
||||
try:
|
||||
from app.services.pp_structure_enhanced import PPStructureEnhanced
|
||||
|
||||
enhanced_processor = PPStructureEnhanced(structure_engine)
|
||||
result = enhanced_processor.analyze_with_full_structure(
|
||||
image_path, output_dir, current_page
|
||||
)
|
||||
|
||||
if result.get('has_parsing_res_list'):
|
||||
logger.info(f"Enhanced PP-StructureV3 analysis successful with {result['total_elements']} elements")
|
||||
logger.info(f"Element types found: {result.get('element_types', {})}")
|
||||
|
||||
# Convert to legacy format for compatibility
|
||||
layout_data = {
|
||||
'elements': result['elements'],
|
||||
'total_elements': result['total_elements'],
|
||||
'reading_order': result['reading_order'],
|
||||
'element_types': result.get('element_types', {}),
|
||||
'enhanced': True
|
||||
}
|
||||
|
||||
# Extract images metadata
|
||||
images_metadata = []
|
||||
for elem in result.get('images', []):
|
||||
images_metadata.append({
|
||||
'element_id': elem['element_id'],
|
||||
'type': 'image',
|
||||
'page': elem['page'],
|
||||
'bbox': elem['bbox']
|
||||
})
|
||||
|
||||
return layout_data, images_metadata
|
||||
else:
|
||||
logger.info("parsing_res_list not available, using standard processing")
|
||||
|
||||
except ImportError:
|
||||
logger.debug("Enhanced PP-StructureV3 module not available, using standard processing")
|
||||
except Exception as e:
|
||||
logger.warning(f"Enhanced processing failed, falling back to standard: {e}")
|
||||
|
||||
# Standard processing (original implementation)
|
||||
logger.info(f"Running standard layout analysis on {image_path.name}")
|
||||
results = structure_engine.predict(str(image_path))
|
||||
|
||||
layout_elements = []
|
||||
@@ -858,20 +913,12 @@ class OCRService:
|
||||
file_path, lang, detect_layout, confidence_threshold, output_dir
|
||||
)
|
||||
|
||||
# Convert OCR result to UnifiedDocument
|
||||
metadata = DocumentMetadata(
|
||||
filename=file_path.name,
|
||||
file_type=file_path.suffix,
|
||||
file_size=file_path.stat().st_size,
|
||||
created_at=start_time,
|
||||
processing_track=ProcessingTrack.OCR,
|
||||
processing_time=(datetime.now() - start_time).total_seconds(),
|
||||
language=lang
|
||||
)
|
||||
|
||||
unified_doc = UnifiedDocumentConverter.from_ocr_result(
|
||||
ocr_result, document_id, metadata
|
||||
# Convert OCR result to UnifiedDocument using the converter
|
||||
processing_time_so_far = (datetime.now() - start_time).total_seconds()
|
||||
unified_doc = self.ocr_to_unified_converter.convert(
|
||||
ocr_result, file_path, processing_time_so_far, lang
|
||||
)
|
||||
unified_doc.document_id = document_id
|
||||
|
||||
# Update processing track metadata
|
||||
unified_doc.metadata.processing_track = (
|
||||
@@ -951,11 +998,13 @@ class OCRService:
|
||||
'processing_time': 0.0,
|
||||
'pages': [],
|
||||
'layout_data': {'elements': []},
|
||||
'images_metadata': []
|
||||
'images_metadata': [],
|
||||
'enhanced_results': [] # For PP-StructureV3 enhanced results
|
||||
}
|
||||
|
||||
total_confidence = 0.0
|
||||
total_regions = 0
|
||||
has_enhanced = False
|
||||
|
||||
for page_num, result in enumerate(results):
|
||||
if result['status'] == 'success':
|
||||
@@ -971,7 +1020,21 @@ class OCRService:
|
||||
|
||||
# Collect layout data
|
||||
if result.get('layout_data'):
|
||||
for elem in result['layout_data'].get('elements', []):
|
||||
layout = result['layout_data']
|
||||
# Check if this is enhanced layout data
|
||||
if layout.get('enhanced'):
|
||||
has_enhanced = True
|
||||
# Store enhanced results separately for converter
|
||||
combined['enhanced_results'].append({
|
||||
'elements': layout.get('elements', []),
|
||||
'reading_order': layout.get('reading_order', []),
|
||||
'element_types': layout.get('element_types', {}),
|
||||
'page': page_num,
|
||||
'width': result.get('ocr_dimensions', {}).get('width', 0),
|
||||
'height': result.get('ocr_dimensions', {}).get('height', 0)
|
||||
})
|
||||
# Always collect elements for backward compatibility
|
||||
for elem in layout.get('elements', []):
|
||||
elem['page'] = page_num
|
||||
combined['layout_data']['elements'].append(elem)
|
||||
|
||||
|
||||
670
backend/app/services/ocr_to_unified_converter.py
Normal file
670
backend/app/services/ocr_to_unified_converter.py
Normal file
@@ -0,0 +1,670 @@
|
||||
"""
|
||||
OCR to UnifiedDocument Converter
|
||||
|
||||
Converts PP-StructureV3 OCR results to UnifiedDocument format, preserving
|
||||
all structure information and metadata.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Any, Union
|
||||
from datetime import datetime
|
||||
import hashlib
|
||||
|
||||
from app.models.unified_document import (
|
||||
UnifiedDocument, DocumentElement, Page, DocumentMetadata,
|
||||
BoundingBox, StyleInfo, TableData, ElementType,
|
||||
ProcessingTrack, TableCell, Dimensions
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class OCRToUnifiedConverter:
|
||||
"""
|
||||
Converter for transforming PP-StructureV3 OCR results to UnifiedDocument format.
|
||||
|
||||
This converter handles:
|
||||
- PP-StructureV3 parsing_res_list results
|
||||
- Markdown fallback results
|
||||
- Multi-page document assembly
|
||||
- Metadata preservation
|
||||
- Structure relationship mapping
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize the converter."""
|
||||
self.element_counter = 0
|
||||
|
||||
def convert(
|
||||
self,
|
||||
ocr_results: Dict[str, Any],
|
||||
file_path: Path,
|
||||
processing_time: float,
|
||||
lang: str = 'ch'
|
||||
) -> UnifiedDocument:
|
||||
"""
|
||||
Convert OCR results to UnifiedDocument.
|
||||
|
||||
Args:
|
||||
ocr_results: Raw OCR results from PP-StructureV3
|
||||
file_path: Original file path
|
||||
processing_time: Time taken for OCR processing
|
||||
lang: Language used for OCR
|
||||
|
||||
Returns:
|
||||
UnifiedDocument with all extracted information
|
||||
"""
|
||||
try:
|
||||
# Create document metadata
|
||||
metadata = self._create_metadata(file_path, processing_time, lang)
|
||||
|
||||
# Extract pages from OCR results
|
||||
pages = self._extract_pages(ocr_results)
|
||||
|
||||
# Create document ID
|
||||
document_id = self._generate_document_id(file_path)
|
||||
|
||||
# Create UnifiedDocument
|
||||
unified_doc = UnifiedDocument(
|
||||
document_id=document_id,
|
||||
metadata=metadata,
|
||||
pages=pages,
|
||||
processing_errors=ocr_results.get('errors', [])
|
||||
)
|
||||
|
||||
# Post-process to establish relationships
|
||||
self._establish_relationships(unified_doc)
|
||||
|
||||
logger.info(f"Successfully converted OCR results to UnifiedDocument: "
|
||||
f"{len(pages)} pages, {self._count_elements(pages)} elements")
|
||||
|
||||
return unified_doc
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error converting OCR results: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
|
||||
# Return minimal document with error
|
||||
return UnifiedDocument(
|
||||
document_id=self._generate_document_id(file_path),
|
||||
metadata=self._create_metadata(file_path, processing_time, lang),
|
||||
pages=[],
|
||||
processing_errors=[{
|
||||
'error': str(e),
|
||||
'type': 'conversion_error',
|
||||
'timestamp': datetime.now().isoformat()
|
||||
}]
|
||||
)
|
||||
|
||||
def _create_metadata(
|
||||
self,
|
||||
file_path: Path,
|
||||
processing_time: float,
|
||||
lang: str
|
||||
) -> DocumentMetadata:
|
||||
"""Create document metadata."""
|
||||
return DocumentMetadata(
|
||||
filename=file_path.name,
|
||||
file_type=file_path.suffix,
|
||||
file_size=file_path.stat().st_size if file_path.exists() else 0,
|
||||
created_at=datetime.now(),
|
||||
processing_track=ProcessingTrack.OCR,
|
||||
processing_time=processing_time,
|
||||
language=lang
|
||||
)
|
||||
|
||||
def _extract_pages(self, ocr_results: Dict[str, Any]) -> List[Page]:
|
||||
"""
|
||||
Extract pages from OCR results.
|
||||
|
||||
Handles both enhanced PP-StructureV3 results (with parsing_res_list)
|
||||
and traditional markdown results.
|
||||
"""
|
||||
pages = []
|
||||
|
||||
# Check if we have enhanced results from PPStructureEnhanced
|
||||
if 'enhanced_results' in ocr_results:
|
||||
pages = self._extract_from_enhanced_results(ocr_results['enhanced_results'])
|
||||
# Check for traditional layout_data structure
|
||||
elif 'layout_data' in ocr_results:
|
||||
pages = self._extract_from_layout_data(ocr_results['layout_data'])
|
||||
# Check for direct PP-StructureV3 results
|
||||
elif 'pages' in ocr_results:
|
||||
pages = self._extract_from_direct_results(ocr_results['pages'])
|
||||
else:
|
||||
logger.warning("No recognized OCR result structure found")
|
||||
|
||||
return pages
|
||||
|
||||
def _extract_from_enhanced_results(
|
||||
self,
|
||||
enhanced_results: List[Dict[str, Any]]
|
||||
) -> List[Page]:
|
||||
"""Extract pages from enhanced PP-StructureV3 results."""
|
||||
pages = []
|
||||
|
||||
for page_idx, page_result in enumerate(enhanced_results):
|
||||
elements = []
|
||||
|
||||
# Process elements from parsing_res_list
|
||||
if 'elements' in page_result:
|
||||
for elem_data in page_result['elements']:
|
||||
element = self._convert_pp3_element(elem_data, page_idx)
|
||||
if element:
|
||||
elements.append(element)
|
||||
|
||||
# Create page
|
||||
page = Page(
|
||||
page_number=page_idx + 1,
|
||||
dimensions=Dimensions(
|
||||
width=page_result.get('width', 0),
|
||||
height=page_result.get('height', 0)
|
||||
),
|
||||
elements=elements,
|
||||
metadata={'reading_order': page_result.get('reading_order', [])}
|
||||
)
|
||||
|
||||
pages.append(page)
|
||||
logger.debug(f"Extracted page {page_idx + 1} with {len(elements)} elements")
|
||||
|
||||
return pages
|
||||
|
||||
def _extract_from_layout_data(
|
||||
self,
|
||||
layout_data: Dict[str, Any]
|
||||
) -> List[Page]:
|
||||
"""Extract pages from traditional layout_data structure."""
|
||||
pages = []
|
||||
|
||||
# Get page dimensions (assuming uniform for all pages)
|
||||
page_width = layout_data.get('page_width', 0)
|
||||
page_height = layout_data.get('page_height', 0)
|
||||
|
||||
# Group elements by page
|
||||
elements_by_page = {}
|
||||
|
||||
# Process text regions
|
||||
for text_region in layout_data.get('text_regions', []):
|
||||
page_num = text_region.get('page', 1)
|
||||
if page_num not in elements_by_page:
|
||||
elements_by_page[page_num] = []
|
||||
|
||||
element = self._convert_text_region(text_region)
|
||||
if element:
|
||||
elements_by_page[page_num].append(element)
|
||||
|
||||
# Process images
|
||||
for img_meta in layout_data.get('images_metadata', []):
|
||||
page_num = img_meta.get('page', 1)
|
||||
if page_num not in elements_by_page:
|
||||
elements_by_page[page_num] = []
|
||||
|
||||
element = self._convert_image_metadata(img_meta)
|
||||
if element:
|
||||
elements_by_page[page_num].append(element)
|
||||
|
||||
# Process tables
|
||||
for table_data in layout_data.get('tables', []):
|
||||
page_num = table_data.get('page', 1)
|
||||
if page_num not in elements_by_page:
|
||||
elements_by_page[page_num] = []
|
||||
|
||||
element = self._convert_table_data(table_data)
|
||||
if element:
|
||||
elements_by_page[page_num].append(element)
|
||||
|
||||
# Create pages
|
||||
max_page = max(elements_by_page.keys()) if elements_by_page else 0
|
||||
for page_num in range(1, max_page + 1):
|
||||
elements = elements_by_page.get(page_num, [])
|
||||
|
||||
# Determine reading order based on position
|
||||
reading_order = self._calculate_reading_order(elements)
|
||||
|
||||
page = Page(
|
||||
page_number=page_num,
|
||||
dimensions=Dimensions(
|
||||
width=page_width,
|
||||
height=page_height
|
||||
),
|
||||
elements=elements,
|
||||
metadata={'reading_order': reading_order}
|
||||
)
|
||||
|
||||
pages.append(page)
|
||||
|
||||
return pages
|
||||
|
||||
def _convert_pp3_element(
|
||||
self,
|
||||
elem_data: Dict[str, Any],
|
||||
page_idx: int
|
||||
) -> Optional[DocumentElement]:
|
||||
"""Convert PP-StructureV3 element to DocumentElement."""
|
||||
try:
|
||||
# Extract bbox
|
||||
bbox_data = elem_data.get('bbox', [0, 0, 0, 0])
|
||||
bbox = BoundingBox(
|
||||
x0=float(bbox_data[0]),
|
||||
y0=float(bbox_data[1]),
|
||||
x1=float(bbox_data[2]),
|
||||
y1=float(bbox_data[3])
|
||||
)
|
||||
|
||||
# Get element type
|
||||
element_type = elem_data.get('type', ElementType.TEXT)
|
||||
if isinstance(element_type, str):
|
||||
# Convert string to ElementType if needed
|
||||
element_type = ElementType[element_type] if element_type in ElementType.__members__ else ElementType.TEXT
|
||||
|
||||
# Prepare content based on element type
|
||||
if element_type == ElementType.TABLE:
|
||||
# For tables, use TableData as content
|
||||
table_data = self._extract_table_data(elem_data)
|
||||
content = table_data if table_data else elem_data.get('content', '')
|
||||
elif element_type in [ElementType.IMAGE, ElementType.FIGURE]:
|
||||
# For images, use metadata dict as content
|
||||
content = {
|
||||
'path': elem_data.get('img_path', ''),
|
||||
'width': elem_data.get('width', 0),
|
||||
'height': elem_data.get('height', 0),
|
||||
'format': elem_data.get('format', 'unknown')
|
||||
}
|
||||
else:
|
||||
content = elem_data.get('content', '')
|
||||
|
||||
# Create element
|
||||
element = DocumentElement(
|
||||
element_id=elem_data.get('element_id', f"elem_{self.element_counter}"),
|
||||
type=element_type,
|
||||
content=content,
|
||||
bbox=bbox,
|
||||
confidence=elem_data.get('confidence', 1.0),
|
||||
metadata=elem_data.get('metadata', {})
|
||||
)
|
||||
|
||||
# Add style info if available
|
||||
if 'style' in elem_data:
|
||||
element.style = self._extract_style_info(elem_data['style'])
|
||||
|
||||
self.element_counter += 1
|
||||
return element
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to convert PP3 element: {e}")
|
||||
return None
|
||||
|
||||
def _convert_text_region(
|
||||
self,
|
||||
text_region: Dict[str, Any]
|
||||
) -> Optional[DocumentElement]:
|
||||
"""Convert text region to DocumentElement."""
|
||||
try:
|
||||
# Extract bbox (handle both formats: [[x1,y1], [x2,y1], [x2,y2], [x1,y2]] or [x0, y0, x1, y1])
|
||||
bbox_data = text_region.get('bbox', [0, 0, 0, 0])
|
||||
|
||||
if isinstance(bbox_data, list) and len(bbox_data) == 4:
|
||||
if isinstance(bbox_data[0], list):
|
||||
# 4-point format: [[x1,y1], [x2,y1], [x2,y2], [x1,y2]]
|
||||
x0 = float(bbox_data[0][0])
|
||||
y0 = float(bbox_data[0][1])
|
||||
x1 = float(bbox_data[2][0])
|
||||
y1 = float(bbox_data[2][1])
|
||||
else:
|
||||
# Simple format: [x0, y0, x1, y1]
|
||||
x0 = float(bbox_data[0])
|
||||
y0 = float(bbox_data[1])
|
||||
x1 = float(bbox_data[2])
|
||||
y1 = float(bbox_data[3])
|
||||
else:
|
||||
x0 = y0 = x1 = y1 = 0
|
||||
|
||||
bbox = BoundingBox(x0=x0, y0=y0, x1=x1, y1=y1)
|
||||
|
||||
element = DocumentElement(
|
||||
element_id=f"text_{self.element_counter}",
|
||||
type=ElementType.TEXT,
|
||||
content=text_region.get('text', ''),
|
||||
bbox=bbox,
|
||||
confidence=text_region.get('confidence', 1.0),
|
||||
metadata={'page': text_region.get('page', 1)}
|
||||
)
|
||||
|
||||
self.element_counter += 1
|
||||
return element
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to convert text region: {e}")
|
||||
return None
|
||||
|
||||
def _convert_image_metadata(
|
||||
self,
|
||||
img_meta: Dict[str, Any]
|
||||
) -> Optional[DocumentElement]:
|
||||
"""Convert image metadata to DocumentElement."""
|
||||
try:
|
||||
# Extract bbox (handle both formats)
|
||||
bbox_data = img_meta.get('bbox', [0, 0, 0, 0])
|
||||
|
||||
if isinstance(bbox_data, list) and len(bbox_data) == 4:
|
||||
if isinstance(bbox_data[0], list):
|
||||
# 4-point format: [[x1,y1], [x2,y1], [x2,y2], [x1,y2]]
|
||||
x0 = float(bbox_data[0][0])
|
||||
y0 = float(bbox_data[0][1])
|
||||
x1 = float(bbox_data[2][0])
|
||||
y1 = float(bbox_data[2][1])
|
||||
else:
|
||||
# Simple format: [x0, y0, x1, y1]
|
||||
x0 = float(bbox_data[0])
|
||||
y0 = float(bbox_data[1])
|
||||
x1 = float(bbox_data[2])
|
||||
y1 = float(bbox_data[3])
|
||||
else:
|
||||
x0 = y0 = x1 = y1 = 0
|
||||
|
||||
bbox = BoundingBox(x0=x0, y0=y0, x1=x1, y1=y1)
|
||||
|
||||
# Create image content dict
|
||||
image_content = {
|
||||
'path': img_meta.get('path', ''),
|
||||
'width': img_meta.get('width', 0),
|
||||
'height': img_meta.get('height', 0),
|
||||
'format': img_meta.get('format', 'unknown')
|
||||
}
|
||||
|
||||
element = DocumentElement(
|
||||
element_id=f"img_{self.element_counter}",
|
||||
type=ElementType.IMAGE,
|
||||
content=image_content,
|
||||
bbox=bbox,
|
||||
metadata={'page': img_meta.get('page', 1)}
|
||||
)
|
||||
|
||||
self.element_counter += 1
|
||||
return element
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to convert image metadata: {e}")
|
||||
return None
|
||||
|
||||
def _convert_table_data(
|
||||
self,
|
||||
table_dict: Dict[str, Any]
|
||||
) -> Optional[DocumentElement]:
|
||||
"""Convert table data to DocumentElement."""
|
||||
try:
|
||||
# Extract bbox
|
||||
bbox_data = table_dict.get('bbox', [0, 0, 0, 0])
|
||||
bbox = BoundingBox(
|
||||
x0=float(bbox_data[0]),
|
||||
y0=float(bbox_data[1]),
|
||||
x1=float(bbox_data[2]),
|
||||
y1=float(bbox_data[3])
|
||||
)
|
||||
|
||||
# Create table data
|
||||
table_data = TableData(
|
||||
rows=table_dict.get('rows', 0),
|
||||
columns=table_dict.get('columns', 0),
|
||||
cells=table_dict.get('cells', []),
|
||||
html=table_dict.get('html', '')
|
||||
)
|
||||
|
||||
element = DocumentElement(
|
||||
element_id=f"table_{self.element_counter}",
|
||||
type=ElementType.TABLE,
|
||||
content=table_data, # Use TableData object as content
|
||||
bbox=bbox,
|
||||
metadata={'page': table_dict.get('page', 1), 'extracted_text': table_dict.get('extracted_text', '')}
|
||||
)
|
||||
|
||||
self.element_counter += 1
|
||||
return element
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to convert table data: {e}")
|
||||
return None
|
||||
|
||||
def _extract_table_data(self, elem_data: Dict) -> Optional[TableData]:
|
||||
"""Extract table data from element."""
|
||||
try:
|
||||
html = elem_data.get('html', '')
|
||||
extracted_text = elem_data.get('extracted_text', '')
|
||||
|
||||
# Try to parse HTML to get rows and columns
|
||||
rows = 0
|
||||
columns = 0
|
||||
cells = []
|
||||
|
||||
if html:
|
||||
# Simple HTML parsing (could be enhanced with BeautifulSoup)
|
||||
rows = html.count('<tr')
|
||||
if rows > 0:
|
||||
# Estimate columns from first row
|
||||
first_row_end = html.find('</tr>')
|
||||
if first_row_end > 0:
|
||||
first_row = html[:first_row_end]
|
||||
columns = first_row.count('<td') + first_row.count('<th')
|
||||
|
||||
return TableData(
|
||||
rows=rows,
|
||||
columns=columns,
|
||||
cells=cells,
|
||||
html=html,
|
||||
extracted_text=extracted_text
|
||||
)
|
||||
except:
|
||||
return None
|
||||
|
||||
|
||||
def _extract_style_info(self, style_data: Dict) -> Optional[StyleInfo]:
|
||||
"""Extract style info from element."""
|
||||
try:
|
||||
return StyleInfo(
|
||||
font_family=style_data.get('font_family'),
|
||||
font_size=style_data.get('font_size'),
|
||||
font_weight=style_data.get('font_weight'),
|
||||
font_style=style_data.get('font_style'),
|
||||
text_color=style_data.get('text_color'),
|
||||
background_color=style_data.get('background_color'),
|
||||
alignment=style_data.get('alignment')
|
||||
)
|
||||
except:
|
||||
return None
|
||||
|
||||
def _calculate_reading_order(self, elements: List[DocumentElement]) -> List[int]:
|
||||
"""Calculate reading order based on element positions."""
|
||||
if not elements:
|
||||
return []
|
||||
|
||||
# Create indexed elements with position
|
||||
indexed_elements = []
|
||||
for i, elem in enumerate(elements):
|
||||
# Use top-left corner for sorting
|
||||
indexed_elements.append((
|
||||
i,
|
||||
elem.bbox.y1, # y coordinate (top to bottom)
|
||||
elem.bbox.x1 # x coordinate (left to right)
|
||||
))
|
||||
|
||||
# Sort by y first (top to bottom), then x (left to right)
|
||||
indexed_elements.sort(key=lambda x: (x[1], x[2]))
|
||||
|
||||
# Return the sorted indices
|
||||
return [idx for idx, _, _ in indexed_elements]
|
||||
|
||||
def _establish_relationships(self, doc: UnifiedDocument):
|
||||
"""
|
||||
Establish relationships between elements.
|
||||
|
||||
This includes:
|
||||
- Linking captions to figures/tables
|
||||
- Grouping list items
|
||||
- Identifying headers and their content
|
||||
"""
|
||||
for page in doc.pages:
|
||||
# Link captions to nearest figure/table
|
||||
self._link_captions(page.elements)
|
||||
|
||||
# Group consecutive list items
|
||||
self._group_list_items(page.elements)
|
||||
|
||||
# Link headers to content
|
||||
self._link_headers(page.elements)
|
||||
|
||||
# Update metadata based on content
|
||||
self._update_metadata(doc)
|
||||
|
||||
def _link_captions(self, elements: List[DocumentElement]):
|
||||
"""Link caption elements to their associated figures/tables."""
|
||||
captions = [e for e in elements if e.type in [ElementType.CAPTION, ElementType.TABLE_CAPTION]]
|
||||
targets = [e for e in elements if e.type in [ElementType.FIGURE, ElementType.TABLE, ElementType.IMAGE]]
|
||||
|
||||
for caption in captions:
|
||||
if not targets:
|
||||
break
|
||||
|
||||
# Find nearest target above the caption
|
||||
best_target = None
|
||||
min_distance = float('inf')
|
||||
|
||||
for target in targets:
|
||||
# Caption should be below the target
|
||||
if target.bbox.y2 <= caption.bbox.y1:
|
||||
distance = caption.bbox.y1 - target.bbox.y2
|
||||
if distance < min_distance:
|
||||
min_distance = distance
|
||||
best_target = target
|
||||
|
||||
if best_target and min_distance < 50: # Within 50 pixels
|
||||
caption.metadata['linked_to'] = best_target.element_id
|
||||
best_target.metadata['caption_id'] = caption.element_id
|
||||
|
||||
def _group_list_items(self, elements: List[DocumentElement]):
|
||||
"""Group consecutive list items."""
|
||||
list_items = [e for e in elements if e.type == ElementType.LIST_ITEM]
|
||||
|
||||
if not list_items:
|
||||
return
|
||||
|
||||
# Sort by position
|
||||
list_items.sort(key=lambda e: (e.bbox.y1, e.bbox.x1))
|
||||
|
||||
# Group consecutive items
|
||||
current_group = []
|
||||
groups = []
|
||||
|
||||
for i, item in enumerate(list_items):
|
||||
if i == 0:
|
||||
current_group = [item]
|
||||
else:
|
||||
prev_item = list_items[i-1]
|
||||
# Check if items are consecutive (similar x position, reasonable y gap)
|
||||
x_aligned = abs(item.bbox.x1 - prev_item.bbox.x1) < 20
|
||||
y_consecutive = (item.bbox.y1 - prev_item.bbox.y2) < 30
|
||||
|
||||
if x_aligned and y_consecutive:
|
||||
current_group.append(item)
|
||||
else:
|
||||
if current_group:
|
||||
groups.append(current_group)
|
||||
current_group = [item]
|
||||
|
||||
if current_group:
|
||||
groups.append(current_group)
|
||||
|
||||
# Mark groups in metadata
|
||||
for group_idx, group in enumerate(groups):
|
||||
group_id = f"list_group_{group_idx}"
|
||||
for item_idx, item in enumerate(group):
|
||||
item.metadata['list_group'] = group_id
|
||||
item.metadata['list_index'] = item_idx
|
||||
|
||||
def _link_headers(self, elements: List[DocumentElement]):
|
||||
"""Link headers to their content sections."""
|
||||
headers = [e for e in elements if e.type in [ElementType.HEADER, ElementType.TITLE]]
|
||||
|
||||
for i, header in enumerate(headers):
|
||||
# Find content between this header and the next
|
||||
next_header_y = float('inf')
|
||||
if i + 1 < len(headers):
|
||||
next_header_y = headers[i + 1].bbox.y1
|
||||
|
||||
# Find all elements between headers
|
||||
content_elements = [
|
||||
e for e in elements
|
||||
if (e.bbox.y1 > header.bbox.y2 and
|
||||
e.bbox.y1 < next_header_y and
|
||||
e.type not in [ElementType.HEADER, ElementType.TITLE])
|
||||
]
|
||||
|
||||
if content_elements:
|
||||
header.metadata['content_elements'] = [e.element_id for e in content_elements]
|
||||
for elem in content_elements:
|
||||
elem.metadata['header_id'] = header.element_id
|
||||
|
||||
def _update_metadata(self, doc: UnifiedDocument):
|
||||
"""Update document metadata based on extracted content."""
|
||||
# For now, just ensure basic metadata is present.
|
||||
# Since DocumentMetadata doesn't have all these fields,
|
||||
# we can store summary data at the document level or in processing_errors
|
||||
pass
|
||||
|
||||
def _generate_document_id(self, file_path: Path) -> str:
|
||||
"""Generate unique document ID."""
|
||||
content = f"{file_path.name}_{datetime.now().isoformat()}"
|
||||
return hashlib.md5(content.encode()).hexdigest()
|
||||
|
||||
def _detect_mime_type(self, file_path: Path) -> str:
|
||||
"""Detect MIME type of file."""
|
||||
try:
|
||||
import magic
|
||||
return magic.from_file(str(file_path), mime=True)
|
||||
except:
|
||||
# Fallback to extension-based detection
|
||||
ext = file_path.suffix.lower()
|
||||
mime_map = {
|
||||
'.pdf': 'application/pdf',
|
||||
'.png': 'image/png',
|
||||
'.jpg': 'image/jpeg',
|
||||
'.jpeg': 'image/jpeg'
|
||||
}
|
||||
return mime_map.get(ext, 'application/octet-stream')
|
||||
|
||||
def _count_elements(self, pages: List[Page]) -> int:
|
||||
"""Count total elements across all pages."""
|
||||
return sum(len(page.elements) for page in pages)
|
||||
|
||||
def _extract_from_direct_results(
|
||||
self,
|
||||
pages_data: List[Dict[str, Any]]
|
||||
) -> List[Page]:
|
||||
"""Extract pages from direct PP-StructureV3 results."""
|
||||
pages = []
|
||||
|
||||
for page_idx, page_data in enumerate(pages_data):
|
||||
elements = []
|
||||
|
||||
# Process each element in the page
|
||||
if 'elements' in page_data:
|
||||
for elem_data in page_data['elements']:
|
||||
element = self._convert_pp3_element(elem_data, page_idx)
|
||||
if element:
|
||||
elements.append(element)
|
||||
|
||||
# Create page
|
||||
page = Page(
|
||||
page_number=page_idx + 1,
|
||||
dimensions=Dimensions(
|
||||
width=page_data.get('width', 0),
|
||||
height=page_data.get('height', 0)
|
||||
),
|
||||
elements=elements,
|
||||
metadata={'reading_order': self._calculate_reading_order(elements)}
|
||||
)
|
||||
|
||||
pages.append(page)
|
||||
|
||||
return pages
|
||||
410
backend/app/services/pp_structure_enhanced.py
Normal file
410
backend/app/services/pp_structure_enhanced.py
Normal file
@@ -0,0 +1,410 @@
|
||||
"""
|
||||
Enhanced PP-StructureV3 processing with full element extraction
|
||||
|
||||
This module provides enhanced PP-StructureV3 processing that extracts all
|
||||
23 element types with their bbox coordinates and reading order.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Tuple, Any
|
||||
import json
|
||||
|
||||
from paddleocr import PPStructureV3
|
||||
from app.models.unified_document import ElementType
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PPStructureEnhanced:
|
||||
"""
|
||||
Enhanced PP-StructureV3 processor that extracts all available element types
|
||||
and structure information from parsing_res_list.
|
||||
"""
|
||||
|
||||
# Mapping from PP-StructureV3 types to our ElementType
|
||||
ELEMENT_TYPE_MAPPING = {
|
||||
'title': ElementType.TITLE,
|
||||
'text': ElementType.TEXT,
|
||||
'paragraph': ElementType.PARAGRAPH,
|
||||
'figure': ElementType.FIGURE,
|
||||
'figure_caption': ElementType.CAPTION,
|
||||
'table': ElementType.TABLE,
|
||||
'table_caption': ElementType.TABLE_CAPTION,
|
||||
'header': ElementType.HEADER,
|
||||
'footer': ElementType.FOOTER,
|
||||
'reference': ElementType.REFERENCE,
|
||||
'equation': ElementType.EQUATION,
|
||||
'formula': ElementType.FORMULA,
|
||||
'list-item': ElementType.LIST_ITEM,
|
||||
'list': ElementType.LIST,
|
||||
'code': ElementType.CODE,
|
||||
'footnote': ElementType.FOOTNOTE,
|
||||
'page-number': ElementType.PAGE_NUMBER,
|
||||
'watermark': ElementType.WATERMARK,
|
||||
'signature': ElementType.SIGNATURE,
|
||||
'stamp': ElementType.STAMP,
|
||||
'logo': ElementType.LOGO,
|
||||
'barcode': ElementType.BARCODE,
|
||||
'qr-code': ElementType.QR_CODE,
|
||||
# Default fallback
|
||||
'image': ElementType.IMAGE,
|
||||
'chart': ElementType.CHART,
|
||||
'diagram': ElementType.DIAGRAM,
|
||||
}
|
||||
|
||||
def __init__(self, structure_engine: PPStructureV3):
|
||||
"""
|
||||
Initialize with existing PP-StructureV3 engine.
|
||||
|
||||
Args:
|
||||
structure_engine: Initialized PPStructureV3 instance
|
||||
"""
|
||||
self.structure_engine = structure_engine
|
||||
|
||||
def analyze_with_full_structure(
|
||||
self,
|
||||
image_path: Path,
|
||||
output_dir: Optional[Path] = None,
|
||||
current_page: int = 0
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Analyze document with full PP-StructureV3 capabilities.
|
||||
|
||||
Args:
|
||||
image_path: Path to image file
|
||||
output_dir: Optional output directory for saving extracted content
|
||||
current_page: Current page number (0-based)
|
||||
|
||||
Returns:
|
||||
Dictionary with complete structure information including:
|
||||
- elements: List of all detected elements with types and bbox
|
||||
- reading_order: Reading order indices
|
||||
- images: Extracted images with metadata
|
||||
- tables: Extracted tables with structure
|
||||
"""
|
||||
try:
|
||||
logger.info(f"Enhanced PP-StructureV3 analysis on {image_path.name}")
|
||||
|
||||
# Perform structure analysis
|
||||
results = self.structure_engine.predict(str(image_path))
|
||||
|
||||
all_elements = []
|
||||
all_images = []
|
||||
all_tables = []
|
||||
|
||||
# Process each page result
|
||||
for page_idx, page_result in enumerate(results):
|
||||
# Try to access parsing_res_list (the complete structure)
|
||||
parsing_res_list = None
|
||||
|
||||
# Method 1: Direct access to json attribute
|
||||
if hasattr(page_result, 'json'):
|
||||
result_json = page_result.json
|
||||
if isinstance(result_json, dict) and 'parsing_res_list' in result_json:
|
||||
parsing_res_list = result_json['parsing_res_list']
|
||||
logger.info(f"Found parsing_res_list with {len(parsing_res_list)} elements")
|
||||
|
||||
# Method 2: Try to access as attribute
|
||||
elif hasattr(page_result, 'parsing_res_list'):
|
||||
parsing_res_list = page_result.parsing_res_list
|
||||
logger.info(f"Found parsing_res_list attribute with {len(parsing_res_list)} elements")
|
||||
|
||||
# Method 3: Check if result has to_dict method
|
||||
elif hasattr(page_result, 'to_dict'):
|
||||
result_dict = page_result.to_dict()
|
||||
if 'parsing_res_list' in result_dict:
|
||||
parsing_res_list = result_dict['parsing_res_list']
|
||||
logger.info(f"Found parsing_res_list in to_dict with {len(parsing_res_list)} elements")
|
||||
|
||||
# Process parsing_res_list if found
|
||||
if parsing_res_list:
|
||||
elements = self._process_parsing_res_list(
|
||||
parsing_res_list, current_page, output_dir
|
||||
)
|
||||
all_elements.extend(elements)
|
||||
|
||||
# Extract tables and images from elements
|
||||
for elem in elements:
|
||||
if elem['type'] == ElementType.TABLE:
|
||||
all_tables.append(elem)
|
||||
elif elem['type'] in [ElementType.IMAGE, ElementType.FIGURE]:
|
||||
all_images.append(elem)
|
||||
else:
|
||||
# Fallback to markdown if parsing_res_list not available
|
||||
logger.warning("parsing_res_list not found, falling back to markdown")
|
||||
elements = self._process_markdown_fallback(
|
||||
page_result, current_page, output_dir
|
||||
)
|
||||
all_elements.extend(elements)
|
||||
|
||||
# Create reading order based on element positions
|
||||
reading_order = self._determine_reading_order(all_elements)
|
||||
|
||||
return {
|
||||
'elements': all_elements,
|
||||
'total_elements': len(all_elements),
|
||||
'reading_order': reading_order,
|
||||
'tables': all_tables,
|
||||
'images': all_images,
|
||||
'element_types': self._count_element_types(all_elements),
|
||||
'has_parsing_res_list': parsing_res_list is not None
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Enhanced PP-StructureV3 analysis error: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
return {
|
||||
'elements': [],
|
||||
'total_elements': 0,
|
||||
'reading_order': [],
|
||||
'tables': [],
|
||||
'images': [],
|
||||
'element_types': {},
|
||||
'has_parsing_res_list': False,
|
||||
'error': str(e)
|
||||
}
|
||||
|
||||
def _process_parsing_res_list(
|
||||
self,
|
||||
parsing_res_list: List[Dict],
|
||||
current_page: int,
|
||||
output_dir: Optional[Path]
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Process parsing_res_list to extract all elements.
|
||||
|
||||
Args:
|
||||
parsing_res_list: List of parsed elements from PP-StructureV3
|
||||
current_page: Current page number
|
||||
output_dir: Optional output directory
|
||||
|
||||
Returns:
|
||||
List of processed elements with normalized structure
|
||||
"""
|
||||
elements = []
|
||||
|
||||
for idx, item in enumerate(parsing_res_list):
|
||||
# Extract element type
|
||||
element_type = item.get('type', 'text').lower()
|
||||
mapped_type = self.ELEMENT_TYPE_MAPPING.get(
|
||||
element_type, ElementType.TEXT
|
||||
)
|
||||
|
||||
# Extract bbox (layout_bbox has the precise coordinates)
|
||||
layout_bbox = item.get('layout_bbox', [])
|
||||
if not layout_bbox and 'bbox' in item:
|
||||
layout_bbox = item['bbox']
|
||||
|
||||
# Ensure bbox has 4 values
|
||||
if len(layout_bbox) >= 4:
|
||||
bbox = layout_bbox[:4] # [x1, y1, x2, y2]
|
||||
else:
|
||||
bbox = [0, 0, 0, 0] # Default if bbox missing
|
||||
|
||||
# Extract content
|
||||
content = item.get('content', '')
|
||||
if not content and 'res' in item:
|
||||
# Some elements have content in 'res' field
|
||||
res = item.get('res', {})
|
||||
if isinstance(res, dict):
|
||||
content = res.get('content', '') or res.get('text', '')
|
||||
elif isinstance(res, str):
|
||||
content = res
|
||||
|
||||
# Create element
|
||||
element = {
|
||||
'element_id': f"pp3_{current_page}_{idx}",
|
||||
'type': mapped_type,
|
||||
'original_type': element_type,
|
||||
'content': content,
|
||||
'page': current_page,
|
||||
'bbox': bbox, # [x1, y1, x2, y2]
|
||||
'index': idx, # Original index in reading order
|
||||
'confidence': item.get('score', 1.0)
|
||||
}
|
||||
|
||||
# Special handling for tables
|
||||
if mapped_type == ElementType.TABLE:
|
||||
# Extract table structure if available
|
||||
if 'res' in item and isinstance(item['res'], dict):
|
||||
html_content = item['res'].get('html', '')
|
||||
if html_content:
|
||||
element['html'] = html_content
|
||||
element['extracted_text'] = self._extract_text_from_html(html_content)
|
||||
|
||||
# Special handling for images/figures
|
||||
elif mapped_type in [ElementType.IMAGE, ElementType.FIGURE]:
|
||||
# Save image if path provided
|
||||
if 'img_path' in item and output_dir:
|
||||
self._save_image(item['img_path'], output_dir, element['element_id'])
|
||||
element['img_path'] = item['img_path']
|
||||
|
||||
# Add any additional metadata
|
||||
if 'metadata' in item:
|
||||
element['metadata'] = item['metadata']
|
||||
|
||||
elements.append(element)
|
||||
logger.debug(f"Processed element {idx}: type={mapped_type}, bbox={bbox}")
|
||||
|
||||
return elements
|
||||
|
||||
def _process_markdown_fallback(
|
||||
self,
|
||||
page_result: Any,
|
||||
current_page: int,
|
||||
output_dir: Optional[Path]
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Fallback to markdown processing if parsing_res_list not available.
|
||||
|
||||
Args:
|
||||
page_result: PP-StructureV3 page result
|
||||
current_page: Current page number
|
||||
output_dir: Optional output directory
|
||||
|
||||
Returns:
|
||||
List of elements extracted from markdown
|
||||
"""
|
||||
elements = []
|
||||
|
||||
# Extract from markdown if available
|
||||
if hasattr(page_result, 'markdown'):
|
||||
markdown_dict = page_result.markdown
|
||||
|
||||
if isinstance(markdown_dict, dict):
|
||||
# Extract markdown texts
|
||||
markdown_texts = markdown_dict.get('markdown_texts', '')
|
||||
if markdown_texts:
|
||||
# Detect if it's a table
|
||||
is_table = '<table' in markdown_texts.lower()
|
||||
|
||||
element = {
|
||||
'element_id': f"md_{current_page}_0",
|
||||
'type': ElementType.TABLE if is_table else ElementType.TEXT,
|
||||
'content': markdown_texts,
|
||||
'page': current_page,
|
||||
'bbox': [0, 0, 0, 0], # No bbox in markdown
|
||||
'index': 0,
|
||||
'from_markdown': True
|
||||
}
|
||||
|
||||
if is_table:
|
||||
element['extracted_text'] = self._extract_text_from_html(markdown_texts)
|
||||
|
||||
elements.append(element)
|
||||
|
||||
# Process images
|
||||
markdown_images = markdown_dict.get('markdown_images', {})
|
||||
for img_idx, (img_path, img_obj) in enumerate(markdown_images.items()):
|
||||
# Save image
|
||||
if output_dir and hasattr(img_obj, 'save'):
|
||||
self._save_pil_image(img_obj, output_dir, f"md_img_{current_page}_{img_idx}")
|
||||
|
||||
# Try to extract bbox from filename
|
||||
bbox = self._extract_bbox_from_filename(img_path)
|
||||
|
||||
element = {
|
||||
'element_id': f"md_img_{current_page}_{img_idx}",
|
||||
'type': ElementType.IMAGE,
|
||||
'content': img_path,
|
||||
'page': current_page,
|
||||
'bbox': bbox,
|
||||
'index': img_idx + 1,
|
||||
'from_markdown': True
|
||||
}
|
||||
elements.append(element)
|
||||
|
||||
return elements
|
||||
|
||||
def _determine_reading_order(self, elements: List[Dict]) -> List[int]:
|
||||
"""
|
||||
Determine reading order based on element positions.
|
||||
|
||||
Args:
|
||||
elements: List of elements with bbox
|
||||
|
||||
Returns:
|
||||
List of indices representing reading order
|
||||
"""
|
||||
if not elements:
|
||||
return []
|
||||
|
||||
# If elements have original indices, use them
|
||||
if all('index' in elem for elem in elements):
|
||||
# Sort by original index
|
||||
indexed_elements = [(i, elem['index']) for i, elem in enumerate(elements)]
|
||||
indexed_elements.sort(key=lambda x: x[1])
|
||||
return [i for i, _ in indexed_elements]
|
||||
|
||||
# Otherwise, sort by position (top to bottom, left to right)
|
||||
indexed_elements = []
|
||||
for i, elem in enumerate(elements):
|
||||
bbox = elem.get('bbox', [0, 0, 0, 0])
|
||||
if len(bbox) >= 2:
|
||||
# Use top-left corner for sorting
|
||||
indexed_elements.append((i, bbox[1], bbox[0])) # (index, y, x)
|
||||
else:
|
||||
indexed_elements.append((i, 0, 0))
|
||||
|
||||
# Sort by y first (top to bottom), then x (left to right)
|
||||
indexed_elements.sort(key=lambda x: (x[1], x[2]))
|
||||
|
||||
return [i for i, _, _ in indexed_elements]
|
||||
|
||||
def _count_element_types(self, elements: List[Dict]) -> Dict[str, int]:
|
||||
"""
|
||||
Count occurrences of each element type.
|
||||
|
||||
Args:
|
||||
elements: List of elements
|
||||
|
||||
Returns:
|
||||
Dictionary with element type counts
|
||||
"""
|
||||
type_counts = {}
|
||||
for elem in elements:
|
||||
elem_type = elem.get('type', ElementType.TEXT)
|
||||
type_counts[elem_type] = type_counts.get(elem_type, 0) + 1
|
||||
return type_counts
|
||||
|
||||
def _extract_text_from_html(self, html: str) -> str:
|
||||
"""Extract plain text from HTML content."""
|
||||
try:
|
||||
from bs4 import BeautifulSoup
|
||||
soup = BeautifulSoup(html, 'html.parser')
|
||||
return soup.get_text(separator=' ', strip=True)
|
||||
except:
|
||||
# Fallback: just remove HTML tags
|
||||
import re
|
||||
text = re.sub(r'<[^>]+>', ' ', html)
|
||||
text = re.sub(r'\s+', ' ', text)
|
||||
return text.strip()
|
||||
|
||||
def _extract_bbox_from_filename(self, filename: str) -> List[int]:
|
||||
"""Extract bbox from filename if it contains coordinate information."""
|
||||
import re
|
||||
match = re.search(r'box_(\d+)_(\d+)_(\d+)_(\d+)', filename)
|
||||
if match:
|
||||
return list(map(int, match.groups()))
|
||||
return [0, 0, 0, 0]
|
||||
|
||||
def _save_image(self, img_path: str, output_dir: Path, element_id: str):
|
||||
"""Save image file to output directory."""
|
||||
try:
|
||||
# Implementation depends on how images are provided
|
||||
pass
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to save image {img_path}: {e}")
|
||||
|
||||
def _save_pil_image(self, img_obj, output_dir: Path, element_id: str):
|
||||
"""Save PIL image object to output directory."""
|
||||
try:
|
||||
img_dir = output_dir / "imgs"
|
||||
img_dir.mkdir(parents=True, exist_ok=True)
|
||||
img_path = img_dir / f"{element_id}.png"
|
||||
img_obj.save(str(img_path))
|
||||
logger.info(f"Saved image to {img_path}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to save PIL image: {e}")
|
||||
Reference in New Issue
Block a user