Few updates

app.py

@@ -1,9 +1,10 @@
 import os
 import base64
 import gradio as gr
-from mistralai import Mistral, ImageURLChunk
+import json
+from mistralai import Mistral, DocumentURLChunk, ImageURLChunk, TextChunk
 from mistralai.models import OCRResponse
-from typing import Union, List, Tuple
+from typing import Union, List, Tuple, Dict
 import requests
 import shutil
 import time
@@ -13,6 +14,11 @@ from tenacity import retry, stop_after_attempt, wait_exponential
 from concurrent.futures import ThreadPoolExecutor
 import socket
 from requests.exceptions import ConnectionError, Timeout
+from pathlib import Path
+from pydantic import BaseModel
+import pycountry
+from enum import Enum
+from PIL import Image
 
 # Constants
 SUPPORTED_IMAGE_TYPES = [".jpg", ".png", ".jpeg"]
@@ -30,6 +36,32 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
+# Language Enum for StructuredOCR
+languages = {lang.alpha_2: lang.name for lang in pycountry.languages if hasattr(lang, 'alpha_2')}
+
+class LanguageMeta(Enum.__class__):
+    def __new__(metacls, cls, bases, classdict):
+        for code, name in languages.items():
+            classdict[name.upper().replace(' ', '_')] = name
+        return super().__new__(metacls, cls, bases, classdict)
+
+class Language(Enum, metaclass=LanguageMeta):
+    pass
+
+class StructuredOCR(BaseModel):
+    file_name: str
+    topics: list[str]
+    languages: list[Language]
+    ocr_contents: dict
+
+    def model_dump_json(self, **kwargs):
+        # Custom JSON serialization to handle Language enums
+        data = self.model_dump(exclude_unset=True, by_alias=True, mode='json')
+        for key, value in data.items():
+            if isinstance(value, list) and all(isinstance(item, Language) for item in value):
+                data[key] = [item.value for item in value]
+        return json.dumps(data, indent=4)
+
 class OCRProcessor:
     def __init__(self, api_key: str):
         if not api_key or not isinstance(api_key, str):
@@ -91,10 +123,12 @@ class OCRProcessor:
     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')
+                encoded = base64.b64encode(image_file.read()).decode('utf-8')
+                logger.info(f"Encoded image {image_path} to base64 (length: {len(encoded)})")
+                return encoded
         except Exception as e:
             logger.error(f"Error encoding image {image_path}: {str(e)}")
-            raise ValueError("Failed to encode image")
+            raise ValueError(f"Failed to encode image: {str(e)}")
 
     @staticmethod
     def _pdf_to_images(pdf_path: str) -> List[Tuple[str, str]]:
@@ -110,10 +144,14 @@ class OCRProcessor:
                     range(pdf_document.page_count)
                 ))
             pdf_document.close()
-            return [data for data in image_data if data]
+            valid_image_data = [(path, encoded) for path, encoded in image_data if path and encoded]
+            if not valid_image_data:
+                raise ValueError("No valid pages converted from PDF")
+            logger.info(f"Converted {len(valid_image_data)} pages to images")
+            return valid_image_data
         except Exception as e:
             logger.error(f"Error converting PDF to images: {str(e)}")
-            return []
+            raise
 
     @staticmethod
     def _convert_page(pdf_path: str, page_num: int) -> Tuple[str, str]:
@@ -132,128 +170,151 @@ class OCRProcessor:
 
     @retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=2, max=10))
     def _call_ocr_api(self, encoded_image: str) -> OCRResponse:
+        if not isinstance(encoded_image, str):
+            raise TypeError(f"Expected encoded_image to be a string, got {type(encoded_image)}")
         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),
+                model="mistral-ocr-latest",
                 include_image_base64=True
             )
             logger.info("OCR API call successful")
+            try:
+                if hasattr(response, 'model_dump_json'):
+                    response_dict = json.loads(response.model_dump_json())
+                else:
+                    response_dict = {k: v for k, v in response.__dict__.items() if isinstance(v, (str, int, float, list, dict))}
+                logger.info(f"Raw OCR response: {json.dumps(response_dict, default=str, indent=4)}")
+            except Exception as log_err:
+                logger.warning(f"Failed to log raw OCR response: {str(log_err)}")
             return response
-        except (ConnectionError, Timeout, socket.error) as e:
+        except (ConnectionError, TimeoutError, socket.error) as e:
             logger.error(f"Network error during OCR API call: {str(e)}")
             raise
