Update app.py

app.py

@@ -1,75 +1,110 @@
 import gradio as gr
-from ultralytics import YOLO
+from ultralytics import YOLO, SAM
 import numpy as np
 import cv2
 
-# Carica il modello locale (già caricato nei Files)
-print("Caricamento modello...")
-model = YOLO('best.pt')
-print("Modello caricato con successo!")
+# Carica modello detection per trovare balloon
+detector = YOLO('best.pt')
 
-def detect_balloons(image, confidence):
+# Carica SAM per segmentazione precisa
+print("Caricamento SAM per segmentazione precisa...")
+segmenter = SAM('sam_b.pt')  # SAM base (più leggero)
+print("Modelli caricati!")
+
+def detect_and_segment_balloons(image, confidence):
     if image is None:
         return None, {"error": "Nessuna immagine"}
     
-    print(f"Rilevamento con confidenza: {confidence}")
+    print(f"1. Rilevamento balloon con confidenza: {confidence}")
     
-    # Detection - rileva SOLO balloon
-    results = model(image, conf=confidence, verbose=False)
+    # Step 1: Rileva balloon con YOLO
+    detection_results = detector(image, conf=confidence, verbose=False)
     
     output = {
         'num_balloons': 0,
         'detections': []
     }
     
-    # Controlla se ci sono detection
-    if results[0].boxes is not None and len(results[0].boxes) > 0:
-        output['num_balloons'] = len(results[0].boxes)
+    if detection_results[0].boxes is None or len(detection_results[0].boxes) == 0:
+        print("Nessun balloon rilevato")
+        return image, output
+    
+    output['num_balloons'] = len(detection_results[0].boxes)
+    print(f"2. Trovati {output['num_balloons']} balloon")
+    
+    # Prepara immagine per visualizzazione
+    annotated = image.copy()
+    
+    # Step 2: Per ogni balloon, usa SAM per segmentazione precisa
+    for i in range(len(detection_results[0].boxes)):
+        box = detection_results[0].boxes.xyxy[i].cpu().numpy()
+        conf = float(detection_results[0].boxes.conf[i].cpu().numpy())
+        
+        print(f"3. Segmentazione balloon {i+1}...")
+        
+        # Usa SAM con il bounding box come prompt
+        seg_results = segmenter(image, bboxes=[box])
         
-        print(f"Trovati {len(results[0].boxes)} balloon")
+        detection_data = {
+            'balloon_id': i + 1,
+            'box': {
+                'x1': int(box[0]),
+                'y1': int(box[1]),
+                'x2': int(box[2]),
+                'y2': int(box[3])
+            },
+            'confidence': round(conf, 3)
+        }
         
-        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())
+        # Estrai la mask precisa da SAM
+        if seg_results[0].masks is not None:
+            mask = seg_results[0].masks.data[0].cpu().numpy()
             
-            detection_data = {
-                'box': {
-                    'x1': int(box[0]),
-                    'y1': int(box[1]),
-                    'x2': int(box[2]),
-                    'y2': int(box[3])
-                },
-                'confidence': round(conf, 3)
-            }
+            # Ridimensiona mask alle dimensioni originali
+            h, w = image.shape[:2]
+            mask_resized = cv2.resize(mask, (w, h))
             
