Update app.py

app.py

@@ -1,25 +1,31 @@
-from transformers import AutoFeatureExtractor, AutoModelForImageClassification
+import streamlit as st
+from transformers import pipeline
 from PIL import Image
-import gradio as gr
+import os
 
-# Load the model and feature extractor
-feature_extractor = AutoFeatureExtractor.from_pretrained("google/vit-base-patch16-224")
-model = AutoModelForImageClassification.from_pretrained("google/vit-base-patch16-224")
+# Set up the OCR pipeline
+@st.cache_resource
+def load_model():
+    return pipeline("image-to-text", model="microsoft/trocr-base-stage1")
 
-# Define prediction function
-def classify_image(image):
-    image = Image.fromarray(image).convert("RGB")
-    inputs = feature_extractor(images=image, return_tensors="pt")
-    outputs = model(**inputs)
-    predicted_class = outputs.logits.argmax(-1).item()
-    return model.config.id2label[predicted_class]
+ocr_model = load_model()
 
-# Create a Gradio app
-app = gr.Interface(
-    fn=classify_image,
-    inputs=gr.Image(type="numpy"),
-    outputs="text",
-    title="Image Classifier"
-)
+# 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.")
 
-app.launch()
+# 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'])