Update app.py

app.py

@@ -1,57 +1,106 @@
 import gradio as gr
-from PIL import Image, ImageOps
+import torch
+import torch.nn as nn
+from torchvision import transforms
+from torchvision.models import resnet18
+from PIL import Image
 import base64
 import io
 
+# ---------------- CONFIG ----------------
+labels = ["Drawings", "Hentai", "Neutral", "Porn", "Sexy"]
+theme_color = "#6C5B7B"
 
-# Funzione di predizione condivisa
-def predict(text, image_b64):
+# ---------------- MODEL ----------------
+class Classifier(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.cnn_layers = resnet18(weights=None)
+        self.fc_layers = nn.Sequential(
+            nn.Linear(1000, 512),
+            nn.Dropout(0.3),
+            nn.Linear(512, 128),
+            nn.ReLU(),
+            nn.Linear(128, 5)
+        )
+
+    def forward(self, x):
+        x = self.cnn_layers(x)
+        x = self.fc_layers(x)
+        return x
+
+preprocess = transforms.Compose([
+    transforms.Resize((224,224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485,0.456,0.406],
+                         std=[0.229,0.224,0.225])
+])
+
+model = Classifier()
+model.load_state_dict(torch.load("classify_nsfw_v3.0.pth", map_location="cpu"))
+model.eval()
+
+# ---------------- FUNZIONE ----------------
+def predict(base64_image):
     """
-    Elabora un testo e un'immagine codificata in base64.
-    - Testo: invertito
-    - Immagine: convertita in scala di grigi e ritrasformata in base64
+    Input: immagine base64
+    Output: label più probabile + probabilità per tutte le classi
     """
-    # Elabora testo
-    out_text = text[::-1] if text else ""
-
-    # Elabora immagine base64
-    out_b64 = None
-    if image_b64:
-        try:
-            # Decodifica da base64 (gestisce anche stringhe tipo "data:image/png;base64,...")
-            img_bytes = base64.b64decode(image_b64.split(",")[-1])
-            img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
-
-            # Conversione in bianco e nero
-            img = ImageOps.grayscale(img)
-
-            # Ricodifica in base64
-            buffer = io.BytesIO()
-            img.save(buffer, format="PNG")
-            out_b64 = "data:image/png;base64," + base64.b64encode(buffer.getvalue()).decode("utf-8")
-        except Exception as e:
-            out_text += f"\n[Errore immagine: {e}]"
-
-    return out_text, out_b64
-
-
-# Interfaccia Gradio
-demo = gr.Interface(
-    fn=predict,
-    inputs=[
-        gr.Textbox(label="Testo"),
-        gr.Textbox(label="Immagine in base64")
-    ],
-    outputs=[
-        gr.Textbox(label="Testo elaborato"),
-        gr.Textbox(label="Immagine in base64")
-    ],
-    title="Demo Base64 API Gradio 3.5",
-    description="Input: testo + immagine base64 → Output: testo invertito + immagine base64 in scala di grigi"
-)
+    try:
+        if base64_image.startswith("data:image"):
+            base64_image = base64_image.split(",",1)[1]
+        img_bytes = base64.b64decode(base64_image)
+        img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
+        img_tensor = preprocess(img).unsqueeze(0)
+
+        with torch.no_grad():
+            pred = torch.nn.functional.softmax(model(img_tensor)[0], dim=0)
+
+        max_label = labels[torch.argmax(pred).item()]
+        probs = {labels[i]: float(pred[i]) for i in range(len(labels))}
 
+        return max_label, probs
 
+    except Exception as e:
+        return f"Error: {str(e)}", {}
+
+def clear_all():
+    return ""
+
+# ---------------- INTERFACCIA ----------------
+with gr.Blocks(title="NSFW Classifier") as demo:
+    gr.HTML(f"""
+    <div style="padding:10px; background:linear-gradient(135deg,#f8f9fa 0%,#e9ecef 100%); border-radius:10px;">
+        <h2 style="color:{theme_color};">🎨 NSFW Image Classifier</h2>
+        <p>Incolla qui l'immagine in base64 per analizzarla.</p>
+    </div>
+    """)
+
+    with gr.Row():
+        with gr.Column(scale=2):
+            base64_input = gr.Textbox(
+                label="📤 Base64 dell'immagine",
+                lines=6,
+                placeholder="Incolla qui la stringa base64 dell'immagine..."
+            )
+            with gr.Row():
+                submit_btn = gr.Button("✨ Analizza", variant="primary")
+                clear_btn = gr.Button("🔄 Pulisci", variant="secondary")
+        with gr.Column(scale=1):
+            label_output = gr.Textbox(label="Classe predetta", interactive=False)
+            result_display = gr.Label(label="Distribuzione probabilità", num_top_classes=5)
+
+    # Eventi
+    submit_btn.click(
+        fn=predict,
+        inputs=base64_input,
+        outputs=[label_output, result_display],
+        api_name="predict"  # <- espone /run/predict
+    )
+    clear_btn.click(fn=clear_all, inputs=None, outputs=base64_input)
+
+# ---------------- LAUNCH ----------------
 if __name__ == "__main__":
-    # In Gradio 3.5 le API sono esposte automaticamente su /api/predict
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(server_name="0.0.0.0", server_port=7860, show_api=True)
+