from doctr.io import DocumentFile
from doctr.models import ocr_predictor
from PIL import Image
import gradio as gr
from transformers import pipeline

# Load the model
predictor = ocr_predictor(pretrained=True)
summarizer = pipeline("summarization", model="facebook/bart-large-cnn")

def read_img(image):
    # save image
    image.save('image.png')
    
    # Read the image
    doc = DocumentFile.from_images(['image.png'])
    # Run the model
    result = predictor(doc)
    # Return the result
    image_text = result.render()
    summary_text = summarizer(image_text, max_length=150, min_length=50, do_sample=False)
    return image_text, summary_text[0].get('summary_text'), result.show(doc)

app = gr.Interface(
    title="Image OCR",
    description="Extract text from images",
    inputs=gr.Image(type="pil", label="Input image"),
    outputs=[gr.Textbox(label="Image Text"), gr.Textbox(label="Summary"), gr.Image(label="Image")],
    fn=read_img
)
    

app.launch(debug=True)