Update app.py

app.py

@@ -1,52 +1,208 @@
 import gradio as gr
-import numpy as np
 import torch
-import cv2
+import numpy as np
 from PIL import Image
-import os
-import urllib.request
-from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
-
-# Modell laden oder herunterladen
-MODEL_URL = "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth"
-MODEL_PATH = "sam_vit_b_01ec64.pth"
-
-if not os.path.exists(MODEL_PATH):
-    print("Modell wird heruntergeladen...")
-    urllib.request.urlretrieve(MODEL_URL, MODEL_PATH)
-    print("Modell heruntergeladen.")
+import cv2
+from transformers import SamModel, SamProcessor
+import json
 
-# Modelltyp
-model_type = "vit_b"
+# SAM Model laden
+print("Lade SAM Model...")
+processor = SamProcessor.from_pretrained("facebook/sam-vit-huge")
+model = SamModel.from_pretrained("facebook/sam-vit-huge")
 device = "cuda" if torch.cuda.is_available() else "cpu"
-sam = sam_model_registry[model_type](checkpoint=MODEL_PATH)
-sam.to(device=device)
+model.to(device)
 
-mask_generator = SamAutomaticMaskGenerator(sam)
+def generate_colors(n):
+    """Generiere verschiedene Farben für Masken"""
+    colors = []
+    for i in range(n):
+        hue = int(180 * i / max(n, 1))
+        color = cv2.cvtColor(np.uint8([[[hue, 255, 255]]]), cv2.COLOR_HSV2RGB)[0][0]
+        colors.append(color.tolist())
+    return colors
 
-def segment_all_objects(image):
+def segment_automatic(image):
+    """Automatische Segmentierung - Hauptobjekt in der Mitte"""
+    if image is None:
+        return None, {"error": "Kein Bild hochgeladen"}
+    
+    # Bild vorbereiten
     image_np = np.array(image)
