rchhibba's picture
creating app.py
68b6746 verified
import gradio as gr
import torch
from transformers import pipeline
from PIL import Image
import numpy as np
import cv2
def process_image(image, effect_type="Gaussian Blur", blur_intensity=15):
"""
Process the image with selected effect
"""
# Resize image to 512x512
image = Image.fromarray(image).resize((512, 512))
if effect_type == "Gaussian Blur":
# Generate segmentation mask
segmenter = pipeline("image-segmentation",
model="openmmlab/upernet-swin-base",
device=0 if torch.cuda.is_available() else -1)
results = segmenter(image)
mask = np.zeros((512, 512), dtype=np.uint8)
for segment in results:
if segment['label'].lower() == 'person':
segment_mask = np.array(segment['mask'])
mask[segment_mask > 0] = 255
# Apply gaussian blur
img_np = np.array(image)
blurred = cv2.GaussianBlur(img_np, (0, 0), blur_intensity)
mask_np = mask / 255.0
mask_np = np.stack([mask_np] * 3, axis=-1)
result = img_np * mask_np + blurred * (1 - mask_np)
return result.astype(np.uint8)
else: # Depth-based blur
# Generate depth map
depth_estimator = pipeline("depth-estimation",
model="Intel/dpt-large",
device=0 if torch.cuda.is_available() else -1)
depth_result = depth_estimator(image)
depth_map = depth_result['predicted_depth']
if torch.is_tensor(depth_map):
depth_map = depth_map.cpu().numpy()
# Apply depth-based blur
img_np = np.array(image)
depth_norm = blur_intensity * (1 - (depth_map - depth_map.min()) /
(depth_map.max() - depth_map.min()))
result = np.zeros_like(img_np)
for sigma in range(int(blur_intensity) + 1):
if sigma == 0:
continue
kernel_size = 2 * int(4 * sigma + 0.5) + 1
mask = (depth_norm >= sigma - 0.5) & (depth_norm < sigma + 0.5)
if not mask.any():
continue
blurred = cv2.GaussianBlur(img_np, (kernel_size, kernel_size), sigma)
result[mask] = blurred[mask]
min_depth_mask = depth_norm > blur_intensity-0.5
result[min_depth_mask] = img_np[min_depth_mask]
return result
# Create Gradio interface
demo = gr.Interface(
fn=process_image,
inputs=[
gr.Image(label="Upload Image", type="numpy"),
gr.Radio(["Gaussian Blur", "Depth-based Blur"], label="Effect Type", value="Gaussian Blur"),
gr.Slider(minimum=1, maximum=30, value=15, label="Blur Intensity")
],
outputs=gr.Image(label="Result"),
title="Image Background Effects",
description="""Upload an image to apply background effects:
1. Gaussian Blur: Blurs the background while keeping the person sharp
2. Depth-based Blur: Applies varying blur based on depth (bokeh effect)""",
examples=[], # You can add example images later
cache_examples=False
)
if __name__ == "__main__":
demo.launch()