File size: 3,323 Bytes
8f8907e
 
c752dca
 
 
8f8907e
 
c752dca
8f8907e
c752dca
8f8907e
c752dca
 
 
 
8f8907e
 
c752dca
8f8907e
 
 
 
c752dca
 
 
 
 
 
8f8907e
c752dca
 
 
 
8f8907e
c752dca
 
 
 
8f8907e
c752dca
 
8f8907e
c752dca
 
8f8907e
c752dca
 
 
 
8f8907e
c752dca
 
 
 
 
 
 
 
 
 
 
 
f48bfcf
c752dca
8f8907e
c752dca
8f8907e
c752dca
8f8907e
 
c752dca
 
8f8907e
c752dca
 
 
 
 
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
import gradio as gr
from PIL import Image, ImageFilter, ImageOps
import numpy as np
import torch
from transformers import SegformerFeatureExtractor, SegformerForSemanticSegmentation, DPTFeatureExtractor, DPTForDepthEstimation
import cv2

# Load segmentation model
seg_model_name = "nvidia/segformer-b1-finetuned-ade-512-512"
seg_feature_extractor = SegformerFeatureExtractor.from_pretrained(seg_model_name)
seg_model = SegformerForSemanticSegmentation.from_pretrained(seg_model_name)

# Load depth estimation model
depth_model_name = "Intel/dpt-hybrid-midas"
depth_feature_extractor = DPTFeatureExtractor.from_pretrained(depth_model_name)
depth_model = DPTForDepthEstimation.from_pretrained(depth_model_name)

# Device configuration
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
seg_model.to(device)
depth_model.to(device)

def process_image(image):
    # Ensure image is in RGB format and resize
    image = ImageOps.exif_transpose(image).resize((512, 512))

    # Perform segmentation
    inputs = seg_feature_extractor(images=image, return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = seg_model(**inputs)
        logits = outputs.logits
        segmentation = torch.argmax(logits, dim=1)[0].cpu().numpy()
    binary_mask = np.where(segmentation > 0, 255, 0).astype(np.uint8)

    # Apply Gaussian Blur to the background
    blurred_background = image.filter(ImageFilter.GaussianBlur(15))
    foreground = Image.fromarray(binary_mask).convert("L").resize(image.size)
    output_blur = Image.composite(image, blurred_background, foreground)

    # Depth estimation for lens blur
    depth_inputs = depth_feature_extractor(images=image, return_tensors="pt").to(device)
    with torch.no_grad():
        depth_outputs = depth_model(**depth_inputs)
        predicted_depth = depth_outputs.predicted_depth.squeeze().cpu().numpy()

    # Normalize depth map
    depth_min, depth_max = predicted_depth.min(), predicted_depth.max()
    normalized_depth = (predicted_depth - depth_min) / (depth_max - depth_min)
    normalized_depth_resized = cv2.resize(normalized_depth, (512, 512))

    # Lens blur using depth map
    blurred_image = np.array(image).astype(np.float32)
    blur_intensity = normalized_depth_resized * 20
    for y in range(image.size[1]):
        for x in range(image.size[0]):
            sigma = blur_intensity[y, x]
            kernel_size = int(2 * sigma + 1)
            if kernel_size > 1:
                patch = image.crop((x - kernel_size//2, y - kernel_size//2, x + kernel_size//2 + 1, y + kernel_size//2 + 1))
                patch = patch.filter(ImageFilter.GaussianBlur(sigma))
                blurred_image[y, x, :] = np.array(patch)[kernel_size//2, kernel_size//2, :]
    lens_blur_image = Image.fromarray(np.clip(blurred_image, 0, 255).astype(np.uint8))

    return image, output_blur, lens_blur_image

iface = gr.Interface(
    fn=process_image,
    inputs=gr.Image(type="pil", label="Upload an Image"),
    outputs=[
        gr.Image(label="Original Image"),
        gr.Image(label="Gaussian Blur Effect"),
        gr.Image(label="Depth-Based Lens Blur Effect")
    ],
    title="Image Blurring with Gaussian and Depth-Based Lens Blur",
    description="Upload an image to see Gaussian blur and depth-based lens blur effects."
)

iface.launch()