Update app.py

app.py

@@ -20,63 +20,44 @@ except Exception as e:
 def classify_image(image):
     """
     Classify an image as 'safe' or 'unsafe' and return probabilities.
-
-    Args:
-        image (PIL.Image.Image): Uploaded image.
-    
-    Returns:
-        dict: Classification results or an error message.
     """
     try:
-        print("Starting image classification...")
-
-        # Validate input
         if image is None:
             raise ValueError("No image provided. Please upload a valid image.")
 
-        # Validate image format
-        if not hasattr(image, "convert"):
-            raise ValueError("Invalid image format. Please upload a valid image (JPEG, PNG, etc.).")
-
         # Define categories
         categories = ["safe", "unsafe"]
 
-        # Process the image with the processor
-        print("Processing the image...")
+        # Process the image
         inputs = processor(text=categories, images=image, return_tensors="pt", padding=True)
-        print(f"Processed inputs: {inputs}")
 
-        # Run inference with the model
-        print("Running model inference...")
+        # Run inference
         outputs = model(**inputs)
-        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"Softmax probabilities: {probs}")
+        # Extract logits and apply softmax
+        logits_per_image = outputs.logits_per_image  # Image-text similarity scores
+        probs = logits_per_image.softmax(dim=1).detach().numpy()  # Convert logits to probabilities
 
         # Extract probabilities for each category
-        safe_prob = probs[0][0].item()
-        unsafe_prob = probs[0][1].item()
+        safe_prob = probs[0][0]  # Safe probability
+        unsafe_prob = probs[0][1]  # Unsafe probability
 
         # 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}")
+        total = safe_prob + unsafe_prob
+        safe_percentage = (safe_prob / total) * 100
+        unsafe_percentage = (unsafe_prob / total) * 100
 
-        # Return results
+        # Return results as percentages
         return {
-            "safe": safe_percentage,
-            "unsafe": unsafe_percentage
+            "safe": round(safe_percentage, 2),  # Rounded to 2 decimal places
+            "unsafe": round(unsafe_percentage, 2)
         }
 
     except Exception as e:
-        print(f"Error during classification: {e}")
         return {"Error": str(e)}
 
 
+
 # Step 3: Set Up Gradio Interface
 iface = gr.Interface(
     fn=classify_image,