Update app.py

app.py

@@ -1,161 +1,299 @@
+import gradio as gr
 import numpy as np
+import torch
 import cv2
 from PIL import Image
+from transformers import SamModel, SamProcessor
+import spaces
+
+# GPU setup
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Model loading
+print("Loading SAM2 model...")
+model_id = "facebook/sam-vit-huge"
+processor = SamProcessor.from_pretrained(model_id)
+model = SamModel.from_pretrained(model_id).to(device)
+model.eval()
+print(f"Model loaded on {device}")
+
+def prepare_image(image):
+    """Convert image to PIL and numpy formats"""
+    if isinstance(image, np.ndarray):
+        image_pil = Image.fromarray(image)
+        image_np = image
+    else:
+        image_pil = image
+        image_np = np.array(image)
+    
+    # Ensure RGB
+    if image_np.ndim == 2:
+        image_np = cv2.cvtColor(image_np, cv2.COLOR_GRAY2RGB)
+    elif image_np.shape[2] == 4:
+        image_np = image_np[:, :, :3]
+    
+    return image_pil, image_np
+
+def refine_mask(mask, kernel_size=5):
+    """Refine mask edges using morphological operations"""
+    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size))
+    mask = cv2.morphologyEx(mask.astype(np.uint8), cv2.MORPH_CLOSE, kernel)
+    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
+    return mask.astype(bool)
 
 def safe_to_cpu_np(tensor):
     """Return numpy array or None if tensor is None"""
     if tensor is None:
         return None
-    return tensor.detach().cpu().numpy()
+    if isinstance(tensor, torch.Tensor):
+        return tensor.detach().cpu().numpy()
+    return np.asarray(tensor)
 
 def postprocess_masks_safe(processor, pred_masks, original_sizes, reshaped_input_sizes):
     """
     Wrapper: versucht verschiedene Formen / Signaturen zu behandeln.
-    Gibt eine Liste von mask arrays (H,W) für das erste Bild zurück.
     """
-    # ensure numpy
-    pred_masks_np = np.asarray(pred_masks)
-    # call processor post_process in a robust way
+    pred_masks_np = safe_to_cpu_np(pred_masks)
+    
     try:
+        # Try the standard post-processing
         masks = processor.image_processor.post_process_masks(
-            pred_masks_np, np.asarray(original_sizes), np.asarray(reshaped_input_sizes)
+            pred_masks_np, 
+            original_sizes, 
+            reshaped_input_sizes,
+            return_tensors="np"
         )
+        return masks
     except Exception as e:
-        # Try alternative signature (some transformers versions accept different args)
-        try:
-            masks = processor.post_process_masks(pred_masks_np, np.asarray(original_sizes), np.asarray(reshaped_input_sizes))
-        except Exception as e2:
-            print("postprocess_masks failed:", e, e2)
-            return []
-    # masks is a list per image - return first (for batch size 1)
-    if isinstance(masks, (list, tuple)) and len(masks) > 0:
-        return masks[0]
-    return []
+        print(f"Postprocess error: {e}")
+        # Return raw masks if post-processing fails
+        return pred_masks_np
 
 def ensure_mask_bool(mask):
-    """normalizes mask to boolean 2D ndarray"""
+    """Normalize mask to boolean 2D ndarray"""
     if mask is None:
         return None
+    
     a = np.asarray(mask)
-    # if mask has channel dim, reduce
-    if a.ndim == 3 and a.shape[0] in (1,):
+    
+    # Handle different shapes
+    while a.ndim > 2:
         a = a[0]
+    
     if a.dtype != bool:
         return (a > 0).astype(bool)
     return a
 
 def pick_best_mask_by_score(masks_list, scores):
-    """Return index of chosen mask; handles missing/short scores"""
+    """Return index of chosen mask"""
     if len(masks_list) == 0:
         return None
+    
     if scores is None or len(scores) == 0:
         return 0
-    scores = np.asarray(scores).flatten()
-    if len(scores) == 1:
+    
+    scores = safe_to_cpu_np(scores)
+    if scores is None:
         return 0
-    # prefer middle mask as before if at least 3
-    if len(scores) >= 3:
-        sorted_idx = np.argsort(scores)
-        return sorted_idx[1]  # middle
+        
+    scores = scores.flatten()
+    
+    # Pick the mask with highest IoU score
     return int(np.argmax(scores))
 
 def keep_largest_component_bool(mask_bool):
     """Return boolean mask with only largest connected component"""
     if mask_bool is None:
         return None
