Update app.py

app.py

@@ -7,7 +7,6 @@ import gradio as gr
 from PIL import Image
 from io import BytesIO
 import pypdfium2 as pdfium
-from pathlib import Path
 
 ENDPOINT = os.environ.get("VLLM_ENDPOINT")
 MODEL = os.environ.get("VLLM_MODEL")
@@ -24,22 +23,22 @@ def image_to_base64(image):
     return base64.b64encode(buffered.getvalue()).decode("utf-8")
 
 
-def render_pdf_page(page, max_resolution=1540, scale=2.77):
+def render_pdf_page(page, max_resolution=1280, scale=2.77):
     width, height = page.get_size()
     pixel_width = width * scale
     pixel_height = height * scale
-    resize_factor = min(1, max_resolution / pixel_width, max_resolution / pixel_height)
+    resize_factor = min(max_resolution / pixel_width, max_resolution / pixel_height)
     target_scale = scale * resize_factor
     return page.render(scale=target_scale, rev_byteorder=True).to_pil()
 
 
-def process_pdf(pdf_path, max_pages=5):
+def process_pdf(pdf_path, num_pages=1):
     pdf = pdfium.PdfDocument(pdf_path)
     total_pages = len(pdf)
-    num_pages = min(total_pages, max_pages)
+    pages_to_process = min(num_pages, total_pages, 5)
     images = []
     
-    for i in range(num_pages):
+    for i in range(pages_to_process):
         page = pdf[i]
         img = render_pdf_page(page)
         images.append(img)
@@ -48,22 +47,7 @@ def process_pdf(pdf_path, max_pages=5):
     return images, total_pages
 
 
-def process_single_page(pdf_path, page_number):
-    pdf = pdfium.PdfDocument(pdf_path)
-    total_pages = len(pdf)
-    
-    if page_number < 1 or page_number > total_pages:
-        pdf.close()
-        return None, total_pages
-    
-    page = pdf[page_number - 1]
-    img = render_pdf_page(page)
-    pdf.close()
-    
-    return img, total_pages
-
-
-def process_input(file_input, temperature, page_number):
+def process_input(file_input, temperature, num_pages):
     if file_input is None:
         yield "Please upload an image or PDF first.", "", "", None
         return
@@ -76,22 +60,15 @@ def process_input(file_input, temperature, page_number):
     
     if file_path.lower().endswith('.pdf'):
         try:
-            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.", "", "", None
-                    return
-                images_to_process = [img]
-                display_image = img
-                page_info = f"Processing page {page_number} of {total_pages}"
-            else:
-                img, total_pages = process_single_page(file_path, 1)
-                if img is None:
-                    yield f"Error: Could not read PDF.", "", "", None
-                    return
-                images_to_process = [img]
-                display_image = img
+            images_to_process, total_pages = process_pdf(file_path, num_pages)
+            if len(images_to_process) == 0:
+                yield "Error: Could not extract pages from PDF.", "", "", None
+                return
+            display_image = images_to_process[0]
+            if len(images_to_process) == 1:
                 page_info = f"Processing page 1 of {total_pages}"
+            else:
+                page_info = f"Processing {len(images_to_process)} pages of {total_pages}"
         except Exception as e:
             yield f"Error processing PDF: {str(e)}", "", "", None
             return
@@ -100,15 +77,11 @@ def process_input(file_input, temperature, page_number):
             img = Image.open(file_path)
             images_to_process = [img]
             display_image = img
+            page_info = "Processing image"
         except Exception as 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.", "", "", None
-            return
-    
     content = [{"type": "text", "text": ""}]
     
     for img in images_to_process:
@@ -144,6 +117,7 @@ def process_input(file_input, temperature, page_number):
         response.raise_for_status()
 
         accumulated_response = ""
+        first_chunk = True
         
         for line in response.iter_lines():
             if line:
@@ -161,7 +135,11 @@ 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, display_image
+                            if first_chunk:
+                                yield accumulated_response, accumulated_response, page_info, display_image
+                                first_chunk = False
+                            else:
+                                yield accumulated_response, accumulated_response, page_info, gr.update()
                 except json.JSONDecodeError:
                     continue
                     
@@ -175,10 +153,10 @@ with gr.Blocks(title="📖 Image/PDF OCR", theme=gr.themes.Soft()) as demo:
         """
         # 📖 Image/PDF to Text Extraction
         **💡 How to use:**
-        1. Upload an image OR a PDF (max 5 pages)
-        2. Click "Extract Text" to process
-        
-        The model will extract and format text from your document.
+        1. Upload an image or PDF
+        2. For PDFs: choose how many pages to process (1-5, default is 1)
+        3. Adjust temperature if needed
+        4. Click "Extract Text"
         """
     )
     
@@ -190,24 +168,24 @@ with gr.Blocks(title="📖 Image/PDF OCR", theme=gr.themes.Soft()) as demo:
                 type="filepath"
             )
             rendered_image = gr.Image(
-                label="📄 Current Page/Image",
+                label="📄 Preview (First Page)",
                 type="pil",
-                height=400,
+                height=600,
                 interactive=False
             )
-            page_number = gr.Number(
-                label="PDF: Page Number (0 = auto first page, or specify 1, 2, 3...)",
-                value=0,
-                minimum=0,
+            num_pages = gr.Slider(
+                minimum=1,
+                maximum=5,
+                value=1,
                 step=1,
-                precision=0
+                label="PDF: Number of Pages to Process",
+                info="Only applies to PDF files (max 5 pages)"
             )
             page_info = gr.Textbox(
-                label="Page Info",
+                label="Processing Info",
                 value="",
                 interactive=False
             )
-            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,
@@ -237,13 +215,13 @@ 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],
+        inputs=[file_input, temperature, num_pages],
         outputs=[output_text, raw_output, page_info, rendered_image]
     )
     
     clear_btn.click(
-        fn=lambda: (None, "", "", 0, "", None),
-        outputs=[file_input, output_text, raw_output, page_number, page_info, rendered_image]
+        fn=lambda: (None, "", "", "", None, 1),
+        outputs=[file_input, output_text, raw_output, page_info, rendered_image, num_pages]
     )