pmd-python/src/eval.py

import math
import json
import cv2
import os
from tqdm import tqdm
import numpy as np
import scipy.io as sio
import matplotlib.pyplot as plt
from scipy.spatial import cKDTree


def point_to_plane_distance(x, y, z, a, b, c):
    z_plane = a * x + b * y + c  
    distances = np.abs(z - z_plane) 
    return distances

def fit_plane(x, y, z):
    A = np.c_[x, y, np.ones(x.shape)]
    C, _, _, _ = np.linalg.lstsq(A, z, rcond=None)  
    return C  


def remove_noise(x, y, z, a, b, c, thresh=0.01):

    distances = point_to_plane_distance(x, y, z, a, b, c)
    mask = distances < thresh  
    mask1 = distances >= thresh
    x_keep = x[mask]
    y_keep = y[mask]
    z_keep = z[mask]
    z_mean = np.mean(z_keep)
    arr_keep = np.column_stack((x_keep,y_keep,z_keep))
    x_fliterd = x[mask1]
    y_fliterd = y[mask1]
    z_fliterd = z[mask1]
    arr_filterd = np.column_stack((x_fliterd,y_fliterd,z_fliterd))
   

    # fig = plt.figure(figsize=(10, 7))
    # ax = fig.add_subplot(111, projection='3d')
    # scatter = ax.scatter(x_keep, y_keep, z_keep, c=z_keep, cmap='viridis', marker='o')

    # cbar = plt.colorbar(scatter, ax=ax, shrink=0.5, aspect=5)
    # cbar.set_label('Z values')

    # ax.set_xlabel('X Axis')
    # ax.set_ylabel('Y Axis')
    # ax.set_zlabel('Z Axis')
    # ax.set_title('Filtered 3D Point Cloud (Outliers Removed)')

    # plt.show()
    return np.vstack([arr_keep, arr_filterd])

# def point_to_line_distance(point, A, B, C):
#     distance = abs(A * point[0] + B * point[1] + C) / np.sqrt(A**2 + B**2)
#     return distance

def get_support_points(x_filtered, y_filtered, z_filtered):
    support_points = []
    center_x = np.mean(x_filtered)
    center_y = np.mean(y_filtered)
    num_regions = 3
    angle_step = 2 * np.pi / num_regions  
    angles = np.arctan2(y_filtered - center_y, x_filtered - center_x) 
    distances = np.sqrt((x_filtered - center_x)**2 + (y_filtered - center_y)**2)  

    for i in range(num_regions):

        angle_min = -np.pi + i * angle_step
        angle_max = -np.pi + (i + 1) * angle_step
        region_mask = (angles >= angle_min) & (angles < angle_min + (2 * np.pi / 360 * 5))
        
        if np.any(region_mask):
            region_distances = distances[region_mask]
            region_idx = np.argmax(region_distances) #farthest points
            filtered_region_x = x_filtered[region_mask]
            filtered_region_y = y_filtered[region_mask]
            filtered_region_z = z_filtered[region_mask]
            point = [filtered_region_x[region_idx], filtered_region_y[region_idx], filtered_region_z[region_idx]]
            support_points.append(point)

    support_points = np.array(support_points)
    return support_points


def get_warpage(x, y, z):
    warpage = 0.0
    a, b, c = fit_plane(x, y, z)
    distances = (a * x + b * y - z + c) / np.sqrt(a**2 + b**2 + 1)
    max_distance = np.max(distances)
    min_distance = np.min(distances)
    warpage = max_distance - min_distance
    return warpage


# 定义函数：计算三维空间中点到直线的距离
def point_to_line_distance(point, line_point1, line_point2):
    line_vec = line_point2 - line_point1
    point_vec = point - line_point1
    distance = np.linalg.norm(np.cross(line_vec, point_vec)) / np.linalg.norm(line_vec)
    return distance


