Update app.py

app.py

@@ -1,299 +1,168 @@
-import gradio as gr
+# app.py - Boostly SAM2 (AutomaticMaskGenerator, CPU friendly defaults)
+import os
+import urllib.request
+import time
+import traceback
+
 import numpy as np
-import torch
-import cv2
 from PIL import Image
-from transformers import SamModel, SamProcessor
-import spaces
+import cv2
+import torch
+import gradio as gr
+
+print("=== Boostly SAM2 Space starting ===", time.asctime())
+
+# ---------------------------
+# Model download / settings
+# ---------------------------
+MODEL_URL = "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth"
+MODEL_PATH = "sam_vit_b.pth"
+MODEL_TYPE = "vit_b"
+
+# Download model if not present
+if not os.path.exists(MODEL_PATH):
+    print("Model not found — downloading (this can take a while)...")
+    urllib.request.urlretrieve(MODEL_URL, MODEL_PATH)
+    print("Model downloaded.")
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print("Device:", device)
+
+# Try to import segment_anything (needs to be in requirements.txt)
+try:
+    from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
+except Exception as e:
+    print("ERROR importing segment_anything:", e)
+    raise
 
-# GPU setup
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# Load model
+sam = sam_model_registry[MODEL_TYPE](checkpoint=MODEL_PATH)
+sam.to(device=device)
 
-# 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}")
+# Automatic mask generator parameters (tweak for CPU)
+mask_generator = SamAutomaticMaskGenerator(
+    sam,
+    points_per_side=16,            # lower -> fewer proposals (faster)
+    pred_iou_thresh=0.86,          # confidence threshold
+    stability_score_thresh=0.7,    # lower -> more masks
+    min_mask_region_area=200       # remove tiny specks
+)
 
-def prepare_image(image):
-    """Convert image to PIL and numpy formats"""
+print("Mask generator ready.")
+
+# ---------------------------
+# Helpers
+# ---------------------------
+def prepare_image(image, max_size=1024):
+    """Convert and resize image to reduce memory on CPU."""
     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
+    if image_pil.mode != "RGB":
+        image_pil = image_pil.convert("RGB")
+    w, h = image_pil.size
+    max_wh = max(w, h)
+    if max_wh > max_size:
+        scale = max_size / max_wh
+        new_w, new_h = int(w * scale), int(h * scale)
+        image_pil = image_pil.resize((new_w, new_h), Image.Resampling.LANCZOS)
+    return image_pil, np.array(image_pil)
 
-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 keep_top_masks(masks, top_k=12):
+    """Sort masks by area and keep top_k largest ones."""
+    if not masks:
+        return masks
+    areas = [np.sum(m["segmentation"]) for m in masks]
+    idx = np.argsort(areas)[::-1]  # descending
+    top_idx = idx[:top_k]
+    return [masks[i] for i in top_idx]
 
-def safe_to_cpu_np(tensor):
-    """Return numpy array or None if tensor is None"""
-    if tensor is None:
-        return None
-    if isinstance(tensor, torch.Tensor):
-        return tensor.detach().cpu().numpy()
-    return np.asarray(tensor)
+def draw_masks_overlay(image_np, masks, draw_labels=True):
+    overlay = image_np.copy()
+    for i, mask in enumerate(masks):
+        m = mask["segmentation"]
+        color = np.random.randint(40, 230, size=(3,))
+        alpha = 0.55
+        overlay[m] = (overlay[m] * (1 - alpha) + color * alpha).astype(np.uint8)
+        # Draw contour
+        contours, _ = cv2.findContours(m.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        cv2.drawContours(overlay, contours, -1, (255,255,255), 2)
+        # Label
+        if draw_labels:
+            ys, xs = np.where(m)
+            if len(xs) and len(ys):
+                cx, cy = int(np.mean(xs)), int(np.mean(ys))
+                cv2.putText(overlay, f"Obj {i+1}", (cx, cy), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255,255,255), 2)
+    return overlay
 
-def postprocess_masks_safe(processor, pred_masks, original_sizes, reshaped_input_sizes):
+# ---------------------------
+# Main function: segment all
+# ---------------------------
+def segment_all_objects(image, max_size=1024, top_k=12, cpu_mode=True):
     """
-    Wrapper: versucht verschiedene Formen / Signaturen zu behandeln.
+    image: PIL or ndarray
+    max_size: resize largest side to this for CPU
+    top_k: keep top_k biggest masks (to reduce noise)
+    cpu_mode: if True, uses conservative params (already set in mask_generator)
     """
