Update app.py

app.py

@@ -12,13 +12,13 @@ print("Model loaded successfully.")
 # Step 2: Define the Inference Function
 def classify_image(image):
     """
-    Classify an image as 'safe' or 'unsafe' with the corresponding percentage.
+    Classify an image as 'safe' or 'unsafe' with probabilities and display as a progress bar.
 
     Args:
         image (PIL.Image.Image): The input image.
     
     Returns:
-        dict: A dictionary containing the main category (safe/unsafe) and its percentage.
+        dict: A dictionary containing probabilities for 'safe' and 'unsafe'.
     """
     # Define the main categories
     main_categories = ["safe", "unsafe"]
@@ -29,23 +29,23 @@ def classify_image(image):
     logits_per_image = outputs.logits_per_image  # Image-text similarity scores
     probs = logits_per_image.softmax(dim=1)  # Convert logits to probabilities
 
-    # Extract the category with the highest probability
+    # Extract the probabilities
     safe_probability = probs[0][0].item() * 100  # Safe percentage
     unsafe_probability = probs[0][1].item() * 100  # Unsafe percentage
 
-    # Determine the main category
-    if safe_probability > unsafe_probability:
-        return {"Category": "safe", "Probability": f"{safe_probability:.2f}%"}
-    else:
-        return {"Category": "unsafe", "Probability": f"{unsafe_probability:.2f}%"}
+    # Return probabilities as a dictionary for display in Gradio's Label component
+    return {
+        "safe": f"{safe_probability:.2f}%",
+        "unsafe": f"{unsafe_probability:.2f}%"
+    }
 
 # Step 3: Set Up Gradio Interface
 iface = gr.Interface(
     fn=classify_image,
     inputs=gr.Image(type="pil"),
-    outputs="json",
+    outputs=gr.Label(label="Output"),
     title="Content Safety Classification",
-    description="Classify images as 'safe' or 'unsafe' with their respective percentage.",
+    description="Upload an image to classify it as 'safe' or 'unsafe' with corresponding probabilities.",
 )
 
 # Step 4: Launch Gradio Interface
@@ -67,3 +67,4 @@ if __name__ == "__main__":
 
 
 
+