def get_mean_line_value(x, y, z, far_point_x, far_point_y, dist_thresh=0.5):
    
    a, b, c = fit_plane(x, y, z)
    x1,y1 = far_point_x, far_point_y
    x2,y2 = np.mean(x), np.mean(y)
    A = -(y2 - y1) if x2 != x1 else 1
    B = x2 - x1 if x2 != x1 else 0
    C = -(x2 * y1 - x1 * y2) if x2 != x1 else - y1
    dist_thresh = 2
    flat_data = np.column_stack((x, y))
    distances = np.abs(A * x + B * y + C) / np.sqrt(A**2 + B**2)
    intersections = flat_data[distances < dist_thresh]
    #print('len(intersections)= ', len(intersections))
    tree = cKDTree(flat_data)
    _, idx = tree.query(intersections, k=1)
    z_values = z[idx]
    if len(z_values) == 0:
        return 0, 0, 0
    else:
        z_plane_values = (a * intersections[:, 0] + b * intersections[:, 1] + c)
        warpage_diff = (z_values - z_plane_values)
        
        line_warpage = max(warpage_diff) - min(warpage_diff)
    
    distances_to_far_point = np.sqrt((intersections[:, 0] - far_point_x)**2 + (intersections[:, 1] - far_point_y)**2)

    # 找到距离最远的点索引
    max_distance_idx = np.argmax(distances_to_far_point)

    # 获取最远点的 x, y 坐标
    far_point_2 = intersections[max_distance_idx]

    # 构建 KD-Tree，便于快速查找最近邻
    tree = cKDTree(np.column_stack((x, y)))

    # 查找 far_point_x, far_point_y 最近的原始点
    _, idx1 = tree.query([far_point_x, far_point_y], k=1)
    far_point_1_z = z[idx1]  # 找到对应的 z 值

    # 查找 far_point_2 最近的原始点
    _, idx2 = tree.query(far_point_2, k=1)
    far_point_2_z = z[idx2]  # 找到对应的 z 值

    # 第一个端点的三维坐标
    far_point_1_3d = np.array([far_point_x, far_point_y, far_point_1_z])

    # 第二个端点的三维坐标
    far_point_2_3d = np.array([far_point_2[0], far_point_2[1], far_point_2_z])

    #print(far_point_1_3d, far_point_2_3d)

    # 查找中心点 (x2, y2) 最近的原始点，确定其 z 值
    _, idx_center = tree.query([round(x2), round(y2)], k=1)
    center_point_z = z[idx_center]  # 根据 (x2, y2) 最近点确定 z 值
    center_point_3d = np.array([round(x2), round(y2), center_point_z])
    
    line_bow = point_to_line_distance(center_point_3d, far_point_1_3d, far_point_2_3d)

    return line_warpage, line_bow, z_values


def get_theta(x, y, z, num_regions=60):
    
    center_x = np.mean(x)
    center_y = np.mean(y)
    angle_step = 2 * np.pi / num_regions  
    angles = np.arctan2(y - center_y, x - center_x) 
    distances = np.sqrt((x - center_x)**2 + (y - center_y)**2)  
    max_line_warpage = 0.0
    max_line_bow = 0.0
    max_warpage_theta = 0.0
    max_bow_theta = 0.0
    max_warpage_z = []
    max_bow_z = []
    for i in range(num_regions):
        angle_min = -np.pi + i * angle_step
        angle_max = -np.pi + (i + 1) * angle_step
        region_mask = (angles >= angle_min) & (angles < angle_max)
        if np.any(region_mask):
            region_distances = distances[region_mask]
            region_idx = np.argmax(region_distances) 
            far_point_x = x[region_mask][region_idx]
            far_point_y = y[region_mask][region_idx]
            mean_line_warpage, mean_line_bow, z_values = get_mean_line_value(x, y, z, far_point_x, far_point_y)
            if max_line_warpage < mean_line_warpage:
                max_line_warpage = mean_line_warpage
                max_warpage_theta = angle_min
                max_warpage_z = z_values
            if max_line_bow < mean_line_bow:
                max_line_bow = mean_line_bow
                max_bow_theta = angle_min
                max_bow_z = z_values
            #print('mean_line_warpage = ', mean_line_warpage)
    if max_warpage_theta < 0:
        max_warpage_theta = max_warpage_theta + np.pi
    if max_bow_theta < 0:
        max_bow_theta = max_bow_theta + np.pi

    return max_line_warpage, max_line_bow, max_warpage_theta, max_bow_theta, max_warpage_z, max_bow_z


