Update app.py

app.py

@@ -1,97 +1,80 @@
-
 import os
-# Ensure Hugging Face runs in PyTorch mode
-os.environ["TRANSFORMERS_NO_TF"] = "1"
+os.environ["TRANSFORMERS_NO_TF"] = "1"  # force PyTorch-only pipelines
 
 import numpy as np
 from PIL import Image, ImageFilter
 import gradio as gr
-
-# Torch is imported lazily so the Space can boot even if torch takes time to install
 import torch
 from transformers import pipeline
 
-# Device selection: GPU if available; otherwise CPU (-1)
+# ---- Config ----
 DEVICE = 0 if torch.cuda.is_available() else -1
+SEG_MODEL = "nvidia/segformer-b0-finetuned-ade-512-512"
+DEPTH_MODEL = "Intel/dpt-hybrid-midas"
+SIZE = (512, 512)
 
-# Model identifiers
-SEG_MODEL_NAME = "nvidia/segformer-b0-finetuned-ade-512-512"
-DEPTH_MODEL_NAME = "Intel/dpt-hybrid-midas"
-TARGET_SIZE = (512, 512)
-
-# Pipelines: loaded once on process start
-seg_pipe = pipeline("image-segmentation", model=SEG_MODEL_NAME, device=DEVICE, framework="pt")
-depth_pipe = pipeline("depth-estimation", model=DEPTH_MODEL_NAME, device=DEVICE, framework="pt")
+# ---- Pipelines (loaded once) ----
+seg_pipe = pipeline("image-segmentation", model=SEG_MODEL, device=DEVICE, framework="pt")
+depth_pipe = pipeline("depth-estimation", model=DEPTH_MODEL, device=DEVICE, framework="pt")
 
-def resize_center_crop(img: Image.Image, size=(512, 512)) -> Image.Image:
-    """Resize with aspect ratio, center-crop to size."""
+# ---- Helpers ----
+def resize_center_crop(img: Image.Image, size=SIZE) -> Image.Image:
     img = img.convert("RGB")
     w, h = img.size
     tw, th = size
-    scale = max(tw / w, th / h)
-    nw, nh = int(round(w * scale)), int(round(h * scale))
+    s = max(tw / w, th / h)
+    nw, nh = int(round(w * s)), int(round(h * s))
     img = img.resize((nw, nh), Image.BICUBIC)
     left, top = (nw - tw) // 2, (nh - th) // 2
     return img.crop((left, top, left + tw, top + th))
 
 def person_mask(img_512: Image.Image) -> Image.Image:
-    """Return a binary mask (L mode) with person=255, background=0; black if no person."""
     results = seg_pipe(img_512)
-    person = None
-    for r in results:
-        if r.get("label", "").lower() == "person":
-            person = r; break
+    person = next((r for r in results if r.get("label", "").lower() == "person"), None)
     if person is None:
-        for r in results:
-            if "person" in r.get("label", "").lower():
-                person = r; break
+        person = next((r for r in results if "person" in r.get("label", "").lower()), None)
     if person is None:
-        return Image.new("L", img_512.size, 0)
+        return Image.new("L", img_512.size, 0)  # no person detected
     m = person["mask"].convert("L")
     m = (np.array(m) > 127).astype(np.uint8) * 255
     return Image.fromarray(m, mode="L")
 
 def gaussian_bg_blur(img_512: Image.Image, sigma: int = 15) -> Image.Image:
-    """Blur only the background (person stays sharp)."""
     m = person_mask(img_512)
     blurred = img_512.filter(ImageFilter.GaussianBlur(radius=int(sigma)))
-    return Image.composite(img_512, blurred, m)
+    return Image.composite(img_512, blurred, m)  # white=person -> keep sharp
 
 def depth_lens_blur(img_512: Image.Image, max_radius: int = 15, keep_subject: bool = True) -> Image.Image:
-    """Apply per-pixel blur proportional to depth (farther = more blur)."""
     out = depth_pipe(img_512)
-    d = out["depth"].resize(TARGET_SIZE, Image.BICUBIC)
+    d = out["depth"].resize(SIZE, Image.BICUBIC)
     dnp = np.array(d).astype(np.float32)
-    d01 = (dnp - dnp.min()) / (dnp.max() - dnp.min() + 1e-8)   # 0..1 (brighter≈closer)
-    far = 1.0 - d01                                            # larger = farther
+    d01 = (dnp - dnp.min()) / (dnp.max() - dnp.min() + 1e-8)  # 0..1
+    far = 1.0 - d01  # larger=farther -> more blur
     if keep_subject:
         m = person_mask(img_512)
         m01 = (np.array(m) > 127).astype(np.float32)
-        far = far * (1.0 - 0.85 * m01)                         # suppress blur on subject
-    # Build blurred pyramid and gather
+        far = far * (1.0 - 0.85 * m01)  # suppress blur on detected subject
     max_radius = int(max(0, min(30, max_radius)))
-    radii = np.arange(max_radius + 1, dtype=np.int32)
     idx = np.clip(np.rint(far * max_radius).astype(np.int32), 0, max_radius)
-    stack = []
-    for r in radii:
-        stack.append(img_512 if r == 0 else img_512.filter(ImageFilter.GaussianBlur(radius=int(r))))
+    stack = [img_512 if r == 0 else img_512.filter(ImageFilter.GaussianBlur(radius=r))
+             for r in range(max_radius + 1)]
     stack_np = np.stack([np.array(im) for im in stack], axis=0)  # [R+1,H,W,3]
     H, W = idx.shape
-    h = np.arange(H)[:, None]
-    w = np.arange(W)[None, :]
+    h = np.arange(H)[:, None]; w = np.arange(W)[None, :]
     out_np = stack_np[idx, h, w]
     return Image.fromarray(out_np.astype(np.uint8))
 
-def run(image: Image.Image, effect: str, sigma: int, max_radius: int, keep_subject: bool):
+def run(image, effect, sigma, max_radius, keep_subject):
     if image is None:
         return None, None
-    img_512 = resize_center_crop(image, TARGET_SIZE)
+    img_512 = resize_center_crop(image, SIZE)
     if effect == "Gaussian Background Blur (subject sharp)":
         out = gaussian_bg_blur(img_512, sigma=int(sigma))
     else:
-        out = depth_lens_blur(img_512, max_radius=int(max_radius), keep_subject=keep_subject)
+        out = depth_lens_blur(img_512, max_radius=int(max_radius), keep_subject=bool(keep_subject))
     return img_512, out
 
+# ---- UI ----
 with gr.Blocks(title="Gaussian & Lens Blur Lab") as demo:
     gr.Markdown("# Gaussian & Lens Blur Lab\nUpload an image and compare effects.")
     with gr.Row():
@@ -99,8 +82,7 @@ with gr.Blocks(title="Gaussian & Lens Blur Lab") as demo:
             in_img = gr.Image(type="pil", label="Upload image")
             effect = gr.Radio(
                 ["Gaussian Background Blur (subject sharp)", "Depth-based Lens Blur"],
-                value="Gaussian Background Blur (subject sharp)",
-                label="Effect"
+                value="Gaussian Background Blur (subject sharp)", label="Effect"
             )
             sigma = gr.Slider(1, 40, value=15, step=1, label="Gaussian sigma")
             max_r = gr.Slider(4, 30, value=15, step=1, label="Max blur radius (lens blur)")