File size: 3,124 Bytes
122fab2 6d9557a 122fab2 6d9557a 122fab2 52261a5 122fab2 f82c801 52261a5 6d9557a 52261a5 6d9557a 52261a5 6d9557a 52261a5 6d9557a 122fab2 6d9557a 52261a5 662e365 6d9557a 52261a5 6d9557a 122fab2 6d9557a 122fab2 52261a5 122fab2 52261a5 c7fc3bf 662e365 122fab2 6d9557a 122fab2 991bd3c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 |
import gradio as gr
import cv2
import numpy as np
import torch
from PIL import Image
from transformers import (
SegformerImageProcessor,
SegformerForSemanticSegmentation,
AutoImageProcessor,
AutoModelForDepthEstimation
)
# Load Segformer model for Gaussian blur
segformer_processor = SegformerImageProcessor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
segformer_model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
# Load Depth-Anything model for lens blur
depth_processor = AutoImageProcessor.from_pretrained("depth-anything/Depth-Anything-V2-Small-hf")
depth_model = AutoModelForDepthEstimation.from_pretrained("depth-anything/Depth-Anything-V2-Small-hf")
def apply_blur(image, blur_type, blur_strength, depth_threshold):
# Convert image to RGB
img = image
if blur_type == "Gaussian":
# Use Segformer for Gaussian blur
pil_image = Image.fromarray(img)
inputs = segformer_processor(images=pil_image, return_tensors="pt")
outputs = segformer_model(**inputs)
logits = outputs.logits
mask = logits[0, 12, :, :].detach().cpu().numpy() > depth_threshold
mask = cv2.resize(mask.astype(np.uint8), (img.shape[1], img.shape[0]))
elif blur_type == "Lens":
# Use Depth-Anything for lens blur
pil_image = Image.fromarray(img)
inputs = depth_processor(images=pil_image, return_tensors="pt")
with torch.no_grad():
outputs = depth_model(**inputs)
predicted_depth = outputs.predicted_depth
prediction = torch.nn.functional.interpolate(
predicted_depth.unsqueeze(1),
size=img.shape[:2],
mode="bicubic",
align_corners=False,
)
mask = prediction[0, 0, :, :].detach().cpu().numpy() > depth_threshold
mask = mask.astype(np.uint8)
# Invert mask using cv2
mask = cv2.bitwise_not(mask)
mask = np.repeat(mask[:, :, np.newaxis], 3, axis=2)
# Apply blur based on selected type
if blur_type == "Gaussian":
blurred_image = cv2.GaussianBlur(img, (0, 0), sigmaX=blur_strength)
elif blur_type == "Lens":
# Simulate lens blur using a larger kernel
kernel_size = int(blur_strength * 2) * 2 + 1
blurred_image = cv2.GaussianBlur(img, (kernel_size, kernel_size), 0)
# Combine blurred and original images using the mask
output = np.where(mask == 255, blurred_image, img)
return output
# Define Gradio interface
iface = gr.Interface(
fn=apply_blur,
inputs=[
gr.Image(label="Input Image"),
gr.Radio(["Gaussian", "Lens"], label="Blur Type", value="Gaussian"),
gr.Slider(1, 30, value=15, step=1, label="Blur Strength"),
gr.Slider(-20, 20, value=-4, step=0.1, label="Depth Threshold")
],
outputs=gr.Image(label="Output Image"),
title="Image Segmentation and Blurring",
description="Upload an image and apply Gaussian or Lens blur to the background using different segmentation models."
)
# Launch the app
iface.launch(share=True)
|