Update app.py

app.py

@@ -5,7 +5,6 @@ Created on Thu Mar 27 13:56:42 2025
 
 @author: perghect
 """
-
 import gradio as gr
 import requests
 import io
@@ -15,6 +14,14 @@ from PIL import Image, ImageFilter
 from torchvision import transforms
 from transformers import AutoModelForImageSegmentation, AutoImageProcessor, AutoModelForDepthEstimation
 
+# Verify kornia and timm are available
+try:
+    import kornia
+    import timm
+    print("kornia and timm are successfully installed!")
+except ImportError as e:
+    print(f"Error: {e}")
+
 # Set device and precision
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 torch.set_float32_matmul_precision('high')
@@ -34,29 +41,32 @@ def load_image_from_link(url: str) -> Image.Image:
 # Gaussian Blur Functions
 def run_rmbg(image: Image.Image, threshold=0.5):
     """Runs the RMBG-2.0 model on the image and returns a binary mask."""
-    image_size = (1024, 1024)
-    transform_image = transforms.Compose([
-        transforms.Resize(image_size),
-        transforms.ToTensor(),
-        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
-    ])
-
-    input_images = transform_image(image).unsqueeze(0).to(device)
-
-    with torch.no_grad():
-        preds = rmbg_model(input_images)
-        if isinstance(preds, list):
-            mask_logits = preds[-1]
-        else:
-            raise ValueError(f"Unexpected output format: {type(preds)}")
-
-    mask_prob = mask_logits.sigmoid().cpu()[0].squeeze()
-    pred_pil = transforms.ToPILImage()(mask_prob)
-    mask_pil = pred_pil.resize(image.size, resample=Image.BILINEAR)
-
-    mask_np = np.array(mask_pil, dtype=np.uint8) / 255.0
-    binary_mask = (mask_np > threshold).astype(np.uint8)
-    return binary_mask
+    try:
+        image_size = (1024, 1024)
+        transform_image = transforms.Compose([
+            transforms.Resize(image_size),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+        ])
+
+        input_images = transform_image(image).unsqueeze(0).to(device)
+
+        with torch.no_grad():
+            preds = rmbg_model(input_images)
+            if isinstance(preds, list):
+                mask_logits = preds[-1]
+            else:
+                raise ValueError(f"Unexpected output format: {type(preds)}")
+
+        mask_prob = mask_logits.sigmoid().cpu()[0].squeeze()
+        pred_pil = transforms.ToPILImage()(mask_prob)
+        mask_pil = pred_pil.resize(image.size, resample=Image.BILINEAR)
+
+        mask_np = np.array(mask_pil, dtype=np.uint8) / 255.0
+        binary_mask = (mask_np > threshold).astype(np.uint8)
+        return binary_mask
+    except Exception as e:
+        raise Exception(f"Error in background removal: {str(e)}")
 
 def apply_background_blur(image: Image.Image, mask: np.ndarray, sigma: float = 15):
     """Applies a Gaussian blur to the background while keeping the foreground sharp."""
@@ -73,28 +83,33 @@ def apply_background_blur(image: Image.Image, mask: np.ndarray, sigma: float = 1
 # Lens Blur Functions
 def run_depth_estimation(image: Image.Image, target_size=(512, 512)):
     """Runs the Depth-Anything-V2-Small model and returns the depth map."""
-    image_resized = image.resize(target_size, resample=Image.BILINEAR)
-    inputs = depth_processor(images=image_resized, return_tensors="pt").to(device)
-
-    with torch.no_grad():
-        outputs = depth_model(**inputs)
-        predicted_depth = outputs.predicted_depth
-
-    prediction = torch.nn.functional.interpolate(
-        predicted_depth.unsqueeze(1),
-        size=image.size[::-1],
-        mode="bicubic",
-        align_corners=False,
-    )
-
-    depth_map = prediction.squeeze().cpu().numpy()
-    depth_max = depth_map.max()
-    depth_min = depth_map.min()
-    depth_map = (depth_map - depth_min) / (depth_max - depth_min)
-    depth_map = 1 - depth_map  # Invert: higher values = farther
-    return depth_map
-
-def apply_depth_based_blur(image: Image.Image, depth_map: np.ndarray, max_radius: float = 30, foreground_percentile: float = 5.0):
+    try:
+        image_resized = image.resize(target_size, resample=Image.BILINEAR)
+        inputs = depth_processor(images=image_resized, return_tensors="pt").to(device)
+
+        with torch.no_grad():
+            outputs = depth_model(**inputs)
+            predicted_depth = outputs.predicted_depth
+
+        prediction = torch.nn.functional.interpolate(
+            predicted_depth.unsqueeze(1),
+            size=image.size[::-1],
+            mode="bicubic",
+            align_corners=False,
+        )
+
+        depth_map = prediction.squeeze().cpu().numpy()
+        depth_max = depth_map.max()
+        depth_min = depth_map.min()
+        if depth_max == depth_min:
+            depth_max = depth_min + 1e-6  # Avoid division by zero
+        depth_map = (depth_map - depth_min) / (depth_max - depth_min)
+        depth_map = 1 - depth_map  # Invert: higher values = farther
+        return depth_map
+    except Exception as e:
+        raise Exception(f"Error in depth estimation: {str(e)}")
+
+def apply_depth_based_blur(image: Image.Image, depth_map: np.ndarray, max_radius: float = 15, foreground_percentile: float = 30):
     """Applies a variable Gaussian blur based on the depth map."""
     image_np = np.array(image)
 
@@ -140,36 +155,76 @@ def apply_depth_based_blur(image: Image.Image, depth_map: np.ndarray, max_radius
     return output_image
 
 # Main Processing Function for Gradio
-def process_image(image, blur_type, sigma=15, max_radius=30, foreground_percentile=5.0):
-    """Processes the image based on the selected blur type and parameters."""
+def process_image(image, blur_type, sigma=15, max_radius=15, foreground_percentile=30):
+    """Processes the image based on the selected blur type."""
     if image is None:
         return None, "Please upload an image."
 
-    image = Image.fromarray(image).convert("RGB")
-
-    if blur_type == "Gaussian Blur":
-        mask = run_rmbg(image, threshold=0.5)
-        output_image = apply_background_blur(image, mask, sigma=sigma)
-        title = f"Gaussian Blur (sigma={sigma})"
-    else:  # Lens Blur
-        depth_map = run_depth_estimation(image, target_size=(512, 512))
-        output_image = apply_depth_based_blur(image, depth_map, max_radius=max_radius, foreground_percentile=foreground_percentile)
-        title = f"Lens Blur (max_radius={max_radius}, foreground_percentile={foreground_percentile})"
+    try:
+        image = Image.fromarray(image).convert("RGB")
+    except Exception as e:
+        return None, f"Error processing image: {str(e)}"
+
+    # Resize image if too large
+    max_size = (1024, 1024)
+    if image.size[0] > max_size[0] or image.size[1] > max_size[1]:
+        image.thumbnail(max_size, Image.Resampling.LANCZOS)
+
+    try:
+        if blur_type == "Gaussian Blur":
+            mask = run_rmbg(image, threshold=0.5)
+            output_image = apply_background_blur(image, mask, sigma=sigma)
+            title = f"Gaussian Blur (sigma={sigma})"
+        else:  # Lens Blur
+            depth_map = run_depth_estimation(image, target_size=(512, 512))
+            output_image = apply_depth_based_blur(image, depth_map, max_radius=max_radius, foreground_percentile=foreground_percentile)
+            title = f"Lens Blur (max_radius={max_radius}, foreground_percentile={foreground_percentile})"
+    except Exception as e:
+        return None, f"Error applying blur: {str(e)}"
 
     return output_image, title
 
-# Gradio Interface
+# Gradio Interface with Conditional Parameter Display
 with gr.Blocks() as demo:
     gr.Markdown("# Image Blur Effects with Gaussian and Lens Blur")
-    gr.Markdown("Upload an image and choose a blur effect. Adjust the parameters to customize the blur intensity.")
+    gr.Markdown("""
+    This app applies blur effects to your images. Follow these steps to use it:
+
+    **Note**: This app is hosted on Hugging Face Spaces’ free tier and may go to "Sleeping" mode after 48 hours of inactivity. If it doesn’t load immediately, please wait a few seconds while it wakes up.
+
+    1. **Upload an Image**: Click the "Upload Image" box to upload an image from your device.
+    2. **Choose a Blur Type**:
+       - **Gaussian Blur**: Applies a uniform blur to the background, keeping the foreground sharp. Adjust the sigma parameter to control blur intensity.
+       - **Lens Blur**: Applies a depth-based blur, simulating a depth-of-field effect (closer objects are sharp, farther objects are blurred). Adjust the max radius and foreground percentile to fine-tune the effect.
+    3. **Adjust Parameters**:
+       - For Gaussian Blur, use the "Gaussian Blur Sigma" slider to control blur intensity (higher values = more blur).
+       - For Lens Blur, use the "Max Blur Radius" slider to control the maximum blur intensity and the "Foreground Percentile" slider to adjust the depth threshold for the foreground.
+    4. **Apply the Blur**: Click the "Apply Blur" button to process the image.
+    5. **View the Result**: The processed image will appear in the "Output Image" box, along with a description of the effect applied.
+
+    **Example**: Try uploading an image with a clear foreground and background (e.g., a person in front of a landscape) to see the effects in action.
+    """)
 
     with gr.Row():
         image_input = gr.Image(label="Upload Image", type="numpy")
         with gr.Column():
             blur_type = gr.Radio(choices=["Gaussian Blur", "Lens Blur"], label="Blur Type", value="Gaussian Blur")
-            sigma = gr.Slider(minimum=1, maximum=50, step=1, value=15, label="Gaussian Blur Sigma (for Gaussian Blur)")
-            max_radius = gr.Slider(minimum=1, maximum=50, step=1, value=15, label="Max Blur Radius (for Lens Blur)")
-            foreground_percentile = gr.Slider(minimum=1, maximum=50, step=1, value=30, label="Foreground Percentile (for Lens Blur)")
+            sigma = gr.Slider(minimum=1, maximum=50, step=1, value=15, label="Gaussian Blur Sigma", visible=True)
+            max_radius = gr.Slider(minimum=1, maximum=50, step=1, value=15, label="Max Blur Radius (for Lens Blur)", visible=False)
+            foreground_percentile = gr.Slider(minimum=1, maximum=50, step=1, value=30, label="Foreground Percentile (for Lens Blur)", visible=False)
+
+    # Update visibility of parameters based on blur type
+    def update_visibility(blur_type):
+        if blur_type == "Gaussian Blur":
+            return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)
+        else:  # Lens Blur
+            return gr.update(visible=False), gr.update(visible=True), gr.update(visible=True)
+
+    blur_type.change(
+        fn=update_visibility,
+        inputs=blur_type,
+        outputs=[sigma, max_radius, foreground_percentile]
+    )
 
     process_button = gr.Button("Apply Blur")
     with gr.Row():