""" Utility functions for OCR processing with Mistral AI. Contains helper functions for working with OCR responses and image handling. """ import json import base64 import io from pathlib import Path from typing import Dict, List, Optional, Union, Any try: from PIL import Image PILLOW_AVAILABLE = True except ImportError: PILLOW_AVAILABLE = False from mistralai import DocumentURLChunk, ImageURLChunk, TextChunk def replace_images_in_markdown(markdown_str: str, images_dict: dict) -> str: """ Replace image placeholders in markdown with base64-encoded images. Args: markdown_str: Markdown text containing image placeholders images_dict: Dictionary mapping image IDs to base64 strings Returns: Markdown text with images replaced by base64 data """ for img_name, base64_str in images_dict.items(): markdown_str = markdown_str.replace( f"![{img_name}]({img_name})", f"![{img_name}]({base64_str})" ) return markdown_str def get_combined_markdown(ocr_response) -> str: """ Combine OCR text and images into a single markdown document. Ensures proper spacing between text and images. Args: ocr_response: Response from OCR processing containing text and images See https://docs.mistral.ai/capabilities/document/ for API reference Returns: Combined markdown string with embedded images """ markdowns: list[str] = [] # Extract images from page for page in ocr_response.pages: image_data = {} for img in page.images: image_data[img.id] = img.image_base64 # Replace image placeholders with actual images page_markdown = replace_images_in_markdown(page.markdown, image_data) # Ensure proper spacing between paragraphs and images # Add extra newlines between paragraphs to improve rendering page_markdown = page_markdown.replace("\n", "\n\n") # Add page separator for multi-page documents markdowns.append(page_markdown) # Join pages with clear separators for multi-page documents return "\n\n---\n\n".join(markdowns) def encode_image_for_api(image_path: Union[str, Path]) -> str: """ Encode an image as base64 for API use. Args: image_path: Path to the image file Returns: Base64 data URL for the image """ # Convert to Path object if string image_file = Path(image_path) if isinstance(image_path, str) else image_path # Verify image exists if not image_file.is_file(): raise FileNotFoundError(f"Image file not found: {image_file}") # Encode image as base64 encoded = base64.b64encode(image_file.read_bytes()).decode() return f"data:image/jpeg;base64,{encoded}" def process_image_with_ocr(client, image_path: Union[str, Path], model: str = "mistral-ocr-latest"): """ Process an image with OCR and return the response. Args: client: Mistral AI client image_path: Path to the image file model: OCR model to use Returns: OCR response object """ # Encode image as base64 base64_data_url = encode_image_for_api(image_path) # Process image with OCR image_response = client.ocr.process( document=ImageURLChunk(image_url=base64_data_url), model=model ) return image_response def ocr_response_to_json(ocr_response, indent: int = 4) -> str: """ Convert OCR response to a formatted JSON string. Args: ocr_response: OCR response object indent: Indentation level for JSON formatting Returns: Formatted JSON string """ # Convert response to JSON response_dict = json.loads(ocr_response.model_dump_json()) return json.dumps(response_dict, indent=indent) def get_combined_markdown_compressed(ocr_response, max_width: int = 1200, quality: int = 92) -> str: """ Combine OCR text and images into a single markdown document with compressed images. Reduces image sizes to improve performance. Args: ocr_response: Response from OCR processing containing text and images max_width: Maximum width to resize images to (preserves aspect ratio) quality: JPEG quality (0-100) for compression Returns: Combined markdown string with embedded compressed images """ if not PILLOW_AVAILABLE: # Fall back to regular method if PIL is not available return get_combined_markdown(ocr_response) markdowns: list[str] = [] # Process each page for page in ocr_response.pages: image_data = {} # Process and compress each image for img in page.images: try: # Decode base64 image img_bytes = base64.b64decode(img.image_base64.split(',')[1] if ',' in img.image_base64 else img.image_base64) # Open with PIL pil_img = Image.open(io.BytesIO(img_bytes)) # Convert to RGB if not already (to ensure CV_8UC3 format) if pil_img.mode != 'RGB': pil_img = pil_img.convert('RGB') # Resize if needed (maintain aspect ratio) original_width, original_height = pil_img.size if original_width > max_width: ratio = max_width / original_width new_height = int(original_height * ratio) pil_img = pil_img.resize((max_width, new_height), Image.LANCZOS) # Convert to bytes with compression buffer = io.BytesIO() format = pil_img.format if pil_img.format else 'JPEG' if format.upper() == 'JPEG' or format.upper() == 'JPG': pil_img.save(buffer, format=format, quality=quality, optimize=True) else: # For non-JPEG formats (PNG, etc.) pil_img.save(buffer, format=format, optimize=True) # Convert back to base64 compressed_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8') mime_type = f"image/{format.lower()}" if format else "image/jpeg" image_data[img.id] = f"data:{mime_type};base64,{compressed_base64}" except Exception as e: # If compression fails, use original image image_data[img.id] = img.image_base64 # Replace image placeholders with compressed images page_markdown = replace_images_in_markdown(page.markdown, image_data) # Ensure proper spacing between paragraphs and images page_markdown = page_markdown.replace("\n", "\n\n") # Add page to list markdowns.append(page_markdown) # Join pages with clear separators return "\n\n---\n\n".join(markdowns) # For display in notebooks try: from IPython.display import Markdown, display def display_ocr_with_images(ocr_response): """ Display OCR response with embedded images in IPython environments. Args: ocr_response: OCR response object """ combined_markdown = get_combined_markdown(ocr_response) display(Markdown(combined_markdown)) except ImportError: # IPython not available pass def create_html_with_images(result_with_pages: dict) -> str: """ Create HTML with embedded images from the OCR result. Args: result_with_pages: OCR result with pages_data containing markdown and images Returns: HTML string with embedded images """ if not result_with_pages.get('has_images', False) or 'pages_data' not in result_with_pages: return "

No images available in the document.

" # Create HTML document html = """ Document with Images
""" # Process each page pages_data = result_with_pages.get('pages_data', []) for page_idx, page in enumerate(pages_data): page_number = page.get('page_number', page_idx + 1) page_markdown = page.get('markdown', '') page_images = page.get('images', []) # Add page header html += f'
\n' if len(pages_data) > 1: html += f'
Page {page_number}
\n' # Process markdown text and replace image references if page_markdown: # Replace image markers with actual images for img in page_images: img_id = img.get('id', '') img_base64 = img.get('image_base64', '') if img_id and img_base64: # Format image tag img_tag = f'Image {img_id}' # Replace markdown image reference with HTML image page_markdown = page_markdown.replace(f'![{img_id}]({img_id})', img_tag) # Convert line breaks to

tags for proper HTML formatting paragraphs = page_markdown.split('\n\n') for paragraph in paragraphs: if paragraph.strip(): # Check if this looks like a header if paragraph.startswith('# '): header_text = paragraph[2:].strip() html += f'

{header_text}

\n' elif paragraph.startswith('## '): header_text = paragraph[3:].strip() html += f'

{header_text}

\n' elif paragraph.startswith('### '): header_text = paragraph[4:].strip() html += f'

{header_text}

\n' else: html += f'

{paragraph}

\n' # Add any images that weren't referenced in the markdown referenced_img_ids = [img.get('id') for img in page_images if img.get('id') in page_markdown] for img in page_images: img_id = img.get('id', '') img_base64 = img.get('image_base64', '') if img_id and img_base64 and img_id not in referenced_img_ids: html += f'Image {img_id}\n' # Close page div html += '
\n' # Close main container and document html += """
""" return html