Update app.py

app.py

@@ -41,10 +41,8 @@ def classify_image(image):
         probs = logits_per_image.softmax(dim=1).detach().numpy()  # Convert logits to probabilities
 
         # Extract probabilities for each category
-        safe_prob = probs[0][0]  # Safe probability
-        unsafe_prob = probs[0][1]  # Unsafe probability
-        #safe_prob = sum(value if categories[i] in safe_categories else 0.0 for i, value in enumerate(probs[0]))
-        #unsafe_prob = sum(value if categories[i] in unsafe_categories else 0.0 for i, value in enumerate(probs[0]))
+        safe_prob = sum(value if categories[i] in safe_categories else 0.0 for i, value in enumerate(probs[0]))
+        unsafe_prob = sum(value if categories[i] in unsafe_categories else 0.0 for i, value in enumerate(probs[0]))
 
         #debug
         for i, value in enumerate(probs[0]):