Update app.py

app.py

@@ -1,13 +1,19 @@
 import gradio as gr
 from transformers import CLIPModel, CLIPProcessor
+from PIL import Image
+import torch
 
 # 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...")
-model = CLIPModel.from_pretrained(model_name, trust_remote_code=True)
-processor = CLIPProcessor.from_pretrained(model_name)
-print("Model loaded successfully.")
+try:
+    model = CLIPModel.from_pretrained(model_name, trust_remote_code=True)
+    processor = CLIPProcessor.from_pretrained(model_name)
+    print("Model loaded successfully.")
+except Exception as e:
+    print(f"Error loading model or processor: {str(e)}")
+    raise
 
 # Step 2: Define the Inference Function
 def classify_image(image):
@@ -18,40 +24,43 @@ def classify_image(image):
         image (PIL.Image.Image): The input image.
     
     Returns:
-        dict: A dictionary containing probabilities for 'safe' and 'unsafe'.
+        dict: A dictionary containing probabilities for 'safe' and 'unsafe' or an error message.
     """
     try:
-        # Debug: Check if the image is loaded
+        # Check if the image is valid
         if image is None:
             raise ValueError("No image provided. Please upload an image.")
+        if not hasattr(image, "convert"):
+            raise ValueError("Uploaded file is not a valid image. Please upload a valid image (JPEG, PNG).")
 
-        # Define the main categories
+        # Define main categories
         main_categories = ["safe", "unsafe"]
 
-        # Process the image with the model processor
+        # Process the image
         print("Processing the image...")
         inputs = processor(text=main_categories, images=image, return_tensors="pt", padding=True)
-        print(f"Inputs processed: {inputs}")
+        print("Inputs processed successfully.")
 
-        # Perform inference using the model
+        # Perform inference
         outputs = model(**inputs)
-        print(f"Model outputs: {outputs}")
+        print("Model inference completed.")
 
-        # Extract probabilities for each category
+        # Calculate probabilities
         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
+        # Extract probabilities for "safe" and "unsafe"
         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 results
         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)}
@@ -60,12 +69,9 @@ def classify_image(image):
 iface = gr.Interface(
     fn=classify_image,
     inputs=gr.Image(type="pil"),
-    outputs=gr.Label(label="Output"),  # Use Gradio's Label component for progress bar display
+    outputs=gr.Label(label="Output"),  # Use Gradio's Label component for user-friendly display
     title="Content Safety Classification",
-    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."
-    ),
+    description="Upload an image to classify it as 'safe' or 'unsafe' with corresponding probabilities.",
 )
 
 # Step 4: Launch Gradio Interface
@@ -89,3 +95,4 @@ if __name__ == "__main__":
 
 
 
+