Update app.py

app.py

@@ -1,4 +1,3 @@
-# app.py
 import gradio as gr
 import torch
 import torch.nn as nn
@@ -10,6 +9,7 @@ import io
 
 # ---------------- CONFIG ----------------
 labels = ["Drawings", "Hentai", "Neutral", "Porn", "Sexy"]
+theme_color = "#6C5B7B"
 
 # ---------------- MODEL ----------------
 class Classifier(nn.Module):
@@ -30,10 +30,10 @@ class Classifier(nn.Module):
         return x
 
 preprocess = transforms.Compose([
-    transforms.Resize((224, 224)),
+    transforms.Resize((224,224)),
     transforms.ToTensor(),
     transforms.Normalize(mean=[0.485,0.456,0.406],
-                         std =[0.229,0.224,0.225])
+                         std=[0.229,0.224,0.225])
 ])
 
 model = Classifier()
@@ -41,62 +41,82 @@ model.load_state_dict(torch.load("classify_nsfw_v3.0.pth", map_location="cpu"))
 model.eval()
 
 # ---------------- FUNZIONE ----------------
-def predict(base64_input: str):
-    if not base64_input:
-        return "Nessun input fornito", {}
-
-    if base64_input.startswith("data:image"):
-        base64_input = base64_input.split(",", 1)[1]
-
-    img_bytes = base64.b64decode(base64_input)
-    img = Image.open(io.BytesIO(img_bytes)).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
-
-# ---------------- HELPER ----------------
-def image_to_base64(img: Image.Image):
-    if img is None:
-        return ""
-    buf = io.BytesIO()
-    img.save(buf, format="JPEG", quality=90)
-    return "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode("utf-8")
+def predict(image_input):
+    """
+    Supporta:
+    - PIL Image (UI web)
+    - stringa base64 (API)
+    """
+    try:
+        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"Error: {str(e)}", {}
+
+def clear_all():
+    return "", ""
 
 # ---------------- INTERFACCIA ----------------
-with gr.Blocks(title="NSFW Classifier") as demo:
-    gr.Markdown("## 🎨 NSFW Image Classifier\nCarica un'immagine o incolla la stringa base64.\n\nAPI standard: **/api/predict**")
+with gr.Blocks(title="NSFW Image 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>Carica un'immagine o incolla la stringa base64 per analizzarla.</p>
+    </div>
+    """)
 
     with gr.Row():
         with gr.Column(scale=2):
+            # Input UI
             img_input = gr.Image(label="📷 Carica immagine", type="pil")
-            base64_input = gr.Textbox(label="📤 Base64 (API)", lines=6)
-
+            base64_input = gr.Textbox(
+                label="📤 Base64 dell'immagine (API)",
+                lines=6,
+                placeholder="Incolla qui la stringa base64..."
+            )
             with gr.Row():
-                analyze_btn = gr.Button("✨ Analizza")
-                clear_btn = gr.Button("🔄 Pulisci")
+                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")
+            label_output = gr.Textbox(label="Classe predetta", interactive=False)
             result_display = gr.Label(label="Distribuzione probabilità", num_top_classes=len(labels))
 
-    # immagine → converte in base64 → textbox
-    img_input.change(fn=image_to_base64, inputs=img_input, outputs=base64_input)
-
-    # unico endpoint API standard
-    analyze_btn.click(fn=predict,
-                      inputs=base64_input,
-                      outputs=[label_output, result_display],
-                      api_name="predict")
-
-    clear_btn.click(fn=lambda: "", inputs=None, outputs=base64_input)
-
-# ---------------- AVVIO ----------------
+    # ---------------- 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])
+
+    # ---------------- 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
+    )
+
+# ---------------- LAUNCH ----------------
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(server_name="0.0.0.0", server_port=7860, show_api=True)