yrosenbloom commited on
Commit
c752dca
·
verified ·
1 Parent(s): 490a2c8

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +54 -58
app.py CHANGED
@@ -1,82 +1,78 @@
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
- # ---- Segmentation ----
28
- seg_inputs = seg_extractor(images=image, return_tensors="pt", do_resize=True, do_normalize=True)
 
29
  with torch.no_grad():
30
- seg_output = seg_model(**seg_inputs.to(device)).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
- # ---- Blur Background ----
36
- image_rgba = image.convert("RGBA")
37
- blurred = image.filter(ImageFilter.GaussianBlur(15)).convert("RGBA")
38
- composite_blur = Image.composite(image_rgba, blurred, foreground_mask)
39
 
40
- # ---- Depth ----
41
- image_np = np.array(image)
42
- depth_inputs = depth_extractor(images=image_np, return_tensors="pt")
43
  with torch.no_grad():
44
- depth_output = depth_model(**depth_inputs.to(device))
45
- predicted_depth = depth_output.predicted_depth.squeeze().cpu().numpy()
46
- normalized_depth = (predicted_depth - predicted_depth.min()) / (predicted_depth.max() - predicted_depth.min())
47
 
48
- # ---- Depth-Based Blur ----
49
- image_np = np.array(image).astype(np.float32)
50
- resized_depth = cv2.resize(normalized_depth, (image_np.shape[1], image_np.shape[0]))
51
- inverted_depth = 1.0 - resized_depth
52
- blur_levels = 4
53
- blurred_variants = []
54
- for i in range(blur_levels):
55
- sigma = i * 3
56
- blurred = cv2.GaussianBlur(image_np, (15, 15), sigmaX=sigma, sigmaY=sigma) if sigma > 0 else image_np.copy()
57
- blurred_variants.append(blurred)
58
 
59
- blur_indices = (inverted_depth * (blur_levels - 1)).astype(np.uint8)
60
- final_blur = np.zeros_like(image_np)
61
- for i in range(blur_levels):
62
- mask = (blur_indices == i)
63
- for c in range(3):
64
- final_blur[:, :, c][mask] = blurred_variants[i][:, :, c][mask]
65
- lens_blur_pil = Image.fromarray(np.clip(final_blur, 0, 255).astype(np.uint8))
 
 
 
 
 
66
 
67
- return image, composite_blur.convert("RGB"), lens_blur_pil
68
 
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()
 
 
1
  import gradio as gr
 
 
2
  from PIL import Image, ImageFilter, ImageOps
3
+ import numpy as np
4
+ import torch
5
+ from transformers import SegformerFeatureExtractor, SegformerForSemanticSegmentation, DPTFeatureExtractor, DPTForDepthEstimation
6
  import cv2
 
 
 
 
7
 
8
+ # Load segmentation model
9
  seg_model_name = "nvidia/segformer-b1-finetuned-ade-512-512"
10
+ seg_feature_extractor = SegformerFeatureExtractor.from_pretrained(seg_model_name)
 
 
11
  seg_model = SegformerForSemanticSegmentation.from_pretrained(seg_model_name)
12
+
13
+ # Load depth estimation model
14
+ depth_model_name = "Intel/dpt-hybrid-midas"
15
+ depth_feature_extractor = DPTFeatureExtractor.from_pretrained(depth_model_name)
16
  depth_model = DPTForDepthEstimation.from_pretrained(depth_model_name)
17
 
18
+ # Device configuration
19
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
20
  seg_model.to(device)
21
  depth_model.to(device)
22
 
23
+ def process_image(image):
24
+ # Ensure image is in RGB format and resize
25
+ image = ImageOps.exif_transpose(image).resize((512, 512))
26
+
27
+ # Perform segmentation
28
+ inputs = seg_feature_extractor(images=image, return_tensors="pt").to(device)
29
  with torch.no_grad():
30
+ outputs = seg_model(**inputs)
31
+ logits = outputs.logits
32
+ segmentation = torch.argmax(logits, dim=1)[0].cpu().numpy()
33
+ binary_mask = np.where(segmentation > 0, 255, 0).astype(np.uint8)
34
 
35
+ # Apply Gaussian Blur to the background
36
+ blurred_background = image.filter(ImageFilter.GaussianBlur(15))
37
+ foreground = Image.fromarray(binary_mask).convert("L").resize(image.size)
38
+ output_blur = Image.composite(image, blurred_background, foreground)
39
 
40
+ # Depth estimation for lens blur
41
+ depth_inputs = depth_feature_extractor(images=image, return_tensors="pt").to(device)
 
42
  with torch.no_grad():
43
+ depth_outputs = depth_model(**depth_inputs)
44
+ predicted_depth = depth_outputs.predicted_depth.squeeze().cpu().numpy()
 
45
 
46
+ # Normalize depth map
47
+ depth_min, depth_max = predicted_depth.min(), predicted_depth.max()
48
+ normalized_depth = (predicted_depth - depth_min) / (depth_max - depth_min)
49
+ normalized_depth_resized = cv2.resize(normalized_depth, (512, 512))
 
 
 
 
 
 
50
 
51
+ # Lens blur using depth map
52
+ blurred_image = np.array(image).astype(np.float32)
53
+ blur_intensity = normalized_depth_resized * 20
54
+ for y in range(image.size[1]):
55
+ for x in range(image.size[0]):
56
+ sigma = blur_intensity[y, x]
57
+ kernel_size = int(2 * sigma + 1)
58
+ if kernel_size > 1:
59
+ patch = image.crop((x - kernel_size//2, y - kernel_size//2, x + kernel_size//2 + 1, y + kernel_size//2 + 1))
60
+ patch = patch.filter(ImageFilter.GaussianBlur(sigma))
61
+ blurred_image[y, x, :] = np.array(patch)[kernel_size//2, kernel_size//2, :]
62
+ lens_blur_image = Image.fromarray(np.clip(blurred_image, 0, 255).astype(np.uint8))
63
 
64
+ return image, output_blur, lens_blur_image
65
 
66
+ iface = gr.Interface(
 
 
67
  fn=process_image,
68
+ inputs=gr.Image(type="pil", label="Upload an Image"),
69
  outputs=[
70
  gr.Image(label="Original Image"),
71
+ gr.Image(label="Gaussian Blur Effect"),
72
+ gr.Image(label="Depth-Based Lens Blur Effect")
73
  ],
74
+ title="Image Blurring with Gaussian and Depth-Based Lens Blur",
75
+ description="Upload an image to see Gaussian blur and depth-based lens blur effects."
76
+ )
77
+
78
+ iface.launch()