Update app.py

app.py

@@ -3,18 +3,36 @@ import torch
 import numpy as np
 from PIL import Image
 import cv2
-from transformers import SamModel, SamProcessor
 import json
 
-# SAM Model laden
-print("Lade SAM Model...")
-processor = SamProcessor.from_pretrained("facebook/sam-vit-huge")
-model = SamModel.from_pretrained("facebook/sam-vit-huge")
+print("🚀 Starting SAM2 App...")
+
+# Device Setup
 device = "cuda" if torch.cuda.is_available() else "cpu"
-model.to(device)
+print(f"📱 Using device: {device}")
+
+# Lazy Loading
+model = None
+processor = None
+
+def load_model():
+    global model, processor
+    if model is None:
+        print("📦 Loading SAM model...")
+        try:
+            from transformers import SamModel, SamProcessor
+            
+            model_name = "facebook/sam-vit-base"
+            processor = SamProcessor.from_pretrained(model_name)
+            model = SamModel.from_pretrained(model_name)
+            model.to(device)
+            print("✅ Model loaded successfully!")
+        except Exception as e:
+            print(f"❌ Error loading model: {e}")
+            raise e
+    return model, processor
 
 def generate_colors(n):
-    """Generiere verschiedene Farben für Masken"""
     colors = []
     for i in range(n):
         hue = int(180 * i / max(n, 1))
@@ -23,72 +41,33 @@ def generate_colors(n):
     return colors
 
 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)
-    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
-
-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)
+    try:
+        model, processor = load_model()
+        
+        if isinstance(image, np.ndarray):
+            image_pil = Image.fromarray(image)
+        else:
+            image_pil = image
+        
+        image_np = np.array(image_pil)
+        h, w = image_np.shape[:2]
+        
+        # Resize wenn zu groß
+        max_size = 1024
+        if max(h, w) > max_size:
+            scale = max_size / max(h, w)
+            new_h, new_w = int(h * scale), int(w * scale)
+            image_pil = image_pil.resize((new_w, new_h))
+            image_np = np.array(image_pil)
+            h, w = new_h, new_w
+        
+        # Mittelpunkt
+        input_points = [[[w//2, h//2]]]
+        
+        inputs = processor(image_pil, input_points=input_points, return_tensors="pt").to(device)
         
         with torch.no_grad():
             outputs = model(**inputs)
@@ -101,108 +80,176 @@ def segment_all_objects(image):
         
         scores = outputs.iou_scores.cpu().numpy()[0]
         best_idx = np.argmax(scores)
+        best_mask = masks[best_idx].numpy().squeeze()
         
-        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
+        overlay = image_np.copy()
+        color = np.array([255, 0, 100])
+        overlay[best_mask] = (overlay[best_mask] * 0.5 + color * 0.5).astype(np.uint8)
+        
+        metadata = {
+            "mode": "automatic",
+            "success": True,
+            "score": float(scores[best_idx]),
+            "image_size": [h, w],
+            "device": device
+        }
+        
+        return Image.fromarray(overlay), metadata
+        
+    except Exception as e:
+        print(f"❌ Error: {e}")
+        return image, {"error": str(e)}
 
-def segment_with_points(image, points_json):
-    """Segmentierung mit benutzerdefinierten Punkten"""
+def segment_multi(image):
     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]])]
+        model, processor = load_model()
         
-        inputs = processor(image, input_points=input_points, return_tensors="pt").to(device)
+        if isinstance(image, np.ndarray):
+            image_pil = Image.fromarray(image)
+        else:
+            image_pil = image
         
-        with torch.no_grad():
-            outputs = model(**inputs)
+        image_np = np.array(image_pil)
+        h, w = image_np.shape[:2]
         
-        masks = processor.image_processor.post_process_masks(
-            outputs.pred_masks.cpu(),
-            inputs["original_sizes"].cpu(),
-            inputs["reshaped_input_sizes"].cpu()
-        )[0]
+        # Resize
+        max_size = 1024
+        if max(h, w) > max_size:
+            scale = max_size / max(h, w)
+            new_h, new_w = int(h * scale), int(w * scale)
+            image_pil = image_pil.resize((new_w, new_h))
+            image_np = np.array(image_pil)
+            h, w = new_h, new_w
         
