Update app.py

Remove pipeline to increase speed

app.py

@@ -1,11 +1,9 @@
-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_pipe = pipeline("image-classification", model="Freepik/nsfw_image_detector", torch_dtype=torch.bfloat16, device="cpu")
+
 classifier_nsfw = NSFWDetector(dtype=torch.bfloat16, device="cpu")
 
 # Define the inference function
@@ -13,7 +11,6 @@ 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]
@@ -25,12 +22,8 @@ def classify_image(image, confidence_level):
     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
+    return nsfw_proba_str, is_nsfw_str
 
 # Create Gradio interface
 demo = gr.Interface(
@@ -44,9 +37,8 @@ demo = gr.Interface(
         )
     ],
     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)
+        gr.Textbox(label="NSFW Categories Scores", lines=3),
+        gr.Textbox(label="NSFW Classification", lines=2),
     ],
     title="NSFW Image Classifier",
     description="Upload an image and select a confidence level to get a prediction using the Freepik/nsfw_image_detector model."