EEE515_Problem2 / app.py
joeWabbit's picture
Update app.py
1218a4a verified
import gradio as gr
import torch
import numpy as np
from transformers import AutoImageProcessor, AutoModelForDepthEstimation
from PIL import Image, ImageFilter
import matplotlib.pyplot as plt
import matplotlib.cm as cm
def compute_depth_map(image: Image.Image, scale_factor: float) -> np.ndarray:
image_processor = AutoImageProcessor.from_pretrained("LiheYoung/depth-anything-large-hf")
model = AutoModelForDepthEstimation.from_pretrained("LiheYoung/depth-anything-large-hf")
inputs = image_processor(images=image, return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs)
predicted_depth = outputs.predicted_depth
prediction = torch.nn.functional.interpolate(
predicted_depth.unsqueeze(1),
size=image.size[::-1],
mode="bicubic",
align_corners=False,
)
depth_min = prediction.min()
depth_max = prediction.max()
depth_vis = (prediction - depth_min) / (depth_max - depth_min + 1e-8)
depth_map = depth_vis.squeeze().cpu().numpy()
depth_map_inverted = 1.0 - depth_map
depth_map_inverted *= scale_factor
return depth_map_inverted
def layered_blur(image: Image.Image, depth_map: np.ndarray, num_layers: int, max_blur: float) -> Image.Image:
blur_radii = np.linspace(0, max_blur, num_layers)
blur_versions = [image.filter(ImageFilter.GaussianBlur(radius)) for radius in blur_radii]
thresholds = np.linspace(0, 1, num_layers + 1)
final_image = blur_versions[-1]
for i in range(num_layers - 1, -1, -1):
mask_array = np.logical_and(
depth_map >= thresholds[i],
depth_map < thresholds[i + 1]
).astype(np.uint8) * 255
mask_image = Image.fromarray(mask_array, mode="L")
final_image = Image.composite(blur_versions[i], final_image, mask_image)
return final_image
def process_depth_blur(uploaded_image, max_blur_value, scale_factor, num_layers):
if not isinstance(uploaded_image, Image.Image):
uploaded_image = Image.open(uploaded_image)
image = uploaded_image.convert("RGB").resize((512, 512))
depth_map = compute_depth_map(image, scale_factor)
final_image = layered_blur(image, depth_map, int(num_layers), max_blur_value)
return final_image
def create_heatmap(depth_map: np.ndarray, intensity: float) -> Image.Image:
normalized = np.clip(depth_map * intensity, 0, 1)
colormap = cm.get_cmap("inferno")
colored = colormap(normalized)
heatmap = (colored[:, :, :3] * 255).astype(np.uint8)
return Image.fromarray(heatmap)
def process_depth_heatmap(uploaded_image, intensity):
if not isinstance(uploaded_image, Image.Image):
uploaded_image = Image.open(uploaded_image)
image = uploaded_image.convert("RGB").resize((512, 512))
depth_map = compute_depth_map(image, scale_factor=1.0)
heatmap_img = create_heatmap(depth_map, intensity)
return heatmap_img
def load_segmentation_model():
global seg_model, seg_device
if "seg_model" not in globals():
from ben2 import BEN_Base # Import BEN2
seg_model = BEN_Base.from_pretrained("PramaLLC/BEN2")
seg_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
seg_model.to(seg_device).eval()
return seg_model, seg_device
def process_segmentation_blur(uploaded_image, seg_blur_radius: float):
if not isinstance(uploaded_image, Image.Image):
uploaded_image = Image.open(uploaded_image)
image = uploaded_image.convert("RGB").resize((512, 512))
seg_model, seg_device = load_segmentation_model()
blurred_image = image.filter(ImageFilter.GaussianBlur(seg_blur_radius))
foreground = seg_model.inference(image, refine_foreground=False)
foreground_rgba = foreground.convert("RGBA")
_, _, _, alpha = foreground_rgba.split()
binary_mask = alpha.point(lambda x: 255 if x > 128 else 0, mode="L")
final_image = Image.composite(image, blurred_image, binary_mask)
return final_image
with gr.Blocks() as demo:
gr.Markdown("Gaussian Blur & Lens Blur Effect")
with gr.Tabs():
with gr.Tab("Gaussian Blur"):
seg_img = gr.Image(type="pil", label="Upload Image")
seg_blur = gr.Slider(5, 30, value=15, step=1, label="Gaussian Blur Radius")
seg_out = gr.Image(label="Gaussian Blurred Image")
seg_button = gr.Button("Process Gaussian Blur")
seg_button.click(process_segmentation_blur, inputs=[seg_img, seg_blur], outputs=seg_out)
with gr.Tab("Depth Heatmap"):
img_input2 = gr.Image(type="pil", label="Upload Image")
intensity_slider = gr.Slider(0.5, 5.0, value=1.0, label="Heatmap Intensity")
heatmap_output = gr.Image(label="Depth Heatmap")
heatmap_button = gr.Button("Generate Depth Heatmap")
heatmap_button.click(
process_depth_heatmap,
inputs=[img_input2, intensity_slider],
outputs=heatmap_output
)
with gr.Tab("Lens Blur"):
img_input = gr.Image(type="pil", label="Upload Image")
blur_slider = gr.Slider(1, 50, value=6, label="Maximum Blur Radius")
scale_slider = gr.Slider(0.1, 2.0, value=0.72, label="Depth Scale Factor")
layers_slider = gr.Slider(2, 10, value=2.91, label="Number of Layers")
blur_output = gr.Image(label="Lens Blur Result")
blur_button = gr.Button("Process Blur")
blur_button.click(
process_depth_blur,
inputs=[img_input, blur_slider, scale_slider, layers_slider],
outputs=blur_output
)
if __name__ == "__main__":
demo.launch(share=True)