-        scores = outputs.iou_scores.cpu().numpy()[0]
-        best_idx = np.argmax(scores)
-        best_mask = masks[best_idx].numpy().squeeze()
+        # Grid von Punkten
+        all_masks = []
+        all_scores = []
+        
+        grid_points = [
+            [w//4, h//4], [w//2, h//4], [3*w//4, h//4],
+            [w//4, h//2], [w//2, h//2], [3*w//4, h//2],
+            [w//4, 3*h//4], [w//2, 3*h//4], [3*w//4, 3*h//4]
+        ]
+        
+        for point in grid_points[:6]:
+            inputs = processor(image_pil, 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])
         
         overlay = image_np.copy()
-        color = [0, 255, 100]  # Grün
-        overlay[best_mask] = overlay[best_mask] * 0.5 + np.array(color) * 0.5
+        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).astype(np.uint8)
         
         metadata = {
-            "mode": "custom_points",
-            "points": input_points[0],
-            "score": float(scores[best_idx]),
-            "success": True
+            "mode": "multi_object",
+            "success": True,
+            "num_masks": len(all_masks),
+            "avg_score": float(np.mean(all_scores)),
+            "device": device
         }
         
-        return Image.fromarray(overlay.astype(np.uint8)), metadata
+        return Image.fromarray(overlay), metadata
         
     except Exception as e:
+        print(f"❌ Error: {e}")
         return image, {"error": str(e)}
 
 # Gradio Interface
-with gr.Blocks(title="SAM2 Segmentierung API") as demo:
+with gr.Blocks(title="SAM2 Segmentation API", theme=gr.themes.Soft()) as demo:
+    
     gr.Markdown("""
     # 🎨 SAM2 Bild Segmentierung
-    ### Keine Training nötig - Zero-Shot Object Segmentation!
+    ### Zero-Shot Object Segmentation - Kein Training nötig!
     """)
     
-    with gr.Tab("🤖 Automatisch (Hauptobjekt)"):
+    with gr.Tab("🤖 Automatisch"):
+        gr.Markdown("Erkennt automatisch das Hauptobjekt in der Bildmitte")
+        
         with gr.Row():
             with gr.Column():
                 input_auto = gr.Image(type="pil", label="Bild hochladen")
-                btn_auto = gr.Button("Objekt erkennen", variant="primary")
+                btn_auto = gr.Button("🚀 Objekt erkennen", variant="primary", size="lg")
+            
             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])
+        btn_auto.click(
+            fn=segment_automatic,
+            inputs=input_auto,
+            outputs=[output_auto, json_auto]
+        )
     
     with gr.Tab("🎯 Mehrere Objekte"):
+        gr.Markdown("Segmentiert mehrere Bereiche im Bild")
+        
         with gr.Row():
             with gr.Column():
                 input_multi = gr.Image(type="pil", label="Bild hochladen")
-                btn_multi = gr.Button("Alle Objekte erkennen", variant="primary")
+                btn_multi = gr.Button("🚀 Alle Objekte erkennen", variant="primary", size="lg")
+            
             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)",
-                    
+        btn_multi.click(
+            fn=segment_multi,
+            inputs=input_multi,
+            outputs=[output_multi, json_multi]
+        )
+    
+    with gr.Tab("📡 API Info"):
+        gr.Markdown("""
+        ### API Nutzung für Lovable/externe Apps
+        
+        **Endpoint:** `https://dein-space-name.hf.space/api/predict`
+        
+        **Beispiel (JavaScript):**
+        ```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
+            })
+        });
+        
+        const result = await response.json();
+        console.log(result.data); // [segmentedImage, metadata]
+        ```
+        
+        **Python:**
+        ```python
+        from gradio_client import Client
+        
+        client = Client("dein-username/dein-space")
+        result = client.predict("image.jpg", api_name="/predict")
+        ```
+        """)
+
+if __name__ == "__main__":
+    print("🌐 Launching Gradio Interface...")
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        show_error=True
+    )