app.py

import gradio as gr
from ultralytics import YOLO
import numpy as np
import cv2

# 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):
    if image is None:
        return None, {"error": "Nessuna immagine"}
    
    # Detection - ora rileva SOLO balloon, non persone
    results = model(image, conf=confidence)
    
    output = {
        'num_balloons': 0,
        'detections': []
    }
    
    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())
            
            # 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
            
            output['detections'].append({
                'box': {
                    'x1': box[0],
                    'y1': box[1],
                    'x2': box[2],
                    'y2': box[3]
                },
                'mask': mask_list,
                'confidence': conf
            })
    
    # Disegna i risultati
    annotated = results[0].plot()
    return annotated, output

# 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()