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

# Set up the OCR pipeline
@st.cache_resource
def load_model():
    return pipeline("image-to-text", model="microsoft/trocr-base-stage1")

ocr_model = load_model()

# Streamlit UI
st.title("Image Data Extractor using Hugging Face")
st.write("Upload an image, and this app will extract text from it using a Hugging Face model.")

# Upload an image
uploaded_image = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])

if uploaded_image:
    # Display the image
    image = Image.open(uploaded_image)
    st.image(image, caption="Uploaded Image", use_column_width=True)

    # Run OCR on the image
    with st.spinner("Extracting text..."):
        text = ocr_model(image.convert("RGB"))
    
    # Display the extracted text
    st.subheader("Extracted Text:")
    st.write(text[0]['generated_text'])