-    masks = mask_generator.generate(image_np)
+    h, w = image_np.shape[:2]
+    
+    # Mittelpunkt als Input (da Objekt zentral sein sollte)
+    input_points = [[[w//2, h//2]]]
+    
+    # SAM Processing
+    inputs = processor(image, input_points=input_points, return_tensors="pt").to(device)
+    
+    with torch.no_grad():
+        outputs = model(**inputs)
+    
+    # Masken extrahieren
+    masks = processor.image_processor.post_process_masks(
+        outputs.pred_masks.cpu(),
+        inputs["original_sizes"].cpu(),
+        inputs["reshaped_input_sizes"].cpu()
+    )[0]
+    
+    # Beste Maske nehmen
+    scores = outputs.iou_scores.cpu().numpy()[0]
+    best_mask_idx = np.argmax(scores)
+    best_mask = masks[best_mask_idx].numpy().squeeze()
+    
+    # Farbige Overlay erstellen
     overlay = image_np.copy()
+    color = [255, 0, 100]  # Pink
+    overlay[best_mask] = overlay[best_mask] * 0.5 + np.array(color) * 0.5
+    
+    # Metadata
+    metadata = {
+        "mode": "automatic",
+        "num_masks": 1,
+        "score": float(scores[best_mask_idx]),
+        "mask_shape": best_mask.shape,
+        "object_detected": True
+    }
+    
+    return Image.fromarray(overlay.astype(np.uint8)), metadata
 
-    for i, mask in enumerate(masks):
-        m = mask["segmentation"]
-        color = np.random.randint(0, 255, size=(3,))
-        overlay[m] = overlay[m] * 0.3 + color * 0.7
-        y, x = np.where(m)
-        if len(x) > 0 and len(y) > 0:
-            cx, cy = int(np.mean(x)), int(np.mean(y))
-            cv2.putText(overlay, f"Obj {i+1}", (cx, cy),
-                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
-
-    return Image.fromarray(overlay.astype(np.uint8))
+def segment_all_objects(image):
+    """Alle Objekte segmentieren - für manuelle Auswahl"""
+    if image is None:
+        return None, {"error": "Kein Bild hochgeladen"}
+    
+    image_np = np.array(image)
+    h, w = image_np.shape[:2]
+    
+    # Grid von Punkten für Segmentierung
+    grid_points = []
+    step = max(h, w) // 8  # 8x8 Grid
+    for y in range(step, h, step):
+        for x in range(step, w, step):
+            grid_points.append([x, y])
+    
+    all_masks = []
+    all_scores = []
+    
+    # Segmentiere jeden Punkt
+    for point in grid_points[:10]:  # Limitiere auf 10 für Performance
+        inputs = processor(image, input_points=[[point]], return_tensors="pt").to(device)
+        
+        with torch.no_grad():
+            outputs = model(**inputs)
+        
+        masks = processor.image_processor.post_process_masks(
+            outputs.pred_masks.cpu(),
+            inputs["original_sizes"].cpu(),
+            inputs["reshaped_input_sizes"].cpu()
+        )[0]
+        
+        scores = outputs.iou_scores.cpu().numpy()[0]
+        best_idx = np.argmax(scores)
+        
+        all_masks.append(masks[best_idx].numpy().squeeze())
+        all_scores.append(scores[best_idx])
+    
+    # Kombiniere alle Masken mit verschiedenen Farben
+    overlay = image_np.copy()
+    colors = generate_colors(len(all_masks))
+    
+    for mask, color in zip(all_masks, colors):
+        overlay[mask] = overlay[mask] * 0.6 + np.array(color) * 0.4
+    
+    metadata = {
+        "mode": "multi_object",
+        "num_masks": len(all_masks),
+        "avg_score": float(np.mean(all_scores)),
+        "masks_data": [
+            {
+                "id": i,
+                "score": float(score),
+                "area": int(mask.sum())
+            } for i, (mask, score) in enumerate(zip(all_masks, all_scores))
+        ]
+    }
+    
+    return Image.fromarray(overlay.astype(np.uint8)), metadata
 
-demo = gr.Interface(
-    fn=segment_all_objects,
-    inputs=gr.Image(type="pil", label="Bild hochladen"),
-    outputs=gr.Image(type="pil", label="Segmentiertes Ergebnis"),
-    title="FishBoost SAM (Meta Original)",
-    description="Segmentiert automatisch alle Objekte im Bild mit Metas offiziellem SAM-Modell."
-)
+def segment_with_points(image, points_json):
+    """Segmentierung mit benutzerdefinierten Punkten"""
+    if image is None:
+        return None, {"error": "Kein Bild hochgeladen"}
+    
+    image_np = np.array(image)
+    
+    try:
+        # Parse JSON input
+        points_data = json.loads(points_json) if isinstance(points_json, str) else points_json
+        input_points = [points_data.get("points", [[image_np.shape[1]//2, image_np.shape[0]//2]])]
+        
+        inputs = processor(image, input_points=input_points, return_tensors="pt").to(device)
+        
+        with torch.no_grad():
+            outputs = model(**inputs)
+        
+        masks = processor.image_processor.post_process_masks(
+            outputs.pred_masks.cpu(),
+            inputs["original_sizes"].cpu(),
+            inputs["reshaped_input_sizes"].cpu()
+        )[0]
+        
+        scores = outputs.iou_scores.cpu().numpy()[0]
+        best_idx = np.argmax(scores)
+        best_mask = masks[best_idx].numpy().squeeze()
+        
+        overlay = image_np.copy()
+        color = [0, 255, 100]  # Grün
+        overlay[best_mask] = overlay[best_mask] * 0.5 + np.array(color) * 0.5
+        
+        metadata = {
+            "mode": "custom_points",
+            "points": input_points[0],
+            "score": float(scores[best_idx]),
+            "success": True
+        }
+        
+        return Image.fromarray(overlay.astype(np.uint8)), metadata
+        
+    except Exception as e:
+        return image, {"error": str(e)}
 
-demo.launch()
+# Gradio Interface
+with gr.Blocks(title="SAM2 Segmentierung API") as demo:
+    gr.Markdown("""
+    # 🎨 SAM2 Bild Segmentierung
+    ### Keine Training nötig - Zero-Shot Object Segmentation!
+    """)
+    
+    with gr.Tab("🤖 Automatisch (Hauptobjekt)"):
+        with gr.Row():
+            with gr.Column():
+                input_auto = gr.Image(type="pil", label="Bild hochladen")
+                btn_auto = gr.Button("Objekt erkennen", variant="primary")
+            with gr.Column():
+                output_auto = gr.Image(label="Segmentiertes Bild")
+                json_auto = gr.JSON(label="Metadata")
+        
+        btn_auto.click(segment_automatic, inputs=input_auto, outputs=[output_auto, json_auto])
+    
+    with gr.Tab("🎯 Mehrere Objekte"):
+        with gr.Row():
+            with gr.Column():
+                input_multi = gr.Image(type="pil", label="Bild hochladen")
+                btn_multi = gr.Button("Alle Objekte erkennen", variant="primary")
+            with gr.Column():
+                output_multi = gr.Image(label="Segmentierte Bereiche")
+                json_multi = gr.JSON(label="Metadata")
+        
+        btn_multi.click(segment_all_objects, inputs=input_multi, outputs=[output_multi, json_multi])
+    
+    with gr.Tab("✋ Custom (mit Punkten)"):
+        with gr.Row():
+            with gr.Column():
+                input_custom = gr.Image(type="pil", label="Bild hochladen")
+                points_input = gr.Textbox(
+                    label="Punkte (JSON)",
+