def get_eval_result(pointcloud_path, json_path, theta_notch, debug):
    data = np.loadtxt(os.path.join(pointcloud_path), delimiter=',')
    
    x = data[:, 0]
    y = data[:, 1]
    z =  data[:, 2]
    
    max_z = max(z)
    min_z = min(z)
    pv = max_z - min_z
    print('pv = ', pv)
    #a, b, c = fit_plane(x, y, z)
    #x_filtered, y_filtered, z_filtered = remove_noise(x, y, z, a, b, c, thresh=0.01)
    x_filtered, y_filtered, z_filtered = x, y, z
    
    support_points = get_support_points(x_filtered, y_filtered, z_filtered)
    a_bow, b_bow, c_bow = fit_plane(support_points[:,0], support_points[:,1], support_points[:,2])

    #center_int = (round(np.mean(x_filtered)), round(np.mean(y_filtered)))
    center_int = (np.mean(x_filtered), np.mean(y_filtered))
    #print('center_int = ', center_int)
    distances = np.sqrt((data[:, 0] - center_int[0]) ** 2 + (data[:, 1] - center_int[1]) ** 2)
    # 找到最近的点
    nearest_index = np.argmin(distances)  # 找到最小距离对应的索引

    # 获取最近点的 z 值
    z_center = data[nearest_index, 2]
    #mask = (data[:, 0] == center_int[0]) & (data[:, 1] == center_int[1])
    #z_center = data[mask, 2][0]

    fit_z_center = a_bow*center_int[0] + b_bow*center_int[1] + c_bow
    bow = (z_center - fit_z_center)
    print('global bow = ', bow)
    
    warpage = get_warpage(x_filtered, y_filtered, z_filtered)
    print('global warpage = ', warpage)

    max_line_warpage, max_line_bow, max_warpage_theta, max_bow_theta, max_warpage_z, max_bow_z = get_theta(x_filtered, y_filtered, z_filtered, 360)

    print('max_warpage_theta = ', max_warpage_theta)
    print('max_line_warpage = ', max_line_warpage)
    print('max_line_bow = ', max_line_bow)
    #print('max_warpage_z = ', max_warpage_z)

    # x2 = int(round(np.mean(x)))  
    # y2 = int(round(np.mean(y)))
    # horizontal_mask = (y == y2)
    # horizontal_z = z[horizontal_mask]  # 水平方向对应的 z 值
    # horizontal_x = x[horizontal_mask]
   
    # # 提取垂直方向的 z 值（即 x == x2 的所有点）
    # vertical_mask = (x == x2)
    # vertical_z = z[vertical_mask]  # 垂直方向对应的 z 值
    # vertical_y = y[vertical_mask]
    # 创建二维坐标数组
    points = np.column_stack((x, y))  # 将 x 和 y 组合成二维坐标点

    # 构建 KD 树
    tree = cKDTree(points)

    # 定义要查询的中心点
    x2 = round(np.mean(x))
    y2 = round(np.mean(y))

    horizontal_indices = np.where(y == y2)[0]  # 找到 y == y_center 的索引
    horizontal_x = x[horizontal_indices]  # 提取 x 值
    horizontal_z = z[horizontal_indices]  # 提取对应的 z 值

    sorted_horizontal_indices = np.argsort(horizontal_x)  # 对 x 排序，返回索引
    #print('sorted_horizontal_indices = ', sorted_horizontal_indices)
    horizontal_x_sorted = horizontal_x[sorted_horizontal_indices]
    horizontal_z_sorted = horizontal_z[sorted_horizontal_indices]

    # 提取经过中心点的垂直方向 (x = x_center) 的所有点
    vertical_indices = np.where(x == x2)[0]  # 找到 x == x2 的索引
    vertical_y = y[vertical_indices]  # 提取 y 值
    vertical_z = z[vertical_indices]  # 提取对应的 z 值

    # 按照 y 从小到大排序
    sorted_vertical_indices = np.argsort(vertical_y)  # 对 y 排序，返回索引
    vertical_y_sorted = vertical_y[sorted_vertical_indices]
    vertical_z_sorted = vertical_z[sorted_vertical_indices]

    # ## 查找水平方向上 (y 接近 y2) 的最近点
    # horizontal_idx = np.where(np.abs(y - y2) == 0)[0]  # 使用一个小容差判断 y 是否接近 y2
    # if len(horizontal_idx) > 0:
    #     horizontal_z = z[horizontal_idx]
    #     horizontal_x = x[horizontal_idx]
    #     #print(f"水平方向最近的点 (x, z): {list(zip(horizontal_x, horizontal_z))}")
    # else:
    #     print("没有找到接近 y2 的水平方向点")

    # # 查找垂直方向上 (x 接近 x2) 的最近点
    # vertical_idx = np.where(np.abs(x - x2) == 0)[0]  # 使用一个小容差判断 x 是否接近 x2
    # if len(vertical_idx) > 0:
    #     vertical_z = z[vertical_idx]
    #     vertical_y = y[vertical_idx]
    #     #print(f"垂直方向最近的点 (y, z): {list(zip(vertical_y, vertical_z))}")
    # else:
    #     print("没有找到接近 x2 的垂直方向点")

    #     # 检查水平和垂直方向是否找到点
    #     if len(horizontal_x) > 0:
    #         right_x = max(horizontal_x)
    #     else:
    #         print("找不到水平方向的点")
    #         right_x = None

    #     if len(vertical_y) > 0:
    #         top_y = max(vertical_y)
    #     else:
    #         print("找不到垂直方向的点")
    #         top_y = None

    top_x = x2
    right_y = y2

    top_x, top_y = x2, max(vertical_y)
    right_x, right_y = max(horizontal_x), y2
    #print(top_x, top_y)
    line_warpage_x, line_bow_x, z_values = get_mean_line_value(x, y, z, right_x, right_y)
    line_warpage_y, line_bow_y, z_values = get_mean_line_value(x, y, z, top_x, top_y)
    if debug:
        # 创建一个图
        plt.figure()

        # 绘制第一个列表
        plt.plot(horizontal_z_sorted.tolist(), label='List 1', color='blue', marker='o')

        # 绘制第二个列表
        plt.plot(vertical_z_sorted.tolist(), label='List 2', color='red', marker='s')

        # 添加标题和标签
        plt.title('Plot of List 1 and List 2')
        plt.xlabel('Index')
        plt.ylabel('Values')

        # 添加图例
        plt.legend()

        # 显示图形
        plt.show()

    with open(json_path, 'w') as file:
        result_data = {
                'x':{
                    'height': horizontal_z_sorted.tolist(),
                    'bow':line_bow_x,
                    'warpage':line_warpage_x
                },
                'y':{
                    'height':vertical_z_sorted.tolist(),
                    'bow':line_bow_y,
                    'warpage':line_warpage_y
                },
                'max_bow':{
                    'angle':max_bow_theta,
                    'height':max_bow_z.tolist(),
                    'max_bow':max_line_bow
                },
                'max_warpage':{
                    'angle':max_warpage_theta,
                    'height':max_warpage_z.tolist(),
                    'max_warpage':max_line_warpage
                },
                'pv': pv,
                'notch_theta':theta_notch
        }
        #print(result_data)
        json.dump(result_data, file, indent=4)

    if debug:
        plt.scatter(x, y, c=z, cmap='viridis')
        plt.colorbar()
        plt.title('3D Scatter Plot of Reconstructed Surface')
        plt.xlabel('X Axis')
        plt.ylabel('Y Axis')

        fig = plt.figure(figsize=(10, 7))
        ax = fig.add_subplot(111, projection='3d')

        ax.scatter(x_filtered, y_filtered, z_filtered, color='green', label='Filtered Data', alpha=0.1)

        ax.scatter(support_points[:, 0], support_points[:, 1], support_points[:, 2], 
                color='red', marker='*', s=200, label='Support Points')
        ax.legend()

        ax.set_xlabel('X Axis')
        ax.set_ylabel('Y Axis')
        ax.set_zlabel('Z Axis')
        ax.set_title('3D Point Cloud with Uniformly Distributed Support Points')

        plt.show()