app.py

import gradio as gr
from PIL import Image, ImageFilter
# import matplotlib.pyplot as plt
import torch
import cv2
import numpy as np
from torchvision import transforms
from transformers import AutoModelForImageSegmentation, DepthProImageProcessorFast, DepthProForDepthEstimation
import requests

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

birefnet = AutoModelForImageSegmentation.from_pretrained('ZhengPeng7/BiRefNet', trust_remote_code=True)
torch.set_float32_matmul_precision(['high', 'highest'][0])
birefnet.to('cuda')
birefnet.eval()
birefnet.half()

def extract_object(image, t1, t2):
    # Data settings
    image_size = (1024, 1024)
    transform_image = transforms.Compose([
        transforms.Resize(image_size),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

    # image = Image.open(imagepath)
    image1 = image.copy()
    input_images = transform_image(image1).unsqueeze(0).to('cuda').half()

    # Prediction
    with torch.no_grad():
        preds = birefnet(input_images)[-1].sigmoid().cpu()
    pred = preds[0].squeeze()
    pred_pil = transforms.ToPILImage()(pred)
    mask = pred_pil.resize(image1.size)
    image1.putalpha(mask)

    blurredBg = cv2.GaussianBlur(np.array(imageResized), (0, 0), sigmaX=15, sigmaY=15)

    mask = np.array(result[1].convert("L"))
    _, maskBinary = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)
    img = cv2.cvtColor(np.array(imageResized), cv2.COLOR_RGB2BGR)
    
    maskInv = cv2.bitwise_not(maskBinary)
    maskInv3 = cv2.cvtColor(maskInv, cv2.COLOR_GRAY2BGR)
    
    foreground = cv2.bitwise_and(img, cv2.bitwise_not(maskInv3))
    background = cv2.bitwise_and(blurredBg, maskInv3)
    finalImg = cv2.add(cv2.cvtColor(foreground, cv2.COLOR_BGR2RGB), background)

# plt.figure(figsize=(15, 5))
    # return image1, mask

# def depth_estimation():
    imageProcessor = DepthProImageProcessorFast.from_pretrained("apple/DepthPro-hf")
    model = DepthProForDepthEstimation.from_pretrained("apple/DepthPro-hf").to(device)
    
    inputs = imageProcessor(images=imageResized, return_tensors="pt").to(device)
    
    with torch.no_grad():
        outputs = model(**inputs)
    
    post_processed_output = imageProcessor.post_process_depth_estimation(
        outputs, target_sizes=[(imageResized.height, imageResized.width)],
    )
    
    field_of_view = post_processed_output[0]["field_of_view"]
    focal_length = post_processed_output[0]["focal_length"]
    depth = post_processed_output[0]["predicted_depth"]
    depth = (depth - depth.min()) / (depth.max() - depth.min())
    depth = depth * 255.
    depth = depth.detach().cpu().numpy()
    # print(depth)
    depthImg = Image.fromarray(depth.astype("uint8"))
    
    # threshold1 = 255 / 20      # ~85
    # threshold2 = 2 * 255 / 3   # ~170

    threshold1 = (t1/10) * 255
    threshold2 = (t2/10) * 255
    
    # Precompute blurred versions for each region
    img_foreground = img.copy()  # No blur for foreground
    img_middleground = cv2.GaussianBlur(img, (0, 0), sigmaX=7, sigmaY=7)
    img_background = cv2.GaussianBlur(img, (0, 0), sigmaX=15, sigmaY=15)
    
    # Create masks for each region (as float arrays for proper blending)
    mask_fg = (depth < threshold1).astype(np.float32)
    mask_mg = ((depth >= threshold1) & (depth < threshold2)).astype(np.float32)
    mask_bg = (depth >= threshold2).astype(np.float32)
    
    # Expand masks to 3 channels (H, W, 3)
    mask_fg = np.stack([mask_fg]*3, axis=-1)
    mask_mg = np.stack([mask_mg]*3, axis=-1)
    mask_bg = np.stack([mask_bg]*3, axis=-1)
    
    # Combine the images using the masks in a vectorized manner.
    final_img = (img_foreground * mask_fg + 
                 img_middleground * mask_mg + 
                 img_background * mask_bg).astype(np.uint8)
    
    # Convert the result back to RGB for display with matplotlib.
    final_img_rgb = cv2.cvtColor(final_img, cv2.COLOR_BGR2RGB)

    return image1, final_img

# Visualization
# plt.axis("off")
# subplots for 3 images: original, segmented, mask

# plt.figure(figsize=(15, 5))

# image = Image.open('/content/drive/MyDrive/eee515-hw3/hw3-q24.jpg')
# #resize the image to 512x512
# imageResized = image.resize((512, 512))

# result = extract_object(birefnet, imageResized)
# plt.subplot(1, 3, 1)
# plt.title("Original Resized Image")
# plt.imshow(imageResized)

# plt.subplot(1, 3, 2)
# plt.title("Segmented Image")
# plt.imshow(result[0])

# plt.subplot(1, 3, 3)
# plt.title("Mask")
# plt.imshow(result[1], cmap="gray")
# plt.show()

# Create a Gradio interface


def build_interface(image1, image2):
    """Build UI for gradio app
    """
    title = "Bokeh and Lens Blur"
    with gr.Blocks(theme=gr.themes.Soft(), title=title, fill_width=True) as interface:
        with gr.Row():
            # with gr.Column(scale=3):
            #     with gr.Group():
            #         input_text_box = gr.Textbox(
            #             value=None,
            #             label="Prompt",
            #             lines=2,
            #         )
            #         # gr.Markdown("### Set the values for Middleground and Background")
            #         # fg = gr.Slider(minimum=0, maximum=99, step=1, value=33, label="Middleground")
            #         # mg = gr.Slider(minimum=0, maximum=99, step=1, value=66, label="Background")
            #     with gr.Row():
            #         submit_button = gr.Button("Submit", variant="primary")
            with gr.Column(scale=3):
                model3d = gr.Model3D(
                    label="Output", height="45em", interactive=False
                )

            with gr.Column(scale=3):
                model3d = gr.Model3D(
                    label="Output", height="45em", interactive=False
                )
    
        submit_button.click(
            handle_text_prompt,
            inputs=[
                input_text_box,
                variance
            ],
            outputs=[
                model3d
            ]
        )
                
    return interface

# demo = gr.Interface(sepia, gr.Image(), "image")

title = "Gaussian Blur Background App"
description = (
    "Upload an image to apply a realistic background blur effect. "
    "The app segments the foreground using RMBG-2.0 and then applies a Gaussian "
    "blur (σ=15) to the background, simulating a video conferencing blur effect."
)

iface = gr.Interface(
    fn=apply_blur_effect,
    inputs=[gr.Image(type="pil", label="Input Image"), gr.Slider(minimum=0, maximum=40, step=1, value=33, label="Middleground"), gr.Slider(minimum=40, maximum=99, step=1, value=66, label="Background")],
    outputs=[gr.Image(type="pil", label="Bokeh Image"), gr.Image(type="pil", label="Lens Blur Image")],
    title=title,
    description=description,
    allow_flagging="never"
)

demo = build_interface()
demo.queue(default_concurrency_limit=1)
demo.launch()