Update app.py

app.py

@@ -34,30 +34,18 @@ def classify_image(image):
         # Run inference
         outputs = model(**inputs)
 
-        # Extract logits
+        # Extract logits and apply softmax
         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}")
+        probs = logits_per_image.softmax(dim=1).detach().numpy()  # Convert logits to probabilities
 
         # Extract probabilities for each category
-        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_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
-
-        # Ensure the sum is exactly 100%
-        print(f"Normalized percentages: Safe={safe_percentage}%, Unsafe={unsafe_percentage}%")
+        safe_prob = probs[0][0]  # Safe probability
+        unsafe_prob = probs[0][1]  # Unsafe probability
+
+        # Return raw probabilities
         return {
-            "safe": round(safe_percentage, 2),  # Rounded to 2 decimal places
-            "unsafe": round(unsafe_percentage, 2)
+            "safe": safe_prob,  # Leave as a fraction (e.g., 0.92)
+            "unsafe": unsafe_prob  # Leave as a fraction (e.g., 0.08)
         }
 
     except Exception as e:
@@ -66,6 +54,7 @@ def classify_image(image):
 
 
 
+
 # Step 3: Set Up Gradio Interface
 iface = gr.Interface(
     fn=classify_image,