-
-    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)
-            
-            if not os.path.exists(pdf_path):
-                raise FileNotFoundError(f"Saved PDF not found at: {pdf_path}")
-            
-            image_data = self._pdf_to_images(pdf_path)
-            if not image_data:
-                raise ValueError("No pages converted from PDF")
-            
-            ocr_results = []
-            image_paths = [path for path, _ in image_data]
-            for i, (_, encoded) in enumerate(image_data):
-                response = self._call_ocr_api(encoded)
-                markdown_with_images = self._get_combined_markdown_with_images(response, image_paths, i)
-                ocr_results.append(markdown_with_images)
-            
-            return "\n\n".join(ocr_results), image_paths
+        except TypeError as e:
+            logger.error(f"TypeError in OCR API call: {str(e)}", exc_info=True)
+            raise
         except Exception as e:
-            return self._handle_error("uploaded PDF processing", e), []
+            logger.error(f"Unexpected error in OCR API call: {str(e)}", exc_info=True)
+            raise
 
-    def ocr_pdf_url(self, pdf_url: str) -> Tuple[str, List[str]]:
-        logger.info(f"Processing PDF URL: {pdf_url}")
+    def _process_pdf_with_ocr(self, pdf_path: str) -> Tuple[str, List[str], List[Dict]]:
         try:
-            file_name = pdf_url.split('/')[-1] or f"pdf_{int(time.time())}.pdf"
-            pdf_path = self._save_uploaded_file(pdf_url, file_name)
-            
-            if not os.path.exists(pdf_path):
-                raise FileNotFoundError(f"Saved PDF not found at: {pdf_path}")
-            
-            image_data = self._pdf_to_images(pdf_path)
-            if not image_data:
-                raise ValueError("No pages converted from PDF")
-            
-            ocr_results = []
-            image_paths = [path for path, _ in image_data]
-            for i, (_, encoded) in enumerate(image_data):
-                response = self._call_ocr_api(encoded)
-                markdown_with_images = self._get_combined_markdown_with_images(response, image_paths, i)
-                ocr_results.append(markdown_with_images)
-            
-            return "\n\n".join(ocr_results), image_paths
+            # Upload PDF and get signed URL
+            uploaded_file = self.client.files.upload(
+                file={"file_name": Path(pdf_path).stem, "content": Path(pdf_path).read_bytes()},
+                purpose="ocr",
+            )
+            signed_url = self.client.files.get_signed_url(file_id=uploaded_file.id, expiry=1).url
+
+            # Process with OCR
+            ocr_response = self.client.ocr.process(
+                document=DocumentURLChunk(document_url=signed_url),
+                model="mistral-ocr-latest",
+                include_image_base64=True
+            )
+            markdown, base64_images = self._get_combined_markdown(ocr_response)
+            json_results = self._convert_to_structured_json(markdown, pdf_path)
+            # Fallback to local images if OCR images are missing
+            image_paths = []
+            if not any(page.images for page in ocr_response.pages):
+                logger.warning("No images found in OCR response; using local images")
+                image_data = self._pdf_to_images(pdf_path)
+                image_paths = [path for path, _ in image_data]
+            else:
+                image_paths = [os.path.join(UPLOAD_FOLDER, f"ocr_page_{i}.png") for i in range(len(ocr_response.pages))]
+                for i, base64_img in enumerate(base64_images):
+                    if base64_img:
+                        try:
+                            img_data = base64.b64decode(base64_img.split(',')[1])
+                            with open(image_paths[i], "wb") as f:
+                                f.write(img_data)
+                            if os.path.exists(image_paths[i]):
+                                logger.info(f"Image {image_paths[i]} saved and exists")
+                            else:
+                                logger.error(f"Image {image_paths[i]} saved but does not exist")
+                        except Exception as e:
+                            logger.error(f"Error saving image {i}: {str(e)}")
+                            image_paths[i] = None
+                image_paths = [path for path in image_paths if path and os.path.exists(path)]
+            logger.info(f"Final image paths: {image_paths}")
+            return markdown, image_paths, json_results
         except Exception as e:
