Update app.py

app.py

@@ -26,7 +26,9 @@ def classify_image(image):
             raise ValueError("No image provided. Please upload a valid image.")
 
         # Define categories
-        categories = ["safe", "unsafe"]
+        unsafe_categories = ["hate", "sexual", "violent", "self-harm"]
+        safe_categories = ["safe", "retail product"]
+        categories = safe_categories + unsafe_categories
 
         # Process the image
         inputs = processor(text=categories, images=image, return_tensors="pt", padding=True)
@@ -39,12 +41,17 @@ 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
+        #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_categoriessafe_categories else 0.0 for i, value in enumerate(probs[0]))
+
+        #debug
+        for i, value in enumerate(probs[0]):
+            print(categories[i], value)
 
         # Return raw probabilities
         return {
-            "unknown" : 0.2,
             "safe": safe_prob,  # Leave as a fraction (e.g., 0.92)
             "unsafe": unsafe_prob  # Leave as a fraction (e.g., 0.08)
         }