Update module_ocr2.py

module_ocr2.py

@@ -11,6 +11,9 @@ import gradio as gr
 import os
 import magic
 
+import pdf2image
+import tempfile
+
 import ocr2 # OCR with software 2.0 models
 
 #
@@ -49,7 +52,46 @@ def process(input_file: str):
         return "Unsupported file type. Please upload a PDF, or an image file."
     return ocr2.process(input_file)
 
+#
+# Preview the document (image or PDF)
+#
+def preview_file(file):
+    if file is None:
+        return None, None
+    
+    file_path = file.name
+    file_extension = file_path.lower().split('.')[-1]
+    
+    if file_extension in ['jpg', 'jpeg', 'png', 'gif', 'bmp']:
+        # For images, return the image directly
+        return file_path, None
     
+    elif file_extension == 'pdf':
+        # For PDFs, convert first page to image using pdf2image
+        try:
+            # Convert only the first page for preview
+            pages = pdf2image.convert_from_path(
+                file_path, 
+                first_page=1, 
+                last_page=1,
+                dpi=150  # Good quality for preview
+            )
+            
+            if pages:
+                # Save the first page as a temporary image
+                with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as tmp_file:
+                    pages[0].save(tmp_file.name, 'PNG')
+                    return tmp_file.name, f"PDF Preview: {os.path.basename(file_path)}"
+            else:
+                return None, "<p>Could not convert PDF to image</p>"
+                
+        except Exception as e:
+            return None, f"<p>Error previewing PDF: {str(e)}</p>"
+    
+    else:
+        return None, f"<p>Preview not available for {file_extension} files</p>"
+
+
 #
 # User interface
 #
@@ -57,7 +99,14 @@ with gr.Blocks() as demo:
 
     # Upload file to process
     with gr.Row():
-        input_file = gr.File(label="Upload a PDF file", scale=1)
+        with gr.Column():
+            input_file = gr.File(
+                label="Upload a PDF or an image file",
+                file_types=[".pdf", ".jpg", ".jpeg", ".png", ".gif", ".bmp"],
+                scale=1)
+            preview_image = gr.Image(label="Preview", show_label=True)
+            preview_text = gr.HTML(label="Status")
+        
         output_text = gr.Textbox(label="OCR output", scale=2)
 
     # Buttons
@@ -78,6 +127,12 @@ with gr.Blocks() as demo:
             cache_examples=False,
             label="Examples"
         )
+    # Update preview when file is uploaded
+    input_file.change(
+        fn=preview_file,
+        inputs=[input_file],
+        outputs=[preview_image, preview_text]
+    )
     
     # Functions
     ocr_btn.click(