Update app.py

app.py

@@ -1,55 +1,103 @@
 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"
 
-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.")
+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 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' with probabilities and display as a progress bar.
+    Classify an image as 'safe' or 'unsafe' and return probabilities.
 
     Args:
-        image (PIL.Image.Image): The input image.
+        image (PIL.Image.Image): Uploaded image.
     
     Returns:
-        dict: A dictionary containing probabilities for 'safe' and 'unsafe'.
+        dict: Classification results or an error message.
     """
-    # 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:
+        print("Starting image classification...")
+
+        # Validate input
+        if image is None:
+            raise ValueError("No image provided. Please upload a valid image.")
+
+        # Validate image format
+        if not hasattr(image, "convert"):
+            raise ValueError("Invalid image format. Please upload a valid image (JPEG, PNG, etc.).")
+
+        # Define categories
+        categories = ["safe", "unsafe"]
+
+        # Process the image with the processor
+        print("Processing the image...")
+        inputs = processor(text=categories, images=image, return_tensors="pt", padding=True)
+        print(f"Processed inputs: {inputs}")
+
+        # Run inference with the model
+        print("Running model inference...")
+        outputs = model(**inputs)
+        print(f"Model outputs: {outputs}")
+
+        # Extract logits and 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"Calculated 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
+
+        # Return results
+        return {
+            "safe": f"{safe_prob:.2f}%",
+            "unsafe": f"{unsafe_prob:.2f}%"
+        }
+
+    except Exception as e:
+        # Log and return detailed error messages
+        print(f"Error during classification: {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"),  # Display probabilities as progress bars
     title="Content Safety Classification",
     description="Upload an image to classify it as 'safe' or 'unsafe' with corresponding probabilities.",
 )
 
-# Step 4: Launch Gradio Interface
+# Step 4: Test Before Launch
 if __name__ == "__main__":
+    print("Testing model locally with a sample image...")
+    try:
+        # Test with a sample image
+        url = "https://huggingface.co/datasets/Narsil/image_dummy/raw/main/parrots.png"
+        test_image = Image.open(requests.get(url, stream=True).raw)
+
+        # Run the classification function
+        print("Running local test...")
+        result = classify_image(test_image)
+        print(f"Local Test Result: {result}")
+    except Exception as e:
+        print(f"Error during local test: {e}")
+
+    # Launch Gradio Interface
+    print("Launching the Gradio interface...")
     iface.launch()
 
 
@@ -79,5 +127,6 @@ if __name__ == "__main__":
 
 
 
+