mul / utils /distance.py

Upload folder using huggingface_hub

90ca518 verified 16 days ago

3.65 kB

"""点云距离与法向一致性指标(倒角 L1/L2、豪斯多夫、|NC| / NC)。
modules.pipeline 与 scripts.eval 共用同一份实现, 保证推理内打印的指标与离线评测同口径。
"""
import numpy as np
from scipy.spatial import cKDTree


def calculate_distances_numpy(pcd1, pcd2):
    """
    [已修改]
    使用 NumPy 高效计算 L2倒角距离、L1倒角距离 和 豪斯多夫距离。
    参数:
    pcd1: np.ndarray, 形状为 (N, 3)，第一个点云。
    pcd2: np.ndarray, 形状为 (M, 3)，第二个点云。
    返回:
    (float, float, float): (L2倒角距离, 豪斯多夫距离, L1倒角距离)
    """
    # --- 计算所有点对之间的差异 ---
    pcd1_expanded = pcd1[:, np.newaxis, :]
    pcd2_expanded = pcd2[np.newaxis, :, :]
    diff = pcd1_expanded - pcd2_expanded

    # --- L2 (欧氏) 距离计算 ---
    dist_matrix_sq = np.sum(diff**2, axis=2)

    # --- L2 倒角距离 (Chamfer Distance L2) ---
    min_dists_1_to_2_l2 = np.sqrt(np.min(dist_matrix_sq, axis=1))
    min_dists_2_to_1_l2 = np.sqrt(np.min(dist_matrix_sq, axis=0))
    chamfer_dist_l2 = np.mean(min_dists_1_to_2_l2) + np.mean(min_dists_2_to_1_l2)

    # --- L2 豪斯多夫距离 (Hausdorff Distance) ---
    hausdorff_dist = max(np.max(min_dists_1_to_2_l2), np.max(min_dists_2_to_1_l2))
    
    # --- L1 (曼哈顿) 距离计算 ---
    dist_matrix_l1 = np.sum(np.abs(diff), axis=2)
    
    # --- L1 倒角距离 (Chamfer Distance L1) ---
    min_dists_1_to_2_l1 = np.min(dist_matrix_l1, axis=1)
    min_dists_2_to_1_l1 = np.min(dist_matrix_l1, axis=0)
    chamfer_dist_l1 = np.mean(min_dists_1_to_2_l1) + np.mean(min_dists_2_to_1_l1)

    return chamfer_dist_l2, hausdorff_dist, chamfer_dist_l1


