Update app.py

app.py

@@ -4,55 +4,70 @@ from transformers import CLIPModel, CLIPProcessor
 # Step 1: Load Fine-Tuned Model from Hugging Face Model Hub
 model_name = "quadranttechnologies/retail-content-safety-clip-finetuned"
 
-print("Loading the fine-tuned model from Hugging Face Model Hub...")
+print("Initializing the application...")
+
 try:
+    print("Loading the model from Hugging Face Model Hub...")
     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 model or processor: {e}")
+    print(f"Error loading the model or processor: {e}")
     raise RuntimeError(f"Failed to load model: {e}")
 
 # Step 2: Define the Inference Function
 def classify_image(image):
     """
-    Classify an image as 'safe' or 'unsafe' and display category and probabilities.
+    Classify an image as 'safe' or 'unsafe' and return probabilities.
 
     Args:
-        image (PIL.Image.Image): Uploaded image.
+        image (PIL.Image.Image): The uploaded image.
     
     Returns:
         str: Predicted category ("safe" or "unsafe").
         dict: Probabilities for "safe" and "unsafe".
     """
     try:
-        # Validate image input
+        print("Starting image classification...")
+
+        # Check if the image is valid
         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.")
 
-        # Define categories
+        # Define main categories
         categories = ["safe", "unsafe"]
+        print(f"Categories: {categories}")
 
         # Process the image
+        print("Processing the image with the processor...")
         inputs = processor(text=categories, images=image, return_tensors="pt", padding=True)
+        print(f"Processed inputs: {inputs}")
+
+        # Perform inference
+        print("Running model inference...")
         outputs = model(**inputs)
+        print(f"Model outputs: {outputs}")
 
-        # Get logits and probabilities
+        # Calculate probabilities
         logits_per_image = outputs.logits_per_image  # Image-text similarity scores
         probs = logits_per_image.softmax(dim=1)  # Convert logits to probabilities
+        print(f"Probabilities: {probs}")
 
-        # Extract probabilities
+        # Extract probabilities for each category
         safe_prob = probs[0][0].item() * 100  # Safe percentage
         unsafe_prob = probs[0][1].item() * 100  # Unsafe percentage
 
         # Determine the predicted category
         predicted_category = "safe" if safe_prob > unsafe_prob else "unsafe"
+        print(f"Predicted category: {predicted_category}")
 
         # Return the predicted category and probabilities
         return predicted_category, {"safe": f"{safe_prob:.2f}%", "unsafe": f"{unsafe_prob:.2f}%"}
 
     except Exception as e:
-        print(f"Error during inference: {e}")
+        print(f"Error during classification: {e}")
         return f"Error: {str(e)}", {}
 
 # Step 3: Set Up Gradio Interface