Update app.py

app.py

@@ -4,7 +4,7 @@ import numpy as np
 from PIL import Image
 import cv2
 
-print("🚀 Starting SAM2 App...")
+print("🚀 Starting SAM2 App v2.0...")
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"📱 Using device: {device}")
@@ -19,31 +19,36 @@ def load_model():
         try:
             from transformers import SamModel, SamProcessor
             
-            model_name = "facebook/sam-vit-base"
+            # UPGRADE: Large Model für bessere Qualität!
+            # Wenn RAM-Error: zurück zu "facebook/sam-vit-base"
+            model_name = "facebook/sam-vit-large"  # ← BESSER!
+            
             processor = SamProcessor.from_pretrained(model_name)
             model = SamModel.from_pretrained(model_name)
             model.to(device)
-            print("✅ Model loaded!")
+            print(f"✅ Model loaded: {model_name}")
         except Exception as e:
             print(f"❌ Model load error: {e}")
-            raise e
+            print("💡 Falling back to base model...")
+            # Fallback
+            model_name = "facebook/sam-vit-base"
+            processor = SamProcessor.from_pretrained(model_name)
+            model = SamModel.from_pretrained(model_name)
+            model.to(device)
     return model, processor
 
 def prepare_image(image, max_size=1024):
-    """Bild vorbereiten und normalisieren"""
     if isinstance(image, np.ndarray):
         image_pil = Image.fromarray(image)
     else:
         image_pil = image
     
-    # RGB sicherstellen
     if image_pil.mode != 'RGB':
         image_pil = image_pil.convert('RGB')
     
     image_np = np.array(image_pil)
     h, w = image_np.shape[:2]
     
-    # Resize
     if max(h, w) > max_size:
         scale = max_size / max(h, w)
         new_h, new_w = int(h * scale), int(w * scale)
@@ -52,133 +57,129 @@ def prepare_image(image, max_size=1024):
     
     return image_pil, image_np
 
-def segment_automatic(image):
-    """Automatische Segmentierung mit robustem Error Handling"""
+def refine_mask(mask, kernel_size=5):
+    """Maske glätten und Löcher füllen"""
+    mask_uint8 = (mask > 0).astype(np.uint8) * 255
+    
+    # Morphological Operations für glattere Kanten
+    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size))
+    
+    # Closing: Löcher füllen
+    mask_closed = cv2.morphologyEx(mask_uint8, cv2.MORPH_CLOSE, kernel)
+    
+    # Opening: Rauschen entfernen
+    mask_refined = cv2.morphologyEx(mask_closed, cv2.MORPH_OPEN, kernel)
+    
+    return mask_refined > 0
+
+def segment_automatic(image, quality="high"):
+    """
+    Automatische Segmentierung mit Multi-Mask Output
+    quality: "high" (langsamer, präziser) oder "fast" (schneller, gröber)
+    """
     
     if image is None:
         return None, {"error": "Kein Bild hochgeladen"}
     
     try:
-        print("🔄 Starting segmentation...")
+        print(f"🔄 Starting segmentation (quality: {quality})...")
         model, processor = load_model()
         
-        # Bild vorbereiten
         image_pil, image_np = prepare_image(image)
         h, w = image_np.shape[:2]
-        print(f"📐 Image size: {w}x{h}")
         
-        # Mittelpunkt
         center_x, center_y = w // 2, h // 2
-        input_points = [[[center_x, center_y]]]
-        input_labels = [[1]]  # 1 = Foreground
-        
-        print(f"🎯 Using point: ({center_x}, {center_y})")
         
-        # Process
+        # WICHTIG: multimask_output=True für bessere Qualität!
         inputs = processor(
             image_pil,
-            input_points=input_points,
-            input_labels=input_labels,
+            input_points=[[[center_x, center_y]]],
+            input_labels=[[1]],
             return_tensors="pt"
         ).to(device)
         
         print("🧠 Running inference...")
         with torch.no_grad():
-            outputs = model(**inputs)
-        
-        print(f"📊 Output shapes:")
-        print(f"   Masks: {outputs.pred_masks.shape}")
-        print(f"   Scores: {outputs.iou_scores.shape}")
+            outputs = model(**inputs, multimask_output=True)  # ← Multi-Mask!
         
-        # Post-process masks
+        # Post-process
         masks = processor.image_processor.post_process_masks(
             outputs.pred_masks.cpu(),
             inputs["original_sizes"].cpu(),
             inputs["reshaped_input_sizes"].cpu()
-        )
-        
-        print(f"✅ Got {len(masks)} mask batches")
-        
-        # Erste Maske nehmen
-        mask_batch = masks[0]
-        print(f"📦 Mask batch shape: {mask_batch.shape}")
+        )[0]
         
