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()
    if image.mode == 'RGBA':
        image = image.convert('RGB')
    image.save(buffered, format="PNG")
    return base64.b64encode(buffered.getvalue()).decode("utf-8")


def render_pdf_page(page, max_resolution=1540, 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)
    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(file_input, temperature, page_number):
    if file_input is None:
        yield "Please upload an image or PDF first.", "", "", None
        return
    
    images_to_process = []
    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:
            if page_number > 0:
                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.", "", "", None
                    return
                images_to_process = [img]
                display_image = img
                page_info = f"Processing page {page_number} of {total_pages}"
            else:
                img, total_pages = process_single_page(file_path, 1)
                if img is None:
                    yield f"Error: Could not read PDF.", "", "", None
                    return
                images_to_process = [img]
                display_image = img
                page_info = f"Processing page 1 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
        except Exception as e:
            yield f"Error opening image: {str(e)}", "", "", None
            return
    
    for img in images_to_process:
        if not isinstance(img, Image.Image):
            yield "Error: Invalid image format.", "", "", None
            return
    
    content = [{"type": "text", "text": ""}]
    
    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
    
    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, display_image
                except json.JSONDecodeError:
                    continue
                    
    except Exception as e:
        error_msg = f"Error: {str(e)}"
        yield error_msg, error_msg, page_info, display_image


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. Click "Extract Text" to process
        
        The model will extract and format text from your document.
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            file_input = gr.File(
                label="🖼️ Upload Image or PDF",
                file_types=[".pdf", ".png", ".jpg", ".jpeg"],
                type="filepath"
            )
            rendered_image = gr.Image(
                label="📄 Current Page/Image",
                type="pil",
                height=400,
                interactive=False
            )
            page_number = gr.Number(
                label="PDF: Page Number (0 = auto first page, or specify 1, 2, 3...)",
                value=0,
                minimum=0,
                step=1,
                precision=0
            )
            page_info = gr.Textbox(
                label="Page Info",
                value="",
                interactive=False
            )
            gr.Markdown("*Upload an image (PNG/JPG) or PDF. For PDF: 0 = page 1 automatically, or specify any page number*")
            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=[file_input, temperature, page_number],
        outputs=[output_text, raw_output, page_info, rendered_image]
    )
    
    clear_btn.click(
        fn=lambda: (None, "", "", 0, "", None),
        outputs=[file_input, output_text, raw_output, page_number, page_info, rendered_image]
    )


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