Update app.py

app.py

@@ -32,68 +32,50 @@ def render_pdf_page(page, max_resolution=1540, scale=2.77):
     return page.render(scale=target_scale, rev_byteorder=True).to_pil()
 
 
-def process_pdf(pdf_path, num_pages=1):
+def process_pdf(pdf_path, page_num=1):
     pdf = pdfium.PdfDocument(pdf_path)
     total_pages = len(pdf)
-    pages_to_process = min(int(num_pages), total_pages, 5)
-    images = []
+    page_idx = min(max(int(page_num) - 1, 0), total_pages - 1)
     
-    for i in range(pages_to_process):
-        page = pdf[i]
-        img = render_pdf_page(page)
-        images.append(img)
+    page = pdf[page_idx]
+    img = render_pdf_page(page)
     
     pdf.close()
-    return images, total_pages
+    return img, total_pages, page_idx + 1
 
 
-def process_input(file_input, temperature, num_pages):
+def process_input(file_input, temperature, page_num):
     if file_input is None:
         yield "Please upload an image or PDF first.", "", "", None
         return
     
-    images_to_process = []
+    image_to_process = None
     page_info = ""
-    display_image = None
     
     file_path = file_input if isinstance(file_input, str) else file_input.name
     
     if file_path.lower().endswith('.pdf'):
         try:
-            images_to_process, total_pages = process_pdf(file_path, int(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}"
+            image_to_process, total_pages, actual_page = process_pdf(file_path, int(page_num))
+            page_info = f"Processing page {actual_page} of {total_pages}"
         except Exception as e:
             yield f"Error processing PDF: {str(e)}", "", "", None
             return
     else:
         try:
-            img = Image.open(file_path)
-            images_to_process = [img]
-            display_image = img
+            image_to_process = Image.open(file_path)
             page_info = "Processing image"
         except Exception as e:
             yield f"Error opening image: {str(e)}", "", "", None
             return
     
-    content = [{"type": "text", "text": "Extract the text from this image."}]
-    
-    for img in images_to_process:
-        try:
-            b64_image = image_to_base64(img)
-            content.append({
-                "type": "image_url",
-                "image_url": {"url": f"data:image/png;base64,{b64_image}"}
-            })
-        except Exception as e:
-            yield f"Error encoding image: {str(e)}", "", "", display_image
-            return
+    content = [
+        {"type": "text", "text": "Extract the text from this image."},
+        {
+            "type": "image_url",
+            "image_url": {"url": f"data:image/png;base64,{image_to_base64(image_to_process)}"}
+        }
+    ]
     
     payload = {
         "model": MODEL,
@@ -131,7 +113,7 @@ def process_input(file_input, temperature, num_pages):
                         if content_delta:
                             accumulated_response += content_delta
                             if first_chunk:
-                                yield accumulated_response, accumulated_response, page_info, display_image
+                                yield accumulated_response, accumulated_response, page_info, image_to_process
                                 first_chunk = False
                             else:
                                 yield accumulated_response, accumulated_response, page_info, gr.update()
@@ -140,7 +122,7 @@ def process_input(file_input, temperature, num_pages):
                     
     except Exception as e:
         error_msg = f"Error: {str(e)}"
-        yield error_msg, error_msg, page_info, display_image
+        yield error_msg, error_msg, page_info, image_to_process
 
 
 with gr.Blocks(title="📖 Image/PDF OCR", theme=gr.themes.Soft()) as demo:
@@ -149,7 +131,7 @@ with gr.Blocks(title="📖 Image/PDF OCR", theme=gr.themes.Soft()) as demo:
 
 **💡 How to use:**
 1. Upload an image or PDF
-2. For PDFs: choose how many pages to process (1-5, default is 1)
+2. For PDFs: select which page to extract (1-20)
 3. Adjust temperature if needed
 4. Click "Extract Text"
 
@@ -171,11 +153,11 @@ with gr.Blocks(title="📖 Image/PDF OCR", theme=gr.themes.Soft()) as demo:
             )
             num_pages = gr.Slider(
                 minimum=1,
-                maximum=5,
+                maximum=20,
                 value=1,
                 step=1,
-                label="PDF: Number of Pages to Process",
-                info="Only applies to PDF files (max 5 pages)"
+                label="PDF: Page Number",
+                info="Select which page to extract"
             )
             page_info = gr.Textbox(
                 label="Processing Info",