Create app.py

app.py

@@ -0,0 +1,158 @@
+import os
+import base64
+from io import BytesIO
+import gradio as gr
+from mistralai import Mistral
+from PIL import Image
+from pathlib import Path
+
+api_key = os.environ.get("MISTRAL")
+client = Mistral(api_key=api_key)
+
+#config
+VALID_DOCUMENT_EXTENSIONS = {".pdf"}
+VALID_IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png",}
+
+def upload_pdf(content, filename):
+    uploaded_file = client.files.upload(
+        file={"file_name": filename, "content": content},
+        purpose="ocr",
+    )
+    signed_url = client.files.get_signed_url(file_id=uploaded_file.id)
+    return signed_url.url
+def process_ocr(document_source):
+    return client.ocr.process(
+        model="mistral-ocr-latest",
+        document=document_source,
+        include_image_base64=True
+    )
+def do_ocr(input_type, url=None, file=None):
+    document_source = None
+    
+    if input_type == "URL":
+        if not url or url.strip() == "":
+            return "Please provide a valid URL.", "", []
+        url_lower = url.lower()
+        if any(url_lower.endswith(ext) for ext in VALID_IMAGE_EXTENSIONS):
+            document_source = {"type": "image_url", "image_url": url.strip()}
+        else:
+            document_source = {"type": "document_url", "document_url": url.strip()}
+
+    elif input_type == "Upload file":
+        if not file:
+            return "Please upload a file.", "", []
+        file_name = file.name.lower()
+        file_extension = os.path.splitext(file_name)[1]
+        if file_extension in VALID_DOCUMENT_EXTENSIONS:
+            with open(file.name, "rb") as f:
+                content = f.read()
+            signed_url = upload_pdf(content, os.path.basename(file_name))
+            document_source = {"type": "document_url", "document_url": signed_url}
+        elif file_extension in VALID_IMAGE_EXTENSIONS:
+            img = Image.open(file)
+            buffered = BytesIO()
+            img.save(buffered, format="PNG")
+            img_str = base64.b64encode(buffered.getvalue()).decode()
+            document_source = {"type": "image_url", "image_url": f"data:image/png;base64,{img_str}"}
+        else:
+            return f"Error: Unsupported file type. Supported types: {', '.join(VALID_DOCUMENT_EXTENSIONS | VALID_IMAGE_EXTENSIONS)}", "", []
+
+    else:
+        return "Invalid input type ", "", []
+
+    ocr_response = process_ocr(document_source)
+    markdown_text = "\n\n".join(page.markdown for page in ocr_response.pages)
+    extracted_text = markdown_text
+    rendered_markdown = markdown_text
+    images = []
+
+    for page in ocr_response.pages:
+        for img in page.images:
+            if img.image_base64:
+                base64_str = img.image_base64
+                if "," in base64_str:
+                    base64_str = base64_str.split(",")[1]
+                img_bytes = base64.b64decode(base64_str)
+                img_pil = Image.open(BytesIO(img_bytes))
+                images.append(img_pil)
+                img_buffer = BytesIO()
+                img_pil.save(img_buffer, format="PNG")
+                img_base64 = base64.b64encode(img_buffer.getvalue()).decode()
+                data_url = f"data:image/png;base64,{img_base64}"
+                rendered_markdown = rendered_markdown.replace(
+                    f"![{img.id}]({img.id})", f"![{img.id}]({data_url})"
+                )
+            else:
+                rendered_markdown += f"\n\n[Image Warning: No base64 data for {img.id}]"
+
+    return extracted_text.strip(), rendered_markdown.strip(), images
+        
+custom_css = """
+    body {font-family: body {font-family: 'Helvetica Neue', Helvetica;}
+    .gr-button {background-color: #4CAF50; color: white; border: none; padding: 10px 20px; border-radius: 5px;}
+    .gr-button:hover {background-color: #45a049;}
+    .gr-textbox {margin-bottom: 15px;}
+    .example-button {background-color: #1E90FF; color: white; border: none; padding: 8px 15px; border-radius: 5px; margin: 5px;}
+    .example-button:hover {background-color: #FF4500;}
+    .tall-radio .gr-radio-item {padding: 15px 0; min-height: 50px; display: flex; align-items: center;}
+    .tall-radio label {font-size: 16px;}
+"""
+with gr.Blocks(
+    title="Mistral OCR Demo",
+    css=custom_css,
+    theme=gr.themes.Soft()
+) as demo:
+    gr.Markdown("<h1 style='text-align: center; color: #333;'>Mistral OCR Demo</h1>")
+    gr.Markdown("<p style='text-align: center; color: #666;'>Extract text and images from PDFs or images using Mistral's latest OCR model. You can also see markdown live.</p>")
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            input_type = gr.Radio(
+                choices=["URL", "Upload file"],
+                label="Input Type",
+                value="URL",
+                elem_classes="tall-radio"
+            )
+            url_input = gr.Textbox(
+                label="Document or Image URL",
+                placeholder="e.g., https://arxiv.org/pdf/2501.12948",
+                visible=True,
+                lines=1
+            )
+            file_input = gr.File(
+                label="Upload PDF or Image",
+                file_types=[".pdf", ".jpg", ".jpeg", ".png"],
+                visible=False
+            )
+            submit_btn = gr.Button("Extract Text and Images")
+            gr.Markdown("### Try These Examples")
+            pdf_example = gr.Button("PDF", elem_classes="example-button")
+            img_example = gr.Button("Image", elem_classes="example-button")
+        with gr.Column(scale=2):
+            cleaned_output = gr.Textbox(label="Extracted Plain Text", lines=10, show_copy_button=True)
+            markdown_output = gr.Markdown(label="Rendered Markdown Text")
+            image_output = gr.Gallery(label="OCR Extracted Images", columns=2, height="auto")
+
+    def update_visibility(choice):
+        return gr.update(visible=(choice == "URL")), gr.update(visible=(choice == "Upload file"))
+    input_type.change(fn=update_visibility, inputs=input_type, outputs=[url_input, file_input])
+    def set_url_and_type(url):
+        return url, "URL"
+
+    pdf_example.click(
+        fn=lambda: set_url_and_type("https://arxiv.org/pdf/2501.12948"),
+        outputs=[url_input, input_type]
+    )
+    img_example.click(
+        fn=lambda: set_url_and_type("https://huggingface.co/datasets/agents-course/course-images/resolve/main/en/unit0/recommended-pace.jpg"),
+        outputs=[url_input, input_type]
+    )
+
+    submit_btn.click(
+        fn=do_ocr,
+        inputs=[input_type, url_input, file_input],
+        outputs=[cleaned_output, markdown_output, image_output]
+    )
+    
+if __name__ == "__main__":
+    demo.launch()