import torch
import torch.nn as nn


class HausdorffDistanceLoss(nn.Module):
    def __init__(self):
        super(HausdorffDistanceLoss, self).__init__()

    def forward(self, pred, target):
        # 扩展为 (B, N, 1, D) 和 (B, 1, M, D)
        pred = pred.unsqueeze(1)  # (N, D) -> (N, 1, D)
        target = target.unsqueeze(0)  # (M, D) -> (1, M, D)

        # 计算两组点之间的 pairwise 距离
        distances = torch.norm(pred - target, dim=2)  # (N, M)

        # Hausdorff 距离的两个方向
        h_a_to_b = distances.min(dim=1)[0].max()  # a -> b 距离
        h_b_to_a = distances.min(dim=0)[0].max()  # b -> a 距离

        return (h_a_to_b + h_b_to_a) / 2  # 返回 Hausdorff 距离


class OptimizedHausdorffDistanceLoss(nn.Module):
    def __init__(self):
        super(OptimizedHausdorffDistanceLoss, self).__init__()

    def forward(self, pred, target):
        """
        pred: (B, N, D) - 预测点集
        target: (B, M, D) - 目标点集
        B: 批量大小
        N: 每个预测点集的点数
        M: 每个目标点集的点数
        D: 点的维度
        """
        # 确保两个张量的数据类型一致
        if pred.dtype != target.dtype:
            target = target.to(pred.dtype)

        # 使用torch.cdist计算距离矩阵
        distances = torch.cdist(pred, target)  # (B, N, M)

        # 从pred到target的最大最小距离
        h_a_to_b = distances.min(dim=1)[0].max(dim=0)[0]  # (B,)
        # 从target到pred的最大最小距离
        h_b_to_a = distances.min(dim=1)[0].max(dim=0)[0]  # (B,)

        # 返回Hausdorff距离作为损失值
        return ((h_a_to_b + h_b_to_a) / 2).mean()  # 返回批量的平均损失


def euclidean_distance_mean(pred, target):
    """
    计算两个点集之间的欧氏距离，并返回平均值。

    参数：
    pred: Tensor (N, D) - 预测点集
    target: Tensor (M, D) - 目标点集

    返回：
    mean_distance: float - 欧氏距离的平均值
    """
    # 计算两个点集之间的 pairwise 欧氏距离
    distances = torch.cdist(pred, target, p=2)  # 使用 p=2 计算欧氏距离

    # 计算 pairwise 距离的平均值
    mean_distance = distances.mean()

    return mean_distance

# 测试部分
if __name__ == "__main__":
    # 创建两个点集，每个点集只有两个点
    pred = torch.tensor([[1.0, 2.0, 3.0], [3.0, 4.0, 5.0]])
    target = torch.tensor([[2.0, 3.0, 3.0], [3.0, 4.0, 5.0]])

    # 初始化损失函数
    loss_fn = HausdorffDistanceLoss()
    loss_fn2 = OptimizedHausdorffDistanceLoss()

    # 计算损失值
    loss = loss_fn(pred, target)
    loss2 = loss_fn2(pred, target)

    print("Hausdorff Distance Loss:", loss.item())
    print("Optimized Hausdorff Distance Loss:", loss2.item())

    # 计算欧氏距离的平均值
    mean_distance = euclidean_distance_mean(pred, target)

    print("Euclidean Distance Mean:", mean_distance.item())

D:\ProgramData\Anaconda3\envs\torch38\python.exe D:\Desktop\Aidyna\test2.py 
Hausdorff Distance Loss: 1.4142135381698608
Optimized Hausdorff Distance Loss: 1.4142135381698608
Euclidean Distance Mean: 1.8319511413574219

进程已结束，退出代码为 0