Update app.py

app.py

@@ -1,26 +1,39 @@
 import gradio as gr
 from transformers import CLIPProcessor, CLIPModel
 from PIL import Image
+import torch
 
 # Load the model and processor
 model = CLIPModel.from_pretrained("geolocal/StreetCLIP")
 processor = CLIPProcessor.from_pretrained("geolocal/StreetCLIP")
 
 def classify_image(image):
-    # Preprocess the image
-    inputs = processor(images=image, return_tensors="pt")
+    # Example labels for classification
+    labels = ["a photo of a cat", "a photo of a dog", "a photo of a car", "a photo of a tree"]
+    
+    # Preprocess the image and text
+    inputs = processor(text=labels, images=image, return_tensors="pt", padding=True)
+    
     # Perform the inference
     outputs = model(**inputs)
+    
     # Postprocess the outputs
     logits_per_image = outputs.logits_per_image  # this is the image-text similarity score
     probs = logits_per_image.softmax(dim=1)  # we can use softmax to get probabilities
-    return probs
+    
+    # Convert the probabilities to a list
+    probs_list = probs.tolist()[0]
+    
+    # Create a dictionary of labels and probabilities
+    result = {label: prob for label, prob in zip(labels, probs_list)}
+    
+    return result
 
 # Define Gradio interface
 iface = gr.Interface(
     fn=classify_image,
     inputs=gr.Image(type="pil"),
-    outputs="text",
+    outputs="label",
     title="Geolocal StreetCLIP Classification",
     description="Upload an image to classify using Geolocal StreetCLIP"
 )