Update app.py

app.py

@@ -41,7 +41,11 @@ def segment_image(image):
         
     # Original dimensions
     original_size = image.size
-    model_image = image.resize((512, 512))
+    
+    # Use higher resolution for better results while staying within model limits
+    # Most models work well with 640x640
+    model_size = (640, 640)
+    model_image = image.resize(model_size, Image.LANCZOS)
     
     # Process image with model
     inputs = seg_processor(images=model_image, return_tensors="pt")
@@ -64,29 +68,30 @@ def segment_image(image):
             if cls_mask.sum() > binary_mask.sum():
                 binary_mask = cls_mask
     
-    # Improve mask with morphological operations
-    mask_small = Image.fromarray((binary_mask * 255).astype(np.uint8))
-    mask_cv = np.array(mask_small)
+    # Convert to uint8 for OpenCV processing
+    mask_cv = (binary_mask * 255).astype(np.uint8)
+    
+    # Apply morphological operations to clean up the mask
     kernel = np.ones((5, 5), np.uint8)
     mask_cv = cv2.morphologyEx(mask_cv, cv2.MORPH_CLOSE, kernel)
     mask_cv = cv2.morphologyEx(mask_cv, cv2.MORPH_OPEN, kernel)
     
-    # Apply Gaussian blur to smooth the edges
-    mask_cv = cv2.GaussianBlur(mask_cv, (9, 9), 0)
+    # Apply Gaussian blur to smooth the edges - less aggressive
+    mask_cv = cv2.GaussianBlur(mask_cv, (7, 7), 0)
     _, mask_cv = cv2.threshold(mask_cv, 128, 255, cv2.THRESH_BINARY)
     
-    # Resize back to original image size
-    mask_small = Image.fromarray(mask_cv)
-    mask_image = mask_small.resize(original_size, Image.BICUBIC)
+    # Resize back to original image size using bicubic interpolation for smoother results
+    mask_pil = Image.fromarray(mask_cv)
+    mask_resized = mask_pil.resize(original_size, Image.LANCZOS)
     
-    # Create binary mask
-    mask_array = np.array(mask_image) > 0
+    # Convert back to numpy
+    mask_array = np.array(mask_resized) > 128
     
-    # Create colored mask for visualization
-    mask_rgb = np.zeros((mask_array.shape[0], mask_array.shape[1], 3), dtype=np.uint8)
-    mask_rgb[:,:,0] = mask_array * 255  # Red channel for visualization
+    # Create visualization of mask (red on black background)
+    mask_viz = np.zeros((mask_array.shape[0], mask_array.shape[1], 3), dtype=np.uint8)
+    mask_viz[:,:,0] = mask_array * 255  # Red channel
     
-    return mask_array, mask_rgb
+    return mask_array, mask_viz
 
 # Function to apply Gaussian blur to background
 def apply_background_blur(image, mask, sigma=15):
@@ -106,9 +111,8 @@ def apply_background_blur(image, mask, sigma=15):
         binary_mask = mask > 0
     
     # Apply Gaussian blur to the entire image
-    blurred = np.zeros_like(image_array)
-    for c in range(3):
-        blurred[:, :, c] = gaussian_filter(image_array[:, :, c], sigma=sigma)
+    # Use OpenCV for better performance on larger images
+    blurred = cv2.GaussianBlur(image_array, (0, 0), sigma)
     
     # Combine original foreground with blurred background
     result = np.copy(blurred)
@@ -128,7 +132,9 @@ def apply_depth_based_blur(image, mask=None, max_sigma=15):
     
     # Original dimensions
     original_size = image.size
-    model_size = (512, 512)
+    
+    # Higher resolution for depth estimation
+    model_size = (640, 640)
     model_image = image.resize(model_size, Image.LANCZOS)
     
     # Process image for depth estimation
@@ -143,43 +149,56 @@ def apply_depth_based_blur(image, mask=None, max_sigma=15):
     depth = predicted_depth.squeeze().cpu().numpy()
     depth_map = (depth - depth.min()) / (depth.max() - depth.min())
     