-            # Aggiungi mask se disponibile (segmentation)
-            if hasattr(results[0], 'masks') and results[0].masks is not None:
-                try:
-                    mask = results[0].masks.data[i].cpu().numpy()
-                    h, w = image.shape[:2]
-                    mask_resized = cv2.resize(mask, (w, h))
-                    # Converti in lista di coordinate dei pixel bianchi (più leggero del full array)
-                    white_pixels = np.argwhere(mask_resized > 0.5).tolist()
-                    detection_data['mask_pixels'] = white_pixels[:1000]  # Limita a 1000 pixel per performance
-                    detection_data['has_mask'] = True
-                except:
-                    detection_data['has_mask'] = False
+            # Converti in mask binaria
+            mask_binary = (mask_resized > 0.5).astype(np.uint8)
+            
+            # Trova contorno del balloon (come bacchetta magica)
+            contours, _ = cv2.findContours(mask_binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+            
+            if len(contours) > 0:
+                # Prendi il contorno più grande
+                largest_contour = max(contours, key=cv2.contourArea)
+                
+                # Converti contorno in lista di punti (poligono)
+                polygon = largest_contour.reshape(-1, 2).tolist()
+                detection_data['polygon'] = polygon
+                detection_data['num_points'] = len(polygon)
+                
+                # Mask come array 2D (0/255) - QUESTO È COME LA BACCHETTA MAGICA
+                mask_uint8 = (mask_binary * 255).astype(np.uint8)
+                detection_data['mask'] = mask_uint8.tolist()
+                
+                # Disegna contorno sull'immagine
+                cv2.drawContours(annotated, [largest_contour], -1, (0, 255, 0), 2)
+                
+                # Aggiungi numero balloon
+                x, y = int(box[0]), int(box[1]) - 10
+                cv2.putText(annotated, f"B{i+1}", (x, y), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
+                
+                detection_data['has_mask'] = True
             else:
                 detection_data['has_mask'] = False
-            
-            output['detections'].append(detection_data)
-    else:
-        print("Nessun balloon rilevato - prova ad abbassare la confidenza")
-    
-    # Disegna i risultati
-    annotated = results[0].plot()
+        else:
+            detection_data['has_mask'] = False
+        
+        output['detections'].append(detection_data)
     
+    print(f"4. Completato! {output['num_balloons']} balloon segmentati")
     return annotated, output
 
 # Interface Gradio
-with gr.Blocks(title="Balloon Detection") as demo:
-    gr.Markdown("# 🎈 Rilevamento Balloon nei Fumetti")
-    gr.Markdown("**Modello addestrato specificamente su balloon di fumetti** (manga, webtoon, western comics)")
+with gr.Blocks(title="Balloon Segmentation") as demo:
+    gr.Markdown("# 🎈 Segmentazione Precisa Balloon (come Bacchetta Magica)")
+    gr.Markdown("**Rileva balloon e crea mask precise pixel-per-pixel**")
     
     with gr.Row():
         with gr.Column():
@@ -77,28 +112,33 @@ with gr.Blocks(title="Balloon Detection") as demo:
             confidence = gr.Slider(
                 minimum=0.1, 
                 maximum=1.0, 
-                value=0.25,  # Più basso per rilevare più balloon
+                value=0.25,
                 step=0.05,
-                label="🎯 Confidenza (più basso = più sensibile)"
+                label="🎯 Confidenza"
             )
-            detect_btn = gr.Button("🔍 Rileva Balloon", variant="primary", size="lg")
+            segment_btn = gr.Button("✂️ Segmenta Balloon", variant="primary", size="lg")
             
         with gr.Column():
-            output_image = gr.Image(label="✅ Risultato")
-            output_json = gr.JSON(label="📊 Dati API")
+            output_image = gr.Image(label="✅ Balloon Segmentati (contorni verdi)")
+            output_json = gr.JSON(label="📊 Dati con Mask Precise")
     
     gr.Markdown("""
-    ### 💡 Suggerimenti:
-    - Se **non rileva balloon**: abbassa la confidenza a 0.15-0.20
-    - Se **rileva troppe cose**: alza la confidenza a 0.4-0.5
-    - Il modello è addestrato su 8000 immagini di fumetti vari
+    ### 📖 Formato Output:
+    - **mask**: Array 2D (0/255) - usa come selezione (come bacchetta magica!)
+    - **polygon**: Lista punti [x, y] del contorno
+    - **box**: Rettangolo {x1, y1, x2, y2}
+    
+    ### 🎯 Come Usare nella Tua App:
+    La **mask** è identica alla selezione della bacchetta magica:
+    - Pixel 255 = dentro il balloon (selezionato)
+    - Pixel 0 = fuori (non selezionato)
     """)
     
-    detect_btn.click(
-        fn=detect_balloons,
+    segment_btn.click(
+        fn=detect_and_segment_balloons,
         inputs=[input_image, confidence],
         outputs=[output_image, output_json]
     )
 
-print("App avviata!")
+print("App di segmentazione avviata!")
 demo.launch()