Update app.py

app.py

@@ -2,16 +2,29 @@
 import streamlit as st
 from transformers import BlipProcessor, BlipForConditionalGeneration
 from PIL import Image
+import requests
+from io import BytesIO
 from googletrans import Translator
 import torch
 
 # Initialize Translator
 translator = Translator()
 
-st.title("🖼️ AI Image Caption Generator - Advanced Version")
-st.write("Upload an image and get multiple captions generated by AI, in your preferred language!")
+st.set_page_config(page_title="AI Image Caption Generator", page_icon="🖼️")
 
-uploaded_file = st.file_uploader("Upload an image", type=["jpg", "png", "jpeg"])
+st.title("🖼️ AI Image Caption Generator")
+st.write("Upload an image or paste a Google Image URL to get multiple captions generated by AI, in your preferred language!")
+
+# Image input options
+upload_option = st.radio("Choose Image Input Method:", ("Upload from Computer", "Paste Image URL"))
+
+uploaded_file = None
+image_url = None
+
+if upload_option == "Upload from Computer":
+    uploaded_file = st.file_uploader("Upload an image", type=["jpg", "png", "jpeg"])
+else:
+    image_url = st.text_input("Paste Image URL (must be direct link ending with .jpg/.png/.jpeg)")
 
 # Language options
 languages = {
@@ -23,39 +36,58 @@ languages = {
     "Arabic": "ar"
 }
 
-selected_language = st.selectbox("Choose Output Language", list(languages.keys()))
+selected_language = st.selectbox("🌐 Choose Output Language", list(languages.keys()))
+
+# Load image based on input
+image = None
 
 if uploaded_file:
-    # Display image
     image = Image.open(uploaded_file).convert('RGB')
-    st.image(image, caption="Uploaded Image", width=300)
+elif image_url:
+    try:
+        response = requests.get(image_url)
+        if response.status_code == 200:
+            image = Image.open(BytesIO(response.content)).convert('RGB')
+        else:
+            st.error("Failed to fetch image. Please check the URL.")
+    except Exception as e:
+        st.error(f"Error fetching image: {e}")
 
-    st.write("Generating captions... please wait ⏳")
+if image:
+    st.image(image, caption="Selected Image", width=300)
 
-    # Resize for model input
-    image = image.resize((384, 384))
+    with st.spinner("Generating captions... please wait ⏳"):
+        # Resize for model input
+        resized_image = image.resize((384, 384))
 
-    processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
-    model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
+        processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+        model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
 
-    inputs = processor(image, return_tensors="pt")
-    
-    # Generate multiple captions
-    out = model.generate(
-        **inputs,
-        num_beams=5,                 # Use beam search
-        num_return_sequences=3,      # Generate 3 different captions
-        max_length=50,
-        early_stopping=True
-    )
+        inputs = processor(resized_image, return_tensors="pt")
 
-    captions = [processor.decode(o, skip_special_tokens=True) for o in out]
+        # Generate multiple captions
+        out = model.generate(
+            **inputs,
+            num_beams=5,
+            num_return_sequences=3,
+            max_length=50,
+            early_stopping=True
+        )
 
-    st.success("📝 Captions Generated!")
-    
-    st.subheader("Here are the captions:")
+        captions = [processor.decode(o, skip_special_tokens=True) for o in out]
+
+    st.success("📝 Captions Generated Successfully!")
+
+    st.subheader(f"Here are the captions in {selected_language}:")
 
     for idx, cap in enumerate(captions):
-        # Translate the caption based on selected language
-        translated_caption = translator.translate(cap, dest=languages[selected_language]).text
-        st.text_input(f"Caption {idx+1} ({selected_language})", value=translated_caption)
+        try:
+            translated_caption = translator.translate(cap, dest=languages[selected_language]).text
+        except Exception as e:
+            translated_caption = f"(Translation Error: Showing English) {cap}"
+
+        st.text_area(f"✏️ Caption {idx+1}", value=translated_caption, height=80)
+
+    st.caption("Tip: You can edit, copy, or download captions easily! ✍️")
+else:
+    st.info("Please upload an image or paste a valid URL to proceed.")