Update app.py

app.py

@@ -65,11 +65,12 @@ def process_single_page(pdf_path, page_number):
 
 def process_input(file_input, temperature, page_number):
     if file_input is None:
-        yield "Please upload an image or PDF first.", "", ""
+        yield "Please upload an image or PDF first.", "", "", None
         return
     
     images_to_process = []
     page_info = ""
+    display_image = None
     
     file_path = file_input if isinstance(file_input, str) else file_input.name
     
@@ -78,30 +79,34 @@ def process_input(file_input, temperature, page_number):
             if page_number > 0:
                 img, total_pages = process_single_page(file_path, page_number)
                 if img is None:
-                    yield f"Error: Page {page_number} does not exist. PDF has {total_pages} pages.", "", ""
+                    yield f"Error: Page {page_number} does not exist. PDF has {total_pages} pages.", "", "", None
                     return
                 images_to_process = [img]
+                display_image = img
                 page_info = f"Processing page {page_number} of {total_pages}"
             else:
-                images_to_process, total_pages = process_pdf(file_path, max_pages=5)
-                if len(images_to_process) == 0:
-                    yield "Error: Could not extract pages from PDF.", "", ""
+                img, total_pages = process_single_page(file_path, 1)
+                if img is None:
+                    yield f"Error: Could not read PDF.", "", "", None
                     return
-                page_info = f"Processing first {len(images_to_process)} pages of {total_pages}"
+                images_to_process = [img]
+                display_image = img
+                page_info = f"Processing page 1 of {total_pages}"
         except Exception as e:
-            yield f"Error processing PDF: {str(e)}", "", ""
+            yield f"Error processing PDF: {str(e)}", "", "", None
             return
     else:
         try:
             img = Image.open(file_path)
             images_to_process = [img]
+            display_image = img
         except Exception as e:
-            yield f"Error opening image: {str(e)}", "", ""
+            yield f"Error opening image: {str(e)}", "", "", None
             return
     
     for img in images_to_process:
         if not isinstance(img, Image.Image):
-            yield "Error: Invalid image format.", "", ""
+            yield "Error: Invalid image format.", "", "", None
             return
     
     content = [{"type": "text", "text": ""}]
@@ -114,7 +119,7 @@ def process_input(file_input, temperature, page_number):
                 "image_url": {"url": f"data:image/png;base64,{b64_image}"}
             })
         except Exception as e:
-            yield f"Error encoding image: {str(e)}", "", ""
+            yield f"Error encoding image: {str(e)}", "", "", display_image
             return
     
     payload = {
@@ -156,13 +161,13 @@ def process_input(file_input, temperature, page_number):
                         content_delta = delta.get('content', '')
                         if content_delta:
                             accumulated_response += content_delta
-                            yield accumulated_response, accumulated_response, page_info
+                            yield accumulated_response, accumulated_response, page_info, display_image
                 except json.JSONDecodeError:
                     continue
                     
     except Exception as e:
         error_msg = f"Error: {str(e)}"
-        yield error_msg, error_msg, page_info
+        yield error_msg, error_msg, page_info, display_image
 
 
 with gr.Blocks(title="📖 Image/PDF OCR", theme=gr.themes.Soft()) as demo:
@@ -185,8 +190,14 @@ with gr.Blocks(title="📖 Image/PDF OCR", theme=gr.themes.Soft()) as demo:
                 file_types=[".pdf", ".png", ".jpg", ".jpeg"],
                 type="filepath"
             )
+            rendered_image = gr.Image(
+                label="📄 Current Page/Image",
+                type="pil",
+                height=400,
+                interactive=False
+            )
             page_number = gr.Number(
-                label="PDF: Page Number (0 = first 5 pages)",
+                label="PDF: Page Number (0 = auto first page, or specify 1, 2, 3...)",
                 value=0,
                 minimum=0,
                 step=1,
@@ -197,7 +208,7 @@ with gr.Blocks(title="📖 Image/PDF OCR", theme=gr.themes.Soft()) as demo:
                 value="",
                 interactive=False
             )
-            gr.Markdown("*Upload an image (PNG/JPG) or PDF. For PDF: 0 = first 5 pages, or specify page number*")
+            gr.Markdown("*Upload an image (PNG/JPG) or PDF. For PDF: 0 = page 1 automatically, or specify any page number*")
             temperature = gr.Slider(
                 minimum=0.1,
                 maximum=1.0,
@@ -228,12 +239,12 @@ with gr.Blocks(title="📖 Image/PDF OCR", theme=gr.themes.Soft()) as demo:
     submit_btn.click(
         fn=process_input,
         inputs=[file_input, temperature, page_number],
-        outputs=[output_text, raw_output, page_info]
+        outputs=[output_text, raw_output, page_info, rendered_image]
     )
     
     clear_btn.click(
-        fn=lambda: (None, "", "", 0, ""),
-        outputs=[file_input, output_text, raw_output, page_number, page_info]
+        fn=lambda: (None, "", "", 0, "", None),
+        outputs=[file_input, output_text, raw_output, page_number, page_info, rendered_image]
     )