Update app.py

app.py

@@ -18,6 +18,8 @@ if not ENDPOINT or not MODEL:
 
 def image_to_base64(image):
     buffered = BytesIO()
+    if image.mode == 'RGBA':
+        image = image.convert('RGB')
     image.save(buffered, format="PNG")
     return base64.b64encode(buffered.getvalue()).decode("utf-8")
 
@@ -61,25 +63,27 @@ def process_single_page(pdf_path, page_number):
     return img, total_pages
 
 
-def process_input(image, pdf_file, temperature, page_number):
-    if image is None and pdf_file is None:
+def process_input(file_input, temperature, page_number):
+    if file_input is None:
         yield "Please upload an image or PDF first.", "", ""
         return
     
     images_to_process = []
     page_info = ""
     
-    if pdf_file is not None:
+    file_path = file_input if isinstance(file_input, str) else file_input.name
+    
+    if file_path.lower().endswith('.pdf'):
         try:
             if page_number > 0:
-                img, total_pages = process_single_page(pdf_file, page_number)
+                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.", "", ""
                     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)
+                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.", "", ""
                     return
@@ -87,17 +91,31 @@ def process_input(image, pdf_file, temperature, page_number):
         except Exception as e:
             yield f"Error processing PDF: {str(e)}", "", ""
             return
-    elif image is not None:
-        images_to_process = [image]
+    else:
+        try:
+            img = Image.open(file_path)
+            images_to_process = [img]
+        except Exception as e:
+            yield f"Error opening image: {str(e)}", "", ""
+            return
+    
+    for img in images_to_process:
+        if not isinstance(img, Image.Image):
+            yield "Error: Invalid image format.", "", ""
+            return
     
     content = [{"type": "text", "text": ""}]
     
     for img in images_to_process:
-        b64_image = image_to_base64(img)
-        content.append({
-            "type": "image_url",
-            "image_url": {"url": f"data:image/png;base64,{b64_image}"}
-        })
+        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)}", "", ""
+            return
     
     payload = {
         "model": MODEL,
@@ -162,19 +180,13 @@ with gr.Blocks(title="📖 Image/PDF OCR", theme=gr.themes.Soft()) as demo:
     
     with gr.Row():
         with gr.Column(scale=1):
-            image_input = gr.Image(
-                type="pil",
-                label="🖼️ Upload Image",
-                sources=["upload", "clipboard"],
-                height=400
-            )
-            pdf_input = gr.File(
-                label="📄 Upload PDF",
-                file_types=[".pdf"],
+            file_input = gr.File(
+                label="🖼️ Upload Image or PDF",
+                file_types=[".pdf", ".png", ".jpg", ".jpeg"],
                 type="filepath"
             )
             page_number = gr.Number(
-                label="Specific Page Number (0 = first 5 pages)",
+                label="PDF: Page Number (0 = first 5 pages)",
                 value=0,
                 minimum=0,
                 step=1,
@@ -185,7 +197,7 @@ with gr.Blocks(title="📖 Image/PDF OCR", theme=gr.themes.Soft()) as demo:
                 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*")
+            gr.Markdown("*Upload an image (PNG/JPG) or PDF. For PDF: 0 = first 5 pages, or specify page number*")
             temperature = gr.Slider(
                 minimum=0.1,
                 maximum=1.0,
@@ -215,13 +227,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, page_number],
+        inputs=[file_input, temperature, page_number],
         outputs=[output_text, raw_output, page_info]
     )
     
     clear_btn.click(
-        fn=lambda: (None, None, "", "", 0, ""),
-        outputs=[image_input, pdf_input, output_text, raw_output, page_number, page_info]
+        fn=lambda: (None, "", "", 0, ""),
+        outputs=[file_input, output_text, raw_output, page_number, page_info]
     )