Refactor app.py to improve OCR functionality and enhance Gradio interface with separate tabs for image and PDF processing; update requirements.txt to include necessary packages.

app.py

@@ -1,18 +1,20 @@
 import gradio as gr
-from PIL import Image
-from transformers import AutoTokenizer, AutoProcessor, AutoModelForImageTextToText
-import torch
 import spaces
+import torch
+from gradio_pdf import PDF
+from pdf2image import convert_from_path
+from PIL import Image
+from transformers import AutoModelForImageTextToText, AutoProcessor, AutoTokenizer
 
 model_path = "nanonets/Nanonets-OCR-s"
 
 # Load model once at startup
 print("Loading Nanonets OCR model...")
 model = AutoModelForImageTextToText.from_pretrained(
-    model_path, 
-    torch_dtype="auto", 
-    device_map="auto", 
-    attn_implementation="flash_attention_2"
+    model_path,
+    torch_dtype="auto",
+    device_map="auto",
+    attn_implementation="flash_attention_2",
 )
 model.eval()
 
@@ -20,40 +22,79 @@ tokenizer = AutoTokenizer.from_pretrained(model_path)
 processor = AutoProcessor.from_pretrained(model_path)
 print("Model loaded successfully!")
 
+
 @spaces.GPU()
 def ocr_image_gradio(image, max_tokens=4096):
     """Process image through Nanonets OCR model for Gradio interface"""
     if image is None:
         return "Please upload an image."
-    
-    try:
-        prompt = """Extract the text from the above document as if you were reading it naturally. Return the tables in html format. Return the equations in LaTeX representation. If there is an image in the document and image caption is not present, add a small description of the image inside the <img></img> tag; otherwise, add the image caption inside <img></img>. Watermarks should be wrapped in brackets. Ex: <watermark>OFFICIAL COPY</watermark>. Page numbers should be wrapped in brackets. Ex: <page_number>14</page_number> or <page_number>9/22</page_number>. Prefer using ☐ and ☑ for check boxes."""
-        
-        # Convert PIL image if needed
-        if not isinstance(image, Image.Image):
-            image = Image.fromarray(image)
-        
-        messages = [
-            {"role": "system", "content": "You are a helpful assistant."},
-            {"role": "user", "content": [
+
+    prompt = """Extract the text from the above document as if you were reading it naturally. Return the tables in html format. Return the equations in LaTeX representation. If there is an image in the document and image caption is not present, add a small description of the image inside the <img></img> tag; otherwise, add the image caption inside <img></img>. Watermarks should be wrapped in brackets. Ex: <watermark>OFFICIAL COPY</watermark>. Page numbers should be wrapped in brackets. Ex: <page_number>14</page_number> or <page_number>9/22</page_number>. Prefer using ☐ and ☑ for check boxes."""
+
+    # Convert PIL image if needed
+    if not isinstance(image, Image.Image):
+        image = Image.fromarray(image)
+
+    messages = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {
+            "role": "user",
+            "content": [
                 {"type": "image", "image": image},
                 {"type": "text", "text": prompt},
-            ]},
+            ],
+        },
+    ]
+
+    text = processor.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+    inputs = processor(text=[text], images=[image], padding=True, return_tensors="pt")
+    inputs = inputs.to(model.device)
+
+    with torch.no_grad():
+        output_ids = model.generate(
+            **inputs,
+            max_new_tokens=max_tokens,
+            do_sample=False,
+            repetition_penalty=1.25,
+        )
+        generated_ids = [
+            output_ids[len(input_ids) :]
+            for input_ids, output_ids in zip(inputs.input_ids, output_ids)
         ]
-        
-        text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-        inputs = processor(text=[text], images=[image], padding=True, return_tensors="pt")
-        inputs = inputs.to(model.device)
-        
-        with torch.no_grad():
-            output_ids = model.generate(**inputs, max_new_tokens=max_tokens, do_sample=False)
-            generated_ids = [output_ids[len(input_ids):] for input_ids, output_ids in zip(inputs.input_ids, output_ids)]
-            
-        output_text = processor.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)
-        return output_text[0]
-        
-    except Exception as e:
-        return f"Error processing image: {str(e)}"
+
+    output_text = processor.batch_decode(
+        generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
+    )
+    return output_text[0]
+
+
+@spaces.GPU()
+def ocr_pdf_gradio(pdf_path, max_tokens=4096, progress=gr.Progress()):
+    """Process each page of a PDF through Nanonets OCR model"""
+    if pdf_path is None:
+        return "Please upload a PDF file."
+
+    # Convert PDF to images
+    progress(0, desc="Converting PDF to images...")
+    pdf_images = convert_from_path(pdf_path)
+
+    # Process each page
+    all_text = []
+    total_pages = len(pdf_images)
+
+    for i, image in enumerate(pdf_images):
+        progress(
+            (i + 1) / total_pages, desc=f"Processing page {i + 1}/{total_pages}..."
+        )
+        page_text = ocr_image_gradio(image, max_tokens)
+        all_text.append(f"--- PAGE {i + 1} ---\n{page_text}\n")
+
+    # Combine results
+    combined_text = "\n".join(all_text)
+    return combined_text
+
 
 # Create Gradio interface
 with gr.Blocks(title="Nanonets OCR Demo") as demo:
