import os import base64 import gradio as gr from mistralai import Mistral, DocumentURLChunk, ImageURLChunk, TextChunk from mistralai.models import OCRResponse from pathlib import Path import pycountry import json import logging from tenacity import retry, stop_after_attempt, wait_exponential import tempfile from typing import Union, Dict, List, Optional, Tuple from contextlib import contextmanager import requests import shutil from concurrent.futures import ThreadPoolExecutor import time # Constants DEFAULT_LANGUAGE = "English" SUPPORTED_IMAGE_TYPES = [".jpg", ".png", ".jpeg"] SUPPORTED_PDF_TYPES = [".pdf"] TEMP_FILE_EXPIRY = 7200 # 2 hours in seconds 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): self.api_key = self._validate_api_key(api_key) self.client = Mistral(api_key=self.api_key) self._validate_client() @staticmethod def _validate_api_key(api_key: str) -> str: if not api_key or not isinstance(api_key, str): raise ValueError("Valid API key must be provided") return api_key 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) # Reset file pointer 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 _encode_image(image_path: str) -> Optional[str]: with open(image_path, "rb") as image_file: return base64.b64encode(image_file.read()).decode('utf-8') @staticmethod def _save_uploaded_file(file_input: Union[str, bytes], filename: str) -> str: file_path = os.path.join(UPLOAD_FOLDER, f"{int(time.time())}_{filename}") if isinstance(file_input, str) and file_input.startswith("http"): response = requests.get(file_input, timeout=10) response.raise_for_status() with open(file_path, 'wb') as f: f.write(response.content) else: with open(file_path, 'wb') as f: if hasattr(file_input, 'read'): shutil.copyfileobj(file_input, f) else: f.write(file_input) return file_path @staticmethod def _pdf_to_images(pdf_path: str) -> List[str]: 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_paths = list(executor.map( lambda i: OCRProcessor._convert_page(pdf_path, i), range(pdf_document.page_count) )) pdf_document.close() return [path for path in image_paths if path] @staticmethod def _convert_page(pdf_path: str, page_num: int) -> Optional[str]: try: pdf_document = fitz.open(pdf_path) page = pdf_document[page_num] pix = page.get_pixmap(dpi=150) # Improved resolution image_path = os.path.join(UPLOAD_FOLDER, f"page_{page_num + 1}_{int(time.time())}.png") pix.save(image_path) pdf_document.close() return image_path except Exception as e: logger.error(f"Error converting page {page_num}: {str(e)}") return None @retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=2, max=10)) def _call_ocr_api(self, document: Union[DocumentURLChunk, ImageURLChunk]) -> OCRResponse: return self.client.ocr.process( model="mistral-ocr-latest", document=document, include_image_base64=True ) @retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=2, max=10)) def _call_chat_complete(self, model: str, messages: List[Dict], **kwargs) -> Dict: return self.client.chat.complete(model=model, messages=messages, **kwargs) def ocr_uploaded_pdf(self, pdf_file: Union[str, bytes]) -> Tuple[str, List[str]]: file_name = getattr(pdf_file, 'name', f"pdf_{int(time.time())}.pdf") logger.info(f"Processing uploaded PDF: {file_name}") try: self._check_file_size(pdf_file) pdf_path = self._save_uploaded_file(pdf_file, file_name) image_paths = self._pdf_to_images(pdf_path) uploaded_file = self.client.files.upload( file={"file_name": pdf_path, "content": open(pdf_path, "rb")}, purpose="ocr" ) signed_url = self.client.files.get_signed_url(file_id=uploaded_file.id, expiry=TEMP_FILE_EXPIRY) response = self._call_ocr_api(DocumentURLChunk(document_url=signed_url.url)) return self._get_combined_markdown(response), image_paths except Exception as e: return self._handle_error("PDF processing", e), [] def ocr_uploaded_image(self, image_file: Union[str, bytes]) -> Tuple[str, str]: file_name = getattr(image_file, 'name', f"image_{int(time.time())}.jpg") logger.info(f"Processing uploaded image: {file_name}") try: self._check_file_size(image_file) image_path = self._save_uploaded_file(image_file, file_name) encoded_image = self._encode_image(image_path) base64_url = f"data:image/jpeg;base64,{encoded_image}" response = self._call_ocr_api(ImageURLChunk(image_url=base64_url)) return self._get_combined_markdown(response), image_path except Exception as e: return self._handle_error("image processing", e), None def document_understanding(self, doc_url: str, question: str) -> str: try: messages = [{"role": "user", "content": [ TextChunk(text=question), DocumentURLChunk(document_url=doc_url) ]}] response = self._call_chat_complete( model="mistral-small-latest", messages=messages, temperature=0.1 ) return response.choices[0].message.content except Exception as e: return self._handle_error("document understanding", e) def structured_ocr(self, image_file: Union[str, bytes]) -> Tuple[str, str]: file_name = getattr(image_file, 'name', f"image_{int(time.time())}.jpg") try: self._check_file_size(image_file) image_path = self._save_uploaded_file(image_file, file_name) encoded_image = self._encode_image(image_path) base64_url = f"data:image/jpeg;base64,{encoded_image}" ocr_response = self._call_ocr_api(ImageURLChunk(image_url=base64_url)) markdown = self._get_combined_markdown(ocr_response) chat_response = self._call_chat_complete( model="pixtral-12b-latest", messages=[{ "role": "user", "content": [ ImageURLChunk(image_url=base64_url), TextChunk(text=( f"This is image's OCR in markdown:\n\n{markdown}\n.\n" "Convert this into a structured JSON response with file_name, topics, languages, and ocr_contents fields" )) ] }], response_format={"type": "json_object"}, temperature=0.1 ) return self._format_structured_response(image_path, json.loads(chat_response.choices[0].message.content)), image_path except Exception as e: return self._handle_error("structured OCR", e), None @staticmethod def _get_combined_markdown(response: OCRResponse) -> str: return "\n\n".join( page.markdown for page in response.pages if page.markdown.strip() ) or "No text detected" @staticmethod def _handle_error(context: str, error: Exception) -> str: logger.error(f"Error in {context}: {str(error)}") return f"**Error in {context}:** {str(error)}" @staticmethod def _format_structured_response(file_path: str, content: Dict) -> str: languages = {lang.alpha_2: lang.name for lang in pycountry.languages if hasattr(lang, 'alpha_2')} content_languages = content.get("languages", [DEFAULT_LANGUAGE]) valid_langs = [l for l in content_languages if l in languages.values()] or [DEFAULT_LANGUAGE] response = { "file_name": Path(file_path).name, "topics": content.get("topics", []), "languages": valid_langs, "ocr_contents": content.get("ocr_contents", {}) } return f"```json\n{json.dumps(response, indent=2, ensure_ascii=False)}\n```" def create_interface(): css = """ .output-markdown {font-size: 14px; max-height: 500px; overflow-y: auto;} .status {color: #666; font-style: italic;} """ with gr.Blocks(title="Mistral OCR App", css=css) as demo: gr.Markdown("# Mistral OCR App\nUpload images or PDFs for OCR processing") with gr.Row(): api_key = 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 successfully" except Exception as e: return None, f"❌ Error: {str(e)}" set_key_btn.click( fn=init_processor, inputs=api_key, outputs=[processor_state, status] ) with gr.Tab("Image OCR"): with gr.Row(): image_input = gr.File( label=f"Upload Image (max {MAX_FILE_SIZE/1024/1024}MB)", file_types=SUPPORTED_IMAGE_TYPES ) image_preview = gr.Image(label="Preview", height=300) image_output = gr.Markdown(label="OCR Result", elem_classes="output-markdown") process_image_btn = gr.Button("Process Image", variant="primary") def process_image(processor, image): if not processor or not image: return "Please set API key and upload an image", None return processor.ocr_uploaded_image(image) process_image_btn.click( fn=process_image, inputs=[processor_state, image_input], outputs=[image_output, image_preview] ) with gr.Tab("PDF OCR"): with gr.Row(): pdf_input = gr.File( label=f"Upload PDF (max {MAX_FILE_SIZE/1024/1024}MB, {MAX_PDF_PAGES} pages)", file_types=SUPPORTED_PDF_TYPES ) pdf_gallery = gr.Gallery(label="PDF Pages", height=300) pdf_output = gr.Markdown(label="OCR Result", elem_classes="output-markdown") process_pdf_btn = gr.Button("Process PDF", variant="primary") def process_pdf(processor, pdf): if not processor or not pdf: return "Please set API key and upload a PDF", [] return processor.ocr_uploaded_pdf(pdf) process_pdf_btn.click( fn=process_pdf, inputs=[processor_state, pdf_input], outputs=[pdf_output, pdf_gallery] ) with gr.Tab("Structured OCR"): structured_input = gr.File( label=f"Upload Image for Structured OCR (max {MAX_FILE_SIZE/1024/1024}MB)", file_types=SUPPORTED_IMAGE_TYPES ) structured_output = gr.Markdown(label="Structured Result", elem_classes="output-markdown") structured_preview = gr.Image(label="Preview", height=300) process_structured_btn = gr.Button("Process Structured OCR", variant="primary") def process_structured(processor, image): if not processor or not image: return "Please set API key and upload an image", None return processor.structured_ocr(image) process_structured_btn.click( fn=process_structured, inputs=[processor_state, structured_input], outputs=[structured_output, structured_preview] ) 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, debug=True, )