Update app.py

app.py

@@ -1,34 +1,35 @@
-import streamlit as st
+import torch
+from transformers import AutoProcessor, AutoModel, VisionEncoderDecoderModel, ViTFeatureExtractor, AutoTokenizer
 from PIL import Image
-import io
+import streamlit as st
 
-st.title("Artisan Product Submission Form")
+# Load the saved model state dictionary
+model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+model.load_state_dict(torch.load("model.pth", map_location=torch.device('cpu')))
 
-uploaded_file = st.file_uploader("Choose a file", type=["png", "jpg", "jpeg"])
- 
-if uploaded_file is not None:
-    # To read file as bytes:
-    bytes_data = uploaded_file.getvalue()
-    st.write("Filename: ", uploaded_file.name)
-    # st.write(bytes_data)  # This will display the raw bytes, typically not useful for users
+# Load the necessary components
+feature_extractor = ViTFeatureExtractor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
 
-    # To display the image
-    image = Image.open(io.BytesIO(bytes_data))
-    st.image(image, caption='Uploaded Image.', use_column_width=True)
+# Function to generate a caption for an image
+@st.cache_resource
+def generate_caption(image):
+    pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values
+    output_ids = model.generate(pixel_values, max_length=100, num_beams=5, early_stopping=True)
+    caption = tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0]
+    return caption
 
-# Creating text input box
+# Streamlit app
+def main():
+    st.title("Image Captioning")
+    uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
 
-st.header("Tell us about your product")
+    if uploaded_file is not None:
+        image = Image.open(uploaded_file)
+        st.image(image, caption="Uploaded Image", use_column_width=True)
 
-# Input fields
-product_type = st.text_input("Type of Product", placeholder="e.g., Handmade Jewelry, Pottery, Painting")
-product_origin = st.text_input("Product Origin", placeholder="e.g., City, Country, Region")
-product_description = st.text_area("Brief Description", placeholder="Provide a brief description of your product")
+        caption = generate_caption(image)
+        st.write(f"Caption: {caption}")
 
-# Submit button
-if st.button("Submit"):
-    st.write("Thank you for your submission!")
-    st.write("### Product Details")
-    st.write(f"**Type of Product:** {product_type}")
-    st.write(f"**Product Origin:** {product_origin}")
-    st.write(f"**Description:** {product_description}")
+if __name__ == "__main__":
+    main()