Update app.py
Browse files
app.py
CHANGED
@@ -0,0 +1,82 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import gradio as gr
|
2 |
+
import torch
|
3 |
+
import numpy as np
|
4 |
+
from PIL import Image, ImageFilter, ImageOps
|
5 |
+
import cv2
|
6 |
+
from transformers import (
|
7 |
+
SegformerFeatureExtractor, SegformerForSemanticSegmentation,
|
8 |
+
DPTFeatureExtractor, DPTForDepthEstimation
|
9 |
+
)
|
10 |
+
|
11 |
+
# Load models
|
12 |
+
seg_model_name = "nvidia/segformer-b1-finetuned-ade-512-512"
|
13 |
+
depth_model_name = "Intel/dpt-hybrid-midas"
|
14 |
+
|
15 |
+
seg_extractor = SegformerFeatureExtractor.from_pretrained(seg_model_name)
|
16 |
+
seg_model = SegformerForSemanticSegmentation.from_pretrained(seg_model_name)
|
17 |
+
depth_extractor = DPTFeatureExtractor.from_pretrained(depth_model_name)
|
18 |
+
depth_model = DPTForDepthEstimation.from_pretrained(depth_model_name)
|
19 |
+
|
20 |
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
21 |
+
seg_model.to(device)
|
22 |
+
depth_model.to(device)
|
23 |
+
|
24 |
+
def process_image(image_pil):
|
25 |
+
image = ImageOps.exif_transpose(image_pil).resize((512, 512)).convert("RGB")
|
26 |
+
|
27 |
+
# ---------- Part 1: Segmentation ----------
|
28 |
+
seg_inputs = seg_extractor(images=image, return_tensors="pt").to(device)
|
29 |
+
with torch.no_grad():
|
30 |
+
seg_output = seg_model(**seg_inputs).logits
|
31 |
+
seg_mask = torch.argmax(seg_output, dim=1)[0].cpu().numpy()
|
32 |
+
binary_mask = np.where(seg_mask > 0, 255, 0).astype(np.uint8)
|
33 |
+
foreground_mask = Image.fromarray(binary_mask).convert("L")
|
34 |
+
|
35 |
+
# ---------- Part 2: Gaussian blur to background ----------
|
36 |
+
blurred_background = image.filter(ImageFilter.GaussianBlur(15))
|
37 |
+
blurred_background = blurred_background.convert("RGBA")
|
38 |
+
image_rgba = image.convert("RGBA")
|
39 |
+
output_blur = Image.composite(image_rgba, blurred_background, foreground_mask)
|
40 |
+
|
41 |
+
# ---------- Part 3: Depth Estimation ----------
|
42 |
+
image_np = np.array(image)
|
43 |
+
depth_inputs = depth_extractor(images=image_np, return_tensors="pt").to(device)
|
44 |
+
with torch.no_grad():
|
45 |
+
depth_output = depth_model(**depth_inputs)
|
46 |
+
predicted_depth = depth_output.predicted_depth.squeeze().cpu().numpy()
|
47 |
+
normalized_depth = (predicted_depth - predicted_depth.min()) / (predicted_depth.max() - predicted_depth.min())
|
48 |
+
|
49 |
+
# ---------- Part 4: Depth-Based Variable Gaussian Blur ----------
|
50 |
+
image_np_float = image_np.astype(np.float32)
|
51 |
+
resized_depth = cv2.resize(normalized_depth, (image_np.shape[1], image_np.shape[0]))
|
52 |
+
inverted_depth = 1.0 - resized_depth
|
53 |
+
total_blur_levels = 4
|
54 |
+
blurred_versions = []
|
55 |
+
for i in range(total_blur_levels):
|
56 |
+
sigma = i * 3
|
57 |
+
blurred = cv2.GaussianBlur(image_np_float, (15, 15), sigmaX=sigma, sigmaY=sigma) if sigma > 0 else image_np_float.copy()
|
58 |
+
blurred_versions.append(blurred)
|
59 |
+
|
60 |
+
blur_indices = (inverted_depth * (total_blur_levels - 1)).astype(np.uint8)
|
61 |
+
final_blurred_np = np.zeros_like(image_np_float)
|
62 |
+
for i in range(total_blur_levels):
|
63 |
+
mask = (blur_indices == i)
|
64 |
+
for c in range(3):
|
65 |
+
final_blurred_np[:, :, c][mask] = blurred_versions[i][:, :, c][mask]
|
66 |
+
depth_blur_img = Image.fromarray(np.clip(final_blurred_np, 0, 255).astype(np.uint8))
|
67 |
+
|
68 |
+
return image, output_blur.convert("RGB"), depth_blur_img
|
69 |
+
|
70 |
+
# Gradio Interface
|
71 |
+
gr.Interface(
|
72 |
+
fn=process_image,
|
73 |
+
inputs=gr.Image(type="pil"),
|
74 |
+
outputs=[
|
75 |
+
gr.Image(label="Original Image"),
|
76 |
+
gr.Image(label="Segmented Gaussian Blur"),
|
77 |
+
gr.Image(label="Depth-Based Lens Blur")
|
78 |
+
],
|
79 |
+
title="Visual Effects Demo: Segmentation & Depth-Based Blur",
|
80 |
+
description="Upload an image to see it segmented with background blur (like Zoom) and depth-based lens blur.",
|
81 |
+
examples=[],
|
82 |
+
).launch()
|