ocr_inside_container.py

"""
DeepSeek-OCR Resume Extraction - runs INSIDE the vLLM container.
Reads pre-converted page images from /workspace/page_images_rotated/
and sends them to the local vLLM server.
"""

import base64
import time
import json
import sys
import os
from io import BytesIO

from PIL import Image
from openai import OpenAI

IMAGES_DIR = "/workspace/page_images"
ROTATED_DIR = "/workspace/page_images_rotated"
OUTPUT_DIR = "/workspace/ocr_output"
VLLM_BASE_URL = "http://localhost:8000/v1"
MODEL_NAME = "deepseek-ai/DeepSeek-OCR"


def rotate_images(src_dir, dst_dir):
    """Rotate all landscape images 90 degrees clockwise to portrait."""
    os.makedirs(dst_dir, exist_ok=True)
    files = sorted([f for f in os.listdir(src_dir) if f.endswith(".png")])
    print(f"Rotating {len(files)} images to portrait orientation...", flush=True)
    for f in files:
        img = Image.open(os.path.join(src_dir, f))
        w, h = img.size
        if w > h:
            img = img.rotate(-90, expand=True)
        img.save(os.path.join(dst_dir, f))
    print(f"  Done rotating.", flush=True)
    return sorted([f for f in os.listdir(dst_dir) if f.endswith(".png")])


def image_to_base64_url(image_path):
    with open(image_path, "rb") as f:
        data = base64.b64encode(f.read()).decode("utf-8")
    return f"data:image/png;base64,{data}"


def ocr_page(client, image_path, page_num):
    base64_url = image_to_base64_url(image_path)

    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image_url", "image_url": {"url": base64_url}},
                {"type": "text", "text": "Free OCR."}
            ]
        }
    ]

    start = time.time()
    try:
        response = client.chat.completions.create(
            model=MODEL_NAME,
            messages=messages,
            max_tokens=4096,
            temperature=0.0,
            extra_body={
                "skip_special_tokens": False,
                "vllm_xargs": {
                    "ngram_size": 30,
                    "window_size": 90,
                    "whitelist_token_ids": [128821, 128822],
                },
            },
        )
        elapsed = time.time() - start
        text = response.choices[0].message.content
        print(f"  Page {page_num}: OCR complete ({elapsed:.1f}s, {len(text)} chars)", flush=True)
        return text
    except Exception as e:
        elapsed = time.time() - start
        print(f"  Page {page_num}: ERROR after {elapsed:.1f}s - {e}", flush=True)
        return f"[ERROR on page {page_num}: {e}]"


def main():
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    # Rotate images to portrait
    image_files = rotate_images(IMAGES_DIR, ROTATED_DIR)
    print(f"Found {len(image_files)} page images", flush=True)

    if not image_files:
        print("No images found! Exiting.")
        sys.exit(1)

    # Connect to vLLM server
    client = OpenAI(api_key="EMPTY", base_url=VLLM_BASE_URL, timeout=3600)
    models = client.models.list()
    print(f"Server ready. Models: {[m.id for m in models.data]}", flush=True)

    # OCR each page
    all_results = []
    print(f"\nStarting OCR on {len(image_files)} pages...", flush=True)

    for i, img_file in enumerate(image_files):
        page_num = i + 1
        img_path = os.path.join(ROTATED_DIR, img_file)
        text = ocr_page(client, img_path, page_num)
        all_results.append({"page": page_num, "text": text})

        # Save individual page result
        page_file = os.path.join(OUTPUT_DIR, f"page_{page_num:03d}.txt")
        with open(page_file, "w", encoding="utf-8") as f:
            f.write(text)

    # Save consolidated output
    consolidated_file = os.path.join(OUTPUT_DIR, "all_resumes_ocr.txt")
    with open(consolidated_file, "w", encoding="utf-8") as f:
        for result in all_results:
            f.write(f"\n{'='*80}\n")
            f.write(f"PAGE {result['page']}\n")
            f.write(f"{'='*80}\n\n")
            f.write(result["text"])
            f.write("\n")

    # Also save as JSON
    json_file = os.path.join(OUTPUT_DIR, "all_resumes_ocr.json")
    with open(json_file, "w", encoding="utf-8") as f:
        json.dump(all_results, f, indent=2, ensure_ascii=False)

    print(f"\nDone! Results saved to:", flush=True)
    print(f"  Individual pages: {OUTPUT_DIR}/page_XXX.txt", flush=True)
    print(f"  Consolidated text: {consolidated_file}", flush=True)
    print(f"  JSON output: {json_file}", flush=True)


if __name__ == "__main__":
    main()