Update app.py

app.py

@@ -1,168 +1 @@
-# app.py - Boostly SAM2 (AutomaticMaskGenerator, CPU friendly defaults)
-import os
-import urllib.request
-import time
-import traceback
-
-import numpy as np
-from PIL import Image
-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
-
-# Load model
-sam = sam_model_registry[MODEL_TYPE](checkpoint=MODEL_PATH)
-sam.to(device=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
-)
-
-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)
-    else:
-        image_pil = image
-    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 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 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
-
-# ---------------------------
-# Main function: segment all
-# ---------------------------
-def segment_all_objects(image, max_size=1024, top_k=12, cpu_mode=True):
-    """
-    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)
-    """
-    try:
-        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()
-
-        # generate masks (this is the heavy call)
-        masks = mask_generator.generate(image_np)
-
-        t1 = time.time()
-        print(f"Generated {len(masks)} masks in {t1-t0:.1f}s")
-
-        # filter: keep largest top_k masks
-        masks = keep_top_masks(masks, top_k=top_k)
-        print(f"Kept top {len(masks)} masks by area")
-
-        # Convert overlay
-        overlay = draw_masks_overlay(image_np, masks, draw_labels=True)
-
-        metadata = {
-            "success": True,
-            "num_masks_generated": len(masks),
-            "generation_time_s": float(t1 - t0),
-            "image_size": [w, h],
-            "device": device
-        }
-        return Image.fromarray(overlay), metadata
-
-    except Exception as e:
-        tb = traceback.format_exc()
-        print("ERROR in segment_all_objects:", e)
-        print(tb)
-        return image, {"error": str(e), "trace": tb}
-
-# ---------------------------
-# 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."
-
-with gr.Blocks(title=title) as demo:
-    gr.Markdown("# 🎨 Boostly - Auto segmentation (SAM)")
-    with gr.Row():
-        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__":
-    print("Launching Gradio app...")
-    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)
+d78a30bf53e5e1fdfc7d05affd3fac335304ed94