+    
     mask_uint8 = (mask_bool.astype(np.uint8) * 255)
     num_labels, labels = cv2.connectedComponents(mask_uint8)
+    
     if num_labels <= 1:
         return mask_bool
+    
+    # Find largest component (excluding background)
     largest = 1
     largest_size = 0
     for i in range(1, num_labels):
-        s = int((labels == i).sum())
-        if s > largest_size:
-            largest_size = s
+        size = int((labels == i).sum())
+        if size > largest_size:
+            largest_size = size
             largest = i
+    
     return (labels == largest)
 
+@spaces.GPU
 def segment_automatic(image, quality="high", filter_largest=True):
+    """Main segmentation function"""
     if image is None:
         return None, {"error": "Kein Bild hochgeladen"}
+    
     try:
-        model, processor = load_model()
+        # Prepare image
         image_pil, image_np = prepare_image(image)
         h, w = image_np.shape[:2]
-        cx, cy = w//2, h//2
-
-        # Build inputs
+        
+        # Use center point as default
+        cx, cy = w // 2, h // 2
+        
+        # Prepare inputs
         inputs = processor(
             image_pil,
             input_points=[[[cx, cy]]],
             input_labels=[[1]],
             return_tensors="pt"
         )
-        inputs = {k: v.to(device) if hasattr(v, "to") else v for k, v in inputs.items()}
-
+        
+        # Move to device
+        inputs = {k: v.to(device) if hasattr(v, 'to') else v for k, v in inputs.items()}
+        
+        # Run model
         with torch.no_grad():
             outputs = model(**inputs, multimask_output=True)
-
-        # safe extraction
-        pred_masks = safe_to_cpu_np(getattr(outputs, "pred_masks", None))
-        iou_scores = safe_to_cpu_np(getattr(outputs, "iou_scores", None))
-
-        # Get original/reshaped sizes from inputs (move to cpu/numpy if tensors)
-        orig_sizes = safe_to_cpu_np(inputs.get("original_sizes", None))
-        reshaped_sizes = safe_to_cpu_np(inputs.get("reshaped_input_sizes", None))
-
-        # postprocess masks robustly
-        masks_post = postprocess_masks_safe(processor, pred_masks, orig_sizes if orig_sizes is not None else [[h, w]], reshaped_sizes if reshaped_sizes is not None else [[h, w]])
-        if not masks_post:
-            return image, {"error": "keine masken erzeugt (postprocess returned empty)"}
-
-        # masks_post is list/array: choose best
-        idx = pick_best_mask_by_score(masks_post, iou_scores)
-        if idx is None:
-            return image, {"error": "keine passende maske gefunden"}
-
-        # normalize mask
-        # masks_post might be array of masks or dict -> handle both
-        mask_candidate = masks_post[idx] if isinstance(masks_post, (list, tuple, np.ndarray)) else masks_post
-        mask_bool = ensure_mask_bool(mask_candidate)
-
-        # filter largest component
+        
+        # Extract predictions
+        pred_masks = outputs.pred_masks.squeeze(0)  # Remove batch dimension
+        iou_scores = outputs.iou_scores.squeeze(0) if hasattr(outputs, 'iou_scores') else None
+        
+        # Convert to numpy
+        pred_masks_np = safe_to_cpu_np(pred_masks)
+        iou_scores_np = safe_to_cpu_np(iou_scores)
+        
+        # Pick best mask
+        if pred_masks_np.ndim == 3:  # Multiple masks
+            idx = pick_best_mask_by_score(pred_masks_np, iou_scores_np)
+            if idx is not None:
+                mask = pred_masks_np[idx]
+            else:
+                mask = pred_masks_np[0]
+        else:
+            mask = pred_masks_np
+        
+        # Ensure boolean mask
+        mask_bool = ensure_mask_bool(mask)
+        
+        # Resize mask to original image size if needed
+        if mask_bool.shape != (h, w):
+            mask_bool = cv2.resize(mask_bool.astype(np.uint8), (w, h), 
+                                   interpolation=cv2.INTER_NEAREST).astype(bool)
+        
+        # Filter largest component
         if filter_largest and mask_bool is not None:
             mask_bool = keep_largest_component_bool(mask_bool)
-
-        # refine edges if requested
+        
+        # Refine edges
         if quality == "high" and mask_bool is not None:
             mask_bool = refine_mask(mask_bool, kernel_size=7)
