ocr.py

import os
import json
from transformers import AutoProcessor, AutoModelForImageTextToText
from PIL import Image
import torch
import re

# Load model & processor
processor = AutoProcessor.from_pretrained("JackChew/Qwen2-VL-2B-OCR")
model = AutoModelForImageTextToText.from_pretrained("JackChew/Qwen2-VL-2B-OCR").to(
    torch.device("cuda" if torch.cuda.is_available() else "cpu")
)

# OCR function
def ocr_image(image, prompt_text):
    conversation = [{
        "role": "user",
        "content": [{"type": "image"}, {"type": "text", "text": prompt_text}]
    }]
    prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
    inputs = processor(text=[prompt], images=[image], padding=True, return_tensors="pt").to(model.device)
    output_ids = model.generate(**inputs, max_new_tokens=2048)
    generated_ids = [output_ids[len(input_ids):] for input_ids, output_ids in zip(inputs.input_ids, output_ids)]
    return processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()

def extract_title_from_text(full_text, debug_filename=""):
    """Extract title from full OCR text using smart heuristics"""
    lines = [line.strip() for line in full_text.split('\n') if line.strip()]
    
    if debug_filename:
        print(f"\nDebug - Analyzing {debug_filename}:")
        for i, line in enumerate(lines[:10]):
            print(f"  {i}: '{line}'")
    
    # Look for title in first several lines
    for i, line in enumerate(lines[:10]):
        # Skip page numbers and very short lines
        if line.isdigit() or len(line) < 2:
            continue
            
        # Skip common non-title elements
        skip_phrases = [
            "frederick thayer", "oakland", "maryland", "published", "forum",
            "to a. s. d.", "your face", "word of god", "when i would"
        ]
        if any(phrase in line.lower() for phrase in skip_phrases):
            continue
            
        # Check for continuation markers
        if "(continued)" in line.lower() or "(cont" in line.lower():
            continue
            
        # Skip parenthetical subtitles for now (we'll add them back later)
        if line.startswith("(") and line.endswith(")"):
            subtitle = line
            continue
            
        # Look for title characteristics
        is_likely_title = False
        
        # All caps or mostly caps (allowing for some lowercase)
        if line.isupper() or (sum(1 for c in line if c.isupper()) > len(line) * 0.6):
            is_likely_title = True
            
        # Title case and reasonable length
        elif line.istitle() and 2 <= len(line) <= 40:
            is_likely_title = True
            
        # Check if it's a short line that's not clearly poem content
        elif (len(line.split()) <= 4 and 
              not line.lower().startswith(("when", "the", "and", "but", "or", "in", "on", "at", "to", "from")) and
              not any(char in line for char in ".,!?;:")):
            is_likely_title = True
            
        if is_likely_title:
            title = line.strip()
            
            # Check if next line is a subtitle in parentheses
            if i + 1 < len(lines):
                next_line = lines[i + 1].strip()
                if next_line.startswith("(") and next_line.endswith(")"):
                    title += f" {next_line}"
                    
            # Clean up spacing (fix OCR issues like "L E G E R D E M A I N")
            if len(title.split()) > 3 and all(len(word) <= 2 for word in title.split() if word.isalpha()):
                title = ''.join(title.split())
                
            if debug_filename:
                print(f"  -> Found title: '{title}'")
            return title
    
    if debug_filename:
        print(f"  -> No title found, using 'Untitled'")
    return "Untitled"

def clean_poem_text(text, title):
    """Remove title and other metadata from poem text"""
    lines = text.split('\n')
    cleaned_lines = []
    
    # Remove title lines from the beginning
    title_words = set(title.lower().replace('(', '').replace(')', '').split())
    skip_count = 0
    
    for i, line in enumerate(lines):
        line_words = set(line.lower().replace('(', '').replace(')', '').split())
        
        # Skip lines that are primarily the title
        if i < 5 and title_words and len(title_words.intersection(line_words)) > len(title_words) * 0.6:
            skip_count += 1
            continue
            
        # Skip metadata lines
        if any(phrase in line.lower() for phrase in ["frederick thayer", "oakland", "maryland", "published"]):
            continue
            
        cleaned_lines.append(line)
    
    return '\n'.join(cleaned_lines).strip()

# Process all images
image_folder = "img"
poems = {}  # Dictionary to group continuation pages

for filename in sorted(os.listdir(image_folder)):
    if filename.lower().endswith((".jpg", ".jpeg", ".png")):
        path = os.path.join(image_folder, filename)
        img = Image.open(path)
        
        print(f"\nProcessing {filename}...")
        
        # Full OCR with better prompt
        poem_text = ocr_image(img, "Transcribe all text from this document exactly as written, preserving line breaks and spacing.")
        
        # Extract title using smart heuristics
        title = extract_title_from_text(poem_text, filename)
        
        # Check if this is a continuation page
        is_continuation = "(continued)" in poem_text.lower() or "(cont" in poem_text.lower()
        
        if is_continuation and poems:
            # Find the most recent poem to continue
            last_poem_key = list(poems.keys())[-1]
            poems[last_poem_key]["text"] += "\n\n" + clean_poem_text(poem_text, title)
            poems[last_poem_key]["pages"].append(filename)
            print(f"  -> Continuation of '{last_poem_key}'")
        else:
            # New poem or first page
            clean_text = clean_poem_text(poem_text, title)
            
            if title in poems:
                # Same title, merge content
                poems[title]["text"] += "\n\n" + clean_text
                poems[title]["pages"].append(filename)
            else:
                # Brand new poem
                poems[title] = {
                    "title": title,
                    "text": clean_text,
                    "pages": [filename]
                }
            print(f"  -> Title: '{title}'")

# Convert to list format for JSON output
ocr_results = []
for poem_data in poems.values():
    ocr_results.append({
        "filename": poem_data["pages"][0],  # First page filename
        "title": poem_data["title"],
        "text": poem_data["text"],
        "pages": poem_data["pages"]  # All pages for this poem
    })

# Save to disk
with open("ocr_output.json", "w", encoding="utf-8") as f:
    json.dump(ocr_results, f, ensure_ascii=False, indent=2)

print(f"\nOCR complete! Processed {len(ocr_results)} poems and saved to ocr_output.json")