def calculate_normal_consistency_score(pcd_gt, normal_gt, pcd_pred, normal_pred):
    """
    【修改版】根据提供的公式计算法向量一致性得分 (Normal Consistency Score)。
    返回两个值：
    1. |NC| (绝对值 NC): np.mean( |n1·n2| )
    2.  NC  (真实值 NC): np.mean( n1·n2 )
    """

    def _calculate_one_way_score(source_pcd, source_normals, target_pcd, target_normals):
        """
        辅助函数：计算从 source 到 target 的单向 NC 得分。
        返回 (mean_abs_dot, mean_real_dot)
        """
        # 1. 计算 source 到 target 的所有点对的距离矩阵
        diff = source_pcd[:, np.newaxis, :] - target_pcd[np.newaxis, :, :]
        dist_sq = np.sum(diff**2, axis=2)

        # 2. 为每个 source 点找到 target 中的最近邻点的索引
        indices = np.argmin(dist_sq, axis=1)

        # 3. 获取这些最近邻点对应的法向量
        corresponding_normals = target_normals[indices]

        # 4. 计算 source 法向量与对应 target 法向量的点积
        dot_products = np.einsum('ij,ij->i', source_normals, corresponding_normals)
        
        # 5. [修改] 返回点积绝对值的平均值 和 真实点积的平均值
        return np.mean(np.abs(dot_products)), np.mean(dot_products)

    # --- Part 1: 计算从 GT -> Pred 的得分 ---
    score_gt_to_pred_abs, score_gt_to_pred_real = _calculate_one_way_score(pcd_gt, normal_gt, pcd_pred, normal_pred)
    
    # --- Part 2: 计算从 Pred -> GT 的得分 ---
    score_pred_to_gt_abs, score_pred_to_gt_real = _calculate_one_way_score(pcd_pred, normal_pred, pcd_gt, normal_gt)
    
    # --- Final Score: 结合两部分 ---
    final_score_abs = (score_gt_to_pred_abs + score_pred_to_gt_abs) / 2.0
    final_score_real = (score_gt_to_pred_real + score_pred_to_gt_real) / 2.0
    
    # [修改] 返回两个分数
    return final_score_abs, final_score_real

	"""点云距离与法向一致性指标(倒角 L1/L2、豪斯多夫、\|NC\| / NC)。

	modules.pipeline 与 scripts.eval 共用同一份实现, 保证推理内打印的指标与离线评测同口径。
	"""
	import numpy as np
	from scipy.spatial import cKDTree


	def calculate_distances_numpy(pcd1, pcd2):
	"""
	[已修改]
	使用 NumPy 高效计算 L2倒角距离、L1倒角距离和豪斯多夫距离。

	参数:
	pcd1: np.ndarray, 形状为 (N, 3)，第一个点云。
	pcd2: np.ndarray, 形状为 (M, 3)，第二个点云。

	返回:
	(float, float, float): (L2倒角距离, 豪斯多夫距离, L1倒角距离)
	"""
	# --- 计算所有点对之间的差异 ---
	pcd1_expanded = pcd1[:, np.newaxis, :]
	pcd2_expanded = pcd2[np.newaxis, :, :]
	diff = pcd1_expanded - pcd2_expanded

	# --- L2 (欧氏) 距离计算 ---
	dist_matrix_sq = np.sum(diff**2, axis=2)

	# --- L2 倒角距离 (Chamfer Distance L2) ---
	min_dists_1_to_2_l2 = np.sqrt(np.min(dist_matrix_sq, axis=1))
	min_dists_2_to_1_l2 = np.sqrt(np.min(dist_matrix_sq, axis=0))
	chamfer_dist_l2 = np.mean(min_dists_1_to_2_l2) + np.mean(min_dists_2_to_1_l2)

	# --- L2 豪斯多夫距离 (Hausdorff Distance) ---
	hausdorff_dist = max(np.max(min_dists_1_to_2_l2), np.max(min_dists_2_to_1_l2))

	# --- L1 (曼哈顿) 距离计算 ---
	dist_matrix_l1 = np.sum(np.abs(diff), axis=2)

	# --- L1 倒角距离 (Chamfer Distance L1) ---
	min_dists_1_to_2_l1 = np.min(dist_matrix_l1, axis=1)
	min_dists_2_to_1_l1 = np.min(dist_matrix_l1, axis=0)
	chamfer_dist_l1 = np.mean(min_dists_1_to_2_l1) + np.mean(min_dists_2_to_1_l1)

	return chamfer_dist_l2, hausdorff_dist, chamfer_dist_l1


	def calculate_normal_consistency_score(pcd_gt, normal_gt, pcd_pred, normal_pred):
	"""
	【修改版】根据提供的公式计算法向量一致性得分 (Normal Consistency Score)。
	返回两个值：
	1. \|NC\| (绝对值 NC): np.mean( \|n1·n2\| )
	2. NC (真实值 NC): np.mean( n1·n2 )
	"""

	def _calculate_one_way_score(source_pcd, source_normals, target_pcd, target_normals):
	"""
	辅助函数：计算从 source 到 target 的单向 NC 得分。
	返回 (mean_abs_dot, mean_real_dot)
	"""
	# 1. 计算 source 到 target 的所有点对的距离矩阵
	diff = source_pcd[:, np.newaxis, :] - target_pcd[np.newaxis, :, :]
	dist_sq = np.sum(diff**2, axis=2)

	# 2. 为每个 source 点找到 target 中的最近邻点的索引
	indices = np.argmin(dist_sq, axis=1)

	# 3. 获取这些最近邻点对应的法向量
	corresponding_normals = target_normals[indices]

	# 4. 计算 source 法向量与对应 target 法向量的点积
	dot_products = np.einsum('ij,ij->i', source_normals, corresponding_normals)

	# 5. [修改] 返回点积绝对值的平均值和真实点积的平均值
	return np.mean(np.abs(dot_products)), np.mean(dot_products)

	# --- Part 1: 计算从 GT -> Pred 的得分 ---
	score_gt_to_pred_abs, score_gt_to_pred_real = _calculate_one_way_score(pcd_gt, normal_gt, pcd_pred, normal_pred)

	# --- Part 2: 计算从 Pred -> GT 的得分 ---
	score_pred_to_gt_abs, score_pred_to_gt_real = _calculate_one_way_score(pcd_pred, normal_pred, pcd_gt, normal_gt)

	# --- Final Score: 结合两部分 ---
	final_score_abs = (score_gt_to_pred_abs + score_pred_to_gt_abs) / 2.0
	final_score_real = (score_gt_to_pred_real + score_pred_to_gt_real) / 2.0

	# [修改] 返回两个分数
	return final_score_abs, final_score_real