-        # Scores extrahieren
         scores = outputs.iou_scores.cpu().numpy()
-        print(f"📊 Scores shape: {scores.shape}")
-        
-        # Beste Maske finden
-        if scores.ndim == 3:
-            # Shape: [batch, num_masks, 1]
-            scores = scores[0, :, 0]
-        elif scores.ndim == 2:
-            # Shape: [batch, num_masks]
-            scores = scores[0]
-        else:
-            # Fallback
+        if scores.ndim > 1:
             scores = scores.flatten()
         
-        print(f"📈 Scores: {scores}")
+        print(f"✅ Got {len(scores)} masks with scores: {scores}")
         
+        # Beste Maske
         best_idx = np.argmax(scores)
         best_score = scores[best_idx]
         
-        # Maske extrahieren
-        if mask_batch.ndim == 4:
-            # Shape: [1, num_masks, H, W]
-            best_mask = mask_batch[0, best_idx].numpy()
-        elif mask_batch.ndim == 3:
-            # Shape: [num_masks, H, W]
-            best_mask = mask_batch[best_idx].numpy()
+        if masks.ndim == 4:
+            best_mask = masks[0, best_idx].numpy()
         else:
-            # Fallback - erste Maske
-            best_mask = mask_batch.numpy().squeeze()
+            best_mask = masks[best_idx].numpy()
         
-        print(f"✅ Best mask shape: {best_mask.shape}, Score: {best_score}")
+        # REFINEMENT: Maske glätten
+        if quality == "high":
+            best_mask_refined = refine_mask(best_mask, kernel_size=7)
+        else:
+            best_mask_refined = best_mask > 0
         
-        # Overlay erstellen
+        # Overlay mit Anti-Aliasing
         overlay = image_np.copy()
-        color = np.array([255, 50, 150])  # Pink
+        color = np.array([255, 50, 150])
+        
+        # Weiche Kanten mit Gaussian Blur
+        mask_float = best_mask_refined.astype(float)
+        if quality == "high":
+            mask_float = cv2.GaussianBlur(mask_float, (5, 5), 0)
         
-        # Maske anwenden
-        mask_bool = best_mask > 0
-        overlay[mask_bool] = (overlay[mask_bool] * 0.4 + color * 0.6).astype(np.uint8)
+        # Alpha Blending
+        for c in range(3):
+            overlay[:, :, c] = (
+                overlay[:, :, c] * (1 - mask_float * 0.6) +
+                color[c] * mask_float * 0.6
+            )
         
-        # Umriss zeichnen
+        # Kontur zeichnen
         contours, _ = cv2.findContours(
-            (best_mask > 0).astype(np.uint8),
+            best_mask_refined.astype(np.uint8),
             cv2.RETR_EXTERNAL,
             cv2.CHAIN_APPROX_SIMPLE
         )
-        cv2.drawContours(overlay, contours, -1, (255, 255, 0), 2)
+        cv2.drawContours(overlay, contours, -1, (255, 255, 0), 3)
         
         metadata = {
             "success": True,
             "mode": "automatic",
+            "quality": quality,
             "score": float(best_score),
+            "all_scores": scores.tolist(),
             "image_size": [w, h],
-            "mask_area": int(np.sum(mask_bool)),
-            "mask_percentage": float(np.sum(mask_bool) / (h * w) * 100),
-            "device": device,
-            "point_used": [center_x, center_y]
+            "mask_area": int(np.sum(best_mask_refined)),
+            "mask_percentage": float(np.sum(best_mask_refined) / (h * w) * 100),
+            "num_contours": len(contours),
+            "device": device
         }
         
         print("✅ Segmentation complete!")
-        return Image.fromarray(overlay), metadata
+        return Image.fromarray(overlay.astype(np.uint8)), metadata
         
     except Exception as e:
         import traceback
-        error_details = traceback.format_exc()
-        print(f"❌ ERROR:\n{error_details}")
-        
-        return image, {
-            "error": str(e),
-            "details": error_details,
-            "suggestion": "Versuche ein anderes Bild oder warte kurz"
-        }
+        print(f"❌ ERROR:\n{traceback.format_exc()}")
+        return image, {"error": str(e)}
 
-def segment_multi(image):
-    """Multi-Object Segmentierung"""
+def segment_multi_dense(image, density="high"):
+    """
+    Multi-Object mit dichterem Grid
+    density: "high" (15x15), "medium" (9x9), "low" (5x5)
+    """
     
     if image is None:
         return None, {"error": "Kein Bild"}
