Update app.py

app.py

@@ -9,6 +9,7 @@ import io
 
 # ---------------- CONFIG ----------------
 labels = ["Drawings", "Hentai", "Neutral", "Porn", "Sexy"]
+theme_color = "#6C5B7B"
 
 # ---------------- MODEL ----------------
 class Classifier(nn.Module):
@@ -38,18 +39,17 @@ model = Classifier()
 model.load_state_dict(torch.load("classify_nsfw_v3.0.pth", map_location="cpu"))
 model.eval()
 
-# ---------------- FUNZIONE ----------------
-def predict(input_data):
+# ---------------- FUNZIONI ----------------
+def predict(base64_input: str):
+    """
+    Unico input: stringa base64 (da API o da UI).
+    """
     try:
-        if isinstance(input_data, str):
-            if input_data.startswith("data:image"):
-                input_data = input_data.split(",",1)[1]
-            img_bytes = base64.b64decode(input_data)
-            img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
-        elif isinstance(input_data, Image.Image):
-            img = input_data.convert("RGB")
-        else:
-            return "Input non valido", {}
+        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():
@@ -63,20 +63,52 @@ def predict(input_data):
     except Exception as e:
         return f"Error: {str(e)}", {}
 
+def image_to_base64(img: Image.Image):
+    """
+    Converte immagine caricata in base64 e la ritorna
+    così finisce nella Textbox e poi viene analizzata.
+    """
+    buffered = io.BytesIO()
+    img.save(buffered, format="JPEG")
+    img_b64 = base64.b64encode(buffered.getvalue()).decode()
+    return "data:image/jpeg;base64," + img_b64
+
+def clear_all():
+    return ""
+
 # ---------------- INTERFACCIA ----------------
-demo = gr.Interface(
-    fn=predict,
-    inputs=[
-        gr.Image(type="pil", label="📷 Carica immagine"),
-        gr.Textbox(label="📤 Oppure Base64", lines=6, placeholder="Incolla base64 qui...")
-    ],
-    outputs=[
-        gr.Textbox(label="Classe predetta"),
-        gr.Label(label="Distribuzione probabilità", num_top_classes=len(labels))
-    ],
-    live=False,
-    allow_flagging="never"
-)
+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.<br>
+        L'API espone <code>/run/predict</code> e accetta <b>solo base64</b>.</p>
+    </div>
+    """)
+
+    with gr.Row():
+        with gr.Column(scale=2):
+            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))
+
+    # Se carico immagine → la converto in base64 → la inserisco nella Textbox
+    img_input.change(fn=image_to_base64, inputs=img_input, outputs=base64_input)
+
+    # Submit → usa sempre base64 come input
+    submit_btn.click(fn=predict, inputs=base64_input, outputs=[label_output, result_display])
+    clear_btn.click(fn=clear_all, inputs=None, outputs=base64_input)
 
 # ---------------- LAUNCH ----------------
 if __name__ == "__main__":