update

app.py

@@ -18,26 +18,25 @@ segformer_model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segfo
 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):
+def apply_blur(image, blur_type, blur_strength, depth_threshold):
     # Convert image to RGB
     img = image
-    
+
     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 = logits[0, 12, :, :].detach().cpu().numpy() > depth_threshold
         mask = cv2.resize(mask.astype(np.uint8), (img.shape[1], img.shape[0]))
-        mask = cv2.bitwise_not(mask)
-        
+
     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
@@ -48,10 +47,12 @@ def apply_blur(image, blur_type, blur_strength):
             mode="bicubic",
             align_corners=False,
         )
-        
-        mask = prediction[0, 0, :, :].detach().cpu().numpy() < 3
+
+        mask = prediction[0, 0, :, :].detach().cpu().numpy() < depth_threshold
         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
@@ -67,12 +68,13 @@ def apply_blur(image, blur_type, blur_strength):
     
     return output
 
+
 # Define Gradio interface
 iface = gr.Interface(
     fn=apply_blur,
     inputs=[
         gr.Image(label="Input Image"),
-        gr.Radio(["Gaussian", "Lens"], label="Blur Type"),
+        gr.Radio(["Gaussian", "Lens"], label="Blur Type", value="Gaussian"),
         gr.Slider(1, 30, value=15, step=1, label="Blur Strength")
     ],
     outputs=gr.Image(label="Output Image"),