@@ -188,25 +189,25 @@ def segment_multi(image):
         image_pil, image_np = prepare_image(image)
         h, w = image_np.shape[:2]
         
-        # Grid von Punkten
-        points = [
-            [w//3, h//3],
-            [2*w//3, h//3],
-            [w//3, 2*h//3],
-            [2*w//3, 2*h//3],
-            [w//2, h//2]
-        ]
+        # Dichteres Grid basierend auf Density
+        if density == "high":
+            grid_size = 5  # 5x5 = 25 Punkte
+        elif density == "medium":
+            grid_size = 4  # 4x4 = 16 Punkte
+        else:
+            grid_size = 3  # 3x3 = 9 Punkte
         
-        overlay = image_np.copy()
-        colors = [
-            [255, 0, 0],    # Rot
-            [0, 255, 0],    # Grün
-            [0, 0, 255],    # Blau
-            [255, 255, 0],  # Gelb
-            [255, 0, 255]   # Magenta
-        ]
+        points = []
+        for i in range(1, grid_size + 1):
+            for j in range(1, grid_size + 1):
+                x = int(w * i / (grid_size + 1))
+                y = int(h * j / (grid_size + 1))
+                points.append([x, y])
         
-        masks_info = []
+        print(f"📍 Using {len(points)} points for segmentation")
+        
+        all_masks = []
+        all_scores = []
         
         for idx, point in enumerate(points):
             inputs = processor(
@@ -217,7 +218,7 @@ def segment_multi(image):
             ).to(device)
             
             with torch.no_grad():
-                outputs = model(**inputs)
+                outputs = model(**inputs, multimask_output=True)
             
             masks = processor.image_processor.post_process_masks(
                 outputs.pred_masks.cpu(),
@@ -225,11 +226,7 @@ def segment_multi(image):
                 inputs["reshaped_input_sizes"].cpu()
             )[0]
             
-            scores = outputs.iou_scores.cpu().numpy()
-            
-            # Beste Maske
-            if scores.ndim > 1:
-                scores = scores.flatten()
+            scores = outputs.iou_scores.cpu().numpy().flatten()
             best_idx = np.argmax(scores)
             
             if masks.ndim == 4:
@@ -237,97 +234,135 @@ def segment_multi(image):
             else:
                 mask = masks[best_idx].numpy()
             
-            mask_bool = mask > 0
-            color = colors[idx % len(colors)]
-            
-            overlay[mask_bool] = (overlay[mask_bool] * 0.5 + np.array(color) * 0.5).astype(np.uint8)
+            # Nur Masken mit gutem Score behalten
+            if scores[best_idx] > 0.7:
+                all_masks.append(refine_mask(mask))
+                all_scores.append(scores[best_idx])
+        
+        print(f"✅ Got {len(all_masks)} quality masks")
+        
+        # Masken kombinieren mit verschiedenen Farben
+        overlay = image_np.copy()
+        
+        # HSV-basierte Farbgenerierung
+        colors = []
+        for i in range(len(all_masks)):
+            hue = int(180 * i / max(len(all_masks), 1))
+            color_hsv = np.uint8([[[hue, 255, 200]]])
+            color_rgb = cv2.cvtColor(color_hsv, cv2.COLOR_HSV2RGB)[0][0]
+            colors.append(color_rgb)
+        
+        # Masken anwenden
+        for mask, color, score in zip(all_masks, colors, all_scores):
+            alpha = 0.4 + (score - 0.7) * 0.2  # Höherer Score = stärkere Farbe
+            overlay[mask] = (
+                overlay[mask] * (1 - alpha) +
+                np.array(color) * alpha
+            ).astype(np.uint8)
             
-            masks_info.append({
-                "id": idx,
-                "point": point,
-                "score": float(scores[best_idx]),
-                "area": int(np.sum(mask_bool))
-            })
+            # Kontur
+            contours, _ = cv2.findContours(
+                mask.astype(np.uint8),
+                cv2.RETR_EXTERNAL,
+                cv2.CHAIN_APPROX_SIMPLE
+            )
+            cv2.drawContours(overlay, contours, -1, color.tolist(), 2)
         
         metadata = {
             "success": True,
-            "mode": "multi_object",
-            "num_regions": len(masks_info),
-            "masks": masks_info
+            "mode": "multi_object_dense",
+            "density": density,
+            "grid_size": f"{grid_size}x{grid_size}",
+            "total_points": len(points),
+            "quality_masks": len(all_masks),
+            "avg_score": float(np.mean(all_scores)) if all_scores else 0,
+            "scores": [float(s) for s in all_scores]
         }
         
         return Image.fromarray(overlay), metadata
         
     except Exception as e:
         import traceback
-        print(f"❌ Multi Error:\n{traceback.format_exc()}")
+        print(f"❌ ERROR:\n{traceback.format_exc()}")
         return image, {"error": str(e)}
 
 # Gradio Interface
-with gr.Blocks(title="SAM2 Segmentation", theme=gr.themes.Soft()) as demo:
+with gr.Blocks(title="SAM2 v2.0", theme=gr.themes.Soft()) as demo:
     
-    gr.Markdown("# 🎨 SAM2 Bild Segmentierung")
+    gr.Markdown("""
+    # 🎨 SAM2 Segmentierung v2.0
+    ### Verbesserte Qualität mit Large Model & Multi-Mask Output
+    """)
     
-    with gr.Tab("🤖 Automatisch"):
+    with gr.Tab("🤖 Automatisch (Verbessert)"):
         with gr.Row():
             with gr.Column():
                 input_auto = gr.Image(type="pil", label="📸 Bild hochladen")
+                quality_radio = gr.Radio(
+                    choices=["high", "fast"],
+                    value="high",
+                    label="Qualität",
+                    info="High = präziser aber langsamer"
+                )
                 btn_auto = gr.Button("🚀 Segmentieren", variant="primary", size="lg")
+                
                 gr.Markdown("""
-                **Tipp:** Objekt sollte zentral im Bild sein.
-                Beste Ergebnisse mit klarem Vordergrund/Hintergrund.
+                **Verbesserungen:**
+                - ✨ Large Model (bessere Details)
+                - 🎯 Multi-Mask Output
+                - 🔄 Kanten-Glättung
+                - 🌊 Weiche Übergänge
                 """)
             
             with gr.Column():
                 output_auto = gr.Image(label="✨ Ergebnis")
-                json_auto = gr.JSON(label="📊 Info")
+                json_auto = gr.JSON(label="📊 Details")
         
         btn_auto.click(
             fn=segment_automatic,
-            inputs=input_auto,
+            inputs=[input_auto, quality_radio],
             outputs=[output_auto, json_auto]
         )
-        
-        gr.Examples(
-            examples=[],
-            inputs=input_auto,
-            label="Beispielbilder (Upload deine eigenen!)"
-        )
     
-    with gr.Tab("🎯 Multi-Region"):
+    with gr.Tab("🎯 Multi-Region (Dicht)"):
         with gr.Row():
             with gr.Column():
                 input_multi = gr.Image(type="pil", label="📸 Bild hochladen")
+                density_radio = gr.Radio(
+                    choices=["high", "medium", "low"],
+                    value="medium",
+                    label="Punkt-Dichte",
+                    info="High = mehr Details, aber langsamer"
+                )
                 btn_multi = gr.Button("🚀 Alle Bereiche", variant="primary", size="lg")
+                
+                gr.Markdown("""
+                **Grid-Größen:**
+                - High: 5x5 = 25 Punkte (für Details wie Flossen!)
+                - Medium: 4x4 = 16 Punkte
+                - Low: 3x3 = 9 Punkte
+                """)
             
             with gr.Column():
                 output_multi = gr.Image(label="✨ Ergebnis")
-                json_multi = gr.JSON(label="📊 Info")
+                json_multi = gr.JSON(label="📊 Details")
         
         btn_multi.click(
-            fn=segment_multi,
-            inputs=input_multi,
+            fn=segment_multi_dense,
+            inputs=[input_multi, density_radio],
             outputs=[output_multi, json_multi]
         )
     
     gr.Markdown("""
     ---
-    ### 📡 API Nutzung
+    ### 📊 Qualitäts-Tipps:
     
-    **Endpoint:** `https://dein-space.hf.space/api/predict`
+    **Beste Ergebnisse:**
+    - 📸 Gute Beleuchtung
+    - 🎯 Objekt zentral & klar erkennbar
+    - 🖼️ Kontrast zwischen Objekt & Hintergrund
+    - 📏 Nicht zu klein im Bild
     
-    ```javascript
-    const response = await fetch('https://dein-space.hf.space/api/predict', {
-        method: 'POST',
-        headers: {'Content-Type': 'application/json'},
-        body: JSON.stringify({
-            data: [imageBase64],
-            fn_index: 0  // 0=Auto, 1=Multi
-        })
-    });
-    ```
-    """)
-
-if __name__ == "__main__":
-    print("🌐 Launching...")
-    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)
+    **Bei fehlenden Details (z.B. Flosse):**
+    - Nutze "Multi-Region" mit "High" Density
+    - Oder