Update app.py

app.py

@@ -12,7 +12,7 @@ print("Model loaded successfully.")
 # Step 2: Define the Inference Function
 def classify_image(image):
     """
-    Classify an image as 'safe' or 'unsafe' with probabilities and display as a progress bar.
+    Classify an image as 'safe' or 'unsafe' with the corresponding percentage.
 
     Args:
         image (PIL.Image.Image): The input image.
@@ -20,32 +20,52 @@ def classify_image(image):
     Returns:
         dict: A dictionary containing probabilities for 'safe' and 'unsafe'.
     """
-    # Define the main categories
-    main_categories = ["safe", "unsafe"]
-
-    # Process the image with the main categories
-    inputs = processor(text=main_categories, images=image, return_tensors="pt", padding=True)
-    outputs = model(**inputs)
-    logits_per_image = outputs.logits_per_image  # Image-text similarity scores
-    probs = logits_per_image.softmax(dim=1)  # Convert logits to probabilities
-
-    # Extract the probabilities
-    safe_probability = probs[0][0].item() * 100  # Safe percentage
-    unsafe_probability = probs[0][1].item() * 100  # Unsafe percentage
-
-    # Return probabilities as a dictionary for display in Gradio's Label component
-    return {
-        "safe": f"{safe_probability:.2f}%",
-        "unsafe": f"{unsafe_probability:.2f}%"
-    }
+    try:
+        # Debug: Check if the image is loaded
+        if image is None:
+            raise ValueError("No image provided. Please upload an image.")
+
+        # Define the main categories
+        main_categories = ["safe", "unsafe"]
+
+        # Process the image with the model processor
+        print("Processing the image...")
+        inputs = processor(text=main_categories, images=image, return_tensors="pt", padding=True)
+        print(f"Inputs processed: {inputs}")
+
+        # Perform inference using the model
+        outputs = model(**inputs)
+        print(f"Model outputs: {outputs}")
+
+        # Extract probabilities for each category
+        logits_per_image = outputs.logits_per_image  # Image-text similarity scores
+        probs = logits_per_image.softmax(dim=1)  # Convert logits to probabilities
+
+        # Safe and unsafe probabilities
+        safe_probability = probs[0][0].item() * 100  # Convert to percentage
+        unsafe_probability = probs[0][1].item() * 100  # Convert to percentage
+
+        print(f"Safe: {safe_probability:.2f}%, Unsafe: {unsafe_probability:.2f}%")
+
+        # Return the results as a dictionary for display in Gradio
+        return {
+            "safe": f"{safe_probability:.2f}%",
+            "unsafe": f"{unsafe_probability:.2f}%"
+        }
+    except Exception as e:
+        print(f"Error during inference: {str(e)}")
+        return {"Error": str(e)}
 
 # Step 3: Set Up Gradio Interface
 iface = gr.Interface(
     fn=classify_image,
     inputs=gr.Image(type="pil"),
-    outputs=gr.Label(label="Output"),
+    outputs=gr.Label(label="Output"),  # Use Gradio's Label component for progress bar display
     title="Content Safety Classification",
-    description="Upload an image to classify it as 'safe' or 'unsafe' with corresponding probabilities.",
+    description=(
+        "Upload an image to classify it as 'safe' or 'unsafe' with corresponding probabilities. "
+        "The model will analyze the image and provide probabilities for each category."
+    ),
 )
 
 # Step 4: Launch Gradio Interface
@@ -68,3 +88,4 @@ if __name__ == "__main__":
 
 
 
+