Upload app.py

app.py

@@ -3,42 +3,27 @@ import numpy as np
 import torch
 from PIL import Image
 import cv2
-import os  # Import the os module
+import requests
 
-from transformers import DPTImageProcessor, DPTForDepthEstimation
+from transformers import pipeline
 
-# Access the token from the Space's secrets
-HF_TOKEN = os.environ.get("HUGGING_FACE_TOKEN")
-if HF_TOKEN is None:
-    raise ValueError("Hugging Face token not found. Make sure you've added it as a secret to your Space.")
-
-# Load the image processor and model
-image_processor = DPTImageProcessor.from_pretrained("Intel/dpt-hybrid-midas", token=HF_TOKEN)
-model = DPTForDepthEstimation.from_pretrained("Intel/dpt-hybrid-midas", low_cpu_mem_usage=True, token=HF_TOKEN)
-model.to('cuda')
-model.eval()
+# Load the depth estimation pipeline
+pipe = pipeline(task="depth-estimation", model="depth-anything/Depth-Anything-V2-Small-hf")
 
 def apply_depth_aware_blur(image, foreground_blur, midground_blur, background_blur, foreground_threshold, midground_lower, midground_upper, background_threshold):
     original_image = Image.fromarray(image).convert("RGB")
     original_image = original_image.resize((512, 512))
     image_np = np.array(original_image)
 
-    inputs = image_processor(images=original_image, return_tensors="pt").to('cuda')
-
-    with torch.no_grad():
-        outputs = model(**inputs)
-        predicted_depth = outputs.predicted_depth
-
-    prediction = torch.nn.functional.interpolate(
-        predicted_depth.unsqueeze(1),
-        size=(512, 512),
-        mode="bicubic",
-        align_corners=False,
-    )
-
-    depth_map = prediction.squeeze().cpu().numpy()
-    normalized_depth_map = (depth_map - np.min(depth_map)) / (np.max(depth_map) - np.min(depth_map))
+    # Inference
+    depth = pipe(original_image)["depth"]
+    depth = np.array(depth) # Convert to numpy array
+    depth = cv2.resize(depth, (512, 512), interpolation=cv2.INTER_CUBIC) # Resize depth map
+    
+    # Normalize the depth map
+    normalized_depth_map = (depth - np.min(depth)) / (np.max(depth) - np.min(depth))
 
+    # Create masks
     foreground_mask = (normalized_depth_map < foreground_threshold).astype(np.uint8) * 255
     midground_mask = ((normalized_depth_map >= midground_lower) & (normalized_depth_map < midground_upper)).astype(np.uint8) * 255
     background_mask = (normalized_depth_map >= background_threshold).astype(np.uint8) * 255
@@ -65,11 +50,10 @@ iface = gr.Interface(
         gr.Slider(0, 1, label="Midground Lower Threshold", default=0.2),
         gr.Slider(0, 1, label="Midground Upper Threshold", default=0.6),
         gr.Slider(0, 1, label="Background Threshold", default=0.6)
-
     ],
     outputs=gr.Image(label="Blurred Image"),
     title="Depth-Aware Lens Blur App",
-    description="Apply depth-based blur to uploaded images. Adjust blur intensity for foreground, midground, and background.",
+    description="Apply depth-based blur to uploaded images using Depth Anything V2. Adjust blur intensity for foreground, midground, and background.",
 )
 
 if __name__ == "__main__":