Create app.py

app.py

@@ -0,0 +1,212 @@
+from unicodedata import normalize
+
+import gradio as gr
+import numpy as np
+import matplotlib.pyplot as plt
+from PIL import Image, ImageFilter
+from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation
+from scipy.ndimage import gaussian_filter
+import torch
+import requests
+from io import BytesIO
+import cv2
+import warnings
+warnings.filterwarnings('ignore')
+from transformers import DPTImageProcessor, DPTForDepthEstimation,  AutoImageProcessor
+
+model_cache = {
+    "seg_name": None, "seg_proc": None, "seg_model": None,
+    "depth_name": None, "depth_proc": None, "depth_model": None
+}
+
+MODEL_CONFIG = {
+    "segmentation": {
+        "Segformer (B0)": "nvidia/segformer-b0-finetuned-ade-512-512",
+        "Segformer (B5)": "nvidia/segformer-b5-finetuned-ade-640-640",
+    },
+    "depth": {
+        "DPT-Large": "Intel/dpt-large",
+        "Facebook-DPT-Dinov2": "facebook/dpt-dinov2-small-nyu",
+    }
+}
+
+def get_seg_model(model_name):
+    global model_cache
+    repo_id = MODEL_CONFIG["segmentation"][model_name]
+    if model_cache["seg_name"] != model_name:
+        print(f"Switching segmentation model to {model_name}...")
+        model_cache["seg_proc"] = SegformerImageProcessor.from_pretrained(repo_id)
+        model_cache["seg_model"] = SegformerForSemanticSegmentation.from_pretrained(repo_id)
+        model_cache["seg_name"] = model_name
+    return model_cache["seg_proc"], model_cache["seg_model"]
+
+def get_depth_model(model_name):
+    global model_cache
+    repo_id = MODEL_CONFIG["depth"][model_name]
+    if model_cache["depth_name"] != model_name:
+        print(f"Switching depth model to {model_name}...")
+        model_cache["depth_proc"] = DPTImageProcessor.from_pretrained(repo_id)
+        model_cache["depth_model"] = DPTForDepthEstimation.from_pretrained(repo_id)
+        model_cache["depth_name"] = model_name
+    return model_cache["depth_proc"], model_cache["depth_model"]
+
+def preprocess_image(image, target_size=512):
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+    if image.mode != 'RGB':
+        image = image.convert('RGB')
+    return image.resize((target_size, target_size), Image.Resampling.LANCZOS)
+
+def segment_human(image, processor, model):
+    inputs = processor(images=image, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model(**inputs)
+    upsampled = torch.nn.functional.interpolate(
+        outputs.logits, size=(512, 512), mode="bilinear", align_corners=False
+    )
+    pred_seg = upsampled.argmax(dim=1)[0].cpu().numpy()
+    # Note: Label 12 is 'person' in ADE20k dataset
+    return (pred_seg == 12).astype(np.uint8) * 255
+
+def apply_background_blur(image, mask, sigma=15):
+    img_array = np.array(image).astype(np.float32)
+    mask_normalized = mask.astype(np.float32) / 255.0
+    mask_smooth = gaussian_filter(mask_normalized, sigma=2)
+    mask_smooth = np.clip(mask_smooth, 0, 1)
+    
+    blurred_array = np.zeros_like(img_array)
+    for i in range(3):
+        blurred_array[:, :, i] = gaussian_filter(img_array[:, :, i], sigma=sigma)
+    
+    mask_3d = np.stack([mask_smooth] * 3, axis=2)
+    result = (img_array * mask_3d + blurred_array * (1 - mask_3d)).astype(np.uint8)
+    return Image.fromarray(result)
+
+def estimate_depth(image, processor, model, invert):
+    inputs = processor(images=image, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model(**inputs)
+    prediction = torch.nn.functional.interpolate(
+        outputs.predicted_depth.unsqueeze(1), size=(512, 512), mode="bicubic", align_corners=False,
+    )
+    depth_map = prediction.squeeze().cpu().numpy()
+    depth_min, depth_max = depth_map.min(), depth_map.max()
+    normalized = (depth_map - depth_min) / (depth_max - depth_min)
+    if invert == True:
+        normalized = 1.0 - normalized
+    return normalized * 15.0
+
+def apply_lens_blur(image, depth_map, max_sigma=15):
+    img_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR).astype(np.float32)
+    
+    # Create blur pyramid
+    num_levels = 10
+    blur_pyramid = []
+    
+    for i in range(num_levels):
+        sigma = (i / (num_levels - 1)) * max_sigma
+        if sigma < 0.5:
+            blur_pyramid.append(img_cv.copy())
+        else:
+            ksize = int(2 * np.ceil(3 * sigma) + 1)
+            if ksize % 2 == 0:
+                ksize += 1
+            blurred = cv2.GaussianBlur(img_cv, (ksize, ksize), sigma)
+            blur_pyramid.append(blurred)
+    
+    # Apply variable blur based on depth
+    depth_norm = depth_map / 15.0
+    output = np.zeros_like(img_cv)
+    
+    depth_scaled = depth_norm * (num_levels - 1)
+    level_low = np.floor(depth_scaled).astype(np.int32)
+    level_high = np.ceil(depth_scaled).astype(np.int32)
+    level_low = np.clip(level_low, 0, num_levels - 1)
+    level_high = np.clip(level_high, 0, num_levels - 1)
+    
+    weight = depth_scaled - level_low
+    weight = np.expand_dims(weight, axis=2)
+    
+    for y in range(img_cv.shape[0]):
+        for x in range(img_cv.shape[1]):
+            ll = level_low[y, x]
+            lh = level_high[y, x]
+            w = weight[y, x, 0]
+            
+            if ll == lh:
+                output[y, x] = blur_pyramid[ll][y, x]
+            else:
+                output[y, x] = (1 - w) * blur_pyramid[ll][y, x] + w * blur_pyramid[lh][y, x]
+    
+    output = np.clip(output, 0, 255).astype(np.uint8)
+    output_rgb = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
+    return Image.fromarray(output_rgb)
+
+def process_gaussian_blur(image, sigma, model_choice):
+    if image is None: return None, "Upload an image!"
+    try:
+        proc, model = get_seg_model(model_choice)
+        img = preprocess_image(image)
+        mask = segment_human(img, proc, model)
+        result = apply_background_blur(img, mask, sigma)
+        return result, f"Applied {model_choice} with σ={sigma}"
+    except Exception as e:
+        return None, f"Error: {str(e)}"
+
+def process_lens_blur(image, max_sigma, model_choice):
+    if image is None: return None, None, "Upload an image!"
+    try:
+        proc, model = get_depth_model(model_choice)
+        if model_choice == "Facebook-DPT-Dinov2":
+            invert = False
+        else:            
+            invert = True
+        img = preprocess_image(image)
+        depth = estimate_depth(img, proc, model, invert)
+        result = apply_lens_blur(img, depth, max_sigma)
+        
+        depth_vis = cv2.applyColorMap(((depth / 15.0) * 255).astype(np.uint8), cv2.COLORMAP_MAGMA)
+        return result, Image.fromarray(cv2.cvtColor(depth_vis, cv2.COLOR_BGR2RGB)), f"Applied {model_choice}"
+    except Exception as e:
+        return None, None, f"Error: {str(e)}"
+
+with gr.Blocks(title="AI Blur Studio", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# AI Blur Studio\nSelect your AI models and adjust blur intensity.")
+    
+    with gr.Tabs():
+        with gr.Tab("📹 Gaussian Background Blur"):
+            with gr.Row():
+                with gr.Column():
+                    gaussian_input = gr.Image(label="Input Image")
+                    seg_model_dropdown = gr.Dropdown(
+                        choices=list(MODEL_CONFIG["segmentation"].keys()),
+                        value=list(MODEL_CONFIG["segmentation"].keys())[0],
+                        label="Segmentation Model"
+                    )
+                    gaussian_sigma = gr.Slider(0, 30, 15, label="Blur σ")
+                    gaussian_btn = gr.Button("Process", variant="primary")
+                with gr.Column():
+                    gaussian_output = gr.Image(label="Result")
+                    gaussian_status = gr.Textbox(label="Status")
+
+        with gr.Tab("📸 Depth-Based Lens Blur"):
+            with gr.Row():
+                with gr.Column():
+                    lens_input = gr.Image(label="Input Image")
+                    depth_model_dropdown = gr.Dropdown(
+                        choices=list(MODEL_CONFIG["depth"].keys()),
+                        value=list(MODEL_CONFIG["depth"].keys())[0],
+                        label="Depth Estimation Model"
+                    )
+                    lens_sigma = gr.Slider(0, 25, 15, label="Max σ")
+                    lens_btn = gr.Button("Process", variant="primary")
+                with gr.Column():
+                    lens_output = gr.Image(label="Blurred Result")
+                    lens_depth = gr.Image(label="Depth Map")
+                    lens_status = gr.Textbox(label="Status")
+
+    gaussian_btn.click(process_gaussian_blur, [gaussian_input, gaussian_sigma, seg_model_dropdown], [gaussian_output, gaussian_status])
+    lens_btn.click(process_lens_blur, [lens_input, lens_sigma, depth_model_dropdown], [lens_output, lens_depth, lens_status])
+
+if __name__ == "__main__":
+    demo.launch(share=True)