123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400 |
- 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 = 1000 * 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()
|