|
import gradio as gr |
|
import numpy as np |
|
import random |
|
import cv2 |
|
import torch |
|
|
|
device = "cuda" if torch.cuda.is_available() else "cpu" |
|
from transformers import pipeline |
|
import requests |
|
from PIL import Image |
|
|
|
|
|
pipe = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True, device = 0) |
|
pipe_depth = pipeline(task="depth-estimation", model="depth-anything/Depth-Anything-V2-Base-hf") |
|
|
|
def lens_blur(pil_img): |
|
pillow_mask = pipe(pil_img, return_mask = True) |
|
mask_arr = np.array(pillow_mask) |
|
mask_bool = mask_arr > 220 |
|
mask_filter = ((mask_arr > 220)*255).astype(np.uint8) |
|
mask_filter_pillow = Image.fromarray(mask_filter.astype(np.uint8)) |
|
|
|
res = cv2.bitwise_and(np.array(pil_img),np.array(pil_img),mask = mask_filter) |
|
depth = pipe_depth(pil_img)["depth"] |
|
depth = np.array(depth) |
|
|
|
blurred_img_sharp = cv2.GaussianBlur(np.array(pil_img),ksize=(5,5),sigmaX=0.1,sigmaY=0.1) |
|
blurred_img_blur = cv2.GaussianBlur(np.array(pil_img),ksize=(5,5),sigmaX=30,sigmaY=30) |
|
|
|
blurred_with_depth_map = np.expand_dims(depth,-1)/255 * blurred_img_sharp + (255 - np.expand_dims(depth,-1))/255 * blurred_img_blur |
|
blurred_with_depth_map = np.clip(blurred_with_depth_map ,0,255) |
|
blurred_with_depth_map = blurred_with_depth_map.astype(np.uint8) |
|
|
|
bg_lens_blur_img = cv2.bitwise_and(blurred_with_depth_map,blurred_with_depth_map,mask = 255 - mask_filter) + res |
|
return Image.fromarray(bg_lens_blur_img) |
|
|
|
|
|
css = """ |
|
#col-container { |
|
margin: 0 auto; |
|
max-width: 640px; |
|
} |
|
""" |
|
|
|
with gr.Blocks(css=css) as demo: |
|
with gr.Column(elem_id="col-container"): |
|
gr.Markdown(" # Lens Blur") |
|
|
|
with gr.Row(): |
|
prompt = gr.Image(type="pil") |
|
|
|
run_button = gr.Button("Run", scale=0, variant="primary") |
|
|
|
result = gr.Image(label="Result", show_label=False) |
|
|
|
gr.on( |
|
triggers=[run_button.click, prompt.upload], |
|
fn=lens_blur, |
|
inputs=[ |
|
prompt, |
|
], |
|
outputs=[result], |
|
) |
|
|
|
if __name__ == "__main__": |
|
demo.launch(share=True) |
|
|