Update app.py

app.py

@@ -49,20 +49,22 @@ def classify_image(image):
         # Run inference with the model
         print("Running model inference...")
         outputs = model(**inputs)
-        print(f"Model outputs: {outputs}")
-
-        # Extract logits and probabilities
         logits_per_image = outputs.logits_per_image  # Image-text similarity scores
+        print(f"Logits per image: {logits_per_image}")
+
+        # Apply softmax to convert logits to probabilities
         probs = logits_per_image.softmax(dim=1)  # Convert logits to probabilities
-        print(f"Calculated probabilities: {probs}")
+        print(f"Softmax probabilities: {probs}")
 
-        # Normalize probabilities to ensure they sum to 100%
+        # Extract probabilities for each category
         safe_prob = probs[0][0].item()
         unsafe_prob = probs[0][1].item()
 
-        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
+        # Normalize probabilities to ensure they sum to 100%
+        total_prob = safe_prob + unsafe_prob
+        safe_percentage = (safe_prob / total_prob) * 100
+        unsafe_percentage = (unsafe_prob / total_prob) * 100
+        print(f"Normalized percentages: safe={safe_percentage}, unsafe={unsafe_percentage}")
 
         # Return results
         return {
@@ -74,6 +76,7 @@ def classify_image(image):
         print(f"Error during classification: {e}")
         return {"Error": str(e)}
 
+
 # Step 3: Set Up Gradio Interface
 iface = gr.Interface(
     fn=classify_image,