Update app.py

app.py

@@ -1,3 +1,4 @@
+# nsfw_app.py
 import gradio as gr
 import torch
 import torch.nn as nn
@@ -6,12 +7,13 @@ from torchvision.models import resnet18
 from PIL import Image
 import base64
 import io
+import traceback
 
 # ---------------- CONFIG ----------------
 labels = ["Drawings", "Hentai", "Neutral", "Porn", "Sexy"]
 theme_color = "#6C5B7B"
 
-# ---------------- MODEL ----------------
+# ---------------- MODEL (stesso file pesi) ----------------
 class Classifier(nn.Module):
     def __init__(self):
         super().__init__()
@@ -30,93 +32,106 @@ 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])
 ])
 
+# Carica pesi (stesso file che usavi)
 model = Classifier()
 model.load_state_dict(torch.load("classify_nsfw_v3.0.pth", map_location="cpu"))
 model.eval()
 
-# ---------------- FUNZIONE ----------------
-def predict(image_input):
+# ---------------- FUNZIONE UNICA predict (accetta SOLO base64) ----------------
+def predict(base64_input: str):
     """
-    Supporta:
-    - PIL Image (UI web)
-    - stringa base64 (API)
+    Unico input dell'API: stringa base64 (es. "data:image/jpeg;base64,...")
+    Ritorna: (label_str, {label:prob})
     """
     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")
+        if not base64_input or not isinstance(base64_input, str):
+            return "Input base64 mancante o non valido", {}
+
+        # rimuovi eventuale prefisso data:image...
+        if base64_input.startswith("data:image"):
+            base64_input = base64_input.split(",", 1)[1]
 
-        img_tensor = preprocess(img).unsqueeze(0)
+        # decodifica base64
+        try:
+            img_bytes = base64.b64decode(base64_input)
+        except Exception as e:
+            return f"Errore decodifica base64: {e}", {}
+
+        # apri immagine
+        try:
+            img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
+        except Exception as e:
+            return f"Errore apertura immagine: {e}", {}
 
+        # preprocess + inferenza
+        img_tensor = preprocess(img).unsqueeze(0)  # 1x3x224x224
         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 Image Classifier") as demo:
+    except Exception:
+        return f"Unhandled error:\n{traceback.format_exc()}", {}
 
+# ---------------- Helper: convert image upload -> base64 ----------------
+def image_to_base64(img: Image.Image):
+    """
+    Converte PIL image in data:image/jpeg;base64,...
+    (usato dall'UI: caricamento immagine -> si popola la textbox base64)
+    """
+    if img is None:
+        return ""
+    buffer = io.BytesIO()
+    img.save(buffer, format="JPEG", quality=90)
+    b64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
+    return "data:image/jpeg;base64," + b64
+
+def clear_box():
+    return ""
+
+# ---------------- UI (Blocks) ----------------
+with gr.Blocks(title="NSFW Image Classifier (base64 single-input)"):
     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 style="padding:12px; background:linear-gradient(135deg,#f8f9fa 0%,#e9ecef 100%); border-radius:8px;">
+        <h2 style="color:{theme_color}; margin:0;">🎨 NSFW Image Classifier</h2>
+        <p style="margin:6px 0 0 0;">Carica un'immagine oppure incolla la base64. L'API accetta solo base64.</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 dell'immagine (API)",
-                lines=6,
-                placeholder="Incolla qui la stringa base64..."
-            )
+            image_input = gr.Image(label="📷 Carica immagine (verrà convertita in base64)", type="pil")
+            base64_input = gr.Textbox(label="📤 Base64 (API) — unico input", lines=6,
+                                     placeholder="Incolla qui la stringa base64 (data:image/..;base64,...)")
             with gr.Row():
-                submit_btn = gr.Button("✨ Analizza", variant="primary")
-                clear_btn = gr.Button("🔄 Pulisci", variant="secondary")
-
+                analyze_btn = gr.Button("✨ Analizza (usa la base64 sopra)")
+                clear_btn = gr.Button("🔄 Pulisci")
         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))
 
-    # ---------------- 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
-    )
+    # quando carichi immagine -> converto e popolo la textbox con la base64
+    image_input.change(fn=image_to_base64, inputs=image_input, outputs=base64_input)
+
+    # quando la base64 cambia -> chiamo predict (questo espone automaticamente l'endpoint API
+    # che accetta solo la textbox base64; gradio mappa l'endpoint in /run/predict in locale)
+    base64_input.change(fn=predict, inputs=base64_input, outputs=[label_output, result_display], api_name="predict")
+
+    # pulsante per analizzare manualmente (usa la base64 contenuta nella textbox)
+    analyze_btn.click(fn=predict, inputs=base64_input, outputs=[label_output, result_display])
+
+    clear_btn.click(fn=clear_box, inputs=None, outputs=base64_input)
 
 # ---------------- LAUNCH ----------------
 if __name__ == "__main__":
+    # show_api=True per vedere il link "View API" nella UI (opzionale)
     demo.launch(server_name="0.0.0.0", server_port=7860, show_api=True)
+