| import gradio as gr | |
| from transformers import pipeline | |
| from PIL import Image, ImageFilter | |
| import numpy as np | |
| segmentation_model = pipeline("image-segmentation", model="nvidia/segformer-b1-finetuned-cityscapes-1024-1024") | |
| depth_estimator = pipeline("depth-estimation", model="Intel/zoedepth-nyu-kitti") | |
| def process_image(input_image, method, blur_intensity): | |
| """ | |
| Process the input image using one of two methods: | |
| 1. Segmentation Blur Model: | |
| - Uses segmentation to extract a foreground mask. | |
| - Applies Gaussian blur to the background. | |
| - Composites the final image. | |
| 2. Monocular Depth Estimation Model: | |
| - Uses depth estimation to generate a depth map. | |
| - Normalizes the depth map to be used as a blending mask. | |
| - Blends a fully blurred version with the original image. | |
| Returns: | |
| - output_image: final composited image. | |
| - mask_image: the mask used (binary for segmentation, normalized depth for depth-based). | |
| """ | |
| input_image = input_image.convert("RGB") | |
| if method == "Segmentation Blur Model": | |
| results = segmentation_model(input_image) | |
| foreground_mask = results[-1]["mask"] | |
| foreground_mask = foreground_mask.convert("L") | |
| binary_mask = foreground_mask.point(lambda p: 255 if p > 128 else 0) | |
| blurred_background = input_image.filter(ImageFilter.GaussianBlur(radius=blur_intensity)) | |
| output_image = Image.composite(input_image, blurred_background, binary_mask) | |
| mask_image = binary_mask | |
| elif method == "Monocular Depth Estimation Model": | |
| depth_results = depth_estimator(input_image) | |
| depth_map = depth_results["depth"] | |
| depth_array = np.array(depth_map).astype(np.float32) | |
| norm = (depth_array - depth_array.min()) / (depth_array.max() - depth_array.min() + 1e-8) | |
| normalized_depth = (norm * 255).astype(np.uint8) | |
| mask_image = Image.fromarray(normalized_depth) | |
| blurred_image = input_image.filter(ImageFilter.GaussianBlur(radius=blur_intensity)) | |
| orig_np = np.array(input_image).astype(np.float32) | |
| blur_np = np.array(blurred_image).astype(np.float32) | |
| alpha = normalized_depth[..., np.newaxis] / 255.0 | |
| blended_np = (1 - alpha) * orig_np + alpha * blur_np | |
| blended_np = np.clip(blended_np, 0, 255).astype(np.uint8) | |
| output_image = Image.fromarray(blended_np) | |
| else: | |
| output_image = input_image | |
| mask_image = input_image.convert("L") | |
| return output_image, mask_image | |
| with gr.Blocks() as demo: | |
| gr.Markdown("## FocusFusion: Segmentation & Depth Blur") | |
| with gr.Row(): | |
| with gr.Column(): | |
| input_image = gr.Image(label="Input Image", type="pil") | |
| method = gr.Radio(label="Processing Method", | |
| choices=["Segmentation Blur Model", "Monocular Depth Estimation Model"], | |
| value="Segmentation Blur Model") | |
| blur_intensity = gr.Slider(label="Blur Intensity (sigma)", | |
| minimum=1, maximum=30, step=1, value=15) | |
| run_button = gr.Button("Process Image") | |
| with gr.Column(): | |
| output_image = gr.Image(label="Output Image") | |
| mask_output = gr.Image(label="Mask") | |
| run_button.click( | |
| fn=process_image, | |
| inputs=[input_image, method, blur_intensity], | |
| outputs=[output_image, mask_output] | |
| ) | |
| demo.launch() |