Update module_ocr.py

module_ocr.py

@@ -7,6 +7,8 @@ Date: 2025-04-06
 
 import gradio as gr
 import ocr
+import pdf2image
+import tempfile
 
 #
 # Process one file
@@ -16,6 +18,45 @@ def process(input_file: str):
     """
     return ocr.process_file(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
@@ -24,7 +65,13 @@ 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 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
@@ -46,6 +93,13 @@ with gr.Blocks() as demo:
             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(
         fn=process,