Upload module_ocr2.py

module_ocr2.py

@@ -0,0 +1,94 @@
+"""
+File: module_mistral_ocr.py
+
+Description: Gradio module to interact the Mistral_OCR model.
+
+Author: Didier Guillevic
+Date: 2025-04-07
+"""
+
+import gradio as gr
+import os
+import magic
+
+import ocr2 # OCR with software 2.0 models
+
+#
+# Get file type: PDF or Image or something else
+#
+def get_file_type(file_path):
+    # Check file extension
+    file_extension = os.path.splitext(file_path)[1].lower()
+
+    # Check MIME type
+    mime = magic.Magic(mime=True)
+    mime_type = mime.from_file(file_path)
+
+    # Determine file type
+    if file_extension == '.pdf' or mime_type == 'application/pdf':
+        return 'PDF'
+    elif file_extension in ['.jpg', '.jpeg', '.png', '.gif'] or mime_type.startswith('image/'):
+        return 'Image'
+    elif file_extension == '.pptx' or mime_type == 'application/vnd.openxmlformats-officedocument.presentationml.presentation':
+        return 'PowerPoint'
+    else:
+        return 'Other'
+
+#
+# Process one file
+#
+def process(input_file: str):
+    """Process given file with OCR using given languages."
+    """
+    file_type = get_file_type(input_file)
+    if file_type == 'PDF':
+        return ocr2.process_pdf(input_file)
+    elif file_type == 'Image':
+        return ocr2.process_image(input_file)
+    else:
+        return "Unsupported file type. Please upload a PDF, or an image file."
+    return ocr2.process(input_file)
+
+    
+#
+# User interface
+#
+with gr.Blocks() as demo:
+
+    # Upload file to process
+    with gr.Row():
+        input_file = gr.File(label="Upload a PDF file", scale=1)
+        output_text = gr.Textbox(label="OCR output", scale=2)
+
+    # Buttons
+    with gr.Row():
+        ocr_btn = gr.Button(value="OCR", variant="primary")
+        clear_btn = gr.Button("Clear", variant="secondary")
+    
+    # Examples
+    with gr.Accordion("Examples", open=False):
+        examples = gr.Examples(
+            [
+                ['./scanned_doc.pdf',],
+            ],
+            inputs=[input_file,],
+            outputs=[output_text,],
+            fn=process,
+            cache_examples=False,
+            label="Examples"
+        )
+    
+    # Functions
+    ocr_btn.click(
+        fn=process,
+        inputs=[input_file,],
+        outputs=[output_text,]
+    )
+    clear_btn.click(
+        fn=lambda : (None, ''),
+        inputs=[],
+        outputs=[input_file, output_text] # input_file, output_text
+    )
+
+if __name__ == '__main__':
+    demo.launch()