-
-        # overlay coloring
+        
+        # Create overlay
         overlay = image_np.copy()
+        
         if mask_bool is None:
-            return image, {"error": "maske ist None nach Verarbeitung"}
-
+            return image, {"error": "Maske ist None nach Verarbeitung"}
+        
+        # Apply colored overlay
         mask_float = mask_bool.astype(float)
         if quality == "high":
-            mask_float = cv2.GaussianBlur(mask_float, (5,5), 0)
-
+            mask_float = cv2.GaussianBlur(mask_float, (5, 5), 0)
+        
         color = np.array([255, 80, 180], dtype=np.uint8)
         for c in range(3):
-            overlay[:,:,c] = (overlay[:,:,c] * (1 - mask_float * 0.65) + color[c] * mask_float * 0.65).astype(np.uint8)
-
-        contours, _ = cv2.findContours((mask_bool.astype(np.uint8)), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-        cv2.drawContours(overlay, contours, -1, (255,255,0), 3)
-
+            overlay[:, :, c] = (overlay[:, :, c] * (1 - mask_float * 0.65) + 
+                               color[c] * mask_float * 0.65).astype(np.uint8)
+        
+        # Draw contours
+        contours, _ = cv2.findContours(mask_bool.astype(np.uint8), 
+                                       cv2.RETR_EXTERNAL, 
+                                       cv2.CHAIN_APPROX_SIMPLE)
+        cv2.drawContours(overlay, contours, -1, (255, 255, 0), 3)
+        
+        # Metadata
         meta = {
             "success": True,
             "mode": "automatic_smart",
             "image_size": [w, h],
             "mask_area": int(mask_bool.sum()),
-            "mask_percentage": float(mask_bool.sum() / (w*h) * 100),
+            "mask_percentage": float(mask_bool.sum() / (w * h) * 100),
             "num_contours": len(contours),
-            "device": device
+            "device": str(device)
         }
+        
         return Image.fromarray(overlay), meta
-
+    
     except Exception as e:
         import traceback
         print(traceback.format_exc())
         return image, {"error": str(e)}
+
+# Gradio Interface
+def segment_image(image, quality, filter_largest):
+    """Gradio wrapper function"""
+    result_image, metadata = segment_automatic(image, quality, filter_largest)
+    
+    # Format metadata for display
+    if "error" in metadata:
+        info_text = f"❌ Fehler: {metadata['error']}"
+    else:
+        info_text = f"""
+✅ Segmentierung erfolgreich!
+📏 Bildgröße: {metadata['image_size'][0]} x {metadata['image_size'][1]}
+🎯 Maskierte Fläche: {metadata['mask_area']:,} Pixel ({metadata['mask_percentage']:.1f}%)
+📍 Anzahl Konturen: {metadata['num_contours']}
+🖥️ Gerät: {metadata['device']}
+"""
+    
+    return result_image, info_text
+
+# Create Gradio interface
+with gr.Blocks(title="SAM2 Segmentierung") as demo:
+    gr.Markdown("""
+    # 🎯 SAM2 Automatische Segmentierung
+    
+    Lädt ein Bild hoch und segmentiert automatisch das zentrale Objekt.
+    """)
+    
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(label="Eingabebild", type="pil")
+            
+            with gr.Row():
+                quality_dropdown = gr.Dropdown(
+                    choices=["standard", "high"],
+                    value="high",
+                    label="Qualität",
+                    info="High = bessere Kanten, aber langsamer"
+                )
+                filter_checkbox = gr.Checkbox(
+                    value=True,
+                    label="Nur größte Komponente",
+                    info="Entfernt kleine isolierte Bereiche"
+                )
+            
+            segment_btn = gr.Button("🚀 Segmentieren", variant="primary")
+        
+        with gr.Column():
+            output_image = gr.Image(label="Segmentiertes Bild")
+            info_text = gr.Textbox(label="Informationen", lines=6)
+    
+    # Examples
+    gr.Examples(
+        examples=[
+            ["example1.jpg", "high", True],
+            ["example2.jpg", "standard", False],
+        ],
+        inputs=[input_image, quality_dropdown, filter_checkbox],
+        outputs=[output_image, info_text],
+        fn=segment_image,
+        cache_examples=False
+    )
+    
+    # Connect button
+    segment_btn.click(
+        fn=segment_image,
+        inputs=[input_image, quality_dropdown, filter_checkbox],
+        outputs=[output_image, info_text]
+    )
+
+if __name__ == "__main__":
+    demo.launch()