-            return self._handle_error("PDF URL processing", e), []
+            return self._handle_error("PDF OCR processing", e), [], []
+
+    def _get_combined_markdown(self, ocr_response: OCRResponse) -> Tuple[str, List[str]]:
+        markdowns = []
+        base64_images = []
+        for i, page in enumerate(ocr_response.pages):
+            image_data = {}
+            for img in page.images:
+                if img.image_base64:
+                    base64_url = f"data:image/png;base64,{img.image_base64}"
+                    image_data[img.id] = base64_url
+                    base64_images.append(base64_url)
+                    logger.info(f"Base64 image {img.id} length: {len(img.image_base64)}")
+                else:
+                    base64_images.append(None)
+            markdown = page.markdown or "No text detected"
+            markdown = replace_images_in_markdown(markdown, image_data)
+            logger.info(f"Page {i} markdown (first 200 chars): {markdown[:200]}...")
+            markdowns.append(markdown)
+        return "\n\n".join(markdowns), base64_images
 
-    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}")
+    def _convert_to_structured_json(self, markdown: str, file_path: str) -> List[Dict]:
         try:
-            self._check_file_size(image_file)
-            image_path = self._save_uploaded_file(image_file, file_name)
-            encoded_image = self._encode_image(image_path)
-            response = self._call_ocr_api(encoded_image)
-            return self._get_combined_markdown_with_images(response), image_path
+            text_only_markdown = re.sub(r'!\[.*?\]\(data:image/[^)]+\)', '', markdown)
+            logger.info(f"Text-only markdown length: {len(text_only_markdown)}")
+            logger.info(f"Text-only markdown content: {text_only_markdown[:200]}...")
+
+            chat_response = self.client.chat.parse(
+                model="pixtral-12b-latest",
+                messages=[
+                    {
+                        "role": "user",
+                        "content": f"Given OCR output from a PDF about African history and artifacts, convert to JSON with file_name, topics (e.g., African Artifacts, Tribal History), languages (e.g., English), and ocr_contents (title and list of items with descriptions and image refs).\n\nOCR Output:\n{text_only_markdown}"
+                    },
+                ],
+                response_format=StructuredOCR,
+                temperature=0
+            )
+            structured_result = chat_response.choices[0].message.parsed
+            json_str = structured_result.model_dump_json()
+            logger.info(f"Structured JSON: {json_str}")
+            return [json.loads(json_str)]
         except Exception as e:
-            return self._handle_error("image processing", e), None
+            logger.error(f"Error converting to structured JSON: {str(e)}", exc_info=True)
+            return [{"error": str(e), "file_name": Path(file_path).stem}]
 
-    @staticmethod
-    def _get_combined_markdown_with_images(response: OCRResponse, image_paths: List[str] = None, page_index: int = None) -> str:
-        markdown_parts = []
-        for i, page in enumerate(response.pages):
-            if page.markdown.strip():
-                markdown = page.markdown
-                logger.info(f"Page {i} markdown: {markdown}")
-                if hasattr(page, 'images') and page.images:
-                    logger.info(f"Found {len(page.images)} images in page {i}")
-                    for img in page.images:
-                        if img.image_base64:
-                            logger.info(f"Replacing image {img.id} with base64")
-                            markdown = markdown.replace(
-                                f"![{img.id}]({img.id})",
-                                f"![{img.id}](data:image/png;base64,{img.image_base64})"
-                            )
-                        else:
-                            logger.warning(f"No base64 data for image {img.id}")
-                            if image_paths and page_index is not None and page_index < len(image_paths):
-                                local_encoded = OCRProcessor._encode_image(image_paths[page_index])
-                                markdown = markdown.replace(
-                                    f"![{img.id}]({img.id})",
-                                    f"![{img.id}](data:image/png;base64,{local_encoded})"
-                                )
-                else:
-                    logger.warning(f"No images found in page {i}")
-                    # Replace known placeholders or append the local image
-                    if image_paths and page_index is not None and page_index < len(image_paths):
-                        local_encoded = OCRProcessor._encode_image(image_paths[page_index])
-                        # Replace placeholders like img-0.jpeg
-                        placeholder = f"img-{i}.jpeg"
-                        if placeholder in markdown:
-                            markdown = markdown.replace(
-                                placeholder,
-                                f"![Page {i} Image](data:image/png;base64,{local_encoded})"
-                            )
-                        else:
-                            # Append the image if no placeholder is found
-                            markdown += f"\n\n![Page {i} Image](data:image/png;base64,{local_encoded})"
-                markdown_parts.append(markdown)
-        return "\n\n".join(markdown_parts) or "No text or images detected"
+    def ocr_uploaded_pdf(self, pdf_file: Union[str, bytes]) -> Tuple[str, List[str], List[Dict]]:
+        file_path = self._save_uploaded_file(pdf_file, getattr(pdf_file, 'name', f"pdf_{int(time.time())}.pdf"))
+        return self._process_pdf_with_ocr(file_path)
+
+    def ocr_pdf_url(self, pdf_url: str) -> Tuple[str, List[str], List[Dict]]:
+        file_path = self._save_uploaded_file(pdf_url, pdf_url.split('/')[-1] or f"pdf_{int(time.time())}.pdf")
+        return self._process_pdf_with_ocr(file_path)
+
+    def ocr_uploaded_image(self, image_file: Union[str, bytes]) -> Tuple[str, str, Dict]:
+        file_path = self._save_uploaded_file(image_file, getattr(image_file, 'name', f"image_{int(time.time())}.jpg"))
+        encoded_image = self._encode_image(file_path)
+        base64_url = f"data:image/png;base64,{encoded_image}"
+        response = self._call_ocr_api(encoded_image)
+        markdown, base64_images = self._get_combined_markdown(response)
+        json_result = self._convert_to_structured_json(markdown, file_path)[0]
+        return markdown, file_path, json_result
 
     @staticmethod
     def _handle_error(context: str, error: Exception) -> str:
