app.py

#!/usr/bin/env python3
import os
import json
import base64
import requests
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")

if not ENDPOINT or not MODEL:
    raise ValueError("VLLM_ENDPOINT and VLLM_MODEL environment variables must be set.")


def image_to_base64(image):
    buffered = BytesIO()
    image.save(buffered, format="PNG")
    return base64.b64encode(buffered.getvalue()).decode("utf-8")


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(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)
    total_pages = len(pdf)
    num_pages = min(total_pages, max_pages)
    images = []
    
    for i in range(num_pages):
        page = pdf[i]
        img = render_pdf_page(page)
        images.append(img)
    
    pdf.close()
    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(image, pdf_file, temperature, page_number):
    if image is None and pdf_file is None:
        yield "Please upload an image or PDF first.", "", ""
        return
    
    images_to_process = []
    page_info = ""
    
    if pdf_file is not None:
        try:
            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)}", "", ""
            return
    elif image is not None:
        images_to_process = [image]
    
    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}"}
        })
    
    payload = {
        "model": MODEL,
        "messages": [
            {
                "role": "user",
                "content": content
            }
        ],
        "temperature": temperature,
        "stream": True
    }

    try:
        response = requests.post(
            ENDPOINT,
            headers={"Content-Type": "application/json"},
            data=json.dumps(payload),
            stream=True
        )
        response.raise_for_status()

        accumulated_response = ""
        
        for line in response.iter_lines():
            if line:
                line = line.decode('utf-8')
                if line.startswith('data: '):
                    line = line[6:]
                    
                if line.strip() == '[DONE]':
                    break
                    
                try:
                    chunk = json.loads(line)
                    if 'choices' in chunk and len(chunk['choices']) > 0:
                        delta = chunk['choices'][0].get('delta', {})
                        content_delta = delta.get('content', '')
                        if content_delta:
                            accumulated_response += content_delta
                            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, page_info


with gr.Blocks(title="📖 Image/PDF OCR", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 📖 Image/PDF to Text Extraction
        **💡 How to use:**
        1. Upload an image OR a PDF (max 5 pages)
        2. Adjust temperature if needed
        3. Click "Extract Text" to process
        
        The model will extract and format text from your document.
        """
    )
    
    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"],
                type="filepath"
            )
            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,
                value=0.2,
                step=0.05,
                label="Temperature"
            )
            submit_btn = gr.Button("Extract Text", variant="primary")
            clear_btn = gr.Button("Clear", variant="secondary")
        
        with gr.Column(scale=2):
            output_text = gr.Markdown(
                label="📄 Extracted Text (Rendered)",
                value="<div style='min-height: 600px; padding: 10px; border: 1px solid #e0e0e0; border-radius: 4px; background-color: #f9f9f9;'><em>Extracted text will appear here...</em></div>",
                height=600
            )
    
    with gr.Row():
        with gr.Column():
            raw_output = gr.Textbox(
                label="Raw Markdown Output",
                placeholder="Raw text will appear here...",
                lines=20,
                max_lines=30,
                show_copy_button=True
            )
    
    submit_btn.click(
        fn=process_input,
        inputs=[image_input, pdf_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]
    )


if __name__ == "__main__":
    demo.launch()