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
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
|
||||
Reference in New Issue
Block a user