updated code

app.py

@@ -3,36 +3,57 @@ import cv2
 import numpy as np
 import torch
 from PIL import Image
-from transformers import AutoImageProcessor, AutoModelForDepthEstimation
+from transformers import (
+    SegformerImageProcessor, 
+    SegformerForSemanticSegmentation,
+    AutoImageProcessor, 
+    AutoModelForDepthEstimation
+)
+
+# Load Segformer model for Gaussian blur
+segformer_processor = SegformerImageProcessor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
+segformer_model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
 
-# Load depth estimation model
-image_processor = AutoImageProcessor.from_pretrained("depth-anything/Depth-Anything-V2-Small-hf")
+# Load Depth-Anything model for lens blur
+depth_processor = AutoImageProcessor.from_pretrained("depth-anything/Depth-Anything-V2-Small-hf")
 depth_model = AutoModelForDepthEstimation.from_pretrained("depth-anything/Depth-Anything-V2-Small-hf")
 
-def apply_blur(image, blur_type, blur_strength, depth_threshold):
+def apply_blur(image, blur_type, blur_strength):
     # Convert image to RGB
     img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
     
-    # Depth estimation
-    pil_image = Image.fromarray(img)
-    inputs = image_processor(images=pil_image, return_tensors="pt")
-    
-    with torch.no_grad():
-        outputs = depth_model(**inputs)
-        predicted_depth = outputs.predicted_depth
-
-    # Interpolate to original size
-    prediction = torch.nn.functional.interpolate(
-        predicted_depth.unsqueeze(1),
-        size=img.shape[:2],
-        mode="bicubic",
-        align_corners=False,
-    )
+    if blur_type == "Gaussian":
+        # Use Segformer for Gaussian blur
+        pil_image = Image.fromarray(img)
+        inputs = segformer_processor(images=pil_image, return_tensors="pt")
+        outputs = segformer_model(**inputs)
+        logits = outputs.logits
+        
+        mask = logits[0, 12, :, :].detach().cpu().numpy() > -4
+        mask = cv2.resize(mask.astype(np.uint8), (img.shape[1], img.shape[0]))
+        
+    elif blur_type == "Lens":
+        # Use Depth-Anything for lens blur
+        pil_image = Image.fromarray(img)
+        inputs = depth_processor(images=pil_image, return_tensors="pt")
+        
+        with torch.no_grad():
+            outputs = depth_model(**inputs)
+            predicted_depth = outputs.predicted_depth
 
-    # Create mask based on depth threshold
-    mask = prediction[0, 0, :, :].detach().cpu().numpy() < depth_threshold
-    mask = mask.astype(np.uint8)
-    mask = np.repeat(mask[:, :, np.newaxis], 3, axis=2) * 255
+        prediction = torch.nn.functional.interpolate(
+            predicted_depth.unsqueeze(1),
+            size=img.shape[:2],
+            mode="bicubic",
+            align_corners=False,
+        )
+        
+        mask = prediction[0, 0, :, :].detach().cpu().numpy() < 3
+        mask = mask.astype(np.uint8)
+    
+    # Invert mask using cv2
+    mask = cv2.bitwise_not(mask)
+    mask = np.repeat(mask[:, :, np.newaxis], 3, axis=2)
 
     # Apply blur based on selected type
     if blur_type == "Gaussian":
@@ -43,7 +64,7 @@ def apply_blur(image, blur_type, blur_strength, depth_threshold):
         blurred_image = cv2.GaussianBlur(img, (kernel_size, kernel_size), 0)
 
     # Combine blurred and original images using the mask
-    output = np.where(mask == 255, img, blurred_image)
+    output = np.where(mask == 255, blurred_image, img)
     
     return output
 
@@ -53,12 +74,11 @@ iface = gr.Interface(
     inputs=[
         gr.Image(label="Input Image"),
         gr.Radio(["Gaussian", "Lens"], label="Blur Type"),
-        gr.Slider(1, 30, value=15, step=1, label="Blur Strength"),
-        gr.Slider(1, 10, value=3, step=0.1, label="Depth Threshold")
+        gr.Slider(1, 30, value=15, step=1, label="Blur Strength")
     ],
     outputs=gr.Image(label="Output Image"),
     title="Image Segmentation and Blurring",
-    description="Upload an image and apply Gaussian or Lens blur to the background based on depth estimation."
+    description="Upload an image and apply Gaussian or Lens blur to the background using different segmentation models."
 )
 
 # Launch the app