Update app.py

app.py

@@ -1,27 +1,54 @@
 import gradio as gr
+import torch
+import torch.nn.functional as F
+from torchvision import transforms
 from PIL import Image
 
-from model_human import predict_human_nsfw
-from model_anime import predict_anime_nsfw
+# Load models
+human_model = torch.load("humanNsfw_Swf.pth", map_location=torch.device("cpu"))
+human_model.eval()
 
-def nsfw_detector(image: Image.Image, model_type: str):
+anime_model = torch.load("animeCartoonNsfw_Sfw.pth", map_location=torch.device("cpu"))
+anime_model.eval()
+
+# Shared preprocessing
+preprocess = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+])
+
+# Prediction functions
+def predict(image, model_type):
     if image is None:
-        return "<div class='result-box'>Pls upload an img</div>"
+        return "<div class='result-box'>Please upload an image.</div>"
+
+    input_tensor = preprocess(image).unsqueeze(0)
 
-    if model_type == "Human":
-        label, confidence = predict_human_nsfw(image)
+    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:
-        label, confidence = predict_anime_nsfw(image)
+        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()
 
-    result = (
-        f"<div class='result-box'>"
-        f"<strong>Model:</strong> {model_type}<br>"
-        f"<strong>Prediction:</strong> {label}<br>"
-        f"<strong>Confidence:</strong> {confidence:.2%}"
-        f"</div>"
-    )
-    return result
+    return f"""
+    <div class='result-box'>
+        <strong>Model:</strong> {model_type}<br>
+        <strong>Prediction:</strong> {label}<br>
+        <strong>Confidence:</strong> {confidence:.2%}
+    </div>
+    """
 
+# Custom glowing style
 custom_css = """
 .result-box {
     background-color: oklch(0.718 0.202 349.761);
@@ -37,30 +64,23 @@ custom_css = """
 .gradio-container { max-width: 900px; margin: auto; }
 """
 
+# Gradio Interface
 with gr.Blocks(css=custom_css) as demo:
     gr.Markdown("## NSFW Detector (Human + Anime/Cartoon)")
     gr.Markdown(
-        "Upload an image and select the appropriate model for detection. "
-        "No data is stored. This is a side project — results may not be fully accurate."
+        "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."
     )
 
     with gr.Row():
         with gr.Column(scale=1):
-            model_choice = gr.Radio(
-                ["Human", "Anime"],
-                label="Select Model Type",
-                value="Human"
-            )
+            model_choice = gr.Radio(["Human", "Anime"], label="Select Model Type", value="Human")
             image_input = gr.Image(type="pil", label="Upload Image")
         with gr.Column(scale=1):
             output_box = gr.HTML("<div class='result-box'>Awaiting input...</div>")
 
-    # Trigger detection when image is uploaded
-    image_input.change(
-        fn=nsfw_detector,
-        inputs=[image_input, model_choice],
-        outputs=output_box
-    )
+    image_input.change(fn=predict, inputs=[image_input, model_choice], outputs=output_box)
+    model_choice.change(fn=predict, inputs=[image_input, model_choice], outputs=output_box)
 
 if __name__ == "__main__":
     demo.launch()