Update app.py
Browse files
app.py
CHANGED
@@ -3,7 +3,6 @@ from transformers import pipeline
|
|
3 |
from PIL import Image, ImageFilter
|
4 |
import numpy as np
|
5 |
|
6 |
-
# Initialize models with fixed choices
|
7 |
segmentation_model = pipeline("image-segmentation", model="nvidia/segformer-b1-finetuned-cityscapes-1024-1024")
|
8 |
depth_estimator = pipeline("depth-estimation", model="Intel/zoedepth-nyu-kitti")
|
9 |
|
@@ -25,47 +24,34 @@ def process_image(input_image, method, blur_intensity):
|
|
25 |
- output_image: final composited image.
|
26 |
- mask_image: the mask used (binary for segmentation, normalized depth for depth-based).
|
27 |
"""
|
28 |
-
# Ensure image is in RGB mode
|
29 |
input_image = input_image.convert("RGB")
|
30 |
|
31 |
if method == "Segmentation Blur Model":
|
32 |
-
# Use segmentation to obtain a foreground mask
|
33 |
results = segmentation_model(input_image)
|
34 |
-
# Assume the last result is the main foreground object
|
35 |
foreground_mask = results[-1]["mask"]
|
36 |
-
# Ensure the mask is grayscale
|
37 |
foreground_mask = foreground_mask.convert("L")
|
38 |
-
# Threshold to create a binary mask
|
39 |
binary_mask = foreground_mask.point(lambda p: 255 if p > 128 else 0)
|
40 |
|
41 |
-
# Blur the background using Gaussian blur
|
42 |
blurred_background = input_image.filter(ImageFilter.GaussianBlur(radius=blur_intensity))
|
43 |
|
44 |
-
# Composite the final image: keep foreground and use blurred background elsewhere
|
45 |
output_image = Image.composite(input_image, blurred_background, binary_mask)
|
46 |
mask_image = binary_mask
|
47 |
|
48 |
elif method == "Monocular Depth Estimation Model":
|
49 |
-
# Generate depth map
|
50 |
depth_results = depth_estimator(input_image)
|
51 |
depth_map = depth_results["depth"]
|
52 |
|
53 |
-
# Convert depth map to numpy array and normalize to [0, 255]
|
54 |
depth_array = np.array(depth_map).astype(np.float32)
|
55 |
norm = (depth_array - depth_array.min()) / (depth_array.max() - depth_array.min() + 1e-8)
|
56 |
normalized_depth = (norm * 255).astype(np.uint8)
|
57 |
mask_image = Image.fromarray(normalized_depth)
|
58 |
|
59 |
-
# Create fully blurred version using Gaussian blur
|
60 |
blurred_image = input_image.filter(ImageFilter.GaussianBlur(radius=blur_intensity))
|
61 |
|
62 |
-
# Convert images to arrays for blending
|
63 |
orig_np = np.array(input_image).astype(np.float32)
|
64 |
blur_np = np.array(blurred_image).astype(np.float32)
|
65 |
-
# Reshape mask for broadcasting
|
66 |
alpha = normalized_depth[..., np.newaxis] / 255.0
|
67 |
|
68 |
-
# Blend pixels: 0 = original; 1 = fully blurred
|
69 |
blended_np = (1 - alpha) * orig_np + alpha * blur_np
|
70 |
blended_np = np.clip(blended_np, 0, 255).astype(np.uint8)
|
71 |
output_image = Image.fromarray(blended_np)
|
@@ -76,7 +62,6 @@ def process_image(input_image, method, blur_intensity):
|
|
76 |
|
77 |
return output_image, mask_image
|
78 |
|
79 |
-
# Build a Gradio interface
|
80 |
with gr.Blocks() as demo:
|
81 |
gr.Markdown("## FocusFusion: Segmentation & Depth Blur")
|
82 |
|
@@ -93,12 +78,10 @@ with gr.Blocks() as demo:
|
|
93 |
output_image = gr.Image(label="Output Image")
|
94 |
mask_output = gr.Image(label="Mask")
|
95 |
|
96 |
-
# Set up event handler
|
97 |
run_button.click(
|
98 |
fn=process_image,
|
99 |
inputs=[input_image, method, blur_intensity],
|
100 |
outputs=[output_image, mask_output]
|
101 |
)
|
102 |
|
103 |
-
# Launch the app
|
104 |
demo.launch()
|
|
|
3 |
from PIL import Image, ImageFilter
|
4 |
import numpy as np
|
5 |
|
|
|
6 |
segmentation_model = pipeline("image-segmentation", model="nvidia/segformer-b1-finetuned-cityscapes-1024-1024")
|
7 |
depth_estimator = pipeline("depth-estimation", model="Intel/zoedepth-nyu-kitti")
|
8 |
|
|
|
24 |
- output_image: final composited image.
|
25 |
- mask_image: the mask used (binary for segmentation, normalized depth for depth-based).
|
26 |
"""
|
|
|
27 |
input_image = input_image.convert("RGB")
|
28 |
|
29 |
if method == "Segmentation Blur Model":
|
|
|
30 |
results = segmentation_model(input_image)
|
|
|
31 |
foreground_mask = results[-1]["mask"]
|
|
|
32 |
foreground_mask = foreground_mask.convert("L")
|
|
|
33 |
binary_mask = foreground_mask.point(lambda p: 255 if p > 128 else 0)
|
34 |
|
|
|
35 |
blurred_background = input_image.filter(ImageFilter.GaussianBlur(radius=blur_intensity))
|
36 |
|
|
|
37 |
output_image = Image.composite(input_image, blurred_background, binary_mask)
|
38 |
mask_image = binary_mask
|
39 |
|
40 |
elif method == "Monocular Depth Estimation Model":
|
|
|
41 |
depth_results = depth_estimator(input_image)
|
42 |
depth_map = depth_results["depth"]
|
43 |
|
|
|
44 |
depth_array = np.array(depth_map).astype(np.float32)
|
45 |
norm = (depth_array - depth_array.min()) / (depth_array.max() - depth_array.min() + 1e-8)
|
46 |
normalized_depth = (norm * 255).astype(np.uint8)
|
47 |
mask_image = Image.fromarray(normalized_depth)
|
48 |
|
|
|
49 |
blurred_image = input_image.filter(ImageFilter.GaussianBlur(radius=blur_intensity))
|
50 |
|
|
|
51 |
orig_np = np.array(input_image).astype(np.float32)
|
52 |
blur_np = np.array(blurred_image).astype(np.float32)
|
|
|
53 |
alpha = normalized_depth[..., np.newaxis] / 255.0
|
54 |
|
|
|
55 |
blended_np = (1 - alpha) * orig_np + alpha * blur_np
|
56 |
blended_np = np.clip(blended_np, 0, 255).astype(np.uint8)
|
57 |
output_image = Image.fromarray(blended_np)
|
|
|
62 |
|
63 |
return output_image, mask_image
|
64 |
|
|
|
65 |
with gr.Blocks() as demo:
|
66 |
gr.Markdown("## FocusFusion: Segmentation & Depth Blur")
|
67 |
|
|
|
78 |
output_image = gr.Image(label="Output Image")
|
79 |
mask_output = gr.Image(label="Mask")
|
80 |
|
|
|
81 |
run_button.click(
|
82 |
fn=process_image,
|
83 |
inputs=[input_image, method, blur_intensity],
|
84 |
outputs=[output_image, mask_output]
|
85 |
)
|
86 |
|
|
|
87 |
demo.launch()
|