yrosenbloom commited on
Commit
8f8907e
·
verified ·
1 Parent(s): 6e7a7c8

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +82 -0
app.py CHANGED
@@ -0,0 +1,82 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ import numpy as np
4
+ from PIL import Image, ImageFilter, ImageOps
5
+ import cv2
6
+ from transformers import (
7
+ SegformerFeatureExtractor, SegformerForSemanticSegmentation,
8
+ DPTFeatureExtractor, DPTForDepthEstimation
9
+ )
10
+
11
+ # Load models
12
+ seg_model_name = "nvidia/segformer-b1-finetuned-ade-512-512"
13
+ depth_model_name = "Intel/dpt-hybrid-midas"
14
+
15
+ seg_extractor = SegformerFeatureExtractor.from_pretrained(seg_model_name)
16
+ seg_model = SegformerForSemanticSegmentation.from_pretrained(seg_model_name)
17
+ depth_extractor = DPTFeatureExtractor.from_pretrained(depth_model_name)
18
+ depth_model = DPTForDepthEstimation.from_pretrained(depth_model_name)
19
+
20
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
21
+ seg_model.to(device)
22
+ depth_model.to(device)
23
+
24
+ def process_image(image_pil):
25
+ image = ImageOps.exif_transpose(image_pil).resize((512, 512)).convert("RGB")
26
+
27
+ # ---------- Part 1: Segmentation ----------
28
+ seg_inputs = seg_extractor(images=image, return_tensors="pt").to(device)
29
+ with torch.no_grad():
30
+ seg_output = seg_model(**seg_inputs).logits
31
+ seg_mask = torch.argmax(seg_output, dim=1)[0].cpu().numpy()
32
+ binary_mask = np.where(seg_mask > 0, 255, 0).astype(np.uint8)
33
+ foreground_mask = Image.fromarray(binary_mask).convert("L")
34
+
35
+ # ---------- Part 2: Gaussian blur to background ----------
36
+ blurred_background = image.filter(ImageFilter.GaussianBlur(15))
37
+ blurred_background = blurred_background.convert("RGBA")
38
+ image_rgba = image.convert("RGBA")
39
+ output_blur = Image.composite(image_rgba, blurred_background, foreground_mask)
40
+
41
+ # ---------- Part 3: Depth Estimation ----------
42
+ image_np = np.array(image)
43
+ depth_inputs = depth_extractor(images=image_np, return_tensors="pt").to(device)
44
+ with torch.no_grad():
45
+ depth_output = depth_model(**depth_inputs)
46
+ predicted_depth = depth_output.predicted_depth.squeeze().cpu().numpy()
47
+ normalized_depth = (predicted_depth - predicted_depth.min()) / (predicted_depth.max() - predicted_depth.min())
48
+
49
+ # ---------- Part 4: Depth-Based Variable Gaussian Blur ----------
50
+ image_np_float = image_np.astype(np.float32)
51
+ resized_depth = cv2.resize(normalized_depth, (image_np.shape[1], image_np.shape[0]))
52
+ inverted_depth = 1.0 - resized_depth
53
+ total_blur_levels = 4
54
+ blurred_versions = []
55
+ for i in range(total_blur_levels):
56
+ sigma = i * 3
57
+ blurred = cv2.GaussianBlur(image_np_float, (15, 15), sigmaX=sigma, sigmaY=sigma) if sigma > 0 else image_np_float.copy()
58
+ blurred_versions.append(blurred)
59
+
60
+ blur_indices = (inverted_depth * (total_blur_levels - 1)).astype(np.uint8)
61
+ final_blurred_np = np.zeros_like(image_np_float)
62
+ for i in range(total_blur_levels):
63
+ mask = (blur_indices == i)
64
+ for c in range(3):
65
+ final_blurred_np[:, :, c][mask] = blurred_versions[i][:, :, c][mask]
66
+ depth_blur_img = Image.fromarray(np.clip(final_blurred_np, 0, 255).astype(np.uint8))
67
+
68
+ return image, output_blur.convert("RGB"), depth_blur_img
69
+
70
+ # Gradio Interface
71
+ gr.Interface(
72
+ fn=process_image,
73
+ inputs=gr.Image(type="pil"),
74
+ outputs=[
75
+ gr.Image(label="Original Image"),
76
+ gr.Image(label="Segmented Gaussian Blur"),
77
+ gr.Image(label="Depth-Based Lens Blur")
78
+ ],
79
+ title="Visual Effects Demo: Segmentation & Depth-Based Blur",
80
+ description="Upload an image to see it segmented with background blur (like Zoom) and depth-based lens blur.",
81
+ examples=[],
82
+ ).launch()