-    # Resize depth map to match image size
-    depth_pil = Image.fromarray(depth_map)
-    depth_map_resized = np.array(depth_pil.resize(model_size, Image.LANCZOS))
+    # Create high-res depth map
+    depth_map_highres = cv2.resize(depth_map, (model_size[0], model_size[1]), interpolation=cv2.INTER_CUBIC)
     
     # Invert depth map (closer objects should be less blurred)
-    inverted_depth_map = 1.0 - depth_map_resized
+    inverted_depth_map = 1.0 - depth_map_highres
     
     # If mask is provided, ensure foreground is not blurred at all
     if mask is not None:
         # Resize mask to match model size
-        mask_pil = Image.fromarray((mask * 255).astype(np.uint8))
+        mask_pil = Image.fromarray((mask.astype(np.uint8) * 255))
         mask_resized = np.array(mask_pil.resize(model_size, Image.LANCZOS)) > 128
         # Set depth map to 0 (no blur) for foreground pixels
         inverted_depth_map = inverted_depth_map * (1 - mask_resized)
     
-    # Apply variable blur based on depth
-    original_array = np.array(model_image)
-    result_array = np.zeros_like(original_array)
+    # Convert to numpy array for processing
+    img_array = np.array(model_image)
     
-    # Apply blur with different intensities based on depth
-    for channel in range(3):
-        # Maximum blur
-        max_blurred = gaussian_filter(original_array[:, :, channel], sigma=max_sigma)
-        # Apply blur based on depth value
-        result_array[:, :, channel] = (1 - inverted_depth_map) * original_array[:, :, channel] + \
-                                      inverted_depth_map * max_blurred
+    # Create a progressive blur effect with multiple levels
+    result = np.copy(img_array)
+    
+    # Apply multiple blur levels for smoother transitions
+    num_levels = 8
+    for i in range(num_levels):
+        # Calculate blur sigma for this level
+        level_sigma = max_sigma * (i + 1) / num_levels
+        
+        # Create a blurred version of the image at this sigma level
+        level_blurred = cv2.GaussianBlur(img_array, (0, 0), level_sigma)
+        
+        # Calculate where to apply this blur level
+        depth_min = i / num_levels
+        depth_max = (i + 1) / num_levels
+        
+        # Create a mask for this depth range
+        level_mask = (inverted_depth_map >= depth_min) & (inverted_depth_map < depth_max)
+        
+        # Apply this blur level
+        for c in range(3):
+            result[:,:,c] = np.where(level_mask, level_blurred[:,:,c], result[:,:,c])
     
-    # Resize back to original image size
-    depth_blur = Image.fromarray(result_array.astype(np.uint8))
-    depth_blur_image = depth_blur.resize(original_size, Image.LANCZOS)
+    # Convert result back to PIL and resize to original dimensions
+    result_pil = Image.fromarray(result.astype(np.uint8))
+    result_resized = result_pil.resize(original_size, Image.LANCZOS)
     
     # Create colored depth map for visualization
     depth_map_colored = plt.cm.viridis(depth_map)[:, :, :3]
     depth_map_viz = Image.fromarray((depth_map_colored * 255).astype(np.uint8))
-    depth_map_image = depth_map_viz.resize(original_size, Image.LANCZOS)
+    depth_map_viz_resized = depth_map_viz.resize(original_size, Image.LANCZOS)
     
-    return np.array(depth_map_image), np.array(depth_blur_image)
+    return np.array(depth_map_viz_resized), np.array(result_resized)
 
 # Main processing function
 def process_image(input_image, blur_type="Gaussian Blur", blur_intensity=15):
@@ -207,7 +226,11 @@ def process_image(input_image, blur_type="Gaussian Blur", blur_intensity=15):
         if blur_type == "Gaussian Blur":
             # Apply regular Gaussian blur
             result = apply_background_blur(pil_img, mask_array, sigma=blur_intensity)
-            depth_viz = np.zeros_like(img)  # Placeholder for depth map
+            # Create placeholder for depth map (black image)
+            depth_viz = np.zeros_like(img)
+            # Add text saying "Depth map not used for Gaussian blur"
+            font = cv2.FONT_HERSHEY_SIMPLEX
+            cv2.putText(depth_viz, "Depth map not used", (50, 50), font, 1, (255, 255, 255), 2)
             
         else:  # "Depth-based Lens Blur"
             # Apply depth-based blur