Update app.py

app.py

@@ -1,26 +1,26 @@
-import os
-import torch
-from torchvision import transforms
-from PIL import Image
-from transformers import CLIPModel, CLIPProcessor
 import gradio as gr
+from transformers import CLIPModel, CLIPProcessor
 
-# Step 1: Ensure Fine-Tuned Model is Available
-fine_tuned_model_path = "fine-tuned-model"
-
-if not os.path.exists(fine_tuned_model_path):
-    raise FileNotFoundError(
-        f"The fine-tuned model is missing. Ensure that the fine-tuned model files are available in the '{fine_tuned_model_path}' directory."
-    )
+# Step 1: Load Fine-Tuned Model from Hugging Face Model Hub
+model_name = "quadranttechnologies/retail-content-safety-clip-finetuned"
 
-# Step 2: Load Fine-Tuned Model
-print("Loading fine-tuned model...")
-model = CLIPModel.from_pretrained(fine_tuned_model_path)
-processor = CLIPProcessor.from_pretrained(fine_tuned_model_path)
-print("Fine-tuned model loaded successfully.")
+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.")
 
-# Step 3: Define Gradio Inference Function
+# Step 2: Define the Inference Function
 def classify_image(image, class_names):
+    """
+    Classify an image as 'safe' or 'unsafe' using the fine-tuned CLIP model.
+
+    Args:
+        image (PIL.Image.Image): The input image.
+        class_names (str): Comma-separated class names (e.g., "safe, unsafe").
+    
+    Returns:
+        dict: A dictionary containing class names and their probabilities.
+    """
     # Split class names from comma-separated input
     labels = [label.strip() for label in class_names.split(",") if label.strip()]
     if not labels:
@@ -29,26 +29,29 @@ def classify_image(image, class_names):
     # Process the image and labels
     inputs = processor(text=labels, images=image, return_tensors="pt", padding=True)
     outputs = model(**inputs)
-    logits_per_image = outputs.logits_per_image
+    logits_per_image = outputs.logits_per_image  # Get image-text similarity scores
     probs = logits_per_image.softmax(dim=1)  # Convert logits to probabilities
 
     # Extract labels with their corresponding probabilities
     result = {label: probs[0][i].item() for i, label in enumerate(labels)}
     return dict(sorted(result.items(), key=lambda item: item[1], reverse=True))
 
-# Step 4: Set Up Gradio Interface
+# Step 3: Set Up Gradio Interface
 iface = gr.Interface(
     fn=classify_image,
     inputs=[
         gr.Image(type="pil"),
-        gr.Textbox(label="Possible class names (comma-separated)", placeholder="e.g., safe, unsafe")
+        gr.Textbox(
+            label="Possible class names (comma-separated)", 
+            placeholder="e.g., safe, unsafe"
+        )
     ],
     outputs=gr.Label(num_top_classes=2),
     title="Content Safety Classification",
-    description="Classify images as 'safe' or 'unsafe' using a fine-tuned CLIP model.",
+    description="Classify images as 'safe' or 'unsafe' using a fine-tuned CLIP model."
 )
 
-# Step 5: Launch Gradio Interface
+# Step 4: Launch Gradio Interface
 if __name__ == "__main__":
     iface.launch()
 
@@ -64,3 +67,4 @@ if __name__ == "__main__":
 
 
 
+