Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 from transformers import CLIPModel, CLIPProcessor
 from PIL import Image
+import requests
 
 # Step 1: Load Fine-Tuned Model from Hugging Face Model Hub
 model_name = "quadranttechnologies/retail-content-safety-clip-finetuned"
@@ -16,7 +17,35 @@ except Exception as e:
     print(f"Error loading the model or processor: {e}")
     raise RuntimeError(f"Failed to load model: {e}")
 
-# Step 2: Define the Inference Function
+# Step 2: Minimal Test Case to Verify Model and Processor
+try:
+    print("Running a minimal test case with the model...")
+
+    # Test Image URL
+    url = "https://huggingface.co/datasets/Narsil/image_dummy/raw/main/parrots.png"
+    image = Image.open(requests.get(url, stream=True).raw)
+
+    # Define test categories
+    test_categories = ["safe", "unsafe"]
+
+    # Process the image
+    test_inputs = processor(text=test_categories, images=image, return_tensors="pt", padding=True)
+    print(f"Test inputs processed: {test_inputs}")
+
+    # Perform inference
+    test_outputs = model(**test_inputs)
+    print(f"Test outputs: {test_outputs}")
+
+    # Check probabilities
+    test_logits = test_outputs.logits_per_image
+    test_probs = test_logits.softmax(dim=1)
+    print(f"Test probabilities: {test_probs}")
+
+except Exception as e:
+    print(f"Error during the minimal test case: {e}")
+    raise RuntimeError(f"Test case failed: {e}")
+
+# Step 3: Define the Inference Function
 def classify_image(image):
     """
     Classify an image as 'safe' or 'unsafe' and return probabilities.
@@ -25,7 +54,7 @@ def classify_image(image):
         image (PIL.Image.Image): Uploaded image.
     
     Returns:
-        str: Predicted category ("safe" or "unsafe").
+        str: Predicted category.
         dict: Probabilities for "safe" and "unsafe".
     """
     try:
@@ -52,13 +81,13 @@ def classify_image(image):
         print(f"Model outputs: {outputs}")
 
         # Calculate probabilities
-        logits_per_image = outputs.logits_per_image  # Image-text similarity scores
-        probs = logits_per_image.softmax(dim=1)  # Convert logits to probabilities
+        logits_per_image = outputs.logits_per_image
+        probs = logits_per_image.softmax(dim=1)
         print(f"Probabilities: {probs}")
 
         # Extract probabilities for each category
-        safe_prob = probs[0][0].item() * 100  # Safe percentage
-        unsafe_prob = probs[0][1].item() * 100  # Unsafe percentage
+        safe_prob = probs[0][0].item() * 100
+        unsafe_prob = probs[0][1].item() * 100
 
         # Determine the predicted category
         predicted_category = "safe" if safe_prob > unsafe_prob else "unsafe"
@@ -71,7 +100,7 @@ def classify_image(image):
         print(f"Error during classification: {e}")
         return f"Error: {str(e)}", {}
 
-# Step 3: Set Up Gradio Interface
+# Step 4: Set Up Gradio Interface
 iface = gr.Interface(
     fn=classify_image,
     inputs=gr.Image(type="pil"),
@@ -83,7 +112,7 @@ iface = gr.Interface(
     description="Upload an image to classify it as 'safe' or 'unsafe' with corresponding probabilities.",
 )
 
-# Step 4: Launch Gradio Interface
+# Step 5: Launch Gradio Interface
 if __name__ == "__main__":
     print("Launching Gradio interface...")
     iface.launch()
@@ -107,5 +136,6 @@ if __name__ == "__main__":
 
 
 
+