Update app.py

app.py

@@ -3,43 +3,46 @@ import torch
 import torch.nn.functional as F
 from torchvision import transforms
 from PIL import Image
+from transformers import ViTForImageClassification, ViTImageProcessor
 
-# Load models
-human_model = torch.load("humanNsfw_Swf.pth", map_location=torch.device("cpu"))
+# Load processor (same for both)
+processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224-in21k")
+
+# Rebuild and load Human NSFW model
+human_model = ViTForImageClassification.from_pretrained(
+    "google/vit-base-patch16-224-in21k", num_labels=2
+)
+human_model.load_state_dict(torch.load("humanNsfw_Swf.pth", map_location="cpu"))
 human_model.eval()
 
-anime_model = torch.load("animeCartoonNsfw_Sfw.pth", map_location=torch.device("cpu"))
+# Rebuild and load Anime NSFW model
+anime_model = ViTForImageClassification.from_pretrained(
+    "google/vit-base-patch16-224-in21k", num_labels=2
+)
+anime_model.load_state_dict(torch.load("animeCartoonNsfw_Sfw.pth", map_location="cpu"))
 anime_model.eval()
 
-# Shared preprocessing
-preprocess = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-])
+# Preprocess function
+def preprocess(image: Image.Image):
+    inputs = processor(images=image, return_tensors="pt")
+    return inputs["pixel_values"]
 
-# Prediction functions
+# Prediction function
 def predict(image, model_type):
     if image is None:
         return "<div class='result-box'>Please upload an image.</div>"
 
-    input_tensor = preprocess(image).unsqueeze(0)
+    inputs = preprocess(image)
+    model = human_model if model_type == "Human" else anime_model
 
     with torch.no_grad():
-        if model_type == "Human":
-            output = human_model(input_tensor)
-        else:
-            output = anime_model(input_tensor)
-
-    if output.shape[-1] == 1:
-        prob = torch.sigmoid(output).item()
-        label = "NSFW" if prob > 0.5 else "SFW"
-        confidence = prob if prob > 0.5 else 1 - prob
-    else:
-        probs = F.softmax(output, dim=1).squeeze()
-        label_index = torch.argmax(probs).item()
-        label = "NSFW" if label_index == 1 else "SFW"
-        confidence = probs[label_index].item()
+        outputs = model(pixel_values=inputs)
+        logits = outputs.logits
+        probs = F.softmax(logits, dim=1)
+        pred_class = torch.argmax(probs, dim=1).item()
+        confidence = probs[0][pred_class].item()
 
+    label = "NSFW" if pred_class == 1 else "SFW"
     return f"""
     <div class='result-box'>
         <strong>Model:</strong> {model_type}<br>
@@ -48,7 +51,7 @@ def predict(image, model_type):
     </div>
     """
 
-# Custom glowing style
+# Custom glow box CSS
 custom_css = """
 .result-box {
     background-color: oklch(0.718 0.202 349.761);
@@ -64,12 +67,12 @@ custom_css = """
 .gradio-container { max-width: 900px; margin: auto; }
 """
 
-# Gradio Interface
+# Gradio UI
 with gr.Blocks(css=custom_css) as demo:
     gr.Markdown("## NSFW Detector (Human + Anime/Cartoon)")
     gr.Markdown(
-        "Upload an image and choose the model. The system will predict whether the content is NSFW or SFW. "
-        "This is a side project. Results may vary. No images are stored."
+        "Upload an image and select the appropriate model. The system will detect whether the content is NSFW or SFW. "
+        "This is a side project. Results are not guaranteed. No images are stored."
     )
 
     with gr.Row():