Update app.py

app.py

@@ -0,0 +1,82 @@
+import gradio as gr
+import torch
+import numpy as np
+from PIL import Image, ImageFilter, ImageOps
+import cv2
+from transformers import (
+    SegformerFeatureExtractor, SegformerForSemanticSegmentation,
+    DPTFeatureExtractor, DPTForDepthEstimation
+)
+
+# Load models
+seg_model_name = "nvidia/segformer-b1-finetuned-ade-512-512"
+depth_model_name = "Intel/dpt-hybrid-midas"
+
+seg_extractor = SegformerFeatureExtractor.from_pretrained(seg_model_name)
+seg_model = SegformerForSemanticSegmentation.from_pretrained(seg_model_name)
+depth_extractor = DPTFeatureExtractor.from_pretrained(depth_model_name)
+depth_model = DPTForDepthEstimation.from_pretrained(depth_model_name)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+seg_model.to(device)
+depth_model.to(device)
+
+def process_image(image_pil):
+    image = ImageOps.exif_transpose(image_pil).resize((512, 512)).convert("RGB")
+
+    # ---------- Part 1: Segmentation ----------
+    seg_inputs = seg_extractor(images=image, return_tensors="pt").to(device)
+    with torch.no_grad():
+        seg_output = seg_model(**seg_inputs).logits
+    seg_mask = torch.argmax(seg_output, dim=1)[0].cpu().numpy()
+    binary_mask = np.where(seg_mask > 0, 255, 0).astype(np.uint8)
+    foreground_mask = Image.fromarray(binary_mask).convert("L")
+
+    # ---------- Part 2: Gaussian blur to background ----------
+    blurred_background = image.filter(ImageFilter.GaussianBlur(15))
+    blurred_background = blurred_background.convert("RGBA")
+    image_rgba = image.convert("RGBA")
+    output_blur = Image.composite(image_rgba, blurred_background, foreground_mask)
+
+    # ---------- Part 3: Depth Estimation ----------
+    image_np = np.array(image)
+    depth_inputs = depth_extractor(images=image_np, return_tensors="pt").to(device)
+    with torch.no_grad():
+        depth_output = depth_model(**depth_inputs)
+    predicted_depth = depth_output.predicted_depth.squeeze().cpu().numpy()
+    normalized_depth = (predicted_depth - predicted_depth.min()) / (predicted_depth.max() - predicted_depth.min())
+
+    # ---------- Part 4: Depth-Based Variable Gaussian Blur ----------
+    image_np_float = image_np.astype(np.float32)
+    resized_depth = cv2.resize(normalized_depth, (image_np.shape[1], image_np.shape[0]))
+    inverted_depth = 1.0 - resized_depth
+    total_blur_levels = 4
+    blurred_versions = []
+    for i in range(total_blur_levels):
+        sigma = i * 3
+        blurred = cv2.GaussianBlur(image_np_float, (15, 15), sigmaX=sigma, sigmaY=sigma) if sigma > 0 else image_np_float.copy()
+        blurred_versions.append(blurred)
+
+    blur_indices = (inverted_depth * (total_blur_levels - 1)).astype(np.uint8)
+    final_blurred_np = np.zeros_like(image_np_float)
+    for i in range(total_blur_levels):
+        mask = (blur_indices == i)
+        for c in range(3):
+            final_blurred_np[:, :, c][mask] = blurred_versions[i][:, :, c][mask]
+    depth_blur_img = Image.fromarray(np.clip(final_blurred_np, 0, 255).astype(np.uint8))
+
+    return image, output_blur.convert("RGB"), depth_blur_img
+
+# Gradio Interface
+gr.Interface(
+    fn=process_image,
+    inputs=gr.Image(type="pil"),
+    outputs=[
+        gr.Image(label="Original Image"),
+        gr.Image(label="Segmented Gaussian Blur"),
+        gr.Image(label="Depth-Based Lens Blur")
+    ],
+    title="Visual Effects Demo: Segmentation & Depth-Based Blur",
+    description="Upload an image to see it segmented with background blur (like Zoom) and depth-based lens blur.",
+    examples=[],
+).launch()