Update app.py
Browse files
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
|
12 |
seg_model_name = "nvidia/segformer-b1-finetuned-ade-512-512"
|
13 |
-
|
14 |
-
|
15 |
-
seg_extractor = SegformerFeatureExtractor.from_pretrained(seg_model_name)
|
16 |
seg_model = SegformerForSemanticSegmentation.from_pretrained(seg_model_name)
|
17 |
-
|
|
|
|
|
|
|
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(
|
25 |
-
image
|
26 |
-
|
27 |
-
|
28 |
-
|
|
|
29 |
with torch.no_grad():
|
30 |
-
|
31 |
-
|
32 |
-
|
33 |
-
|
34 |
|
35 |
-
#
|
36 |
-
|
37 |
-
|
38 |
-
|
39 |
|
40 |
-
#
|
41 |
-
|
42 |
-
depth_inputs = depth_extractor(images=image_np, return_tensors="pt")
|
43 |
with torch.no_grad():
|
44 |
-
|
45 |
-
|
46 |
-
normalized_depth = (predicted_depth - predicted_depth.min()) / (predicted_depth.max() - predicted_depth.min())
|
47 |
|
48 |
-
#
|
49 |
-
|
50 |
-
|
51 |
-
|
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 |
-
|
60 |
-
|
61 |
-
|
62 |
-
|
63 |
-
for
|
64 |
-
|
65 |
-
|
|
|
|
|
|
|
|
|
|
|
66 |
|
67 |
-
return image,
|
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="
|
77 |
-
gr.Image(label="Depth-Based Lens Blur")
|
78 |
],
|
79 |
-
title="
|
80 |
-
description="Upload an image to see
|
81 |
-
|
82 |
-
|
|
|
|
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()
|