import streamlit as st
from transformers import pipeline
from PIL import Image

# Set the title of the app
st.title("Image-to-Text Converter using Donut")

# Description of the app
st.write("Upload an image to extract text using the Donut model (naver-clova-ix/donut-base).")

# Create a file uploader for image files
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

# Initialize the pipeline
@st.cache_resource(show_spinner=False)
def load_pipeline():
    return pipeline("image-to-text", model="naver-clova-ix/donut-base")

pipe = load_pipeline()

if uploaded_file is not None:
    try:
        # Open the image file and convert to RGB (if necessary)
        image = Image.open(uploaded_file).convert("RGB")
        st.image(image, caption="Uploaded Image", use_column_width=True)
        
        # Process the image through the pipeline
        result = pipe(image)
        
        # Extract generated text from the result list
        generated_text = result[0].get("generated_text", "No text generated.")
        
        st.subheader("Extracted Text")
        st.text_area("Result", generated_text, height=200)
    except Exception as e:
        st.error(f"An error occurred: {e}")