Update app.py

app.py

@@ -1,78 +1,76 @@
 import gradio as gr
-from PIL import Image, ImageFilter, ImageOps
+from PIL import Image, ImageFilter
 import numpy as np
 import torch
-from transformers import SegformerFeatureExtractor, SegformerForSemanticSegmentation, DPTFeatureExtractor, DPTForDepthEstimation
+from transformers import (
+    SegformerFeatureExtractor, SegformerForSemanticSegmentation,
+    DPTFeatureExtractor, DPTForDepthEstimation
+)
 import cv2
+import os, json
 
-# Load segmentation model
+# ————————————————
+# load segmentation model
 seg_model_name = "nvidia/segformer-b1-finetuned-ade-512-512"
-seg_feature_extractor = SegformerFeatureExtractor.from_pretrained(seg_model_name)
+seg_fe = SegformerFeatureExtractor.from_pretrained(seg_model_name)
 seg_model = SegformerForSemanticSegmentation.from_pretrained(seg_model_name)
 
-# Load depth estimation model
+# load depth model
 depth_model_name = "Intel/dpt-hybrid-midas"
-depth_feature_extractor = DPTFeatureExtractor.from_pretrained(depth_model_name)
+depth_fe = DPTFeatureExtractor.from_pretrained(depth_model_name)
 depth_model = DPTForDepthEstimation.from_pretrained(depth_model_name)
 
-# Device configuration
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 seg_model.to(device)
 depth_model.to(device)
 
-def process_image(image):
-    # Ensure image is in RGB format and resize
-    image = image.resize((512, 512))
-
-    # Perform segmentation
-    inputs = seg_feature_extractor(images=image, return_tensors="pt").to(device)
+def process_image(image: Image.Image):
+    # 1) prep
+    image = image.convert("RGB").resize((512,512))
+    
+    # 2) segmentation → binary mask
+    seg_inputs = seg_fe(images=image, return_tensors="pt").to(device)
     with torch.no_grad():
-        outputs = seg_model(**inputs)
-        logits = outputs.logits
-        segmentation = torch.argmax(logits, dim=1)[0].cpu().numpy()
-    binary_mask = np.where(segmentation > 0, 255, 0).astype(np.uint8)
+        seg_logits = seg_model(**seg_inputs).logits
+    seg_map = torch.argmax(seg_logits, dim=1)[0].cpu().numpy()
+    mask = (seg_map > 0).astype(np.uint8) * 255
+    mask = Image.fromarray(mask).resize((512,512))
 
-    # Apply Gaussian Blur to the background
-    blurred_background = image.filter(ImageFilter.GaussianBlur(15))
-    foreground = Image.fromarray(binary_mask).convert("L").resize(image.size)
-    output_blur = Image.composite(image, blurred_background, foreground)
+    # 3) gaussian-blur background
+    bg_blur = image.filter(ImageFilter.GaussianBlur(15))
+    output_blur = Image.composite(image, bg_blur, mask)
 
-    # Depth estimation for lens blur
-    depth_inputs = depth_feature_extractor(images=image, return_tensors="pt").to(device)
+    # 4) depth estimation
+    depth_inputs = depth_fe(images=image, return_tensors="pt").to(device)
     with torch.no_grad():
-        depth_outputs = depth_model(**depth_inputs)
-        predicted_depth = depth_outputs.predicted_depth.squeeze().cpu().numpy()
-
-    # Normalize depth map
-    depth_min, depth_max = predicted_depth.min(), predicted_depth.max()
-    normalized_depth = (predicted_depth - depth_min) / (depth_max - depth_min)
-    normalized_depth_resized = cv2.resize(normalized_depth, (512, 512))
+        depth_pred = depth_model(**depth_inputs).predicted_depth.squeeze().cpu().numpy()
+    # normalize & resize
+    dmin, dmax = depth_pred.min(), depth_pred.max()
+    depth_norm = (depth_pred - dmin) / (dmax - dmin + 1e-8)
+    depth_norm = cv2.resize(depth_norm, (512,512))
 
-    # Lens blur using depth map
-    blurred_image = np.array(image).astype(np.float32)
-    blur_intensity = normalized_depth_resized * 20
-    for y in range(image.size[1]):
-        for x in range(image.size[0]):
-            sigma = blur_intensity[y, x]
-            kernel_size = int(2 * sigma + 1)
-            if kernel_size > 1:
-                patch = image.crop((x - kernel_size//2, y - kernel_size//2, x + kernel_size//2 + 1, y + kernel_size//2 + 1))
-                patch = patch.filter(ImageFilter.GaussianBlur(sigma))
-                blurred_image[y, x, :] = np.array(patch)[kernel_size//2, kernel_size//2, :]
-    lens_blur_image = Image.fromarray(np.clip(blurred_image, 0, 255).astype(np.uint8))
+    # 5) vectorized depth-based blur
+    img_np = np.array(image).astype(np.float32)
+    # two extremes
+    near_blur = cv2.GaussianBlur(img_np, (21,21), 5)
+    far_blur  = cv2.GaussianBlur(img_np, (81,81), 20)
+    alpha = depth_norm[...,None]
+    combined = near_blur * (1 - alpha) + far_blur * alpha
+    lens_blur = Image.fromarray(np.clip(combined,0,255).astype(np.uint8))
 
-    return image, output_blur, lens_blur_image
+    return image, output_blur, lens_blur
 
 iface = gr.Interface(
     fn=process_image,
-    inputs=gr.Image(type="pil", label="Upload an Image"),
+    inputs=gr.Image(type="pil", label="Upload Image"),
     outputs=[
-        gr.Image(label="Original Image"),
-        gr.Image(label="Gaussian Blur Effect"),
-        gr.Image(label="Depth-Based Lens Blur Effect")
+      gr.Image(type="pil", label="Original"),
+      gr.Image(type="pil", label="Gaussian Blur"),
+      gr.Image(type="pil", label="Depth-Based Lens Blur"),
     ],
-    title="Image Blurring with Gaussian and Depth-Based Lens Blur",
-    description="Upload an image to see Gaussian blur and depth-based lens blur effects."
+    title="Image Blurring with CLAHE + Depth-Based Blur",
+    description="Upload a selfie to see background blur and depth-based lens blur."
 )
 
-iface.launch()
+if __name__ == "__main__":
+    iface.launch(share=True)