import gradio as gr
from PIL import Image, ImageEnhance, ImageFilter
import torchvision.transforms as T
import torch
from transformers import AutoModel, AutoTokenizer

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Running on: {device.upper()}")

print("Loading model, please wait...")
model = AutoModel.from_pretrained(
    "baidu/Qianfan-OCR",
    trust_remote_code=True,
    torch_dtype=torch.float32
).to(device).eval()

tokenizer = AutoTokenizer.from_pretrained(
    "baidu/Qianfan-OCR",
    trust_remote_code=True
)
print("Model loaded!")

def enhance_image(image):
    image = image.convert("L").convert("RGB")
    image = ImageEnhance.Contrast(image).enhance(2.5)
    image = ImageEnhance.Sharpness(image).enhance(2.0)
    image = ImageEnhance.Brightness(image).enhance(1.3)
    image = image.filter(ImageFilter.SHARPEN)
    return image

def preprocess_image(image):
    transform = T.Compose([
        T.Resize((448, 448)),
        T.ToTensor(),
        T.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        )
    ])
    return transform(image.convert("RGB")).unsqueeze(0).to(device, torch.float32)

def extract_text(image):
    if image is None:
        return "Please upload an image."
    try:
        enhanced = enhance_image(image)
        pixel_values = preprocess_image(enhanced)
        response = model.chat(
            tokenizer,
            pixel_values,
            question="Please extract all text from this document.",
            generation_config={"max_new_tokens": 2048}
        )
        return response
    except Exception as e:
        return f"Error: {str(e)}"

app = gr.Interface(
    fn=extract_text,
    inputs=gr.Image(type="pil", label="Upload Document or Image"),
    outputs=gr.Textbox(label="Extracted Text", lines=20),
    title="🏛️ Museum Document OCR",
    description="Upload any document image to automatically extract all text from it.",
    theme=gr.themes.Soft()
)

app.launch()