@@ -77,47 +118,83 @@ with gr.Blocks(title="Nanonets OCR Demo") as demo:
         </div>
     </div>
     """)
-    
-    with gr.Row():
-        with gr.Column(scale=1):
-            image_input = gr.Image(
-                label="Upload Document Image",
-                type="pil",
-                height=400
-            )
-            max_tokens_slider = gr.Slider(
-                minimum=1024,
-                maximum=8192,
-                value=4096,
-                step=512,
-                label="Max Tokens",
-                info="Maximum number of tokens to generate"
-            )
-            extract_btn = gr.Button("Extract Text", variant="primary", size="lg")
-            
-        with gr.Column(scale=2):
-            output_text = gr.Textbox(
-                label="Extracted Text",
-                lines=20,
-                show_copy_button=True,
-                placeholder="Extracted text will appear here..."
-            )
-    
-    # Event handlers
-    extract_btn.click(
+
+    with gr.Tabs() as tabs:
+        # Image tab
+        with gr.TabItem("Image OCR"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    image_input = gr.Image(
+                        label="Upload Document Image", type="pil", height=400
+                    )
+                    image_max_tokens = gr.Slider(
+                        minimum=1024,
+                        maximum=8192,
+                        value=4096,
+                        step=512,
+                        label="Max Tokens",
+                        info="Maximum number of tokens to generate",
+                    )
+                    image_extract_btn = gr.Button(
+                        "Extract Text", variant="primary", size="lg"
+                    )
+
+                with gr.Column(scale=2):
+                    image_output_text = gr.Textbox(
+                        label="Extracted Text",
+                        lines=20,
+                        show_copy_button=True,
+                        placeholder="Extracted text will appear here...",
+                    )
+
+        # PDF tab
+        with gr.TabItem("PDF OCR"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    pdf_input = PDF(label="Upload PDF Document", height=400)
+                    pdf_max_tokens = gr.Slider(
+                        minimum=1024,
+                        maximum=8192,
+                        value=4096,
+                        step=512,
+                        label="Max Tokens per Page",
+                        info="Maximum number of tokens to generate for each page",
+                    )
+                    pdf_extract_btn = gr.Button(
+                        "Extract PDF Text", variant="primary", size="lg"
+                    )
+
+                with gr.Column(scale=2):
+                    pdf_output_text = gr.Textbox(
+                        label="Extracted Text (All Pages)",
+                        lines=20,
+                        show_copy_button=True,
+                        placeholder="Extracted text will appear here...",
+                    )
+
+    # Event handlers for Image tab
+    image_extract_btn.click(
         fn=ocr_image_gradio,
-        inputs=[image_input, max_tokens_slider],
-        outputs=output_text,
-        show_progress=True
+        inputs=[image_input, image_max_tokens],
+        outputs=image_output_text,
+        show_progress=True,
     )
-    
+
     image_input.change(
         fn=ocr_image_gradio,
-        inputs=[image_input, max_tokens_slider],
-        outputs=output_text,
-        show_progress=True
+        inputs=[image_input, image_max_tokens],
+        outputs=image_output_text,
+        show_progress=True,
+    )
+
+    # Event handlers for PDF tab
+    pdf_extract_btn.click(
+        fn=ocr_pdf_gradio,
+        inputs=[pdf_input, pdf_max_tokens],
+        outputs=pdf_output_text,
+        show_progress=True,
     )
-    
+
     # Add model information section
     with gr.Accordion("About Nanonets-OCR-s", open=False):
         gr.Markdown("""

requirements.txt

@@ -4,3 +4,5 @@ torchvision
 accelerate
 spaces
 https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.0.8/flash_attn-2.7.4.post1+cu126torch2.7-cp310-cp310-linux_x86_64.whl
+pdf2image
+gradio-pdf