Update app.py

app.py

@@ -1,26 +1,38 @@
 import streamlit as st
-from transformers import AutoFeatureExtractor, AutoModelForImageCaptioning
+import torch
+from transformers import AutoFeatureExtractor, AutoModelForSequenceClassification, AutoTokenizer
 from PIL import Image
-import requests
-from io import BytesIO
 
-st.title("Image Captioning App")
-
-feature_extractor = AutoFeatureExtractor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
-model = AutoModelForImageCaptioning.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+# Load the pretrained model and tokenizer
+model_name = "nlpconnect/vit-gpt2-image-captioning"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSequenceClassification.from_pretrained(model_name)
 
+# Define a function to generate captions from an image
 def generate_caption(image):
-    inputs = feature_extractor(images=image, return_tensors="pt")
-    caption = model.generate(**inputs)
-    return caption[0]
+    inputs = tokenizer(image, return_tensors="pt")
+    with torch.no_grad():
+        logits = model(**inputs).logits
+    caption = tokenizer.decode(logits.argmax(1)[0], skip_special_tokens=True)
+    return caption
 
-uploaded_image = st.file_uploader("Upload an image", type=["jpg", "png", "jpeg"])
+def main():
+    st.title("Image to Text Captioning")
 
-if uploaded_image:
-    image = Image.open(uploaded_image)
-    st.image(image, caption="Uploaded Image", use_column_width=True)
+    with st.form("my_form"):
+        uploaded_file = st.file_uploader("Choose an image file", type=["jpg", "jpeg", "png"])
 
-    caption = generate_caption(image)
+        if uploaded_file is not None:
+            # Display the uploaded image
+            image = Image.open(uploaded_file)
+            st.image(image, caption="Uploaded Image", use_column_width=True)
+        
+        clicked = st.form_submit_button("Generate Caption")
+        if clicked:
+            if "image" in locals():
+                caption = generate_caption(image)
+                st.subheader("Generated Caption:")
+                st.write(caption)
 
-    st.subheader("Generated Caption:")
-    st.write(caption)
+if __name__ == "__main__":
+    main()