Update app.py

app.py

@@ -4,32 +4,53 @@ from PIL import Image
 import requests
 from io import BytesIO
 
-# Load the pipeline
+# Initialize the pipeline
 image_to_text = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")
 
-st.title('Upload the pictures you would like to purchase')
+st.title('Image Captioning Application')
 
-# Upload an image
-uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+# Function to load images from URL
+def load_image_from_url(url):
+    try:
+        response = requests.get(url)
+        img = Image.open(BytesIO(response.content))
+        return img
+    except Exception as e:
+        st.error(f"Error loading image from URL: {e}")
+        return None
 
-if uploaded_file is not None:
-    # Convert the file to an image
-    image = Image.open(uploaded_file)
-    st.image(image, caption='Uploaded Image', use_column_width=True)
-    
-    # Save image to a buffer
-    buf = BytesIO()
-    image.save(buf, format='JPEG')
-    byte_img = buf.getvalue()
+# User option to select input type: Upload or URL
+input_type = st.radio("Select input type:", ("Upload Image", "Image URL"))
 
-    # You might need to adjust this part depending on your model's needs
-    # For demonstration, assuming we need a URL - typically you'd need a model that accepts image bytes directly
-    # Example: result = your_model_that_accepts_bytes(byte_img)
+if input_type == "Upload Image":
+    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+    if uploaded_file is not None:
+        image = Image.open(uploaded_file)
+        st.image(image, caption='Uploaded Image', use_column_width=True)
+elif input_type == "Image URL":
+    image_url = st.text_input("Enter the image URL here:", "")
+    if image_url:
+        image = load_image_from_url(image_url)
+        if image:
+            st.image(image, caption='Image from URL', use_column_width=True)
 
-    # For this example, we just display a placeholder text
-    st.write("It seems like you want...")
-
-    # If your model accepts image bytes directly, you can call it like:
-    # result = image_to_text(byte_img)
-    # And then display the result
-    # st.write(result)
+# Generate caption button
+if st.button('Generate Caption'):
+    if not image:
+        st.warning("Please upload an image or enter an image URL.")
+    else:
+        with st.spinner("Generating caption..."):
+            # Process the image and generate caption
+            if input_type == "Upload Image":
+                # Save the uploaded image to a temporary file to pass its path to the model
+                with open("temp_image.jpg", "wb") as f:
+                    f.write(uploaded_file.getbuffer())
+                result = image_to_text("temp_image.jpg")
+            elif input_type == "Image URL" and image_url:
+                result = image_to_text(image_url)
+                
+            if result:
+                generated_text = result[0]['generated_text']
+                st.success(f'Generated Caption: {generated_text}')
+            else:
+                st.error("Failed to generate caption.")