Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import numpy as np
-from PIL import Image
+from PIL import Image, ImageFilter
 from transformers import pipeline
 
 try:
@@ -8,17 +8,110 @@ try:
 except ImportError:
     cv2 = None
 
-# Load depth estimation model once
+# --- Models ---
 depth_pipe = pipeline("depth-estimation", model="depth-anything/Depth-Anything-V2-Large-hf")
 
-# Global state
+# A robust, widely available panoptic segmentation model in transformers
+try:
+    seg_pipe = pipeline("image-segmentation", model="facebook/detr-resnet-50-panoptic")
+except Exception:
+    seg_pipe = None  # We'll handle gracefully
+
+# --- Global state ---
 current_original_image = None
 current_depth_norm = None
 current_depth_map_pil = None
+current_near_softmask = None   # float32 0..1, soft foreground mask
+current_edge_band = None       # float32 0..1, where to snap edges
+
+
+# ---------- Utilities ----------
+def _np_from_pil_gray(img_pil):
+    return np.array(img_pil.convert("L"), dtype=np.float32) / 255.0
+
+def _ensure_float01(x):
+    x = x.astype(np.float32)
+    if x.max() > 1.0 or x.min() < 0.0:
+        x = (x - x.min()) / max(x.max() - x.min(), 1e-6)
+    return x
+
+def _soften_mask(mask01, sigma_px=2.0):
+    """Gaussian blur softening for a binary mask in [0,1]."""
+    if cv2 is not None:
+        return _ensure_float01(cv2.GaussianBlur(mask01, (0, 0), sigmaX=max(sigma_px, 1e-6)))
+    # PIL fallback
+    pil = Image.fromarray((mask01 * 255).astype(np.uint8), mode="L").filter(ImageFilter.GaussianBlur(radius=float(sigma_px)))
+    return np.array(pil, dtype=np.float32) / 255.0
+
+def _edge_band_from_mask(mask01, band_px=6):
+    """Thin band around mask boundary where we enforce snapping."""
+    m = (mask01 > 0.5).astype(np.uint8)
+    if cv2 is not None:
+        k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
+        edge = cv2.morphologyEx(m, cv2.MORPH_GRADIENT, k)  # 1px outline
+        if band_px > 1:
+            edge = cv2.dilate(edge, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (int(band_px), int(band_px))))
+        return edge.astype(np.float32)
+    # PIL fallback: crude edge + dilation
+    pil = Image.fromarray(m * 255, mode="L").filter(ImageFilter.FIND_EDGES)
+    arr = (np.array(pil) > 0).astype(np.uint8)
+    # simple max filter dilation
+    for _ in range(max(int(band_px // 2), 1)):
+        arr = np.maximum.reduce([
+            np.pad(arr, ((1,0),(0,0)))[:-1,:],
+            np.pad(arr, ((0,1),(0,0)))[1:,:],
+            np.pad(arr, ((0,0),(1,0)))[:, :-1],
+            np.pad(arr, ((0,0),(0,1)))[:, 1:],
+            arr
+        ])
+    return arr.astype(np.float32)
+
+def _pick_nearest_segment_mask(segments, depth01, invert_depth=False):
+    """
+    Given pipeline outputs (each with a 'mask' PIL) and normalized depth map,
+    choose the segment with smallest median *distance to camera*.
+    If invert_depth=True, we treat 1-depth as distance.
+    """
+    if not segments:
+        return None
+
+    h, w = depth01.shape
+    best_mask = None
+    best_score = None
+
+    # Interpret "nearness": if invert_depth True, higher values mean near
+    near_map = (1.0 - depth01) if invert_depth else depth01
+    # If your model outputs "larger = farther", set invert_depth=True via UI
+
+    for seg in segments:
+        m = np.array(seg["mask"].resize((w, h)).convert("L")) > 127
+        if m.sum() < 50:  # ignore tiny segments
+            continue
+        # Higher median(nearness) = nearer
+        score = float(np.median(near_map[m]))
+        if (best_score is None) or (score > best_score):
+            best_score = score
+            best_mask = m.astype(np.float32)
+
+    if best_mask is None:
+        return None
+    return _ensure_float01(best_mask)
 
+def _depth_only_near_mask(depth01, near_percent=15, invert_depth=False):
+    """
+    Fallback: use top N% nearest pixels (by depth) as a binary mask.
+    near_percent in [1..40] typical.
+    """
+    flat = depth01.flatten()
+    if invert_depth:
+        flat = 1.0 - flat
+    q = np.quantile(flat, 1.0 - (near_percent / 100.0))
+    m = ((1.0 - depth01) if invert_depth else depth01) >= q
+    return m.astype(np.float32)
 
+
+# ---------- Preprocess depth ----------
 def preprocess_depth(depth_norm, smoothing_radius):
-    """Smooth the depth map using bilateral filtering if radius > 0 and cv2 is available."""
     if smoothing_radius > 0 and cv2 is not None:
         depth_uint8 = (depth_norm * 255.0).astype(np.uint8)
         sigma = max(smoothing_radius * 10.0, 1.0)
@@ -27,44 +120,64 @@ def preprocess_depth(depth_norm, smoothing_radius):
     return depth_norm
 
 
+# ---------- Effect ----------
 def apply_effect(threshold, depth_scale, feather, red_brightness, blue_brightness, gamma,
-                 black_level_percent, white_level_percent, smoothing_percent):
+                 black_level_percent, white_level_percent, smoothing_percent,
+                 use_segmentation, edge_snap_strength, edge_band_px,
+                 invert_depth, near_percent):
     """
-    Apply chromostereopsis effect using adjustable parameters.
-    threshold: percentage [0,100] controlling blend midpoint.
-    depth_scale: percentage [0,100] controlling steepness of logistic curve.
-    feather: percentage [0,100] affecting the smoothness of the transition.
-    red_brightness, blue_brightness: percentages [0,100] controlling channel intensities.
-    gamma: percentage [0,100] mapped to gamma range [0.1, 3.0].
-    black_level_percent, white_level_percent: percentages mapped to 0..255 levels.
-    smoothing_percent: percentage [0,100] mapped to bilateral filter radius.
+    Adds segmentation-assisted edge snapping:
+      - use_segmentation: run / use segmentation-assisted mask if available
+      - edge_snap_strength: 0..100 weight of snapping in edge band
+      - edge_band_px: band width around boundary to force crisp transition
+      - invert_depth: flip near/far interpretation if needed
+      - near_percent: fallback mask when segmentation is off/failed
     """
     global current_original_image, current_depth_norm
+    global current_near_softmask, current_edge_band
+
     if current_original_image is None or current_depth_norm is None:
         return None
 
-    # Levels adjustment
+    # Levels adjustment (same as yours)
     black_level = black_level_percent * 2.55
     white_level = white_level_percent * 2.55
     gray = np.array(current_original_image.convert("L"), dtype=np.float32)
     denom = max(white_level - black_level, 1e-6)
-    adjusted_gray = (gray - black_level) / denom
-    adjusted_gray = np.clip(adjusted_gray, 0.0, 1.0)
+    adjusted_gray = np.clip((gray - black_level) / denom, 0.0, 1.0)
 
-    # Gamma correction
+    # Gamma
     gamma_val = 0.1 + (gamma / 100.0) * 2.9
     adjusted_gray = np.clip(adjusted_gray ** gamma_val, 0.0, 1.0)
 
-    # Smooth depth map
+    # Depth smoothing
     smoothing_radius = smoothing_percent / 10.0
     depth_smoothed = preprocess_depth(current_depth_norm, smoothing_radius)
+    depth_for_blend = (1.0 - depth_smoothed) if invert_depth else depth_smoothed
 
-    # Compute blend factor using logistic function
+    # Logistic blend from depth
     threshold_norm = threshold / 100.0
     steepness = max(depth_scale, 1e-3)
     feather_norm = feather / 100.0
     steepness_adj = steepness / (feather_norm * 10.0 + 1.0)
-    blend = 1.0 / (1.0 + np.exp(-steepness_adj * (depth_smoothed - threshold_norm)))
+    blend = 1.0 / (1.0 + np.exp(-steepness_adj * (depth_for_blend - threshold_norm)))
+
+    # Edge snapping: mix blend with (soft) near-object mask ONLY within edge band
+    snap_w = np.clip(edge_snap_strength / 100.0, 0.0, 1.0)
+
+    if current_near_softmask is None:
+        # Build a fallback mask so snapping can still help
+        fallback = _depth_only_near_mask(current_depth_norm, near_percent=near_percent, invert_depth=invert_depth)
+        current_near_softmask_local = _soften_mask(fallback, sigma_px=max(edge_band_px / 2.0, 1.0))
+        current_edge_band_local = _edge_band_from_mask(fallback, band_px=max(int(edge_band_px), 1))
+    else:
+        current_near_softmask_local = current_near_softmask
+        current_edge_band_local = current_edge_band
+
+    if snap_w > 0.0:
+        # Per-pixel alpha only near edges
+        per_pixel_alpha = snap_w * _ensure_float01(current_edge_band_local)
+        blend = (1.0 - per_pixel_alpha) * blend + per_pixel_alpha * _ensure_float01(current_near_softmask_local)
 
     # Map brightness to factors (0-2)
     red_factor = red_brightness / 50.0
@@ -81,11 +194,43 @@ def apply_effect(threshold, depth_scale, feather, red_brightness, blue_brightnes
     output[..., 0] = red_img
     output[..., 1] = 0
     output[..., 2] = blue_img
-
     return Image.fromarray(output, mode="RGB")
 
 
-def generate_depth_map(input_image):
+# ---------- Pipeline steps ----------
+def _compute_segmentation_assist(img_pil, depth01, invert_depth, edge_band_px, near_percent):
+    """
+    Build current_near_softmask and current_edge_band using segmentation if enabled;
+    else fallback to depth-only near mask.
+    """
+    global current_near_softmask, current_edge_band
+
+    h, w = depth01.shape
+    near_mask = None
+
+    if seg_pipe is not None:
+        try:
+            segs = seg_pipe(img_pil)
+            # Some models return dict with 'segments'; normalize to list
+            if isinstance(segs, dict) and "segments_info" in segs and "segmentation" in segs:
+                # Panoptic map; fall back to simple depth-based mask
+                segs = []  # transformers panoptic map formats vary; keep generic path below
+            # segs should be a list of dicts, each with a 'mask' PIL
+            candidates = [s for s in segs if isinstance(s.get("mask", None), Image.Image)]
+            near_mask = _pick_nearest_segment_mask(candidates, depth01, invert_depth=invert_depth)
+        except Exception:
+            near_mask = None
+
+    if near_mask is None:
+        # Fallback: depth-only near mask
+        near_mask = _depth_only_near_mask(depth01, near_percent=near_percent, invert_depth=invert_depth)
+
+    # Build soft mask + edge band
+    current_near_softmask = _soften_mask(near_mask, sigma_px=max(edge_band_px / 2.0, 1.0))
+    current_edge_band = _edge_band_from_mask(near_mask, band_px=max(int(edge_band_px), 1))
+
+
+def generate_depth_map(input_image, use_segmentation, edge_band_px, invert_depth, near_percent):
     """Generate normalized depth map and initial effect image."""
     global current_original_image, current_depth_norm, current_depth_map_pil
     if input_image is None:
@@ -93,8 +238,9 @@ def generate_depth_map(input_image):
         current_depth_norm = None
         current_depth_map_pil = None
         return None, None
+
     current_original_image = input_image
-    # Run depth estimation
+    # Depth estimation
     result = depth_pipe(input_image)
     depth = np.array(result["depth"], dtype=np.float32)
     depth -= depth.min()
@@ -103,6 +249,11 @@ def generate_depth_map(input_image):
         depth /= max_val
     current_depth_norm = depth
     current_depth_map_pil = Image.fromarray((depth * 255.0).astype(np.uint8), mode="L")
+
+    # Build segmentation assist (or fallback)
+    if use_segmentation or True:  # we compute once so UI changes can immediately work
+        _compute_segmentation_assist(input_image, current_depth_norm, invert_depth, edge_band_px, near_percent)
+
     # Default effect parameters
     effect = apply_effect(
         threshold=50,
@@ -114,13 +265,20 @@ def generate_depth_map(input_image):
         black_level_percent=0,
         white_level_percent=100,
         smoothing_percent=0,
+        use_segmentation=use_segmentation,
+        edge_snap_strength=60,
+        edge_band_px=edge_band_px,
+        invert_depth=invert_depth,
+        near_percent=near_percent,
     )
     return current_depth_map_pil.convert("RGB"), effect
 
 
 def update_effect(threshold, depth_scale, feather, red_brightness, blue_brightness,
-                  gamma, black_level, white_level, smoothing):
-    """Update the effect when any slider changes."""
+                  gamma, black_level, white_level, smoothing,
+                  use_segmentation, edge_snap_strength, edge_band_px,
+                  invert_depth, near_percent):
+    """Update the effect when any control changes."""
     return apply_effect(
         threshold=threshold,
         depth_scale=depth_scale,
@@ -131,54 +289,81 @@ def update_effect(threshold, depth_scale, feather, red_brightness, blue_brightne
         black_level_percent=black_level,
         white_level_percent=white_level,
         smoothing_percent=smoothing,
+        use_segmentation=use_segmentation,
+        edge_snap_strength=edge_snap_strength,
+        edge_band_px=edge_band_px,
+        invert_depth=invert_depth,
+        near_percent=near_percent,
     )
 
 
 def clear_results():
     """Reset global state and clear outputs."""
     global current_original_image, current_depth_norm, current_depth_map_pil
+    global current_near_softmask, current_edge_band
     current_original_image = None
     current_depth_norm = None
     current_depth_map_pil = None
+    current_near_softmask = None
+    current_edge_band = None
     return None, None
 
 
+# ---------- UI ----------
 with gr.Blocks(title="ChromoStereoizer Enhanced", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# ChromoStereoizer Enhanced")
+    gr.Markdown("# ChromoStereoizer Enhanced (Segmentation-Assisted)")
     with gr.Row():
         with gr.Column(scale=1):
             input_image = gr.Image(label="Upload Image", type="pil", height=400)
+            with gr.Accordion("Segmentation & Mask Options", open=False):
+                use_segmentation = gr.Checkbox(value=True, label="Use segmentation-assisted edge snapping (falls back to depth-only)")
+                edge_band_px = gr.Slider(1, 20, value=6, step=1, label="Edge Band Width (px)")
+                invert_depth = gr.Checkbox(value=False, label="Invert depth (toggle if near/far feels flipped)")
+                near_percent = gr.Slider(1, 40, value=15, step=1, label="Fallback: top N% nearest pixels")
             generate_btn = gr.Button("Generate Depth Map", variant="primary", size="lg")
+
         with gr.Column(scale=1):
             gr.Markdown("**Depth Map**")
             depth_output = gr.Image(type="pil", height=400, interactive=False, show_download_button=True, show_label=False)
             gr.Markdown("**ChromoStereoizer Result**")
             chromo_output = gr.Image(type="pil", height=400, interactive=False, show_download_button=True, show_label=False)
             gr.Markdown("## Controls")
-            threshold_slider = gr.Slider(minimum=0, maximum=100, value=50, step=1, label="Threshold (%)")
-            depth_scale_slider = gr.Slider(minimum=0, maximum=100, value=50, step=1, label="Depth Scale (Steepness)")
-            feather_slider = gr.Slider(minimum=0, maximum=100, value=10, step=1, label="Feather (%)")
-            red_slider = gr.Slider(minimum=0, maximum=100, value=50, step=1, label="Red Brightness")
-            blue_slider = gr.Slider(minimum=0, maximum=100, value=50, step=1, label="Blue Brightness")
-            gamma_slider = gr.Slider(minimum=0, maximum=100, value=50, step=1, label="Gamma")
-            black_slider = gr.Slider(minimum=0, maximum=100, value=0, step=1, label="Black Level (%)")
-            white_slider = gr.Slider(minimum=0, maximum=100, value=100, step=1, label="White Level (%)")
-            smoothing_slider = gr.Slider(minimum=0, maximum=100, value=0, step=1, label="Smoothing (%)")
+            threshold_slider = gr.Slider(0, 100, value=50, step=1, label="Threshold (%)")
+            depth_scale_slider = gr.Slider(0, 100, value=50, step=1, label="Depth Scale (Steepness)")
+            feather_slider = gr.Slider(0, 100, value=10, step=1, label="Feather (%)")
+            red_slider = gr.Slider(0, 100, value=50, step=1, label="Red Brightness")
+            blue_slider = gr.Slider(0, 100, value=50, step=1, label="Blue Brightness")
+            gamma_slider = gr.Slider(0, 100, value=50, step=1, label="Gamma")
+            black_slider = gr.Slider(0, 100, value=0, step=1, label="Black Level (%)")
+            white_slider = gr.Slider(0, 100, value=100, step=1, label="White Level (%)")
+            smoothing_slider = gr.Slider(0, 100, value=0, step=1, label="Depth Smoothing (%)")
+            edge_snap_strength = gr.Slider(0, 100, value=60, step=1, label="Edge Snap Strength (%)")
             clear_btn = gr.Button("Clear", variant="secondary")
-    # Event bindings
+
+    # Events
     generate_btn.click(
         fn=generate_depth_map,
-        inputs=[input_image],
+        inputs=[input_image, use_segmentation, edge_band_px, invert_depth, near_percent],
         outputs=[depth_output, chromo_output],
         show_progress=True,
     )
-    for slider in [threshold_slider, depth_scale_slider, feather_slider, red_slider, blue_slider, gamma_slider, black_slider, white_slider, smoothing_slider]:
-        slider.change(
+
+    for ctrl in [
+        threshold_slider, depth_scale_slider, feather_slider, red_slider, blue_slider,
+        gamma_slider, black_slider, white_slider, smoothing_slider,
+        use_segmentation, edge_snap_strength, edge_band_px, invert_depth, near_percent
+    ]:
+        ctrl.change(
             fn=update_effect,
-            inputs=[threshold_slider, depth_scale_slider, feather_slider, red_slider, blue_slider, gamma_slider, black_slider, white_slider, smoothing_slider],
+            inputs=[
+                threshold_slider, depth_scale_slider, feather_slider, red_slider, blue_slider,
+                gamma_slider, black_slider, white_slider, smoothing_slider,
+                use_segmentation, edge_snap_strength, edge_band_px, invert_depth, near_percent
+            ],
             outputs=chromo_output,
             show_progress=False,
         )
+
     clear_btn.click(
         fn=clear_results,
         inputs=[],