Update app.py

app.py

@@ -15,10 +15,9 @@ 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
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer, pipeline
 from detoxify import Detoxify
 from PIL import Image
-from tensorflow.keras.models import load_model
 import uvicorn
 from datetime import datetime, timedelta
 from collections import defaultdict, deque
@@ -91,24 +90,10 @@ model.eval()
 
 detoxify_model = Detoxify('multilingual')
 
-teachable_machine_url = "https://teachablemachine.withgoogle.com/models/gJOADmf_u/"
-model_url = teachable_machine_url + "model.json"
-weights_url = teachable_machine_url + "weights.bin"
-
-model_path = os.path.join(os.getcwd(), "teachable_machine_model")
-os.makedirs(model_path, exist_ok=True)
-
-if not os.path.exists(os.path.join(model_path, "model.json")):
-    response = requests.get(model_url)
-    with open(os.path.join(model_path, "model.json"), "wb") as f:
-        f.write(response.content)
-
-if not os.path.exists(os.path.join(model_path, "weights.bin")):
-    response = requests.get(weights_url)
-    with open(os.path.join(model_path, "weights.bin"), "wb") as f:
-        f.write(response.content)
-
-image_model = load_model(model_path)
+# Use a Hugging Face pipeline for NSFW image detection
+print("Loading NSFW image classification model...")
+image_classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection")
+print("NSFW image classification model loaded.")
 
 MODERATION_SYSTEM_PROMPT = (
     "You are a multilingual content moderation classifier. "
@@ -320,22 +305,23 @@ def classify_text_with_detoxify(text):
 def classify_image(image_data):
     try:
         img = Image.open(io.BytesIO(image_data)).convert("RGB")
-        img = img.resize((224, 224))
-        img_array = np.array(img) / 255.0
-        img_array = np.expand_dims(img_array, axis=0)
+        results = image_classifier(img)
+        
+        # Extract the top result
+        top_result = results[0]
+        label = top_result['label']
+        score = top_result['score']
         
-        predictions = image_model.predict(img_array)
-        class_idx = np.argmax(predictions[0])
-        classes = ["nothing", "nsfw"]
-        class_name = classes[class_idx]
-        confidence = float(predictions[0][class_idx])
+        # Map the label: 'normal' -> 's', 'nsfw' -> 'u'
+        classification = 'u' if label == 'nsfw' else 's'
+        nsfw_score = score if label == 'nsfw' else 1.0 - score
         
         return {
-            "classification": "u" if class_name == "nsfw" else "s",
-            "label": "NSFW" if class_name == "nsfw" else "SFW",
-            "description": "Content may contain inappropriate or harmful material." if class_name == "nsfw" else "Content appears to be safe and appropriate.",
-            "confidence": confidence,
-            "nsfw_score": confidence if class_name == "nsfw" else 1.0 - confidence
+            "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": score,
+            "nsfw_score": nsfw_score
         }
     except Exception as e:
         return {