Update app.py

app.py

@@ -2,20 +2,30 @@ import gradio as gr
 from transformers import CLIPModel, CLIPProcessor
 from PIL import Image
 
-# Load the model and processor
+# Load Model and Processor
 model_name = "quadranttechnologies/retail-content-safety-clip-finetuned"
-print("Loading the model and processor...")
+print("Initializing the model and processor...")
 
 try:
     model = CLIPModel.from_pretrained(model_name, trust_remote_code=True)
     processor = CLIPProcessor.from_pretrained(model_name)
     print("Model and processor loaded successfully.")
 except Exception as e:
-    print(f"Error loading the model or processor: {e}")
-    raise RuntimeError(f"Failed to load model: {e}")
+    print(f"Error loading model or processor: {e}")
+    raise RuntimeError(f"Failed to load the model: {e}")
 
-# Define the inference function
+
+# Inference Function
 def classify_image(image):
+    """
+    Classifies an image as 'safe' or 'unsafe' using the CLIP model.
+
+    Args:
+        image (PIL.Image.Image): Uploaded image.
+
+    Returns:
+        Tuple: Predicted category and probabilities for "safe" and "unsafe".
+    """
     try:
         print("Starting image classification...")
 
@@ -23,52 +33,56 @@ def classify_image(image):
         if image is None:
             raise ValueError("No image provided. Please upload a valid image.")
         if not hasattr(image, "convert"):
-            raise ValueError("Uploaded file is not a valid image format.")
+            raise ValueError("Uploaded file is not a valid image format (JPEG, PNG, etc.).")
 
-        # Define categories
+        # Define classification categories
         categories = ["safe", "unsafe"]
-        print(f"Categories: {categories}")
+        print(f"Using categories: {categories}")
 
-        # Process the image
+        # Process image using CLIPProcessor
         inputs = processor(text=categories, images=image, return_tensors="pt", padding=True)
-        print(f"Processed inputs: {inputs}")
+        print("Image processed successfully.")
 
         # Perform inference
         outputs = model(**inputs)
-        print(f"Model outputs: {outputs}")
+        print("Inference completed successfully.")
 
         # Calculate probabilities
         logits_per_image = outputs.logits_per_image
         probs = logits_per_image.softmax(dim=1)
         print(f"Probabilities: {probs}")
 
-        # Extract probabilities
+        # Extract probabilities for each category
         safe_prob = probs[0][0].item() * 100
         unsafe_prob = probs[0][1].item() * 100
         print(f"Safe: {safe_prob:.2f}%, Unsafe: {unsafe_prob:.2f}%")
 
         # Determine the predicted category
         predicted_category = "safe" if safe_prob > unsafe_prob else "unsafe"
+        print(f"Predicted Category: {predicted_category}")
+
+        # Return category and probabilities
         return predicted_category, {"safe": f"{safe_prob:.2f}%", "unsafe": f"{unsafe_prob:.2f}%"}
 
     except Exception as e:
         print(f"Error during classification: {e}")
-        return f"Error: {str(e)}", {}
+        return "Error", {"safe": "N/A", "unsafe": "N/A"}
+
 
-# Gradio interface
+# Gradio Interface
 iface = gr.Interface(
     fn=classify_image,
-    inputs=gr.Image(type="pil"),
+    inputs=gr.Image(type="pil"),  # Accept image input
     outputs=[
-        gr.Textbox(label="Predicted Category"),
-        gr.Label(label="Probabilities"),
+        gr.Textbox(label="Predicted Category"),  # Predicted category
+        gr.Label(label="Probabilities"),        # Probabilities as progress bars
     ],
     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 probabilities.",
 )
 
 if __name__ == "__main__":
-    print("Launching Gradio interface...")
+    print("Launching the Gradio interface...")
     iface.launch()
 
 
@@ -93,5 +107,6 @@ if __name__ == "__main__":
 
 
 
+