-    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, 
-            original_sizes, 
-            reshaped_input_sizes,
-            return_tensors="np"
-        )
-        return masks
-    except Exception as e:
-        print(f"Postprocess error: {e}")
-        # Return raw masks if post-processing fails
-        return pred_masks_np
+        if image is None:
+            return None, {"error": "Kein Bild hochgeladen"}
+        # prepare image
+        image_pil, image_np = prepare_image(image, max_size=max_size)
+        h, w = image_np.shape[:2]
+        print(f"Starting mask generation — image size: {w}x{h}, device: {device}")
+        t0 = time.time()
 
-def ensure_mask_bool(mask):
-    """Normalize mask to boolean 2D ndarray"""
-    if mask is None:
-        return None
-    
-    a = np.asarray(mask)
-    
-    # Handle different shapes
-    while a.ndim > 2:
-        a = a[0]
-    
-    if a.dtype != bool:
-        return (a > 0).astype(bool)
-    return a
+        # generate masks (this is the heavy call)
+        masks = mask_generator.generate(image_np)
 
-def pick_best_mask_by_score(masks_list, scores):
-    """Return index of chosen mask"""
-    if len(masks_list) == 0:
-        return None
-    
-    if scores is None or len(scores) == 0:
-        return 0
-    
-    scores = safe_to_cpu_np(scores)
-    if scores is None:
-        return 0
-        
-    scores = scores.flatten()
-    
-    # Pick the mask with highest IoU score
-    return int(np.argmax(scores))
+        t1 = time.time()
+        print(f"Generated {len(masks)} masks in {t1-t0:.1f}s")
 
-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):
-        size = int((labels == i).sum())
-        if size > largest_size:
-            largest_size = size
-            largest = i
-    
-    return (labels == largest)
+        # filter: keep largest top_k masks
+        masks = keep_top_masks(masks, top_k=top_k)
+        print(f"Kept top {len(masks)} masks by area")
 
-@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:
-        # Prepare image
-        image_pil, image_np = prepare_image(image)
-        h, w = image_np.shape[:2]
-        
-        # 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"
-        )
-        
-        # 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)
-        
-        # 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 quality == "high" and mask_bool is not None:
-            mask_bool = refine_mask(mask_bool, kernel_size=7)
-        
-        # Create overlay
-        overlay = image_np.copy()
-        
-        if mask_bool is None:
-            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)
-        
-        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)
-        
-        # 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 = {
+        # Convert overlay
+        overlay = draw_masks_overlay(image_np, masks, draw_labels=True)
+
+        metadata = {
             "success": True,
-            "mode": "automatic_smart",
+            "num_masks_generated": len(masks),
+            "generation_time_s": float(t1 - t0),
             "image_size": [w, h],
-            "mask_area": int(mask_bool.sum()),
-            "mask_percentage": float(mask_bool.sum() / (w * h) * 100),
-            "num_contours": len(contours),
-            "device": str(device)
+            "device": device
         }
-        
-        return Image.fromarray(overlay), meta
-    
+        return Image.fromarray(overlay), metadata
+
     except Exception as e:
-        import traceback
-        print(traceback.format_exc())
-        return image, {"error": str(e)}
+        tb = traceback.format_exc()
+        print("ERROR in segment_all_objects:", e)
+        print(tb)
+        return image, {"error": str(e), "trace": tb}
 
-# 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
+# ---------------------------
+# Gradio UI
+# ---------------------------
+title = "Boostly SAM2 - AutomaticMaskGenerator"
+desc = "Upload an image — SAM AutomaticMaskGenerator will propose many segmentations. On CPU it's slower; use small images for tests."
 
-# 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.Blocks(title=title) as demo:
+    gr.Markdown("# 🎨 Boostly - Auto segmentation (SAM)")
     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]
-    )
+        with gr.Column(scale=1):
+            inp = gr.Image(type="pil", label="Upload image")
+            max_size = gr.Slider(minimum=256, maximum=1600, value=1024, step=128, label="Max image side (resize)")
+            top_k = gr.Slider(minimum=1, maximum=40, value=12, step=1, label="Keep top K masks (largest)")
+            btn = gr.Button("Segment (Auto)")
+        with gr.Column(scale=1):
+            out_img = gr.Image(label="Segmented overlay")
+            out_meta = gr.JSON(label="Metadata")
+
+    btn.click(fn=segment_all_objects, inputs=[inp, max_size, top_k, gr.State(True)], outputs=[out_img, out_meta])
 
 if __name__ == "__main__":
-    demo.launch()
+    print("Launching Gradio app...")
+    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)