Spaces:
Sleeping
Sleeping
| import asyncio | |
| import base64 | |
| import io | |
| import logging | |
| import os | |
| from pathlib import Path | |
| import tempfile | |
| from typing import Any, Dict, List, Optional, Tuple, Union | |
| # Optional PIL import for image processing | |
| try: | |
| from PIL import Image | |
| PIL_AVAILABLE = True | |
| except ImportError: | |
| Image = None | |
| PIL_AVAILABLE = False | |
| logger = logging.getLogger(__name__) | |
| logger.warning("PIL (Pillow) not available - image processing features will be limited") | |
| try: | |
| import PyPDF2 | |
| except ImportError: | |
| PyPDF2 = None | |
| # Optional numpy import | |
| try: | |
| import numpy as np | |
| NUMPY_AVAILABLE = True | |
| except ImportError: | |
| np = None | |
| NUMPY_AVAILABLE = False | |
| # LLM Service Integration | |
| try: | |
| from core.llm_service import LLMService | |
| LLM_SERVICE_AVAILABLE = True | |
| except ImportError: | |
| LLM_SERVICE_AVAILABLE = False | |
| logger = logging.getLogger(__name__) | |
| class PDFOCRService: | |
| """ | |
| Enhanced PDF processing service with OCR capabilities and fallback mechanisms. | |
| Supports both searchable PDFs and scanned/image-based PDFs. | |
| """ | |
| def __init__( | |
| self, | |
| tesseract_path: Optional[str] = None, | |
| easyocr_languages: List[str] = None, | |
| tenant_id: str = "default", | |
| ): | |
| """ | |
| Initialize the PDF OCR service. | |
| Args: | |
| tesseract_path: Path to tesseract executable (if not in PATH) | |
| easyocr_languages: List of languages for EasyOCR (default: ['en']) | |
| tenant_id: Tenant ID for metered AI operations | |
| """ | |
| self.tesseract_path = tesseract_path | |
| self.easyocr_languages = easyocr_languages or ["en"] | |
| self.tenant_id = tenant_id | |
| # Use unified LLMService for AI vision/OCR | |
| self.llm_service = None | |
| if LLM_SERVICE_AVAILABLE: | |
| self.llm_service = LLMService(tenant_id=tenant_id) | |
| logger.info(f"PDF OCR Service initialized with LLMService for tenant: {tenant_id}") | |
| # Initialize OCR readers | |
| self._init_ocr_readers() | |
| # Service availability flags | |
| self.service_status = self._check_service_availability() | |
| logger.info(f"PDF OCR Service initialized - Status: {self.service_status}") | |
| def _init_ocr_readers(self): | |
| """Initialize OCR readers based on available libraries.""" | |
| self.ocr_readers = {} | |
| # Docling (highest priority - advanced document understanding with OCR) | |
| if DOCLING_AVAILABLE: | |
| try: | |
| self.ocr_readers["docling"] = get_docling_processor() | |
| logger.info("Docling document processor initialized (highest priority)") | |
| except Exception as e: | |
| logger.warning(f"Failed to initialize Docling: {e}") | |
| # Tesseract OCR | |
| if TESSERACT_AVAILABLE: | |
| try: | |
| if self.tesseract_path: | |
| pytesseract.pytesseract.tesseract_cmd = self.tesseract_path | |
| self.ocr_readers["tesseract"] = pytesseract | |
| logger.info("Tesseract OCR initialized") | |
| except Exception as e: | |
| logger.warning(f"Failed to initialize Tesseract: {e}") | |
| # EasyOCR | |
| if EASYOCR_AVAILABLE: | |
| try: | |
| self.ocr_readers["easyocr"] = easyocr.Reader(self.easyocr_languages) | |
| logger.info("EasyOCR initialized") | |
| except Exception as e: | |
| logger.warning(f"Failed to initialize EasyOCR: {e}") | |
| # AI Vision (for advanced image comprehension) - Using unified LLMService | |
| if self.llm_service: | |
| self.ocr_readers["ai_vision"] = self.llm_service | |
| logger.info("AI Vision (LLMService) initialized") | |
| def _check_service_availability(self) -> Dict[str, bool]: | |
| """Check availability of different OCR services.""" | |
| status = { | |
| "basic_pdf": True, # Always available (PyPDF2) | |
| "docling": "docling" in self.ocr_readers, | |
| "tesseract": "tesseract" in self.ocr_readers, | |
| "easyocr": "easyocr" in self.ocr_readers, | |
| "openai_vision": "ai_vision" in self.ocr_readers, | |
| "fallback_available": len(self.ocr_readers) > 0, | |
| "byok_integrated": self.llm_service is not None, | |
| } | |
| return status | |
| async def process_pdf( | |
| self, | |
| pdf_data: Union[bytes, str, Path], | |
| use_ocr: bool = True, | |
| extract_images: bool = True, | |
| use_advanced_comprehension: bool = False, | |
| fallback_strategy: str = "cascade", | |
| ) -> Dict[str, Any]: | |
| """ | |
| Process PDF with optional OCR and image comprehension. | |
| Args: | |
| pdf_data: PDF file as bytes, file path, or Path object | |
| use_ocr: Whether to use OCR for scanned PDFs | |
| extract_images: Whether to extract and process images | |
| use_advanced_comprehension: Whether to use AI for image understanding | |
| fallback_strategy: "cascade" (try best first) or "parallel" (try all) | |
| Returns: | |
| Dictionary with extracted text, metadata, and processing results | |
| """ | |
| try: | |
| # Convert input to bytes if it's a file path | |
| if isinstance(pdf_data, (str, Path)): | |
| with open(pdf_data, "rb") as f: | |
| pdf_data = f.read() | |
| # Step 1: Try basic text extraction first | |
| basic_result = await self._extract_basic_text(pdf_data) | |
| # Step 2: Check if we need OCR (low text content or specific request) | |
| needs_ocr = use_ocr and ( | |
| basic_result["text_ratio"] < 0.1 | |
| or len(basic_result["extracted_text"].strip()) < 100 | |
| ) | |
| # Step 3: Process with OCR if needed | |
| ocr_result = None | |
| if needs_ocr: | |
| ocr_result = await self._process_with_ocr( | |
| pdf_data, fallback_strategy, use_advanced_comprehension | |
| ) | |
| # Step 4: Extract and process images if requested | |
| image_results = None | |
| if extract_images: | |
| image_results = await self._extract_and_process_images( | |
| pdf_data, use_advanced_comprehension | |
| ) | |
| # Combine results | |
| final_result = self._combine_results( | |
| basic_result, ocr_result, image_results, needs_ocr | |
| ) | |
| return final_result | |
| except Exception as e: | |
| logger.error(f"PDF processing failed: {e}") | |
| return self._create_error_result(str(e)) | |
| async def _extract_basic_text(self, pdf_data: bytes) -> Dict[str, Any]: | |
| """Extract text using basic PyPDF2 method.""" | |
| try: | |
| pdf_file = io.BytesIO(pdf_data) | |
| pdf_reader = PyPDF2.PdfReader(pdf_file) | |
| text_content = [] | |
| total_chars = 0 | |
| page_count = len(pdf_reader.pages) | |
| for page_num, page in enumerate(pdf_reader.pages): | |
| page_text = page.extract_text() | |
| text_content.append( | |
| { | |
| "page": page_num + 1, | |
| "text": page_text, | |
| "char_count": len(page_text), | |
| } | |
| ) | |
| total_chars += len(page_text) | |
| # Calculate text ratio (rough estimate of searchable content) | |
| text_ratio = min(total_chars / (page_count * 1000), 1.0) # Normalize | |
| return { | |
| "method": "basic_pdf", | |
| "extracted_text": "\n".join([p["text"] for p in text_content]), | |
| "page_texts": text_content, | |
| "page_count": page_count, | |
| "total_chars": total_chars, | |
| "text_ratio": text_ratio, | |
| "success": True, | |
| } | |
| except Exception as e: | |
| logger.warning(f"Basic text extraction failed: {e}") | |
| return { | |
| "method": "basic_pdf", | |
| "extracted_text": "", | |
| "page_texts": [], | |
| "page_count": 0, | |
| "total_chars": 0, | |
| "text_ratio": 0.0, | |
| "success": False, | |
| "error": str(e), | |
| } | |
| async def _process_with_ocr( | |
| self, pdf_data: bytes, fallback_strategy: str, use_advanced_comprehension: bool | |
| ) -> Dict[str, Any]: | |
| """Process PDF using OCR with fallback strategy.""" | |
| methods_tried = [] | |
| best_result = None | |
| ocr_methods = self._get_available_ocr_methods(use_advanced_comprehension) | |
| if fallback_strategy == "cascade": | |
| # Try methods in order of preference | |
| for method_name in ocr_methods: | |
| try: | |
| logger.info(f"Trying OCR method: {method_name}") | |
| result = await self._run_ocr_method(method_name, pdf_data) | |
| methods_tried.append(method_name) | |
| if result["success"] and result["total_chars"] > 0: | |
| best_result = result | |
| break # Found good result, stop trying | |
| except Exception as e: | |
| logger.warning(f"OCR method {method_name} failed: {e}") | |
| methods_tried.append(f"{method_name}_failed") | |
| elif fallback_strategy == "parallel": | |
| # Try all methods and pick the best | |
| results = [] | |
| for method_name in ocr_methods: | |
| try: | |
| result = await self._run_ocr_method(method_name, pdf_data) | |
| methods_tried.append(method_name) | |
| if result["success"]: | |
| results.append(result) | |
| except Exception as e: | |
| logger.warning(f"OCR method {method_name} failed: {e}") | |
| methods_tried.append(f"{method_name}_failed") | |
| if results: | |
| # Pick result with most text | |
| best_result = max(results, key=lambda x: x["total_chars"]) | |
| return { | |
| "best_result": best_result, | |
| "methods_tried": methods_tried, | |
| "success": best_result is not None, | |
| } | |
| def _get_available_ocr_methods(self, use_advanced_comprehension: bool) -> List[str]: | |
| """Get available OCR methods in priority order.""" | |
| methods = [] | |
| # Use BYOK optimization if available | |
| if self.use_byok and use_advanced_comprehension: | |
| try: | |
| optimal_provider = self.byok_manager.get_optimal_provider( | |
| "image_comprehension" | |
| ) | |
| if optimal_provider == "openai" and "openai" in self.ocr_readers: | |
| methods.append("openai_vision") | |
| logger.info( | |
| f"BYOK selected {optimal_provider} for image comprehension" | |
| ) | |
| except Exception as e: | |
| logger.warning(f"BYOK optimization failed: {e}") | |
| # Docling first (highest priority - best OCR with layout analysis) | |
| if "docling" in self.ocr_readers: | |
| methods.insert(0, "docling") | |
| # Fallback to default logic if BYOK not available or failed | |
| if not methods and use_advanced_comprehension and "openai" in self.ocr_readers: | |
| methods.append("openai_vision") | |
| # Then standard OCR methods | |
| if "easyocr" in self.ocr_readers: | |
| methods.append("easyocr") | |
| if "tesseract" in self.ocr_readers: | |
| methods.append("tesseract") | |
| return methods | |
| async def _run_ocr_method( | |
| self, method_name: str, pdf_data: bytes | |
| ) -> Dict[str, Any]: | |
| """Run specific OCR method on PDF.""" | |
| if method_name == "docling": | |
| return await self._ocr_with_docling(pdf_data) | |
| elif method_name == "tesseract": | |
| return await self._ocr_with_tesseract(pdf_data) | |
| elif method_name == "easyocr": | |
| return await self._ocr_with_easyocr(pdf_data) | |
| elif method_name in ["openai_vision", "ai_vision"]: | |
| return await self._ocr_with_ai_vision(pdf_data) | |
| else: | |
| raise ValueError(f"Unknown OCR method: {method_name}") | |
| async def _ocr_with_docling(self, pdf_data: bytes) -> Dict[str, Any]: | |
| """Extract text using Docling with advanced OCR and layout analysis.""" | |
| if "docling" not in self.ocr_readers: | |
| raise RuntimeError("Docling not available") | |
| try: | |
| processor = self.ocr_readers["docling"] | |
| result = await processor.process_pdf(pdf_data, use_ocr=True) | |
| if result.get("success"): | |
| return { | |
| "method": "docling", | |
| "extracted_text": result.get("extracted_text", ""), | |
| "page_texts": result.get("page_texts", []), | |
| "page_count": result.get("page_count", 0), | |
| "total_chars": result.get("total_chars", 0), | |
| "tables": result.get("tables", []), | |
| "success": True, | |
| } | |
| else: | |
| raise RuntimeError(result.get("error", "Docling processing failed")) | |
| except Exception as e: | |
| logger.error(f"Docling OCR failed: {e}") | |
| return { | |
| "method": "docling", | |
| "extracted_text": "", | |
| "page_texts": [], | |
| "page_count": 0, | |
| "total_chars": 0, | |
| "success": False, | |
| "error": str(e), | |
| } | |
| async def _ocr_with_tesseract(self, pdf_data: bytes) -> Dict[str, Any]: | |
| """Extract text using Tesseract OCR.""" | |
| if "tesseract" not in self.ocr_readers: | |
| raise RuntimeError("Tesseract not available") | |
| try: | |
| # Convert PDF to images for OCR | |
| images = await self._pdf_to_images(pdf_data) | |
| text_content = [] | |
| total_chars = 0 | |
| for page_num, image in enumerate(images): | |
| # Convert PIL image to format tesseract expects | |
| text = pytesseract.image_to_string(image) | |
| text_content.append( | |
| {"page": page_num + 1, "text": text, "char_count": len(text)} | |
| ) | |
| total_chars += len(text) | |
| return { | |
| "method": "tesseract", | |
| "extracted_text": "\n".join([p["text"] for p in text_content]), | |
| "page_texts": text_content, | |
| "page_count": len(images), | |
| "total_chars": total_chars, | |
| "success": True, | |
| } | |
| except Exception as e: | |
| logger.error(f"Tesseract OCR failed: {e}") | |
| return { | |
| "method": "tesseract", | |
| "extracted_text": "", | |
| "page_texts": [], | |
| "page_count": 0, | |
| "total_chars": 0, | |
| "success": False, | |
| "error": str(e), | |
| } | |
| async def _ocr_with_easyocr(self, pdf_data: bytes) -> Dict[str, Any]: | |
| """Extract text using EasyOCR.""" | |
| if "easyocr" not in self.ocr_readers: | |
| raise RuntimeError("EasyOCR not available") | |
| try: | |
| images = await self._pdf_to_images(pdf_data) | |
| text_content = [] | |
| total_chars = 0 | |
| for page_num, image in enumerate(images): | |
| # Convert PIL image to numpy array (if numpy is available) | |
| if not NUMPY_AVAILABLE: | |
| logger.error("NumPy is required for EasyOCR but not available") | |
| raise ImportError("NumPy is required for EasyOCR") | |
| image_np = np.array(image) | |
| # Run OCR | |
| results = self.ocr_readers["easyocr"].readtext(image_np) | |
| # Combine text from all detections | |
| page_text = " ".join([result[1] for result in results]) | |
| text_content.append( | |
| { | |
| "page": page_num + 1, | |
| "text": page_text, | |
| "char_count": len(page_text), | |
| } | |
| ) | |
| total_chars += len(page_text) | |
| return { | |
| "method": "easyocr", | |
| "extracted_text": "\n".join([p["text"] for p in text_content]), | |
| "page_texts": text_content, | |
| "page_count": len(images), | |
| "total_chars": total_chars, | |
| "success": True, | |
| } | |
| except Exception as e: | |
| logger.error(f"EasyOCR failed: {e}") | |
| return { | |
| "method": "easyocr", | |
| "extracted_text": "", | |
| "page_texts": [], | |
| "page_count": 0, | |
| "total_chars": 0, | |
| "success": False, | |
| "error": str(e), | |
| } | |
| async def _ocr_with_ai_vision(self, pdf_data: bytes) -> Dict[str, Any]: | |
| """Extract text and comprehend images using AI Vision via unified LLMService.""" | |
| if not self.llm_service: | |
| raise RuntimeError("AI Vision (LLMService) not available") | |
| import base64 | |
| try: | |
| images = await self._pdf_to_images(pdf_data) | |
| text_content = [] | |
| total_chars = 0 | |
| image_descriptions = [] | |
| for page_num, image in enumerate(images): | |
| # Convert PIL image to bytes for API | |
| img_byte_arr = io.BytesIO() | |
| image.save(img_byte_arr, format="PNG") | |
| img_byte_arr = img_byte_arr.getvalue() | |
| # Use LLMService for multimodal vision processing | |
| response_data = await self.llm_service.generate_completion( | |
| messages=[ | |
| { | |
| "role": "user", | |
| "content": [ | |
| { | |
| "type": "text", | |
| "text": "Extract all text from this image and describe any visual elements that might be important for understanding the document.", | |
| }, | |
| { | |
| "type": "image_url", | |
| "image_url": { | |
| "url": f"data:image/png;base64,{base64.b64encode(img_byte_arr).decode('utf-8')}" | |
| }, | |
| }, | |
| ], | |
| } | |
| ], | |
| model="auto", # Multimodal capability resolved internally | |
| tenant_id=self.tenant_id | |
| ) | |
| if response_data.get("success"): | |
| page_text = response_data.get("content", "").strip() | |
| else: | |
| logger.warning(f"AI Vision failed for page {page_num + 1}: {response_data.get('error')}") | |
| page_text = "" | |
| text_content.append( | |
| { | |
| "page": page_num + 1, | |
| "text": page_text, | |
| "char_count": len(page_text), | |
| } | |
| ) | |
| total_chars += len(page_text) | |
| image_descriptions.append( | |
| {"page": page_num + 1, "description": page_text} | |
| ) | |
| return { | |
| "method": "openai_vision", # Keeping method name for consistency in logs/UI | |
| "extracted_text": "\n".join([p["text"] for p in text_content]), | |
| "page_texts": text_content, | |
| "page_count": len(images), | |
| "total_chars": total_chars, | |
| "image_descriptions": image_descriptions, | |
| "success": True, | |
| } | |
| except Exception as e: | |
| logger.error(f"Unified AI Vision failed: {e}") | |
| return { | |
| "method": "openai_vision", | |
| "extracted_text": "", | |
| "page_texts": [], | |
| "page_count": 0, | |
| "total_chars": 0, | |
| "success": False, | |
| "error": str(e), | |
| } | |
| # BYOK Integration Methods | |
| def _get_openai_api_key(self) -> Optional[str]: | |
| """Get OpenAI API key from BYOK system or fallback.""" | |
| if self.use_byok and self.byok_manager: | |
| try: | |
| # Try to get API key from BYOK system | |
| api_key = self.byok_manager.get_api_key("openai") | |
| if api_key: | |
| return api_key | |
| except Exception as e: | |
| logger.warning(f"Failed to get OpenAI key from BYOK: {e}") | |
| # Fallback to constructor parameter or environment variable | |
| if self.openai_api_key: | |
| return self.openai_api_key | |
| return os.getenv("OPENAI_API_KEY") | |
| async def _optimize_provider_selection( | |
| self, use_advanced_comprehension: bool, fallback_strategy: str | |
| ) -> Dict[str, Any]: | |
| """Optimize provider selection using BYOK system.""" | |
| if not self.use_byok or not self.byok_manager: | |
| return {"optimized": False, "reason": "BYOK not available"} | |
| try: | |
| # Determine task type based on requirements | |
| if use_advanced_comprehension: | |
| task_type = "image_comprehension" | |
| else: | |
| task_type = "pdf_ocr" | |
| # Get optimal provider | |
| optimal_provider = self.byok_manager.get_optimal_provider(task_type) | |
| return { | |
| "optimized": True, | |
| "task_type": task_type, | |
| "optimal_provider": optimal_provider, | |
| "fallback_strategy": fallback_strategy, | |
| "available_providers": list(self.ocr_readers.keys()), | |
| } | |
| except Exception as e: | |
| logger.error(f"Provider optimization failed: {e}") | |
| return {"optimized": False, "error": str(e)} | |
| async def _track_byok_usage( | |
| self, ocr_result: Dict[str, Any], use_advanced_comprehension: bool | |
| ): | |
| """Track usage with BYOK system.""" | |
| if not self.use_byok or not self.byok_manager: | |
| return | |
| try: | |
| # Determine which provider was used | |
| best_method = ocr_result.get("best_result", {}).get("method", "") | |
| provider_id = self._map_method_to_provider(best_method) | |
| if not provider_id: | |
| return | |
| # Estimate tokens used (rough calculation) | |
| total_chars = ocr_result.get("best_result", {}).get("total_chars", 0) | |
| estimated_tokens = max( | |
| total_chars // 4, 100 | |
| ) # Rough estimate: ~4 chars per token | |
| # Track successful usage | |
| self.byok_manager.track_usage( | |
| provider_id=provider_id, success=True, tokens_used=estimated_tokens | |
| ) | |
| logger.debug( | |
| f"Tracked BYOK usage: {provider_id}, {estimated_tokens} tokens" | |
| ) | |
| except Exception as e: | |
| logger.warning(f"Failed to track BYOK usage: {e}") | |
| def _map_method_to_provider(self, method: str) -> Optional[str]: | |
| """Map OCR method to BYOK provider ID.""" | |
| method_to_provider = { | |
| "openai_vision": "openai", | |
| "tesseract": "openai", # Tesseract doesn't have BYOK provider, map to default | |
| "easyocr": "openai", # EasyOCR doesn't have BYOK provider, map to default | |
| "basic_pdf": None, # No BYOK tracking for basic extraction | |
| } | |
| return method_to_provider.get(method) | |
| async def _pdf_to_images(self, pdf_data: bytes) -> List[Image.Image]: | |
| """Convert PDF to list of PIL Images.""" | |
| try: | |
| # Try using pdf2image if available (best quality) | |
| try: | |
| from pdf2image import convert_from_bytes | |
| logger.debug("Using pdf2image for PDF to image conversion") | |
| # Convert PDF to list of images at 200 DPI for OCR quality | |
| images = await asyncio.to_thread( | |
| convert_from_bytes, pdf_data, dpi=200, fmt="jpeg" | |
| ) | |
| logger.info(f"Converted {len(images)} pages using pdf2image") | |
| return images | |
| except ImportError: | |
| logger.warning( | |
| "pdf2image not available, using fallback method. " | |
| "Install with: pip install pdf2image" | |
| ) | |
| # Fallback: Try to render PDF pages using PyMuPDF (fitz) if available | |
| try: | |
| import fitz | |
| logger.debug("Using PyMuPDF (fitz) for PDF to image conversion") | |
| pdf_document = fitz.open(stream=pdf_data, filetype="pdf") | |
| images = [] | |
| for page_num in range(pdf_document.page_count): | |
| page = pdf_document[page_num] | |
| # Render page to pixmap (zoom=2 for better quality) | |
| pix = page.get_pixmap(matrix=fitz.Matrix(2, 2)) | |
| img_data = pix.tobytes("jpeg") | |
| img = Image.open(io.BytesIO(img_data)) | |
| images.append(img) | |
| pdf_document.close() | |
| logger.info(f"Converted {len(images)} pages using PyMuPDF") | |
| return images | |
| except ImportError: | |
| logger.warning( | |
| "PyMuPDF not available. Install with: pip install PyMuPDF" | |
| ) | |
| # Final fallback: Create placeholder images with page info | |
| pdf_file = io.BytesIO(pdf_data) | |
| pdf_reader = PyPDF2.PdfReader(pdf_file) | |
| page_count = len(pdf_reader.pages) | |
| logger.warning( | |
| f"Using placeholder images for {page_count} pages. " | |
| "Install pdf2image or PyMuPDF for proper conversion." | |
| ) | |
| images = [] | |
| for i, page in enumerate(pdf_reader.pages): | |
| # Try to extract page dimensions | |
| try: | |
| mediabox = page.mediabox | |
| width = int(mediabox.width) | |
| height = int(mediabox.height) | |
| # Limit max size to avoid memory issues | |
| width = min(width, 2000) | |
| height = min(height, 2000) | |
| except (AttributeError, ValueError, TypeError) as e: | |
| logger.debug(f"Could not extract page dimensions, using defaults: {e}") | |
| width, height = 800, 1000 | |
| except Exception as e: | |
| logger.warning(f"Unexpected error extracting page dimensions: {e}", exc_info=True) | |
| width, height = 800, 1000 | |
| # Create a white image with extracted text overlay if possible | |
| img = Image.new("RGB", (width, height), color="white") | |
| # Try to extract text and add to image (basic rendering) | |
| try: | |
| text = page.extract_text() | |
| if text and text.strip(): | |
| from PIL import ImageDraw, ImageFont | |
| draw = ImageDraw.Draw(img) | |
| # Use default font | |
| try: | |
| font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 12) | |
| except (IOError, OSError) as e: | |
| logger.debug(f"Custom font not available, using default: {e}") | |
| font = ImageFont.load_default() | |
| except Exception as e: | |
| logger.warning(f"Unexpected error loading font: {e}", exc_info=True) | |
| font = ImageFont.load_default() | |
| # Draw text (first 500 chars to avoid overflow) | |
| lines = text[:500].split("\n") | |
| y_offset = 20 | |
| for line in lines[:30]: # Max 30 lines | |
| if line.strip(): | |
| draw.text((20, y_offset), line, fill="black", font=font) | |
| y_offset += 20 | |
| except Exception as e: | |
| logger.debug(f"Could not add text to placeholder image: {e}") | |
| images.append(img) | |
| return images | |
| except Exception as e: | |
| logger.error(f"PDF to image conversion failed: {e}") | |
| return [] | |
| async def _extract_and_process_images( | |
| self, pdf_data: bytes, use_advanced_comprehension: bool | |
| ) -> Dict[str, Any]: | |
| """Extract and process images from PDF.""" | |
| try: | |
| images_found = 0 | |
| image_descriptions = [] | |
| # Try using PyMuPDF (fitz) which has excellent image extraction | |
| try: | |
| import fitz | |
| logger.debug("Using PyMuPDF for image extraction") | |
| pdf_document = fitz.open(stream=pdf_data, filetype="pdf") | |
| for page_num in range(pdf_document.page_count): | |
| page = pdf_document[page_num] | |
| image_list = page.get_images(full=True) | |
| for img_index, img in enumerate(image_list): | |
| xref = img[0] | |
| base_image = pdf_document.extract_image(xref) | |
| if base_image: | |
| images_found += 1 | |
| image_info = { | |
| "page": page_num + 1, | |
| "index": img_index, | |
| "format": base_image.get("ext", "unknown"), | |
| "width": base_image.get("width", 0), | |
| "height": base_image.get("height", 0), | |
| "size_bytes": len(base_image.get("image", b"")), | |
| } | |
| # Basic description based on dimensions | |
| if base_image.get("width", 0) > 500: | |
| image_info["description"] = "Large image (possibly photo or chart)" | |
| elif base_image.get("width", 0) > 200: | |
| image_info["description"] = "Medium image (possibly icon or diagram)" | |
| else: | |
| image_info["description"] = "Small image (possibly icon or bullet point)" | |
| image_descriptions.append(image_info) | |
| # Advanced comprehension if requested and BYOK available | |
| if use_advanced_comprehension and self.use_byok and self.byok_manager: | |
| try: | |
| # Save image to temp file for processing | |
| import tempfile | |
| with tempfile.NamedTemporaryFile( | |
| delete=False, suffix=f".{base_image.get('ext', 'png')}" | |
| ) as tmp: | |
| tmp.write(base_image["image"]) | |
| tmp_path = tmp.name | |
| # Use vision model to describe image | |
| from PIL import Image as PILImage | |
| import base64 | |
| import io | |
| img_pil = PILImage.open(tmp_path) | |
| # Convert PIL image to base64 for vision API | |
| buffered = io.BytesIO() | |
| img_pil.save(buffered, format="PNG") | |
| img_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8') | |
| # Get vision description using BYOK handler | |
| try: | |
| byok_handler = self.byok_manager.get_handler( | |
| tenant_id="default", # System-level operation | |
| db=None | |
| ) | |
| # Use coordinated vision description | |
| vision_description = await byok_handler._get_coordinated_vision_description( | |
| image_payload=img_base64, | |
| tenant_plan="free", | |
| is_managed=True | |
| ) | |
| if vision_description: | |
| image_info["ai_description"] = vision_description | |
| logger.info(f"Generated AI description for image on page {page_num + 1}") | |
| except Exception as vision_error: | |
| logger.warning(f"Vision API call failed: {vision_error}") | |
| # Fall back to basic description (already set above) | |
| # Clean up temp file | |
| os.unlink(tmp_path) | |
| except Exception as e: | |
| logger.debug(f"Advanced image comprehension failed: {e}") | |
| pdf_document.close() | |
| except ImportError: | |
| logger.warning("PyMuPDF not available for image extraction") | |
| # Fallback: Try using PyPDF2 to count images | |
| try: | |
| pdf_file = io.BytesIO(pdf_data) | |
| pdf_reader = PyPDF2.PdfReader(pdf_file) | |
| for page_num, page in enumerate(pdf_reader.pages): | |
| if "/XObject" in page["/Resources"]: | |
| xObject = page["/Resources"]["/XObject"].get_object() | |
| for obj in xObject: | |
| if xObject[obj]["/Subtype"] == "/Image": | |
| images_found += 1 | |
| image_descriptions.append( | |
| { | |
| "page": page_num + 1, | |
| "description": "Image detected (limited info without PyMuPDF)", | |
| } | |
| ) | |
| except Exception as e: | |
| logger.debug(f"PyPDF2 image extraction failed: {e}") | |
| logger.info(f"Extracted {images_found} images from PDF") | |
| return { | |
| "images_found": images_found, | |
| "image_descriptions": image_descriptions, | |
| "success": True, | |
| } | |
| except Exception as e: | |
| logger.error(f"Image extraction failed: {e}") | |
| return { | |
| "images_found": 0, | |
| "image_descriptions": [], | |
| "success": False, | |
| "error": str(e), | |
| } | |
| def _combine_results( | |
| self, | |
| basic_result: Dict[str, Any], | |
| ocr_result: Optional[Dict[str, Any]], | |
| image_results: Optional[Dict[str, Any]], | |
| used_ocr: bool, | |
| ) -> Dict[str, Any]: | |
| """Combine results from different processing methods.""" | |
| # Determine which text to use | |
| if used_ocr and ocr_result and ocr_result["success"]: | |
| best_text_result = ocr_result["best_result"] | |
| else: | |
| best_text_result = basic_result | |
| # Combine all information | |
| combined_result = { | |
| "processing_summary": { | |
| "used_ocr": used_ocr, | |
| "ocr_methods_tried": ocr_result["methods_tried"] if ocr_result else [], | |
| "best_method": best_text_result["method"], | |
| "total_pages": best_text_result["page_count"], | |
| "total_characters": best_text_result["total_chars"], | |
| }, | |
| "extracted_content": { | |
| "text": best_text_result["extracted_text"], | |
| "page_breakdown": best_text_result["page_texts"], | |
| "images": image_results or {}, | |
| }, | |
| "service_status": self.service_status, | |
| "success": basic_result["success"] or (ocr_result and ocr_result["success"]) | |
| if ocr_result | |
| else basic_result["success"], | |
| } | |