-        logger.error(f"Error in {context}: {str(error)}")
+        logger.error(f"Error in {context}: {str(error)}", exc_info=True)
         return f"**Error in {context}:** {str(error)}"
 
+def replace_images_in_markdown(markdown_str: str, images_dict: dict) -> str:
+    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 create_interface():
     css = """
     .output-markdown {font-size: 14px; max-height: 500px; overflow-y: auto;}
@@ -291,17 +352,19 @@ def create_interface():
                 )
                 image_preview = gr.Image(label="Preview", height=300)
             image_output = gr.Markdown(label="OCR Result", elem_classes="output-markdown")
+            image_json_output = gr.JSON(label="Structured JSON Output")
             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)
+                    return "Please set API key and upload an image", None, {}
+                markdown, image_path, json_data = processor.ocr_uploaded_image(image)
+                return markdown, image_path, json_data
 
             process_image_btn.click(
                 fn=process_image,
                 inputs=[processor_state, image_input],
-                outputs=[image_output, image_preview]
+                outputs=[image_output, image_preview, image_json_output]
             )
 
         with gr.Tab("PDF OCR"):
@@ -317,24 +380,32 @@ def create_interface():
                     )
                 pdf_gallery = gr.Gallery(label="PDF Pages", height=300)
             pdf_output = gr.Markdown(label="OCR Result", elem_classes="output-markdown")
+            pdf_json_output = gr.JSON(label="Structured JSON Output")
             process_pdf_btn = gr.Button("Process PDF", variant="primary")
 
             def process_pdf(processor, pdf_file, pdf_url):
                 if not processor:
-                    return "Please set API key first", []
+                    return "Please set API key first", [], {}
                 logger.info(f"Received inputs - PDF file: {pdf_file}, PDF URL: {pdf_url}")
                 if pdf_file is not None and hasattr(pdf_file, 'name'):
                     logger.info(f"Processing as uploaded PDF: {pdf_file.name}")
-                    return processor.ocr_uploaded_pdf(pdf_file)
+                    markdown, image_paths, json_data = processor.ocr_uploaded_pdf(pdf_file)
                 elif pdf_url and pdf_url.strip():
                     logger.info(f"Processing as PDF URL: {pdf_url}")
-                    return processor.ocr_pdf_url(pdf_url)
-                return "Please upload a PDF or provide a valid URL", []
+                    markdown, image_paths, json_data = processor.ocr_pdf_url(pdf_url)
+                else:
+                    return "Please upload a PDF or provide a valid URL", [], {}
+                # Fallback to display images if markdown rendering fails
+                image_components = []
+                for path in image_paths:
+                    if path and os.path.exists(path):
+                        image_components.append(gr.Image(path, label=f"Page Image"))
+                return markdown, image_paths, json_data, gr.Column(*image_components) if image_components else gr.Markdown("No images available")
 
             process_pdf_btn.click(
                 fn=process_pdf,
                 inputs=[processor_state, pdf_input, pdf_url_input],
-                outputs=[pdf_output, pdf_gallery]
+                outputs=[pdf_output, pdf_gallery, pdf_json_output, gr.Column()]
             )
 
     return demo