Update app.py

app.py

@@ -1,22 +1,36 @@
+import gradio as gr
 import torch
 import numpy as np
 from PIL import Image, ImageFilter
-import gradio as gr
+import matplotlib.pyplot as plt
 from torchvision import transforms
-from transformers import (
-    AutoModelForImageSegmentation,
-    AutoProcessor,
-    AutoModelForDepthEstimation,
-)
-
-
+from transformers import AutoProcessor, AutoModelForImageSegmentation, AutoModelForDepthEstimation
 
 def load_segmentation_model():
     model_name = "ZhengPeng7/BiRefNet"
     model = AutoModelForImageSegmentation.from_pretrained(model_name, trust_remote_code=True)
     return model
 
-def segment_image(input_tensor, model):
+def load_depth_model():
+    model_name = "depth-anything/Depth-Anything-V2-Metric-Indoor-Base-hf"
+    processor = AutoProcessor.from_pretrained(model_name)
+    model = AutoModelForDepthEstimation.from_pretrained(model_name)
+    return processor, model
+
+def process_segmentation_image(image):
+    transform = transforms.Compose([
+        transforms.Resize((512, 512)),
+        transforms.ToTensor(),
+    ])
+    input_tensor = transform(image).unsqueeze(0)
+    return image, input_tensor
+
+def process_depth_image(image, processor):
+    image = image.resize((512, 512))
+    inputs = processor(images=image, return_tensors="pt")
+    return image, inputs
+
+def segment_image(image, input_tensor, model):
     with torch.no_grad():
         outputs = model(input_tensor)
         output_tensor = outputs[0] if isinstance(outputs, list) else outputs.logits
@@ -24,54 +38,50 @@ def segment_image(input_tensor, model):
         mask = (mask > 0.5).astype(np.uint8) * 255
     return mask
 
-def load_depth_model():
-    model_name = "depth-anything/Depth-Anything-V2-Metric-Indoor-Base-hf"
-    processor = AutoProcessor.from_pretrained(model_name)
-    model = AutoModelForDepthEstimation.from_pretrained(model_name)
-    return processor, model
-
 def estimate_depth(inputs, model):
     with torch.no_grad():
         outputs = model(**inputs)
-    return outputs.predicted_depth.squeeze().cpu().numpy()
+    depth_map = outputs.predicted_depth.squeeze().cpu().numpy()
+    return depth_map
 
 def normalize_depth_map(depth_map):
-    return (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min())
+    min_val = np.min(depth_map)
+    max_val = np.max(depth_map)
+    normalized_depth = (depth_map - min_val) / (max_val - min_val)
+    return normalized_depth
+
+def apply_blur(image, mask):
+    mask_pil = Image.fromarray(mask).resize(image.size, Image.BILINEAR)
+    blurred_background = image.filter(ImageFilter.GaussianBlur(15))
+    final_image = Image.composite(image, blurred_background, mask_pil)
+    return final_image
 
 def apply_depth_based_blur(image, depth_map):
     normalized_depth = normalize_depth_map(depth_map)
+    image = image.resize((512, 512))
     blurred_image = image.copy()
-
-    for y in range(0, image.height, 20):
-        for x in range(0, image.width, 20):
+    for y in range(image.height):
+        for x in range(image.width):
             depth_value = float(normalized_depth[y, x])
-            blur_radius = depth_value * 20
-
-            box = (max(x-10, 0), max(y-10, 0), min(x+10, image.width), min(y+10, image.height))
-            cropped = image.crop(box)
-            blurred_region = cropped.filter(ImageFilter.GaussianBlur(blur_radius))
-            blurred_image.paste(blurred_region, box)
-
+            blur_radius = max(0, depth_value * 20)
+            cropped_region = image.crop((max(x-10, 0), max(y-10, 0), min(x+10, image.width), min(y+10, image.height)))
+            blurred_region = cropped_region.filter(ImageFilter.GaussianBlur(blur_radius))
+            blurred_image.paste(blurred_region, (max(x-10, 0), max(y-10, 0)))
     return blurred_image
 
 def process_image_pipeline(image):
-    original_image = image.convert("RGB").resize((512, 512))
-
-    # Segmentation
-    seg_transform = transforms.Compose([transforms.Resize((512, 512)), transforms.ToTensor()])
-    input_tensor = seg_transform(original_image).unsqueeze(0)
-    seg_model = load_segmentation_model()
-    mask = segment_image(input_tensor, seg_model)
-
-    # Depth Estimation
+    segmentation_model = load_segmentation_model()
     depth_processor, depth_model = load_depth_model()
-    depth_inputs = depth_processor(images=original_image, return_tensors="pt")
-    depth_map = estimate_depth(depth_inputs, depth_model)
-
-    # Depth-based Blur
-    blurred_image = apply_depth_based_blur(original_image, depth_map)
-
-    return original_image, Image.fromarray(mask), Image.fromarray(np.uint8(depth_map / depth_map.max() * 255)), blurred_image
+    
+    _, input_tensor = process_segmentation_image(image)
+    _, inputs = process_depth_image(image, depth_processor)
+    
+    segmentation_mask = segment_image(image, input_tensor, segmentation_model)
+    depth_map = estimate_depth(inputs, depth_model)
+    blurred_image = apply_depth_based_blur(image, depth_map)
+    gaussian_blur_image = apply_blur(image, segmentation_mask)
+    
+    return image, Image.fromarray(segmentation_mask), Image.fromarray((depth_map / np.max(depth_map) * 255).astype(np.uint8)), blurred_image, gaussian_blur_image
 
 iface = gr.Interface(
     fn=process_image_pipeline,
@@ -80,11 +90,12 @@ iface = gr.Interface(
         gr.Image(label="Original Image"),
         gr.Image(label="Segmentation Mask"),
         gr.Image(label="Depth Map"),
-        gr.Image(label="Depth-based Blurred Image")
+        gr.Image(label="Depth-based Blurred Image"),
+        gr.Image(label="Gaussian Blur Image")
     ],
     title="Segmentation and Depth-Based Image Processing",
-    description="Upload an image to get segmentation mask, depth map, and depth-based blur effect.",
+    description="Upload an image to get segmentation mask, depth map, depth-based blur effect, and Gaussian blur effect."
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()