Update app.py

app.py

@@ -3,31 +3,51 @@ from PIL import Image, ImageFilter
 import numpy as np
 import cv2
 import torch
-from transformers import DPTFeatureExtractor, DPTForDepthEstimation
+from transformers import SegformerFeatureExtractor, SegformerForSemanticSegmentation, DPTFeatureExtractor, DPTForDepthEstimation
 
-# Load model and feature extractor
-feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large")
-model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large")
+# Load models
+segformer_extractor = SegformerFeatureExtractor.from_pretrained("nvidia/segformer-b1-finetuned-ade-512-512")
+segformer_model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b1-finetuned-ade-512-512")
+dpt_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large")
+dpt_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large")
 
-# Gaussian Blur function
-def apply_gaussian_blur(image, blur_radius):
-    return image.filter(ImageFilter.GaussianBlur(blur_radius))
+# Gaussian Blur Background Function
+def gaussian_blur_background(image):
+    # Preprocess image for segmentation
+    inputs = segformer_extractor(images=image, return_tensors="pt")
+    outputs = segformer_model(**inputs)
+    logits = outputs.logits
+    segmentation = torch.argmax(logits, dim=1)[0].numpy()
 
-# Lens Blur function
-def apply_lens_blur(image):
-    # Get depth map
-    inputs = feature_extractor(images=image, return_tensors="pt")
+    # Create a binary mask for 'person' class (class index 12)
+    human_mask = (segmentation == 12).astype(np.uint8) * 255
+    human_mask_image = Image.fromarray(human_mask).resize(image.size)
+
+    # Apply Gaussian blur to the entire image
+    blurred_background = image.filter(ImageFilter.GaussianBlur(15))
+
+    # Composite the original image with blurred background using the mask
+    composite_image = Image.composite(image, blurred_background, human_mask_image)
+    return composite_image
+
+# Depth-Based Lens Blur Function
+def lens_blur(image):
+    # Preprocess image for depth estimation
+    inputs = dpt_extractor(images=image, return_tensors="pt")
     with torch.no_grad():
-        outputs = model(**inputs)
+        outputs = dpt_model(**inputs)
         depth_map = outputs.predicted_depth.squeeze().cpu().numpy()
+
+    # Normalize depth map to range [0, 15] and invert for blur intensity
     depth_map = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min()) * 15
+    depth_map = 15 - depth_map
     depth_map_resized = cv2.resize(depth_map, (image.width, image.height))
-    depth_map_resized = 15 - depth_map_resized
 
-    # Convert to OpenCV format
+    # Convert image to OpenCV format
     image_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
     blurred_image = np.zeros_like(image_cv, dtype=np.float32)
 
+    # Apply variable blur based on depth
     for blur_radius in range(1, 16):
         blurred_layer = cv2.GaussianBlur(image_cv, (0, 0), sigmaX=blur_radius)
         mask = ((depth_map_resized >= (blur_radius - 1)) & (depth_map_resized < blur_radius)).astype(np.float32)
@@ -35,30 +55,26 @@ def apply_lens_blur(image):
         blurred_image += blurred_layer * mask
 
     blurred_image = np.clip(blurred_image, 0, 255).astype(np.uint8)
-    return Image.fromarray(cv2.cvtColor(blurred_image, cv2.COLOR_BGR2RGB))
+    blurred_image_pil = Image.fromarray(cv2.cvtColor(blurred_image, cv2.COLOR_BGR2RGB))
+    return blurred_image_pil
 
-# Gradio app interface
-def process_image(image, effect, blur_radius):
-    if effect == "Gaussian Blur":
-        return apply_gaussian_blur(image, blur_radius)
+# Gradio Interface
+def process_image(image, effect):
+    if effect == "Gaussian Blur Background":
+        return gaussian_blur_background(image)
     elif effect == "Lens Blur":
-        return apply_lens_blur(image)
-    else:
-        return image
+        return lens_blur(image)
 
-# Gradio Interface
 with gr.Blocks() as demo:
-    gr.Markdown("# Gaussian and Lens Blur Effects")
+    gr.Markdown("# BlurEffectsApp: Gaussian and Lens Blur Effects")
     with gr.Row():
         with gr.Column():
-            uploaded_image = gr.Image(type="pil")
-            effect = gr.Radio(["Gaussian Blur", "Lens Blur"], value="Gaussian Blur", label="Effect")
-            blur_radius = gr.Slider(1, 15, value=5, step=1, label="Blur Radius (for Gaussian Blur)")
-            submit_button = gr.Button("Apply Effect")
+            uploaded_image = gr.Image(type="pil", label="Upload an Image")
+            effect = gr.Radio(["Gaussian Blur Background", "Lens Blur"], label="Choose Effect")
+            process_button = gr.Button("Apply Effect")
         with gr.Column():
             output_image = gr.Image(type="pil", label="Processed Image")
+    
+    process_button.click(process_image, inputs=[uploaded_image, effect], outputs=output_image)
 
-    submit_button.click(process_image, inputs=[uploaded_image, effect, blur_radius], outputs=output_image)
-
-# Launch the app
 demo.launch()