import streamlit as st
from PIL import Image
from transformers import TrOCRProcessor, VisionEncoderDecoderModel
import requests
import io

# Load the OCR model and processor
model_name = "microsoft/trocr-base-stage1"
processor = TrOCRProcessor.from_pretrained(model_name)
model = VisionEncoderDecoderModel.from_pretrained(model_name)

# Streamlit app title
st.title("OCR with TrOCR")

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

if uploaded_image is not None:
    # Open and display the uploaded image
    image = Image.open(uploaded_image)
    st.image(image, caption="Uploaded Image", use_column_width=True)

    # Convert image to suitable format
    inputs = processor(images=image, return_tensors="pt")

    # Perform OCR
    with torch.no_grad():
        outputs = model.generate(**inputs)

    # Decode the generated text
    text = processor.decode(outputs[0], skip_special_tokens=True)

    # Display the OCR result
    st.write("Extracted Text:")
    st.text(text)