Update app.py

app.py

@@ -24,7 +24,7 @@ image_transform = transforms.Compose([
 
 def load_image_from_url(url):
     try:
-        response = requests.get(url, timeout=10)
+        response = requests.get(url.strip(), timeout=10)
         response.raise_for_status()
         return Image.open(BytesIO(response.content)).convert("RGB")
     except Exception:
@@ -34,7 +34,7 @@ def load_image_from_url(url):
 def predict_tags(image_url, threshold=0.5):
     image = load_image_from_url(image_url)
     if image is None:
-        return [], "Could not load image from the provided URL."
+        return None, "Could not load image."
 
     image_tensor = image_transform(image).unsqueeze(0).to(device)
     with torch.no_grad():
@@ -49,25 +49,36 @@ def predict_tags(image_url, threshold=0.5):
     results.sort(key=lambda x: x[1], reverse=True)
     return results, None
 
-def gradio_predict(url, threshold):
-    tags, error = predict_tags(url, threshold)
-    if error:
-        return {"error": error}
-    if not tags:
-        return {"error": "No tags above threshold."}
-    
-    top_tag, top_score = tags[0]
-    return {"tag_name": top_tag, "tag_score": round(top_score, 4)}
+def gradio_predict(urls, threshold):
+    url_list = [u.strip() for u in urls.split(",") if u.strip()]
+    output = []
+
+    for url in url_list:
+        tags, error = predict_tags(url, threshold)
+        if error or not tags:
+            output.append({
+                "image_url": url,
+                "error": error or "No tags above threshold."
+            })
+        else:
+            top_tag, top_score = tags[0]
+            output.append({
+                "image_url": url,
+                "tag_name": top_tag,
+                "tag_score": round(top_score, 4)
+            })
+
+    return str(output)  # Return as string for textbox display
 
 demo = gr.Interface(
     fn=gradio_predict,
     inputs=[
-        gr.Textbox(label="Image URL"),
+        gr.Textbox(label="Image URL(s) (comma-separated)"),
         gr.Slider(0, 1, value=0.5, step=0.01, label="Threshold"),
     ],
     outputs=gr.Textbox(label="Tags"),
-    title="Image Tagging with ViT",
-    description="Paste an image URL and get predicted tags using thelabel/240903-image-tagging model.",
+    title="Batch Image Tagging with ViT",
+    description="Paste one or more image URLs separated by commas to get predicted tags using thelabel/240903-image-tagging model.",
 )
 
 if __name__ == "__main__":