import gradio as gr
from paddleocr import PaddleOCR
from transformers import AutoModelForCausalLM, AutoTokenizer

# Initialize PaddleOCR
ocr = PaddleOCR(use_angle_cls=True, lang='en')

# Load GPT-NeoX model and tokenizer
model_name = "EleutherAI/gpt-neox-20b"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

def ocr_and_correct(text):
    # Function to correct text using GPT-NeoX
    inputs = tokenizer(text, return_tensors="pt")
    outputs = model.generate(**inputs)
    corrected_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return corrected_text

def process_image(image):
    # Perform OCR on the image
    result = ocr.ocr(image, cls=True)
    
    # Extract and join the text from OCR results
    extracted_text = " ".join([line[1][0] for line in result])
    
    # Correct the extracted text
    corrected_text = ocr_and_correct(extracted_text)
    
    return corrected_text

# Gradio interface
iface = gr.Interface(fn=process_image, inputs="image", outputs="text")

iface.launch()