Files
OCR/backend/app/services/pp_structure_enhanced.py
egg ea0dd7456c feat: implement layout preprocessing backend
Backend implementation for add-layout-preprocessing proposal:
- Add LayoutPreprocessingService with CLAHE, sharpen, binarize
- Add auto-detection: analyze_image_quality() for contrast/edge metrics
- Integrate preprocessing into OCR pipeline (analyze_layout)
- Add Preview API: POST /api/v2/tasks/{id}/preview/preprocessing
- Add config options: layout_preprocessing_mode, thresholds
- Add schemas: PreprocessingConfig, PreprocessingPreviewResponse

Preprocessing only affects layout detection input.
Original images preserved for element extraction.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>
2025-11-27 15:17:20 +08:00

626 lines
26 KiB
Python

"""
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
import gc
# Optional torch import for additional GPU memory management
try:
import torch
TORCH_AVAILABLE = True
except ImportError:
TORCH_AVAILABLE = False
import paddle
from paddleocr import PPStructureV3
from PIL import Image
import numpy as np
from app.models.unified_document import ElementType
from app.core.config import settings
from app.services.memory_manager import prediction_context
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,
'paragraph_title': ElementType.TITLE, # PP-StructureV3 block_label
'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,
preprocessed_image: Optional[Image.Image] = None
) -> Dict[str, Any]:
"""
Analyze document with full PP-StructureV3 capabilities.
Args:
image_path: Path to original image file (used for cropping extracted images)
output_dir: Optional output directory for saving extracted content
current_page: Current page number (0-based)
preprocessed_image: Optional preprocessed PIL Image for layout detection.
If provided, this is used for PP-Structure prediction,
but original image_path is still used for cropping images.
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}")
if preprocessed_image:
logger.info("Using preprocessed image for layout detection")
# Perform structure analysis with semaphore control
# This prevents OOM errors from multiple simultaneous predictions
with prediction_context(timeout=settings.service_acquire_timeout_seconds) as acquired:
if not acquired:
logger.error("Failed to acquire prediction slot (timeout), returning empty result")
return {
'has_parsing_res_list': False,
'elements': [],
'total_elements': 0,
'images': [],
'tables': [],
'element_types': {},
'error': 'Prediction slot timeout'
}
# Use preprocessed image if provided, otherwise use original path
if preprocessed_image is not None:
# Convert PIL to numpy array (BGR format for PP-Structure)
predict_input = np.array(preprocessed_image)
if len(predict_input.shape) == 3 and predict_input.shape[2] == 3:
# Convert RGB to BGR
predict_input = predict_input[:, :, ::-1]
results = self.structure_engine.predict(predict_input)
else:
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 (check both top-level and res)
if hasattr(page_result, 'json'):
result_json = page_result.json
if isinstance(result_json, dict):
# Check top-level
if 'parsing_res_list' in result_json:
parsing_res_list = result_json['parsing_res_list']
logger.info(f"Found parsing_res_list at top level with {len(parsing_res_list)} elements")
# Check inside 'res' (new structure in paddlex)
elif 'res' in result_json and isinstance(result_json['res'], dict):
if 'parsing_res_list' in result_json['res']:
parsing_res_list = result_json['res']['parsing_res_list']
logger.info(f"Found parsing_res_list inside 'res' with {len(parsing_res_list)} elements")
# Method 2: Try direct dict access (LayoutParsingResultV2 inherits from dict)
elif isinstance(page_result, dict):
if 'parsing_res_list' in page_result:
parsing_res_list = page_result['parsing_res_list']
logger.info(f"Found parsing_res_list via dict access with {len(parsing_res_list)} elements")
elif 'res' in page_result and isinstance(page_result['res'], dict):
if 'parsing_res_list' in page_result['res']:
parsing_res_list = page_result['res']['parsing_res_list']
logger.info(f"Found parsing_res_list inside page_result['res'] with {len(parsing_res_list)} elements")
# Method 3: 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 4: 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")
elif 'res' in result_dict and isinstance(result_dict['res'], dict):
if 'parsing_res_list' in result_dict['res']:
parsing_res_list = result_dict['res']['parsing_res_list']
logger.info(f"Found parsing_res_list in to_dict['res'] 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, image_path
)
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()
# Clean up GPU memory on error
try:
if TORCH_AVAILABLE and torch.cuda.is_available():
torch.cuda.empty_cache()
torch.cuda.synchronize()
if paddle.device.is_compiled_with_cuda():
paddle.device.cuda.empty_cache()
gc.collect()
except:
pass # Ignore cleanup errors
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],
source_image_path: Optional[Path] = None
) -> 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
source_image_path: Path to source image for cropping image regions
Returns:
List of processed elements with normalized structure
"""
elements = []
for idx, item in enumerate(parsing_res_list):
# Debug: log the structure of the first item
if idx == 0:
logger.info(f"First parsing_res_list item structure: {list(item.keys()) if isinstance(item, dict) else type(item)}")
logger.info(f"First parsing_res_list item sample: {str(item)[:500]}")
# Extract element type (check both 'type' and 'block_label')
element_type = item.get('type', '') or item.get('block_label', 'text')
element_type = element_type.lower()
mapped_type = self.ELEMENT_TYPE_MAPPING.get(
element_type, ElementType.TEXT
)
# Extract bbox (check multiple possible keys)
layout_bbox = (
item.get('layout_bbox', []) or
item.get('block_bbox', []) or
item.get('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
logger.warning(f"Element {idx} has invalid bbox: {layout_bbox}")
# Extract content (check multiple possible keys)
content = (
item.get('content', '') or
item.get('block_content', '') or
''
)
# Additional fallback for content in 'res' field
if not content and 'res' in item:
res = item.get('res', {})
if isinstance(res, dict):
content = res.get('content', '') or res.get('text', '')
elif isinstance(res, str):
content = res
# Content-based HTML table detection: PP-StructureV3 sometimes
# classifies tables as 'text' but returns HTML table content
html_table_content = None
if content and '<table' in content.lower():
if mapped_type == ElementType.TEXT or element_type == 'text':
logger.info(f"Element {idx}: Detected HTML table content in 'text' type, reclassifying to TABLE")
mapped_type = ElementType.TABLE
html_table_content = content # Store for later use
# 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:
# Use HTML content from content-based detection or extract from 'res'
html_content = html_table_content # From content-based detection
if not html_content and '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:
saved_path = self._save_image(item['img_path'], output_dir, element['element_id'])
if saved_path:
element['saved_path'] = saved_path
element['img_path'] = item['img_path'] # Keep original for reference
else:
logger.warning(f"Failed to save image for element {element['element_id']}")
# Crop image from source if no img_path but source image is available
elif source_image_path and output_dir and bbox != [0, 0, 0, 0]:
cropped_path = self._crop_and_save_image(
source_image_path, bbox, output_dir, element['element_id']
)
if cropped_path:
element['saved_path'] = cropped_path
element['img_path'] = cropped_path
logger.info(f"Cropped and saved image region for {element['element_id']}")
else:
logger.warning(f"Failed to crop image for element {element['element_id']}")
# 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) -> Optional[str]:
"""Save image file to output directory and return relative path.
Args:
img_path: Path to image file or image data
output_dir: Base output directory for results
element_id: Unique identifier for the element
Returns:
Relative path to saved image, or None if save failed
"""
import shutil
import numpy as np
from PIL import Image
try:
# Create imgs subdirectory
img_dir = output_dir / "imgs"
img_dir.mkdir(parents=True, exist_ok=True)
# Determine output file path
dst_path = img_dir / f"{element_id}.png"
relative_path = f"imgs/{element_id}.png"
# Handle different input types
if isinstance(img_path, str):
src_path = Path(img_path)
if src_path.exists() and src_path.is_file():
# Copy existing file
shutil.copy2(src_path, dst_path)
logger.info(f"Copied image from {src_path} to {dst_path}")
else:
logger.warning(f"Image file not found: {img_path}")
return None
elif isinstance(img_path, np.ndarray):
# Save numpy array as image
Image.fromarray(img_path).save(dst_path)
logger.info(f"Saved numpy array image to {dst_path}")
else:
logger.warning(f"Unknown image type: {type(img_path)}")
return None
# Return relative path for reference
return relative_path
except Exception as e:
logger.error(f"Failed to save image for element {element_id}: {e}")
return None
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}")
def _crop_and_save_image(
self,
source_image_path: Path,
bbox: List[float],
output_dir: Path,
element_id: str
) -> Optional[str]:
"""
Crop image region from source image and save to output directory.
Args:
source_image_path: Path to the source image
bbox: Bounding box [x1, y1, x2, y2]
output_dir: Output directory for saving cropped image
element_id: Element ID for naming
Returns:
Relative filename (not full path) to saved image, consistent with
Direct Track which stores "filename.png" that gets joined with
result_dir by pdf_generator_service.
"""
try:
from PIL import Image
# Open source image
with Image.open(source_image_path) as img:
# Ensure bbox values are integers
x1, y1, x2, y2 = [int(v) for v in bbox[:4]]
# Validate bbox
img_width, img_height = img.size
x1 = max(0, min(x1, img_width))
x2 = max(0, min(x2, img_width))
y1 = max(0, min(y1, img_height))
y2 = max(0, min(y2, img_height))
if x2 <= x1 or y2 <= y1:
logger.warning(f"Invalid bbox for cropping: {bbox}")
return None
# Crop the region
cropped = img.crop((x1, y1, x2, y2))
# Save directly to output directory (no subdirectory)
# Consistent with Direct Track which saves to output_dir directly
image_filename = f"{element_id}.png"
img_path = output_dir / image_filename
cropped.save(str(img_path), "PNG")
# Return just the filename (relative to result_dir)
# PDF generator will join with result_dir to get full path
logger.info(f"Cropped image saved: {img_path} ({x2-x1}x{y2-y1} pixels)")
return image_filename
except Exception as e:
logger.error(f"Failed to crop and save image for {element_id}: {e}")
return None