Update app.py

Improve app look and capabilities

app.py

@@ -1,24 +1,55 @@
 from transformers import pipeline
 from PIL import Image
 import gradio as gr
+from nsfw_image_detector import NSFWDetector
+import torch
 
 # Load your model from the Hub
-classifier = pipeline("image-classification", model="Freepik/nsfw_image_detector")
+classifier_pipe = pipeline("image-classification", model="Freepik/nsfw_image_detector")
+classifier_nsfw = NSFWDetector(dtype=torch.bfloat16, device="cpu")
 
 # Define the inference function
-def classify_image(image):
-    result = classifier(image)
-    # Format the result into a string with label and score (2 decimals)
-    formatted_result = "\n".join([f"{entry['label']}: {entry['score']:.2f}" for entry in result])
-    return formatted_result
+def classify_image(image, confidence_level):
+    # Get predictions from both models
+    result_nsfw_proba = classifier_nsfw.predict_proba(image)
+    is_nsfw_method = result_nsfw_proba[0][confidence_level] >= 0.5
+    result_pipe = classifier_pipe(image)
+    
+    # Format NSFW probability scores
+    proba_dict = result_nsfw_proba[0]
+    nsfw_proba_str = "NSFW Probability Scores:\n"
+    for level, score in proba_dict.items():
+        nsfw_proba_str += f"{level.value.title()}: {score:.4f}\n"
+    
+    # Format NSFW classification
+    is_nsfw_str = f"NSFW Classification ({confidence_level.title()}):\n"
+    is_nsfw_str += "🔴 True" if is_nsfw_method else "🟢 False"
+    
+    # Format pipeline results
+    pipe_str = "Pipeline Results:\n"
+    for result in result_pipe:
+        pipe_str += f"{result['label']}: {result['score']:.4f}\n"
+    
+    return nsfw_proba_str, is_nsfw_str, pipe_str
 
 # Create Gradio interface
 demo = gr.Interface(
     fn=classify_image,
-    inputs=gr.Image(type="pil", label="Upload an image"),
-    outputs=gr.Textbox(label="Prediction"),
+    inputs=[
+        gr.Image(type="pil", label="Upload an image"),
+        gr.Dropdown(
+            choices=["low", "medium", "high"],
+            value="medium",
+            label="Low is the most restrictive, high is the least restrictive"
+        )
+    ],
+    outputs=[
+        gr.Textbox(label="NSFW Probability Scores (recommended)", lines=3),
+        gr.Textbox(label="NSFW Classification (recommended)", lines=2),
+        gr.Textbox(label="Pipeline Results (not recommended, specially in production)", lines=3)
+    ],
     title="NSFW Image Classifier",
-    description="Upload an image and get a prediction using the Freepik/nsfw_image_detector model."
+    description="Upload an image and select a confidence level to get a prediction using the Freepik/nsfw_image_detector model."
 )
 
 # Launch app