Update app.py

app.py

@@ -15,7 +15,7 @@ from fastapi.templating import Jinja2Templates
 from pydantic import BaseModel, Field
 from dotenv import load_dotenv
 from huggingface_hub import snapshot_download
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer, AutoModelForImageClassification, ViTImageProcessor
 from detoxify import Detoxify
 from PIL import Image
 import uvicorn
@@ -90,9 +90,10 @@ model.eval()
 
 detoxify_model = Detoxify('multilingual')
 
-# Use the Falconsai NSFW image detection model
+# Load the NSFW image detection model and processor directly
 print("Loading NSFW image classification model...")
-nsfw_classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection")
+nsfw_model = AutoModelForImageClassification.from_pretrained("Falconsai/nsfw_image_detection")
+nsfw_processor = ViTImageProcessor.from_pretrained('Falconsai/nsfw_image_detection')
 print("NSFW image classification model loaded.")
 
 MODERATION_SYSTEM_PROMPT = (
@@ -305,24 +306,33 @@ def classify_text_with_detoxify(text):
 def classify_image(image_data):
     try:
         img = Image.open(io.BytesIO(image_data)).convert("RGB")
-        # Use the Falconsai NSFW detector
-        results = nsfw_classifier(img)
         
-        # Extract the NSFW score
-        nsfw_score = 0.0
-        for result in results:
-            if result['label'].lower() == 'nsfw':
-                nsfw_score = result['score']
-                break
-        
-        # Use a threshold of 0.5 for NSFW detection
-        classification = 'u' if nsfw_score > 0.5 else 's'
+        # Use the model and processor directly as shown in the example
+        with torch.no_grad():
+            inputs = nsfw_processor(images=img, return_tensors="pt")
+            outputs = nsfw_model(**inputs)
+            logits = outputs.logits
+            
+            # Get the predicted label
+            predicted_label = logits.argmax(-1).item()
+            label = nsfw_model.config.id2label[predicted_label]
+            
+            # Get the confidence score
+            confidence = torch.softmax(logits, dim=-1)[0][predicted_label].item()
+            
+            # Convert to our classification system
+            if label.lower() == "nsfw":
+                classification = "u"
+                nsfw_score = confidence
+            else:  # normal
+                classification = "s"
+                nsfw_score = 1.0 - confidence
         
         return {
             "classification": classification,
             "label": "NSFW" if classification == 'u' else "SFW",
             "description": "Content may contain inappropriate or harmful material." if classification == 'u' else "Content appears to be safe and appropriate.",
-            "confidence": nsfw_score,
+            "confidence": confidence,
             "nsfw_score": nsfw_score
         }
     except Exception as e: