Update app.py

app.py

@@ -1,447 +1,161 @@
-import gradio as gr
-import torch
 import numpy as np
-from PIL import Image
 import cv2
-import io
-import time
-import traceback
-
-print("🚀 Starting SAM2 Boostly - Stable CPU-friendly Edition")
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"📱 Using device: {device}")
+from PIL import Image
 
-model = None
-processor = None
+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()
 
-# ---- Helper: robustes post_process_masks, kompatibel mit Transformers-Versionen ----
-def safe_post_process_masks(processor, outputs, inputs):
+def postprocess_masks_safe(processor, pred_masks, original_sizes, reshaped_input_sizes):
     """
-    Versucht mehrere Varianten von post_process_masks aufzurufen.
-    Liefert eine Liste von mask-arrays passend zur originalen Größe.
+    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
     try:
-        # Neue API (manche Transformers-Versionen)
-        masks = processor.post_process_masks(
-            outputs=outputs,
-            original_sizes=inputs["original_sizes"].cpu(),
-            reshaped_input_sizes=inputs["reshaped_input_sizes"].cpu()
+        masks = processor.image_processor.post_process_masks(
+            pred_masks_np, np.asarray(original_sizes), np.asarray(reshaped_input_sizes)
         )
-        return masks
-    except TypeError:
-        try:
-            # Alternative: processor.image_processor.post_process_masks(pred_masks, original_sizes, reshaped_input_sizes)
-            pred = outputs.pred_masks.cpu()
-            masks = processor.image_processor.post_process_masks(
-                pred,
-                inputs["original_sizes"].cpu(),
-                inputs["reshaped_input_sizes"].cpu()
-            )
-            return masks
-        except Exception:
-            # Fallback: einfache Verarbeitung - nehme pred_masks (resize zur originalgröße)
-            try:
-                pred = outputs.pred_masks.cpu().numpy()
-                # pred shape handling: try common shapes
-                # pred: [batch, num_masks, H, W] or [num_masks, H, W]
-                if pred.ndim == 4:
-                    batch_masks = []
-                    for b in range(pred.shape[0]):
-                        mask_list = []
-                        for m in range(pred.shape[1]):
-                            mask = pred[b, m]
-                            mask_list.append(mask)
-                        batch_masks.append(mask_list)
-                    # Convert to expected structure: list of masks for first image
-                    # We will wrap them similar to new API: [[mask1_resized, ...]]
-                    resized = []
-                    orig_size = inputs["original_sizes"][0].cpu().tolist()
-                    for mask in batch_masks[0]:
-                        mask_img = Image.fromarray((mask * 255).astype(np.uint8))
-                        mask_img = mask_img.resize((orig_size[1], orig_size[0]), Image.Resampling.NEAREST)
-                        resized.append(np.array(mask_img) > 127)
-                    return [ [resized] ]
-                elif pred.ndim == 3:
-                    # single image, multiple masks
-                    orig_size = inputs["original_sizes"][0].cpu().tolist()
-                    resized = []
-                    for m in range(pred.shape[0]):
-                        mask = pred[m]
-                        mask_img = Image.fromarray((mask * 255).astype(np.uint8))
-                        mask_img = mask_img.resize((orig_size[1], orig_size[0]), Image.Resampling.NEAREST)
-                        resized.append(np.array(mask_img) > 127)
-                    return [ [resized] ]
-            except Exception as e:
-                print("⚠️ safe_post_process_masks fallback failed:", e)
-                traceback.print_exc()
-            return None
-
-# ---- Modell laden (einmal) ----
-def load_model():
-    global model, processor
-    if model is None:
-        print("📦 Loading SAM model (base recommended for CPU)...")
+    except Exception as e:
+        # Try alternative signature (some transformers versions accept different args)
         try:
-            from transformers import SamModel, SamProcessor
-            model_name = "facebook/sam-vit-base"  # DEFAULT -> base
-            processor = SamProcessor.from_pretrained(model_name)
-            model = SamModel.from_pretrained(model_name)
-            model.to(device)
-            print(f"✅ Model loaded: {model_name}")
-        except Exception as e:
-            print("❌ Error loading transformers SAM-model:", e)
-            traceback.print_exc()
-            raise e
-    return model, processor
-
-# ---- Bild vorbereiten ----
-def prepare_image(image, max_size=1024):
-    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')
+            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 []
+
+def ensure_mask_bool(mask):
+    """normalizes 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,):
+        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"""
+    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:
+        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
+    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
+    largest = 1
+    largest_size = 0
+    for i in range(1, num_labels):
+        s = int((labels == i).sum())
+        if s > largest_size:
+            largest_size = s
+            largest = i
+    return (labels == largest)
 
-    image_np = np.array(image_pil)
-    h, w = image_np.shape[:2]
-
-    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.Resampling.LANCZOS)
-        image_np = np.array(image_pil)
-
-    return image_pil, image_np
-
-# ---- Masken-Raffinement ----
-def refine_mask(mask, kernel_size=5):
-    mask_uint8 = (mask > 0).astype(np.uint8) * 255
-    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size))
-    mask_closed = cv2.morphologyEx(mask_uint8, cv2.MORPH_CLOSE, kernel)
-    mask_refined = cv2.morphologyEx(mask_closed, cv2.MORPH_OPEN, kernel)
-    return mask_refined > 0
-
-# ---- SMART Automatic segmentation (einzelobjekt) ----
 def segment_automatic(image, quality="high", filter_largest=True):
     if image is None:
         return None, {"error": "Kein Bild hochgeladen"}
-
-    start = time.time()
     try:
         model, processor = load_model()
         image_pil, image_np = prepare_image(image)
         h, w = image_np.shape[:2]
-        center_x, center_y = w // 2, h // 2
-
-        # Device-based multimask decision: only use multimask on GPU
-        use_multimask = True if device == "cuda" else False
+        cx, cy = w//2, h//2
 
+        # Build inputs
         inputs = processor(
             image_pil,
-            input_points=[[[center_x, center_y]]],
+            input_points=[[[cx, cy]]],
             input_labels=[[1]],
             return_tensors="pt"
-        ).to(device)
+        )
+        inputs = {k: v.to(device) if hasattr(v, "to") else v for k, v in inputs.items()}
 
         with torch.no_grad():
-            outputs = model(**inputs, multimask_output=use_multimask)
+            outputs = model(**inputs, multimask_output=True)
 
-        masks_result = safe_post_process_masks(processor, outputs, inputs)
-        if masks_result is None:
-            print("⚠️ No masks returned by post_process_masks; using basic threshold fallback.")
-            # fallback: try to create single mask from pred_masks
-            try:
-                pred = outputs.pred_masks.cpu().numpy()
-                if pred.ndim == 4:
-                    mask = pred[0, 0] > 0.5
-                elif pred.ndim == 3:
-                    mask = pred[0] > 0.5
-                else:
-                    mask = pred > 0.5
-                masks = [mask]
-            except Exception as e:
-                print("Fallback also failed:", e)
-                return image, {"error": "Mask processing failed"}
-        else:
-            # masks_result is like [[ [mask1_arr, mask2_arr, ... ] ]]
-            # take first batch
-            batch_masks = masks_result[0]
-            # Take flattened mask arrays
-            if isinstance(batch_masks, list) and len(batch_masks) > 0 and isinstance(batch_masks[0], list):
-                masks = batch_masks[0]
-            else:
-                # Some shapes: try to flatten
-                masks = []
-                for item in batch_masks:
-                    if isinstance(item, np.ndarray):
-                        masks.append(item)
-        # Scores handling (best effort)
-        try:
-            scores = outputs.iou_scores.cpu().numpy().flatten()
-        except Exception:
-            scores = np.array([0.0] * len(masks))
+        # safe extraction
+        pred_masks = safe_to_cpu_np(getattr(outputs, "pred_masks", None))
+        iou_scores = safe_to_cpu_np(getattr(outputs, "iou_scores", None))
 
-        # Choose mask: prefer middle mask if multiple and not high GPU
-        if len(scores) >= 3:
-            sorted_indices = np.argsort(scores)
-            best_idx = sorted_indices[1] if len(sorted_indices) > 1 else sorted_indices[-1]
-        elif len(scores) >= 1:
-            best_idx = int(np.argmax(scores))
-        else:
-            best_idx = 0
+        # 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))
 
-        selected_mask = masks[best_idx] if isinstance(masks, list) and len(masks) > best_idx else masks[0]
+        # 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)"}
 
-        # If mask is float 0..1, convert to boolean
-        selected_mask = (np.array(selected_mask) > 0.5)
+        # 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"}
 
-        # Filter keep largest component
-        final_mask = selected_mask
-        if filter_largest:
-            num_labels, labels = cv2.connectedComponents(final_mask.astype(np.uint8))
-            if num_labels > 1:
-                largest_component = 0
-                largest_size = 0
-                for i in range(1, num_labels):
-                    size = int(np.sum(labels == i))
-                    if size > largest_size:
-                        largest_size = size
-                        largest_component = i
-                final_mask = (labels == largest_component)
+        # 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)
 
-        if quality == "high":
-            final_mask = refine_mask(final_mask, kernel_size=7)
+        # filter largest component
+        if filter_largest and mask_bool is not None:
+            mask_bool = keep_largest_component_bool(mask_bool)
+
+        # refine edges if requested
+        if quality == "high" and mask_bool is not None:
+            mask_bool = refine_mask(mask_bool, kernel_size=7)
 
-        # create overlay
+        # overlay coloring
         overlay = image_np.copy()
-        color = np.array([255, 80, 180], dtype=np.uint8)
-        mask_float = final_mask.astype(float)
+        if mask_bool is None:
+            return image, {"error": "maske ist None nach Verarbeitung"}
+
+        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)
+            overlay[:,:,c] = (overlay[:,:,c] * (1 - mask_float * 0.65) + color[c] * mask_float * 0.65).astype(np.uint8)
 
-        # contours
-        contours, _ = cv2.findContours(final_mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-        cv2.drawContours(overlay, contours, -1, (255, 255, 0), 3)
+        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",
-            "quality": quality,
-            "filtered": filter_largest,
-            "selected_mask_score": float(scores[best_idx]) if len(scores) > 0 else 0.0,
-            "all_scores": scores.tolist() if hasattr(scores, "tolist") else [],
             "image_size": [w, h],
-            "mask_area": int(np.sum(final_mask)),
-            "mask_percentage": float(np.sum(final_mask) / (h * w) * 100),
+            "mask_area": int(mask_bool.sum()),
+            "mask_percentage": float(mask_bool.sum() / (w*h) * 100),
             "num_contours": len(contours),
-            "device": device,
-            "time_s": round(time.time() - start, 2)
-        }
-
-        return Image.fromarray(overlay.astype(np.uint8)), metadata
-
-    except Exception as e:
-        print("❌ ERROR in segment_automatic:", e)
-        traceback.print_exc()
-        return image, {"error": str(e)}
-
-# ---- Multi-region CPU-optimized (keine Änderung notwendig, nur small improvements) ----
-def segment_multi_dense(image, density="low"):
-    if image is None:
-        return None, {"error": "Kein Bild"}
-
-    start = time.time()
-    try:
-        model, processor = load_model()
-        image_pil, image_np = prepare_image(image)
-        h, w = image_np.shape[:2]
-
-        if density == "high":
-            grid_size = 4
-        elif density == "medium":
-            grid_size = 3
-        else:
-            grid_size = 2
-
-        points = []
-        for i in range(1, grid_size + 1):
-            for j in range(1, grid_size + 1):
-                x = int(w * i / (grid_size + 1))
-                y = int(h * j / (grid_size + 1))
-                points.append([x, y])
-
-        all_masks = []
-        all_scores = []
-
-        # use multimask only on GPU to avoid memory issues
-        use_multimask = True if device == "cuda" else False
-
-        for idx, point in enumerate(points):
-            inputs = processor(
-                image_pil,
-                input_points=[[point]],
-                input_labels=[[1]],
-                return_tensors="pt"
-            ).to(device)
-
-            with torch.no_grad():
-                outputs = model(**inputs, multimask_output=use_multimask)
-
-            masks_result = safe_post_process_masks(processor, outputs, inputs)
-            if masks_result is None:
-                continue
-
-            # flatten masks_result similar to above
-            try:
-                batch_masks = masks_result[0]
-                if isinstance(batch_masks[0], list):
-                    found_masks = batch_masks[0]
-                else:
-                    # if a list of arrays
-                    found_masks = [m for m in batch_masks]
-            except Exception:
-                found_masks = []
-
-            # try to get scores
-            try:
-                scores = outputs.iou_scores.cpu().numpy().flatten()
-            except Exception:
-                scores = np.array([0.0] * len(found_masks))
-
-            for m_idx, mask_arr in enumerate(found_masks):
-                mask_bool = (np.array(mask_arr) > 0.5)
-                best_score = float(scores[m_idx]) if m_idx < len(scores) else 0.0
-                if best_score > 0.5:
-                    all_masks.append(mask_bool)
-                    all_scores.append(best_score)
-
-        print(f"✅ Got {len(all_masks)} quality masks")
-
-        overlay = image_np.copy()
-        colors = []
-        for i in range(len(all_masks)):
-            hue = int(180 * i / max(len(all_masks), 1))
-            color_hsv = np.uint8([[[hue, 255, 200]]])
-            color_rgb = cv2.cvtColor(color_hsv, cv2.COLOR_HSV2RGB)[0][0]
-            colors.append(color_rgb)
-
-        for mask, color, score in zip(all_masks, colors, all_scores):
-            alpha = 0.5
-            overlay[mask] = (
-                overlay[mask] * (1 - alpha) +
-                np.array(color) * alpha
-            ).astype(np.uint8)
-
-            contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-            cv2.drawContours(overlay, contours, -1, color.tolist(), 2)
-
-        metadata = {
-            "success": True,
-            "mode": "multi_object_fast",
-            "density": density,
-            "grid_size": f"{grid_size}x{grid_size}",
-            "total_points": len(points),
-            "quality_masks": len(all_masks),
-            "avg_score": float(np.mean(all_scores)) if all_scores else 0,
-            "scores": [float(s) for s in all_scores],
-            "time_s": round(time.time() - start, 2)
+            "device": device
         }
-
-        return Image.fromarray(overlay), metadata
+        return Image.fromarray(overlay), meta
 
     except Exception as e:
-        print("❌ ERROR in segment_multi_dense:", e)
-        traceback.print_exc()
+        import traceback
+        print(traceback.format_exc())
         return image, {"error": str(e)}
-
-# ---- Gradio UI ----
-demo = gr.Blocks(title="SAM2 Boostly (Stable)", theme=gr.themes.Soft())
-
-with demo:
-    gr.Markdown("# 🎨 SAM2 Segmentierung - Boostly (Stable, CPU-friendly)")
-    gr.Markdown("### ⚡ Smart Filtering & CPU Optimized")
-
-    with gr.Tab("🤖 Automatisch SMART"):
-        gr.Markdown("**Intelligente Filterung** - Nur das Hauptobjekt, keine Jacke/Hintergrund!")
-        with gr.Row():
-            with gr.Column():
-                input_auto = gr.Image(type="pil", label="📸 Bild hochladen")
-                quality_radio = gr.Radio(
-                    choices=["high", "fast"],
-                    value="fast",  # default to fast for CPU
-                    label="⚙️ Qualität",
-                    info="High = glattere Kanten (langsamer)"
-                )
-                filter_checkbox = gr.Checkbox(
-                    value=True,
-                    label="🔍 Nur größtes Objekt",
-                    info="Filtert Jacke, Hintergrund, kleine Teile"
-                )
-                btn_auto = gr.Button("🚀 Segmentieren", variant="primary", size="lg")
-
-                gr.Markdown("""
-                **✨ Smart Features:**
-                - Mittlere Maske (vermeidet Over-Segmentation)
-                - Connected Components Filtering
-                - Entfernt automatisch Jacke/Hintergrund
-                - ⚡ Schnell: empfohlen für CPU = fast / low
-                """)
-
-            with gr.Column():
-                output_auto = gr.Image(label="✨ Segmentiertes Bild")
-                json_auto = gr.JSON(label="📊 Metadata")
-
-        btn_auto.click(
-            fn=segment_automatic,
-            inputs=[input_auto, quality_radio, filter_checkbox],
-            outputs=[output_auto, json_auto]
-        )
-
-    with gr.Tab("🎯 Multi-Region (Schnell)"):
-        gr.Markdown("**CPU-Optimiert** - Reduzierte Punkte für schnellere Verarbeitung")
-        with gr.Row():
-            with gr.Column():
-                input_multi = gr.Image(type="pil", label="📸 Bild hochladen")
-                density_radio = gr.Radio(
-                    choices=["high", "medium", "low"],
-                    value="low",
-                    label="📊 Punkt-Dichte",
-                    info="Start with LOW for tests!"
-                )
-                btn_multi = gr.Button("🎯 Segmentieren", variant="primary", size="lg")
-
-                gr.Markdown("""
-                **⚡ Optimierte Grid-Größen:**
-                - 💨 Low: 2x2 = 4 Punkte (schnell) ← Empfohlen für CPU
-                - ⚡ Medium: 3x3 = 9 Punkte
-                - 🔥 High: 4x4 = 16 Punkte (langsamer)
-                """)
-
-            with gr.Column():
-                output_multi = gr.Image(label="✨ Segmentiertes Bild")
-                json_multi = gr.JSON(label="📊 Metadata")
-
-        btn_multi.click(
-            fn=segment_multi_dense,
-            inputs=[input_multi, density_radio],
-            outputs=[output_multi, json_multi]
-        )
-
-    with gr.Tab("📡 API Documentation"):
-        gr.Markdown("### 🔗 API Endpoint")
-        gr.Code("https://EnginDev-Boostly.hf.space/api/predict", label="Base URL")
-        gr.Markdown("### ⚙️ Notes")
-        gr.Markdown("- Default model: facebook/sam-vit-base (CPU friendly)\n- Use low density for tests\n- If you have GPU, change device to cuda and adjust density")
-
-if __name__ == "__main__":
-    print("🌐 Launching Boostly SAM2 (Stable)...")
-    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)