import os import base64 import gradio as gr import requests import shutil import time import pymupdf as fitz import logging from mistralai import Mistral, ImageURLChunk from mistralai.models import OCRResponse from typing import Union, List, Tuple, Optional, Dict from tenacity import retry, stop_after_attempt, wait_exponential from concurrent.futures import ThreadPoolExecutor import tempfile # Constants SUPPORTED_IMAGE_TYPES = [".jpg", ".png", ".jpeg"] SUPPORTED_PDF_TYPES = [".pdf"] UPLOAD_FOLDER = "./uploads" MAX_FILE_SIZE = 50 * 1024 * 1024 # 50MB MAX_PDF_PAGES = 50 # Configuration os.makedirs(UPLOAD_FOLDER, exist_ok=True) logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', handlers=[logging.StreamHandler()] ) logger = logging.getLogger(__name__) class OCRProcessor: def __init__(self, api_key: str): if not api_key or not isinstance(api_key, str): raise ValueError("Valid API key must be provided") self.client = Mistral(api_key=api_key) self._validate_client() def _validate_client(self) -> None: try: models = self.client.models.list() if not models: raise ValueError("No models available") except Exception as e: raise ValueError(f"API key validation failed: {str(e)}") @staticmethod def _check_file_size(file_input: Union[str, bytes]) -> None: if isinstance(file_input, str) and os.path.exists(file_input): size = os.path.getsize(file_input) elif hasattr(file_input, 'read'): size = len(file_input.read()) file_input.seek(0) else: size = len(file_input) if size > MAX_FILE_SIZE: raise ValueError(f"File size exceeds {MAX_FILE_SIZE/1024/1024}MB limit") @staticmethod def _save_uploaded_file(file_input: Union[str, bytes], filename: str) -> str: clean_filename = os.path.basename(filename).replace(os.sep, "_") file_path = os.path.join(UPLOAD_FOLDER, f"{int(time.time())}_{clean_filename}") try: if isinstance(file_input, str) and file_input.startswith("http"): response = requests.get(file_input, timeout=30) response.raise_for_status() with open(file_path, 'wb') as f: f.write(response.content) elif isinstance(file_input, str) and os.path.exists(file_input): shutil.copy2(file_input, file_path) else: with open(file_path, 'wb') as f: if hasattr(file_input, 'read'): shutil.copyfileobj(file_input, f) else: f.write(file_input) if not os.path.exists(file_path): raise FileNotFoundError(f"Failed to save file at {file_path}") return file_path except Exception as e: logger.error(f"Error saving file {filename}: {str(e)}") raise @staticmethod def _encode_image(image_path: str) -> str: try: with open(image_path, "rb") as image_file: return base64.b64encode(image_file.read()).decode('utf-8') except Exception as e: logger.error(f"Error encoding image {image_path}: {str(e)}") raise ValueError(f"Failed to encode image: {str(e)}") @staticmethod def _pdf_to_images(pdf_path: str) -> List[Tuple[str, str]]: try: pdf_document = fitz.open(pdf_path) if pdf_document.page_count > MAX_PDF_PAGES: pdf_document.close() raise ValueError(f"PDF exceeds maximum page limit of {MAX_PDF_PAGES}") with ThreadPoolExecutor() as executor: image_data = list(executor.map( lambda i: OCRProcessor._convert_page(pdf_path, i), range(pdf_document.page_count) )) pdf_document.close() return [data for data in image_data if data] except Exception as e: logger.error(f"Error converting PDF to images: {str(e)}") return [] @staticmethod def _convert_page(pdf_path: str, page_num: int) -> Tuple[str, str]: try: pdf_document = fitz.open(pdf_path) page = pdf_document[page_num] pix = page.get_pixmap(dpi=150) image_path = os.path.join(UPLOAD_FOLDER, f"page_{page_num + 1}_{int(time.time())}.png") pix.save(image_path) encoded = OCRProcessor._encode_image(image_path) pdf_document.close() return image_path, encoded except Exception as e: logger.error(f"Error converting page {page_num}: {str(e)}") return None, None @retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=2, max=10)) def _call_ocr_api(self, encoded_image: str) -> OCRResponse: base64_url = f"data:image/png;base64,{encoded_image}" try: logger.info("Calling OCR API") response = self.client.ocr.process( model="mistral-ocr-latest", document=ImageURLChunk(image_url=base64_url), include_image_base64=True ) return response except Exception as e: logger.error(f"OCR API call failed: {str(e)}") raise def process_file(self, file: gr.File) -> Tuple[str, str, List[str]]: """Process uploaded file (image or PDF).""" if not file: return "## No file provided", "", [] file_name = file.name self._check_file_size(file) file_path = self._save_uploaded_file(file, file_name) if file_name.lower().endswith(tuple(SUPPORTED_IMAGE_TYPES)): encoded_image = self._encode_image(file_path) response = self._call_ocr_api(encoded_image) markdown = self._combine_markdown(response) return markdown, file_path, [file_path] elif file_name.lower().endswith('.pdf'): image_data = self._pdf_to_images(file_path) if not image_data: return "## No pages converted from PDF", file_path, [] ocr_results = [] image_paths = [path for path, _ in image_data] for _, encoded in image_data: response = self._call_ocr_api(encoded) markdown = self._combine_markdown(response) ocr_results.append(markdown) return "\n\n".join(ocr_results), file_path, image_paths return "## Unsupported file type", file_path, [] def process_url(self, url: str) -> Tuple[str, str, List[str]]: """Process URL (image or PDF).""" if not url: return "## No URL provided", "", [] file_name = url.split('/')[-1] or f"file_{int(time.time())}" file_path = self._save_uploaded_file(url, file_name) if file_name.lower().endswith(tuple(SUPPORTED_IMAGE_TYPES)): encoded_image = self._encode_image(file_path) response = self._call_ocr_api(encoded_image) markdown = self._combine_markdown(response) return markdown, url, [file_path] elif file_name.lower().endswith('.pdf'): image_data = self._pdf_to_images(file_path) if not image_data: return "## No pages converted from PDF", url, [] ocr_results = [] image_paths = [path for path, _ in image_data] for _, encoded in image_data: response = self._call_ocr_api(encoded) markdown = self._combine_markdown(response) ocr_results.append(markdown) return "\n\n".join(ocr_results), url, image_paths return "## Unsupported URL content type", url, [] @staticmethod def _combine_markdown(response: OCRResponse) -> str: """Combine markdown from OCR response.""" markdown_parts = [] for page in response.pages: if not page.markdown.strip(): continue markdown = page.markdown if hasattr(page, 'images') and page.images: for img in page.images: if img.image_base64: markdown = markdown.replace( f"![{img.id}]({img.id})", f"![{img.id}](data:image/png;base64,{img.image_base64})" ) markdown_parts.append(markdown) return "\n\n".join(markdown_parts) or "## No text detected" def create_interface(): css = """ .output-markdown {font-size: 14px; max-height: 500px; overflow-y: auto;} .status {color: #666; font-style: italic;} .preview {max-height: 300px;} """ with gr.Blocks(title="Mistral OCR Demo", css=css) as demo: gr.Markdown("# Mistral OCR Demo") gr.Markdown(f""" Process PDFs and images (max {MAX_FILE_SIZE/1024/1024}MB, {MAX_PDF_PAGES} pages for PDFs) via upload or URL. View previews and OCR results with embedded images. Learn more at [Mistral OCR](https://mistral.ai/news/mistral-ocr). """) # API Key Setup with gr.Row(): api_key_input = gr.Textbox(label="Mistral API Key", type="password", placeholder="Enter your API key") set_key_btn = gr.Button("Set API Key", variant="primary") processor_state = gr.State() status = gr.Markdown("Please enter API key", elem_classes="status") def init_processor(key): try: processor = OCRProcessor(key) return processor, "✅ API key validated" except Exception as e: return None, f"❌ Error: {str(e)}" set_key_btn.click(fn=init_processor, inputs=api_key_input, outputs=[processor_state, status]) # File Upload Tab with gr.Tab("Upload File"): with gr.Row(): file_input = gr.File(label="Upload PDF/Image", file_types=SUPPORTED_IMAGE_TYPES + SUPPORTED_PDF_TYPES) file_preview = gr.Gallery(label="Preview", elem_classes="preview") file_output = gr.Markdown(label="OCR Result", elem_classes="output-markdown") file_raw_output = gr.Textbox(label="Raw File Path") file_button = gr.Button("Process", variant="primary") def update_file_preview(file): return [file.name] if file else [] file_input.change(fn=update_file_preview, inputs=file_input, outputs=file_preview) file_button.click( fn=lambda p, f: p.process_file(f) if p else ("## Set API key first", "", []), inputs=[processor_state, file_input], outputs=[file_output, file_raw_output, file_preview] ) # URL Tab with gr.Tab("URL Input"): with gr.Row(): url_input = gr.Textbox(label="URL to PDF/Image") url_preview = gr.Gallery(label="Preview", elem_classes="preview") url_output = gr.Markdown(label="OCR Result", elem_classes="output-markdown") url_raw_output = gr.Textbox(label="Raw URL") url_button = gr.Button("Process", variant="primary") def update_url_preview(url): if not url: return [] try: temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.tmp') response = requests.get(url, timeout=10) temp_file.write(response.content) temp_file.close() return [temp_file.name] except Exception as e: logger.error(f"URL preview error: {str(e)}") return [] url_input.change(fn=update_url_preview, inputs=url_input, outputs=url_preview) url_button.click( fn=lambda p, u: p.process_url(u) if p else ("## Set API key first", "", []), inputs=[processor_state, url_input], outputs=[url_output, url_raw_output, url_preview] ) # Examples gr.Examples( examples=[], inputs=[file_input, url_input] ) return demo if __name__ == "__main__": os.environ['START_TIME'] = time.strftime('%Y-%m-%d %H:%M:%S') print(f"===== Application Startup at {os.environ['START_TIME']} =====") create_interface().launch(share=True, max_threads=1)