Update app.py

app.py

@@ -1,4 +1,5 @@
 import gradio as gr
+import gradio as gr
 from transformers import CLIPModel, CLIPProcessor
 
 # Step 1: Load Fine-Tuned Model from Hugging Face Model Hub
@@ -10,45 +11,59 @@ processor = CLIPProcessor.from_pretrained(model_name)
 print("Model loaded successfully.")
 
 # Step 2: Define the Inference Function
-def classify_image(image, class_names):
+def classify_image(image):
     """
-    Classify an image as 'safe' or 'unsafe' using the fine-tuned CLIP model.
+    Classify an image as 'safe' or 'unsafe' with probabilities and subcategories.
 
     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.
+        dict: A dictionary containing main categories (safe/unsafe) 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:
-        return {"Error": "Please enter at least one valid class name."}
-
-    # 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  # 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))
+    # Define the predefined categories
+    main_categories = ["safe", "unsafe"]
+    safe_subcategories = ["retail product", "other safe content"]
+    unsafe_subcategories = ["harmful", "violent", "sexual", "self harm"]
+
+    # Process the image with the main categories
+    main_inputs = processor(text=main_categories, images=image, return_tensors="pt", padding=True)
+    main_outputs = model(**main_inputs)
+    logits_per_image = main_outputs.logits_per_image  # Image-text similarity scores
+    main_probs = logits_per_image.softmax(dim=1)  # Convert logits to probabilities
+
+    # Determine the main category
+    main_result = {main_categories[i]: main_probs[0][i].item() for i in range(len(main_categories))}
+    main_category = max(main_result, key=main_result.get)  # Either "safe" or "unsafe"
+
+    # Process the image with subcategories based on the main category
+    subcategories = safe_subcategories if main_category == "safe" else unsafe_subcategories
+    sub_inputs = processor(text=subcategories, images=image, return_tensors="pt", padding=True)
+    sub_outputs = model(**sub_inputs)
+    sub_logits = sub_outputs.logits_per_image
+    sub_probs = sub_logits.softmax(dim=1)  # Convert logits to probabilities
+
+    # Create a structured result
+    result = {
+        "Main Category": main_category,
+        "Main Probabilities": main_result,
+        "Subcategory Probabilities": {
+            subcategories[i]: sub_probs[0][i].item() for i in range(len(subcategories))
+        }
+    }
+    return result
 
 # 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"
-        )
-    ],
-    outputs=gr.Label(num_top_classes=2),
-    title="Content Safety Classification",
-    description="Classify images as 'safe' or 'unsafe' using a fine-tuned CLIP model."
+    inputs=gr.Image(type="pil"),
+    outputs="json",
+    title="Enhanced Content Safety Classification",
+    description=(
+        "Classify images as 'safe' or 'unsafe' using a fine-tuned CLIP model. "
+        "For 'safe', identify subcategories such as 'retail product'. "
+        "For 'unsafe', identify subcategories such as 'harmful', 'violent', 'sexual', or 'self harm'."
+    ),
 )
 
 # Step 4: Launch Gradio Interface
@@ -68,3 +83,4 @@ if __name__ == "__main__":
 
 
 
+