Update app.py

app.py

@@ -1,64 +1,68 @@
 import gradio as gr
 from ultralytics import YOLO
 import numpy as np
-from PIL import Image
 import cv2
 
-# Carica il modello YOLOv8 segmentation
-model = YOLO('yolov8n-seg.pt')  # Usa il modello nano per velocità su CPU
+# Scarica e carica il modello specifico per balloon
+# Questo modello è addestrato su 8000 immagini di fumetti
+model = YOLO('https://huggingface.co/ogkalu/comic-speech-bubble-detector-yolov8m/resolve/main/best.pt')
 
-def detect_balloons(image, confidence=0.5):
-    """
-    Rileva balloon nell'immagine e restituisce le mask
-    """
-    # Esegui prediction
+def detect_balloons(image, confidence):
+    if image is None:
+        return None, {"error": "Nessuna immagine"}
+    
+    # Detection - ora rileva SOLO balloon, non persone
     results = model(image, conf=confidence)
     
-    # Estrai risultati
     output = {
-        'masks': [],
-        'boxes': [],
-        'polygons': []
+        'num_balloons': 0,
+        'detections': []
     }
     
-    if results[0].masks is not None:
-        for i, mask in enumerate(results[0].masks):
-            # Mask bitmap (array 2D)
-            mask_array = mask.data[0].cpu().numpy()
-            output['masks'].append(mask_array.tolist())
+    if results[0].boxes is not None and len(results[0].boxes) > 0:
+        output['num_balloons'] = len(results[0].boxes)
+        
+        for i in range(len(results[0].boxes)):
+            box = results[0].boxes.xyxy[i].cpu().numpy().tolist()
+            conf = float(results[0].boxes.conf[i].cpu().numpy())
             
-            # Bounding box
-            box = results[0].boxes.xyxy[i].cpu().numpy()
-            output['boxes'].append(box.tolist())
+            # Se il modello ha le mask (segmentation)
+            if results[0].masks is not None:
+                mask = results[0].masks.data[i].cpu().numpy()
+                h, w = image.shape[:2]
+                mask_resized = cv2.resize(mask, (w, h))
+                mask_list = (mask_resized * 255).astype(np.uint8).tolist()
+            else:
+                mask_list = None
             
-            # Poligono (coordinate x,y)
-            polygon = mask.xy[0]
-            output['polygons'].append(polygon.tolist())
-    
-    # Disegna i risultati sull'immagine
-    annotated_image = results[0].plot()
+            output['detections'].append({
+                'box': {
+                    'x1': box[0],
+                    'y1': box[1],
+                    'x2': box[2],
+                    'y2': box[3]
+                },
+                'mask': mask_list,
+                'confidence': conf
+            })
     
-    return annotated_image, output
+    # Disegna i risultati
+    annotated = results[0].plot()
+    return annotated, output
 
-# Crea interfaccia Gradio
-with gr.Blocks() as demo:
-    gr.Markdown("# YOLOv8 Balloon Segmentation")
-    gr.Markdown("Carica un'immagine di fumetto per rilevare automaticamente i balloon")
-    
-    with gr.Row():
-        with gr.Column():
-            input_image = gr.Image(type="pil", label="Immagine Input")
-            confidence_slider = gr.Slider(0.1, 1.0, value=0.5, label="Confidence Threshold")
-            detect_btn = gr.Button("Rileva Balloon")
-        
-        with gr.Column():
-            output_image = gr.Image(label="Risultato")
-            output_json = gr.JSON(label="Dati Segmentazione (masks, boxes, polygons)")
-    
-    detect_btn.click(
-        fn=detect_balloons,
-        inputs=[input_image, confidence_slider],
-        outputs=[output_image, output_json]
-    )
+# Interface Gradio
+demo = gr.Interface(
+    fn=detect_balloons,
+    inputs=[
+        gr.Image(type="numpy", label="📷 Immagine Fumetto"),
+        gr.Slider(0.1, 1.0, value=0.3, label="🎯 Confidenza (più basso = più balloon)")
+    ],
+    outputs=[
+        gr.Image(label="✅ Balloon Rilevati"),
+        gr.JSON(label="📊 Dati JSON")
+    ],
+    title="🎈 Balloon Detection - Solo Nuvolette",
+    description="Rileva SOLO i balloon (nuvolette di dialogo) nei fumetti, non le persone"
+)
 
 demo.launch()