Update app.py

app.py

@@ -1,46 +1,122 @@
 import gradio as gr
-import base64
-from io import BytesIO
+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"
+
+# ---------------- 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, len(labels))
+        )
+
+    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 di analisi immagine base64
-def analyze_base64(b64_string: str):
+# ---------------- FUNZIONE ----------------
+def predict(image_input):
+    """
+    Supporta:
+    - PIL Image (UI web)
+    - stringa base64 (API)
+    """
     try:
-        # Decodifica
-        if b64_string.startswith("data:image"):
-            b64_string = b64_string.split(",")[1]
-
-        image_data = base64.b64decode(b64_string)
-        img = Image.open(BytesIO(image_data))
-
-        # Qui metti il tuo modello / logica di analisi
-        # Per esempio: restituisco dimensioni e formato
-        result = {
-            "width": img.width,
-            "height": img.height,
-            "format": img.format
-        }
-        return f"Analisi completata: {result}"
+        if isinstance(image_input, str):
+            if image_input.startswith("data:image"):
+                image_input = image_input.split(",",1)[1]
+            img_bytes = base64.b64decode(image_input)
+            img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
+        else:
+            img = image_input.convert("RGB")
+
+        img_tensor = preprocess(img).unsqueeze(0)
+
+        with torch.no_grad():
+            logits = model(img_tensor)
+            probs = torch.nn.functional.softmax(logits[0], dim=0)
+
+        probs_dict = {labels[i]: float(probs[i]) for i in range(len(labels))}
+        max_label = max(probs_dict, key=probs_dict.get)
+
+        return max_label, probs_dict
+
     except Exception as e:
-        return f"Errore: {str(e)}"
+        return f"Error: {str(e)}", {}
+
+def clear_all():
+    return "", ""
 
-# Se carichi da web → converto subito in base64
-def file_to_base64(img: Image.Image):
-    buffered = BytesIO()
-    img.save(buffered, format="PNG")
-    return "data:image/png;base64," + base64.b64encode(buffered.getvalue()).decode()
+# ---------------- INTERFACCIA ----------------
+with gr.Blocks(title="NSFW Image Classifier") as demo:
 
-with gr.Blocks() as demo:
-    gr.Markdown("## Analisi Immagini via Base64")
+    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>Carica un'immagine o incolla la stringa base64 per analizzarla.</p>
+    </div>
+    """)
 
     with gr.Row():
-        img_input = gr.Image(type="pil", label="Carica immagine (verrà convertita in base64)")
-        b64_input = gr.Textbox(label="📤 Base64 dell'immagine (API)", lines=6)
+        with gr.Column(scale=2):
+            # Input UI
+            img_input = gr.Image(label="📷 Carica immagine", type="pil")
+            base64_input = gr.Textbox(
+                label="📤 Base64 dell'immagine (API)",
+                lines=6,
+                placeholder="Incolla qui la stringa base64..."
+            )
+            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=len(labels))
 
-    output = gr.Textbox(label="Risultato analisi")
+    # ---------------- Eventi UI ----------------
+    submit_btn.click(
+        fn=predict,
+        inputs=[img_input],
+        outputs=[label_output, result_display]
+    )
+    clear_btn.click(fn=clear_all, inputs=None, outputs=[img_input, base64_input])
 
-    img_input.change(fn=file_to_base64, inputs=img_input, outputs=b64_input)
-    b64_input.submit(fn=analyze_base64, inputs=b64_input, outputs=output)
+    # ---------------- Pulsante invisibile per API base64 ----------------
+    api_button = gr.Button(visible=False)
+    api_button.click(
+        fn=predict,
+        inputs=[base64_input],
+        outputs=[label_output, result_display],
+        api_name="predict"  # espone /run/predict
+    )
 
-# Avvio server
-demo.launch(server_name="0.0.0.0", server_port=7860)
+# ---------------- LAUNCH ----------------
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860, show_api=True)