Update app.py

app.py

@@ -34,20 +34,27 @@ def classify_image(image):
         # Run inference
         outputs = model(**inputs)
 
-        # Extract logits and apply softmax
-        logits_per_image = outputs.logits_per_image  # Image-text similarity scores
-        probs = logits_per_image.softmax(dim=1).detach().numpy()  # Convert logits to probabilities
+        # Extract logits
+        logits_per_image = outputs.logits_per_image  # Shape: [1, 2]
+        print(f"Logits: {logits_per_image}")
+
+        # Apply softmax to logits to get probabilities
+        probs = logits_per_image.softmax(dim=1)  # Shape: [1, 2]
+        print(f"Softmax probabilities: {probs}")
 
         # Extract probabilities for each category
-        safe_prob = probs[0][0]  # Safe probability
-        unsafe_prob = probs[0][1]  # Unsafe probability
+        safe_prob = probs[0][0].item()  # Extract 'safe' probability
+        unsafe_prob = probs[0][1].item()  # Extract 'unsafe' probability
+        print(f"Safe probability: {safe_prob}, Unsafe probability: {unsafe_prob}")
 
         # Normalize probabilities to ensure they sum to 100%
-        total = safe_prob + unsafe_prob
-        safe_percentage = (safe_prob / total) * 100
-        unsafe_percentage = (unsafe_prob / total) * 100
+        total_prob = safe_prob + unsafe_prob
+        print(f"Total probability before normalization: {total_prob}")
+        safe_percentage = (safe_prob / total_prob) * 100
+        unsafe_percentage = (unsafe_prob / total_prob) * 100
 
-        # Return results as percentages
+        # Ensure the sum is exactly 100%
+        print(f"Normalized percentages: Safe={safe_percentage}%, Unsafe={unsafe_percentage}%")
         return {
             "safe": round(safe_percentage, 2),  # Rounded to 2 decimal places
             "unsafe": round(unsafe_percentage, 2)
@@ -58,6 +65,7 @@ def classify_image(image):
 
 
 
+
 # Step 3: Set Up Gradio Interface
 iface = gr.Interface(
     fn=classify_image,