Update app.py

app.py

@@ -56,14 +56,18 @@ def classify_image(image):
         probs = logits_per_image.softmax(dim=1)  # Convert logits to probabilities
         print(f"Calculated probabilities: {probs}")
 
-        # Extract probabilities for each category
-        safe_prob = probs[0][0].item() * 100  # Safe percentage
-        unsafe_prob = probs[0][1].item() * 100  # Unsafe percentage
+        # Normalize probabilities to ensure they sum to 100%
+        safe_prob = probs[0][0].item()
+        unsafe_prob = probs[0][1].item()
 
-        # Format data for visualization
+        total_prob = safe_prob + unsafe_prob  # Sum of probabilities
+        safe_percentage = (safe_prob / total_prob) * 100  # Normalize
+        unsafe_percentage = (unsafe_prob / total_prob) * 100  # Normalize
+
+        # Return results
         return {
-            "safe": safe_prob,
-            "unsafe": unsafe_prob
+            "safe": safe_percentage,
+            "unsafe": unsafe_percentage
         }
 
     except Exception as e: