MrGo2/dots-ocr

1import gradio as gr2from transformers import AutoModel, AutoTokenizer3from PIL import Image4import torch5 6# Load model and tokenizer7model_name = "rednote-hilab/dots.ocr"8print("Loading dots.ocr model...")9tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)10model = AutoModel.from_pretrained(11    model_name,12    trust_remote_code=True,13    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,14    device_map="auto"15)16print("Model loaded!")17 18def process_image(image):19    """Process image with dots.ocr"""20    if image is None:21        return "Please upload an image"22 23    try:24        # Run OCR25        result = model.generate(image, tokenizer)26        return result27    except Exception as e:28        return f"Error: {str(e)}"29 30# Create Gradio interface31iface = gr.Interface(32    fn=process_image,33    inputs=gr.Image(type="pil", label="Upload Image"),34    outputs=gr.Textbox(label="OCR Result", lines=10),35    title="dots.ocr - Multilingual OCR",36    description="Upload an image to extract text using dots.ocr. Supports 100+ languages, tables, formulas, and complex layouts.",37    examples=[38        ["examples/example1.jpg"],39        ["examples/example2.png"]40    ],41    theme="soft"42)43 44if __name__ == "__main__":45    iface.launch()