Update app.py

app.py

@@ -1,9 +1,7 @@
 import gradio as gr
-import torch
 from transformers import AutoProcessor, AutoModelForZeroShotImageClassification
 from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
 from PIL import Image
-import requests
 from datasets import load_dataset 
 
 # Load your fine-tuned model and dataset
@@ -17,14 +15,6 @@ labels = list(set(dataset['train']['label']))
 label2id = {label: i for i, label in enumerate(labels)}
 id2label = {i: label for label, i in label2id.items()}
 
-# Define transformations for input images
-transform = Compose([
-    Resize((224, 224)),
-    CenterCrop(224),
-    ToTensor(),
-    Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-])
-
 # Function to classify image using CLIP model
 def classify_image(image):
     # Preprocess the image
@@ -34,13 +24,18 @@ def classify_image(image):
     # Run inference
     outputs = model(**inputs)
 
-    # Get predicted label
-    predicted_label_id = torch.argmax(outputs, dim=1).item()
-    print(predicted_label_id)
+    # Extract logits and apply softmax
+    logits_per_image = outputs.logits_per_image  # logits_per_image is a tensor with shape [1, num_labels]
+    probs = logits_per_image[0].softmax(dim=0)  # Take the softmax across the labels
+
+    # Get predicted label id and score
+    predicted_label_id = probs.argmax().item()
+
     predicted_label = id2label[predicted_label_id]
 
     return predicted_label
 
+
 # Gradio interface
 iface = gr.Interface(
     fn=classify_image,
@@ -51,4 +46,4 @@ iface = gr.Interface(
 )
 
 # Launch the Gradio interface
-iface.launch()
+iface.launch()