Update app.py

app.py

@@ -1,82 +1,39 @@
 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
+# Load the model and processor
 model_name = "quadranttechnologies/retail-content-safety-clip-finetuned"
-
 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: 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}")
+    raise
 
-    # 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
+# Define the inference function
 def classify_image(image):
-    """
-    Classify an image as 'safe' or 'unsafe' and return probabilities.
-
-    Args:
-        image (PIL.Image.Image): Uploaded image.
-    
-    Returns:
-        str: Predicted category.
-        dict: Probabilities for "safe" and "unsafe".
-    """
     try:
         print("Starting image classification...")
 
-        # Check if the image is valid
+        # Validate image input
         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 main categories
+        # Define 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}")
 
@@ -85,34 +42,31 @@ def classify_image(image):
         probs = logits_per_image.softmax(dim=1)
         print(f"Probabilities: {probs}")
 
-        # Extract probabilities for each category
+        # Extract probabilities
         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 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 classification: {e}")
         return f"Error: {str(e)}", {}
 
-# Step 4: Set Up Gradio Interface
+# Gradio interface
 iface = gr.Interface(
     fn=classify_image,
     inputs=gr.Image(type="pil"),
     outputs=[
-        gr.Textbox(label="Predicted Category"),  # Display the predicted category prominently
-        gr.Label(label="Probabilities"),        # Display probabilities with a progress bar
+        gr.Textbox(label="Predicted Category"),
+        gr.Label(label="Probabilities"),
     ],
     title="Content Safety Classification",
     description="Upload an image to classify it as 'safe' or 'unsafe' with corresponding probabilities.",
 )
 
-# Step 5: Launch Gradio Interface
 if __name__ == "__main__":
     print("Launching Gradio interface...")
     iface.launch()
@@ -137,5 +91,6 @@ if __name__ == "__main__":
 
 
 
+