src/test/try_depthanything.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from .depth_anything.dpt import DepthAnything
from PIL import Image
import os
import os.path as osp
import matplotlib.pyplot as plt
import numpy as np

class DepthAnythingWrapper(nn.Module):
    """
    A wrapper module for DepthAnything model with frozen parameters.
    It normalizes input using UniMatch-style normalization, resizes to multiples of 14,
    performs depth prediction, and resizes output back to original resolution.

    Args:
        encoder (str): one of 'vitl', 'vitb', 'vits'.
        checkpoint_path (str): path to the pretrained checkpoint (.pth file).
    """
    def __init__(self, encoder: str, checkpoint_path: str):
        super().__init__()
        # model configurations
        model_configs = {
            'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]},
            'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768]},
            'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]}
        }
        assert encoder in model_configs, f"Unsupported encoder: {encoder}"

        # load DepthAnything model
        self.depth_model = DepthAnything(model_configs[encoder])
        state_dict = torch.load(checkpoint_path, map_location='cpu')
        self.depth_model.load_state_dict(state_dict)
        self.depth_model.eval()

        # freeze parameters
        for param in self.depth_model.parameters():
            param.requires_grad = False

    def normalize_images(self, images: torch.Tensor) -> torch.Tensor:
        """
        Normalize images to match the pretrained UniMatch model.

        Args:
            images (torch.Tensor): input tensor of shape (B, V, C, H, W) or (B, C, H, W)
        Returns:
            torch.Tensor: normalized tensor with same shape as input.
        """
        # Determine extra dims before channel dim
        extras = images.dim() - 3
        shape = [1] * extras + [3, 1, 1]
        mean = torch.tensor([0.485, 0.456, 0.406], device=images.device).reshape(*shape)
        std  = torch.tensor([0.229, 0.224, 0.225], device=images.device).reshape(*shape)
        return (images - mean) / std
    

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass of the wrapper.

        Args:
            x (torch.Tensor): input tensor of shape (B, C, H, W) or (B, V, C, H, W), values in [0, 1].
        Returns:
            torch.Tensor: depth map of shape matching spatial dims of input, with channel=1.
        """
        # apply UniMatch-style normalization
        x = self.normalize_images(x)

        # if extra view dimension, merge
        merged = x
        merged_shape = merged.shape
        if merged.dim() == 5:
            B, V, C, H, W = merged_shape
            merged = merged.view(B * V, C, H, W)
        else:
            B, C, H, W = merged_shape

        # compute multiples of 14
        h14 = (H // 14) * 14
        w14 = (W // 14) * 14
        # resize to multiples of 14
        x_resized = F.interpolate(merged, size=(h14, w14), mode='bilinear', align_corners=True)

        # depth prediction
        with torch.no_grad():
            d = self.depth_model(x_resized)
            if d.dim() == 3:
                d = d.unsqueeze(1)
        # resize back to original
        d_final = F.interpolate(d, size=(H, W), mode='bilinear', align_corners=True)

        # normalize and invert depth
        depth_min = d_final.min()
        depth_max = d_final.max()
        
        #镜像
        d = depth_min + depth_max - d_final
        
        # 7. Clip to physical range [0.5, 200] meters
        d = torch.clamp(d, 0.5, 200.0)
        
        # un-merge view dimension
        if merged.dim() == 5:
            d = d.view(B, V, 1, H, W)
        return d


def visualize_depth(depth, output_path=None, cmap='viridis', vmin=None, vmax=None):
    """
    可视化深度图
    
    参数:
    depth: numpy数组 - 深度图数据 (H, W)
    output_path: str - 保存路径 (如果不提供则显示但不保存)
    cmap: str - 使用的色彩映射 ('jet', 'viridis', 'plasma'等)
    vmin, vmax: float - 颜色映射范围
    """
    # 确保深度图为2D数组
    if depth.ndim > 2:
        depth = depth.squeeze()
    
    plt.figure(figsize=(10, 8))
    
    # 设置颜色范围
    if vmin is None: vmin = np.nanmin(depth)
    if vmax is None: vmax = np.nanmax(depth)
    
    # 创建伪彩色图
    plt.imshow(depth, cmap=cmap, vmin=vmin, vmax=vmax)
    
    # 添加颜色条
    cbar = plt.colorbar()
    cbar.set_label('Depth (meters)', fontsize=12)
    
    # 设置标题和标签
    plt.title('Depth Map Visualization', fontsize=14)
    plt.axis('off')  # 不显示坐标轴
    
    # 保存或显示
    if output_path:
        os.makedirs(osp.dirname(output_path), exist_ok=True)
        plt.savefig(output_path, bbox_inches='tight', dpi=200)
        plt.close()
        print(f"深度图已保存至: {output_path}")
    else:
        plt.show()

def test_depth_wrapper():
    # 配置
    encoder = 'vitb'
    checkpoint_path = 'pretrained/pretrained_weights/depth_anything_v2_vitb.pth'
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    image_path = "/mnt/pfs/users/chaojun.ni/wangweijie_mnt/yeqing/BEV-Splat/checkpoints/dl3dv-256x448-depthsplat-base-randview2-6/images/dl3dv_14eb48a50e37df548894ab6d8cd628a21dae14bbe6c462e894616fc5962e6c49/color/000028_gt.png"
    
    # 实例化并移动到 device
    model = DepthAnythingWrapper(encoder, checkpoint_path).to(device)
    model.eval()

    # 1. 加载图像
    image = Image.open(image_path).convert('RGB')
    original_size = image.size  # 保存原始尺寸用于后续比较
    
    # 手动实现预处理 - 替换 torchvision.transforms
    # 1. 调整大小
    target_size = (448, 256)
    resized_image = image.resize(target_size, Image.BILINEAR)
    
    # 2. 转换为numpy数组并归一化到[0,1]
    image_array = np.array(resized_image).astype(np.float32) / 255.0
    
    # 3. 转换为PyTorch张量并调整维度顺序 (H, W, C) -> (C, H, W)
    image_tensor = torch.from_numpy(image_array).permute(2, 0, 1)
    
    # 4. 添加批次维度
    x = image_tensor.unsqueeze(0).to(device)
    
    print(f"输入张量尺寸: {x.shape} (设备: {x.device})")
    
    # 前向推理
    with torch.no_grad():
        depth = model(x)

    # 检查输出
    print(f"输入形状:  {x.shape}")
    print(f"输出形状:  {depth.shape}  (应为 [B,1,H,W])")

    # 打印深度图统计信息
    depth_np = depth.squeeze().cpu().numpy()  # 移除批次和通道维度 -> (H, W)
    
    print("\n深度图统计:")
    print(f"最小值: {depth_np.min():.4f} m, 最大值: {depth_np.max():.4f} m")
    print(f"均值:   {depth_np.mean():.4f} m, 标准差: {depth_np.std():.4f}")
    
    # =================== 新增深度图可视化 ===================
    # 1. 保存原始图像
    input_dir = './visualization'
    os.makedirs(input_dir, exist_ok=True)
    input_image_path = osp.join(input_dir, 'input_image.jpg')
    image.save(input_image_path)
    print(f"原始图像已保存至: {input_image_path}")
    
    # 2. 保存灰度深度图
    # 归一化深度图到0-255范围
    depth_min = depth_np.min()
    depth_max = depth_np.max()
    depth_normalized = (depth_np - depth_min) / (depth_max - depth_min + 1e-6)
    depth_grayscale = (depth_normalized * 255).astype(np.uint8)
    depth_pil = Image.fromarray(depth_grayscale)
    
    # 恢复原始尺寸（如果需要）
    depth_pil = depth_pil.resize(original_size, Image.NEAREST)
    
    # 保存灰度图
    depth_grayscale_path = osp.join(input_dir, 'depth_grayscale.jpg')
    depth_pil.save(depth_grayscale_path)
    print(f"灰度深度图已保存至: {depth_grayscale_path}")
    
   # 3. 创建伪彩色深度图 (Jet色彩映射)
    depth_colored_path = osp.join(input_dir, 'depth_colormap.png')

    # 创建无坐标轴的伪彩色图
    plt.figure(figsize=(10, 8))
    plt.imshow(depth_np, cmap='jet', vmin=depth_np.min(), vmax=depth_np.max())
    plt.axis('off')  # 关闭坐标轴
    plt.subplots_adjust(left=0, right=1, top=1, bottom=0)  # 移除所有边距
    plt.savefig(depth_colored_path, bbox_inches='tight', pad_inches=0, dpi=200)
    plt.close()
    print(f"伪彩色深度图已保存至: {depth_colored_path}")
    
    # 4. 显示深度图与原始图像的对比
    fig, axes = plt.subplots(1, 2, figsize=(15, 6))
    
    # 原始图像
    axes[0].imshow(image)
    axes[0].set_title('Original Image', fontsize=12)
    axes[0].axis('off')
    
    # 伪彩色深度图
    depth_colored = plt.imread(depth_colored_path)
    axes[1].imshow(depth_colored)
    axes[1].set_title('Depth Map (Jet Colormap)', fontsize=12)
    axes[1].axis('off')
    
    # 保存对比图
    comparison_path = osp.join(input_dir, 'depth_comparison.png')
    plt.savefig(comparison_path, bbox_inches='tight', dpi=150)
    plt.close()
    
    print(f"图像对比图已保存至: {comparison_path}")
    print("\n可视化已完成，结果保存在:", input_dir)

if __name__ == '__main__':
    test_depth_wrapper()


# from .depth_anything.dpt import DepthAnything
# import torch
# import torch.nn.functional as F

# model_configs = {
#     'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]},
#     'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768]},
#     'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]}
# }

# encoder = 'vitb'
# checkpoint_path = f'/mnt/pfs/users/wangweijie/yeqing/BEV-Splat/src/depth_anything/checkpoints/depth_anything_{encoder}14.pth'

# # 1. 创建模型并加载权重
# depth_anything = DepthAnything(model_configs[encoder])
# state_dict = torch.load(checkpoint_path, map_location='cpu')
# depth_anything.load_state_dict(state_dict)
# depth_anything.eval()

# # 2. 构造随机输入并做标准化
# test_input = torch.randn(6, 256, 448, 3).float()  # (B, H, W, C)
# test_input = test_input.permute(0, 3, 1, 2)        # -> (B, C, H, W)
# test_input = test_input / 255.0
# mean = torch.tensor([0.485, 0.456, 0.406]).view(1,3,1,1)
# std  = torch.tensor([0.229, 0.224, 0.225]).view(1,3,1,1)
# test_input = (test_input - mean) / std

# # 3. GPU 加速
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# depth_anything = depth_anything.to(device)
# test_input = test_input.to(device)

# # —— 新增部分 —— #
# # 4. 记录原始 H/W
# ori_h, ori_w = test_input.shape[-2:]          # 256, 448

# # 5. 计算向下取整到 14 的倍数
# resize_h = (ori_h // 14) * 14  # 252
# resize_w = (ori_w // 14) * 14  # 448

# # 6. 对输入做插值
# test_input_resized = F.interpolate(
#     test_input,
#     size=(resize_h, resize_w),
#     mode='bilinear',
#     align_corners=True
# )

# # 7. 推理
# with torch.no_grad():
#     depth_pred_resized = depth_anything(test_input_resized)  # (6,1,252,448)
#     if depth_pred_resized.dim() == 3:
#         depth_pred_resized = depth_pred_resized.unsqueeze(1)  # -> (B,1,H,W)

# # 8. 将预测结果插值回原始大小
# depth_pred = F.interpolate(
#     depth_pred_resized,
#     size=(ori_h, ori_w),
#     mode='bilinear',
#     align_corners=True
# )  # (6,1,256,448)

# # —— 恢复后续流程 —— #
# print("预测深度图形状:", depth_pred.shape)  # (6,1,256,448)

# depth_maps = depth_pred.squeeze(1).cpu().numpy()  # (6,256,448)
# print("第一幅深度图统计:")
# print(f"最小值: {depth_maps[0].min():.4f}, 最大值: {depth_maps[0].max():.4f}")
# print(f"均值:   {depth_maps[0].mean():.4f}, 标准差: {depth_maps[0].std():.4f}")