""" 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 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 ) -> 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 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' } 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 ) 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] ) -> 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): # 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 # 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: 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']}") # 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 = ' 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}")