Assignment1 / app.py
CR7CAD's picture
Update app.py
e1ee436 verified
raw
history blame
1.23 kB
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}")