distill-any-depth

app.py

@@ -1,32 +1,42 @@
 import gradio as gr
 import torch
 from PIL import Image
-import cv2
 import numpy as np
 from distillanydepth.modeling.archs.dam.dam import DepthAnything
-from distillanydepth.utils.image_util import colorize_depth_maps
+from distillanydepth.utils.image_util import chw2hwc, colorize_depth_maps
 from distillanydepth.midas.transforms import Resize, NormalizeImage, PrepareForNet
 from torchvision.transforms import Compose
-import os
+import cv2
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file  # 导入 safetensors 库
 
-# Helper function to load model (same as your original code)
+# Helper function to load model from Hugging Face
 def load_model_by_name(arch_name, checkpoint_path, device):
+    model = None
     if arch_name == 'depthanything':
-        if '.safetensors' in checkpoint_path:
-            model = DepthAnything.from_pretrained(os.path.dirname(checkpoint_path)).to(device)
-        else:
-            raise NotImplementedError("Model architecture not implemented.")
+        # 使用 safetensors 加载模型权重
+        model_weights = load_file(checkpoint_path)  # safetensors 加载方式
+        
+        # 初始化模型
+        model = DepthAnything(checkpoint_path=None).to(device)
+        model.load_state_dict(model_weights)  # 将加载的权重应用到模型
+
+        model = model.to(device)  # 确保模型在正确的设备上
     else:
         raise NotImplementedError(f"Unknown architecture: {arch_name}")
     return model
 
-# Image processing function (same as your original code, modified for Gradio)
+# Image processing function
 def process_image(image, model, device):
+    if model is None:
+        return None
+    
     # Preprocess the image
     image_np = np.array(image)[..., ::-1] / 255
+    
     transform = Compose([
-        Resize(512, 512, resize_target=None, keep_aspect_ratio=False, ensure_multiple_of=32, image_interpolation_method=cv2.INTER_CUBIC),
-        NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+        Resize(756, 756, resize_target=False, keep_aspect_ratio=True, ensure_multiple_of=14, resize_method='lower_bound', image_interpolation_method=cv2.INTER_CUBIC),
+        NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
         PrepareForNet()
     ])
     
@@ -35,31 +45,74 @@ def process_image(image, model, device):
     
     with torch.no_grad():  # Disable autograd since we don't need gradients on CPU
         pred_disp, _ = model(image_tensor)
-        pred_disp_np = pred_disp.cpu().detach().numpy()[0, :, :, :].transpose(1, 2, 0)
-        pred_disp = (pred_disp_np - pred_disp_np.min()) / (pred_disp_np.max() - pred_disp_np.min())
 
+    # Ensure the depth map is in the correct shape before colorization
+    pred_disp_np = pred_disp.cpu().detach().numpy()[0, 0, :, :]  # Remove extra singleton dimensions
+    
+    # Normalize depth map
+    pred_disp = (pred_disp_np - pred_disp_np.min()) / (pred_disp_np.max() - pred_disp_np.min())
+    
     # Colorize depth map
-    cmap = "Spectral_r"  # Default colormap for relative depth
-    depth_colored = colorize_depth_maps(pred_disp[None, ...], 0, 1, cmap=cmap).squeeze()
+    cmap = "Spectral_r"
+    depth_colored = colorize_depth_maps(pred_disp[None, ..., None], 0, 1, cmap=cmap).squeeze()  # Ensure correct dimension
+    
+    # Convert to uint8 for image display
     depth_colored = (depth_colored * 255).astype(np.uint8)
     
-    depth_image = Image.fromarray(depth_colored)
+    # Convert to HWC format (height, width, channels)
+    depth_colored_hwc = chw2hwc(depth_colored)
+    
+    # Resize to match the original image dimensions (height, width)
+    h, w = image_np.shape[:2]
+    depth_colored_hwc = cv2.resize(depth_colored_hwc, (w, h), cv2.INTER_LINEAR)
+    
+    # Convert to a PIL image
+    depth_image = Image.fromarray(depth_colored_hwc)
     return depth_image
 
 # Gradio interface function
 def gradio_interface(image):
-    # Set device to CPU explicitly
-    device = torch.device("cpu")  # Force using CPU
-    model = load_model_by_name("depthanything", "your_checkpoint_path_here", device)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    model_kwargs = dict(
+        vitb=dict(
+            encoder='vitb',
+            features=128,
+            out_channels=[96, 192, 384, 768],
+        ),
+        vitl=dict(
+            encoder="vitl", 
+            features=256, 
+            out_channels=[256, 512, 1024, 1024], 
+            use_bn=False, 
+            use_clstoken=False, 
+            max_depth=150.0, 
+            mode='disparity',
+            pretrain_type='dinov2',
+            del_mask_token=False
+        )
+    )
+    # Load model
+    model = DepthAnything(**model_kwargs['vitl']).to(device)
+    checkpoint_path = hf_hub_download(repo_id=f"xingyang1/Distill-Any-Depth", filename=f"large/model.safetensors", repo_type="model")
 
+    # 使用 safetensors 加载模型权重
+    model_weights = load_file(checkpoint_path)  # safetensors 加载方式
+    model.load_state_dict(model_weights)
+    model = model.to(device)  # 确保模型在正确的设备上
+    
+    if model is None:
+        return None
+    
     # Process image and return output
-    return process_image(image, model, device)
+    depth_image = process_image(image, model, device)
+    return depth_image
 
 # Create Gradio interface
 iface = gr.Interface(
     fn=gradio_interface,
     inputs=gr.Image(type="pil"),  # Only image input, no mode selection
-    outputs=gr.Image(type="pil"),
+    outputs=gr.Image(type="pil"),  # Only depth image output, no debug info
     title="Depth Estimation Demo",
     description="Upload an image to see the depth estimation results."
 )