Spaces:
Sleeping
Sleeping
import gradio as gr | |
import numpy as np | |
import torch | |
import cv2 | |
from PIL import Image | |
from transformers import ( | |
AutoImageProcessor, | |
SegformerForSemanticSegmentation, | |
DPTImageProcessor, | |
DPTForDepthEstimation | |
) | |
# Load models | |
seg_processor = AutoImageProcessor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512") | |
seg_model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512") | |
depth_processor = DPTImageProcessor.from_pretrained("Intel/dpt-hybrid-midas") | |
depth_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-hybrid-midas") | |
# Segmentation + Gaussian Blur | |
def segmentation_blur(image): | |
inputs = seg_processor(images=image, return_tensors="pt") | |
with torch.no_grad(): | |
outputs = seg_model(**inputs) | |
pred = outputs.logits.argmax(dim=1)[0].cpu().numpy() | |
mask = (pred == 12).astype(np.uint8) * 255 # Person class | |
mask_resized = cv2.resize(mask, image.size, interpolation=cv2.INTER_NEAREST) | |
image_np = np.array(image) | |
mask_3ch = np.stack([mask_resized] * 3, axis=-1) // 255 | |
blurred = cv2.GaussianBlur(image_np, (0, 0), sigmaX=15) | |
combined = np.where(mask_3ch == 1, image_np, blurred) | |
return Image.fromarray(combined) | |
# Depth-based Variable Blur | |
def apply_depth_based_blur(image_np, depth_map_norm, max_blur=25): | |
h, w = depth_map_norm.shape | |
blurred_image = np.zeros_like(image_np) | |
image_np = image_np.astype(np.float32) | |
for i in range(1, max_blur+1, 2): | |
mask = ((depth_map_norm >= (i / max_blur)) & (depth_map_norm < ((i + 2) / max_blur))).astype(np.uint8) | |
if np.sum(mask) == 0: | |
continue | |
blurred = cv2.GaussianBlur(image_np, (i, i), 0) | |
mask_3ch = np.stack([mask]*3, axis=-1) | |
blurred_image = np.where(mask_3ch == 1, blurred, blurred_image) | |
sharpest_mask = (depth_map_norm < 0.1).astype(np.uint8) | |
sharpest_mask_3ch = np.stack([sharpest_mask]*3, axis=-1) | |
final_image = np.where(sharpest_mask_3ch == 1, image_np, blurred_image) | |
return final_image.astype(np.uint8) | |
def depth_blur(image): | |
inputs = depth_processor(images=image, return_tensors="pt") | |
with torch.no_grad(): | |
outputs = depth_model(**inputs) | |
depth_map = outputs.predicted_depth[0].cpu().numpy() | |
depth_map = cv2.resize(depth_map, image.size) | |
norm_depth = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min()) | |
image_np = np.array(image) | |
blurred_image = apply_depth_based_blur(image_np, norm_depth) | |
return Image.fromarray(blurred_image) | |
# Gradio interface | |
def process(image, mode): | |
if mode == "Segmentation Blur": | |
return segmentation_blur(image) | |
else: | |
return depth_blur(image) | |
demo = gr.Interface( | |
fn=process, | |
inputs=[ | |
gr.Image(type="pil", label="Upload Image"), | |
gr.Radio(["Segmentation Blur", "Depth Blur"], value="Segmentation Blur", label="Blur Mode") | |
], | |
outputs=gr.Image(label="Output"), | |
title="Image Blur Effects (Segmentation & Depth)", | |
description="Upload an image and apply either Zoom-style background blur or DSLR-style depth blur." | |
) | |
demo.launch() |