Update app.py

app.py

@@ -22,18 +22,19 @@ 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):
     pdf = pdfium.PdfDocument(pdf_path)
-    num_pages = min(len(pdf), max_pages)
+    total_pages = len(pdf)
+    num_pages = min(total_pages, max_pages)
     images = []
     
     for i in range(num_pages):
@@ -42,24 +43,49 @@ def process_pdf(pdf_path, max_pages=5):
         images.append(img)
     
     pdf.close()
-    return images
+    return images, total_pages
 
 
-def process_input(image, pdf_file, temperature):
+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(image, pdf_file, temperature, page_number):
     if image is None and pdf_file is None:
-        yield "Please upload an image or PDF first.", ""
+        yield "Please upload an image or PDF first.", "", ""
         return
     
     images_to_process = []
+    page_info = ""
     
     if pdf_file is not None:
         try:
-            images_to_process = process_pdf(pdf_file, max_pages=5)
-            if len(images_to_process) == 0:
-                yield "Error: Could not extract pages from PDF.", ""
-                return
+            if page_number > 0:
+                img, total_pages = process_single_page(pdf_file, page_number)
+                if img is None:
+                    yield f"Error: Page {page_number} does not exist. PDF has {total_pages} pages.", "", ""
+                    return
+                images_to_process = [img]
+                page_info = f"Processing page {page_number} of {total_pages}"
+            else:
+                images_to_process, total_pages = process_pdf(pdf_file, max_pages=5)
+                if len(images_to_process) == 0:
+                    yield "Error: Could not extract pages from PDF.", "", ""
+                    return
+                page_info = f"Processing first {len(images_to_process)} pages of {total_pages}"
         except Exception as e:
-            yield f"Error processing PDF: {str(e)}", ""
+            yield f"Error processing PDF: {str(e)}", "", ""
             return
     elif image is not None:
         images_to_process = [image]
@@ -112,13 +138,13 @@ def process_input(image, pdf_file, temperature):
                         content_delta = delta.get('content', '')
                         if content_delta:
                             accumulated_response += content_delta
-                            yield accumulated_response, accumulated_response
+                            yield accumulated_response, accumulated_response, page_info
                 except json.JSONDecodeError:
                     continue
                     
     except Exception as e:
         error_msg = f"Error: {str(e)}"
-        yield error_msg, error_msg
+        yield error_msg, error_msg, page_info
 
 
 with gr.Blocks(title="📖 Image/PDF OCR", theme=gr.themes.Soft()) as demo:
@@ -143,11 +169,23 @@ with gr.Blocks(title="📖 Image/PDF OCR", theme=gr.themes.Soft()) as demo:
                 height=400
             )
             pdf_input = gr.File(
-                label="📄 Upload PDF (max 5 pages)",
+                label="📄 Upload PDF",
                 file_types=[".pdf"],
                 type="filepath"
             )
-            gr.Markdown("*Upload either an image or PDF, not both*")
+            page_number = gr.Number(
+                label="Specific Page Number (0 = first 5 pages)",
+                value=0,
+                minimum=0,
+                step=1,
+                precision=0
+            )
+            page_info = gr.Textbox(
+                label="Page Info",
+                value="",
+                interactive=False
+            )
+            gr.Markdown("*Upload either an image or PDF. For PDF: set page number to 0 for first 5 pages, or specify a page*")
             temperature = gr.Slider(
                 minimum=0.1,
                 maximum=1.0,
@@ -177,13 +215,13 @@ with gr.Blocks(title="📖 Image/PDF OCR", theme=gr.themes.Soft()) as demo:
     
     submit_btn.click(
         fn=process_input,
-        inputs=[image_input, pdf_input, temperature],
-        outputs=[output_text, raw_output]
+        inputs=[image_input, pdf_input, temperature, page_number],
+        outputs=[output_text, raw_output, page_info]
     )
     
     clear_btn.click(
-        fn=lambda: (None, None, "", ""),
-        outputs=[image_input, pdf_input, output_text, raw_output]
+        fn=lambda: (None, None, "", "", 0, ""),
+        outputs=[image_input, pdf_input, output_text, raw_output, page_number, page_info]
     )