import numpy as np import json import math import os import matplotlib.pyplot as plt from scipy import interpolate from typing import List, Tuple, Dict, Any def calculate_polygon_centroid(points: List[List[int]]) -> Tuple[float, float]: """计算多边形的形心""" if len(points) < 3: return np.mean(points, axis=0) # 使用鞋带公式计算多边形面积和形心 x = [p[0] for p in points] y = [p[1] for p in points] area = 0.0 centroid_x = 0.0 centroid_y = 0.0 n = len(points) for i in range(n): j = (i + 1) % n cross_product = x[i] * y[j] - x[j] * y[i] area += cross_product centroid_x += (x[i] + x[j]) * cross_product centroid_y += (y[i] + y[j]) * cross_product area *= 0.5 if area == 0: return np.mean(points, axis=0) centroid_x /= (6 * area) centroid_y /= (6 * area) return centroid_x, centroid_y def get_edge_angle(point: List[float], centroid: Tuple[float, float]) -> float: """计算点相对于形心的角度（从y轴正向开始）""" dx = point[0] - centroid[0] dy = point[1] - centroid[1] # 计算角度（弧度），从y轴正向开始，顺时针方向 angle = math.atan2(dx, dy) if angle < 0: angle += 2 * math.pi return angle def sort_points_by_angle(points: List[List[int]], centroid: Tuple[float, float]) -> List[List[int]]: """按角度排序点（从y轴正向开始）""" points_with_angles = [(point, get_edge_angle(point, centroid)) for point in points] points_with_angles.sort(key=lambda x: x[1]) return [point for point, angle in points_with_angles] def interpolate_points_with_spline(visible_edges: List[List[int]], num_points: int, k: int = 3) -> List[List[float]]: """使用样条插值进行等间距采样""" if len(visible_edges) < 3: return visible_edges # 提取x和y坐标 x = np.array([p[0] for p in visible_edges]) y = np.array([p[1] for p in visible_edges]) # 确保多边形闭合 x = np.r_[x, x[0]] y = np.r_[y, y[0]] # 使用B样条曲线进行拟合 try: tck, u = interpolate.splprep([x, y], k=k, s=0, per=True) # 生成等间距参数 u_new = np.linspace(0, 1, num_points) # 计算样条曲线上的点 xi, yi = interpolate.splev(u_new, tck) # 组合成点列表 sampled_points = [[float(xi[i]), float(yi[i])] for i in range(len(xi))] return sampled_points except Exception as e: print(f"样条插值失败，使用线性插值: {e}") return interpolate_points_linear(visible_edges, num_points) def interpolate_points_linear(visible_edges: List[List[int]], num_points: int) -> List[List[float]]: """使用线性插值进行等间距采样（备用方法）""" if len(visible_edges) < 3: return visible_edges # 确保多边形闭合 if visible_edges[0] != visible_edges[-1]: closed_points = visible_edges + [visible_edges[0]] else: closed_points = visible_edges # 计算周长 total_length = 0 segment_lengths = [] for i in range(len(closed_points) - 1): p1 = closed_points[i] p2 = closed_points[i + 1] length = math.sqrt((p2[0] - p1[0]) ** 2 + (p2[1] - p1[1]) ** 2) segment_lengths.append(length) total_length += length # 等间距采样 sampled_points = [] step = total_length / num_points for i in range(num_points): target_dist = i * step current_dist = 0 for j in range(len(segment_lengths)): if current_dist + segment_lengths[j] >= target_dist: # 在当前线段内插值 p1 = closed_points[j] p2 = closed_points[j + 1] t = (target_dist - current_dist) / segment_lengths[j] x = p1[0] + t * (p2[0] - p1[0]) y = p1[1] + t * (p2[1] - p1[1]) sampled_points.append([x, y]) break else: current_dist += segment_lengths[j] return sampled_points def sample_by_angle_original(visible_edges: List[List[int]], centroid: Tuple[float, float], num_points: int): """原始版本的等角度采样方法""" sampled_points = [] angles = [] # for i in range(37, num_points): for i in range(num_points): angle = 2 * math.pi * i / num_points angles.append(angle) # 找到与当前角度射线相交的边界点 intersections = [] n = len(visible_edges) for j in range(n): p1 = visible_edges[j] p2 = visible_edges[(j + 1) % n] # 计算线段与射线的交点 intersection = find_ray_segment_intersection(centroid, angle, p1, p2) if intersection: intersections.append(intersection) if intersections: # 选择距离形心最远的交点 farthest_point = max(intersections, key=lambda p: math.sqrt((p[0] - centroid[0]) ** 2 + (p[1] - centroid[1]) ** 2)) sampled_points.append(farthest_point) else: # sampled_points.append(sampled_points[-1]) sampled_points.append([centroid[0], centroid[1]]) # 以形心填补 return sampled_points, angles def find_ray_segment_intersection(origin: Tuple[float, float], angle: float, p1: List[float], p2: List[float]) -> List[float]: """找到射线与线段的交点""" ox, oy = origin dx, dy = math.sin(angle), math.cos(angle) # y轴正向对应角度0 # 线段参数 x1, y1 = p1 x2, y2 = p2 # 射线参数方程: (ox, oy) + t*(dx, dy) # 线段参数方程: (x1, y1) + u*(x2-x1, y2-y1) denominator = dx * (y2 - y1) - dy * (x2 - x1) if abs(denominator) < 1e-10: return None # 平行或重合 t = ((x1 - ox) * (y2 - y1) - (y1 - oy) * (x2 - x1)) / denominator u = ((x1 - ox) * dy - (y1 - oy) * dx) / denominator if t >= 0 and 0 <= u <= 1: return [ox + t * dx, oy + t * dy] return None def assign_labels(sampled_points: List[List[float]], # full_centroid: List[float], occluded_edges: List[List[int]], unoccluded_edges: List[List[int]], visible_edges: List[List[int]]) -> List[int]: """为采样点分配遮挡标签""" labels = [] # 创建KD树用于最近邻搜索（简化版本） def find_nearest_point(query_point, point_list): min_dist = float('inf') nearest_point = None for point in point_list: dist = math.sqrt((query_point[0] - point[0]) ** 2 + (query_point[1] - point[1]) ** 2) if dist < min_dist: min_dist = dist nearest_point = point return nearest_point, min_dist for point in sampled_points: # 找到visible_edges中最近的点 nearest_visible, dist_visible = find_nearest_point(point, visible_edges) # 检查该点在occluded还是unoccluded中 nearest_occluded, dist_occluded = find_nearest_point(nearest_visible, occluded_edges) nearest_unoccluded, dist_unoccluded = find_nearest_point(nearest_visible, unoccluded_edges) # 如果到occluded点的距离更近，则标签为1（被遮挡） if dist_occluded < dist_unoccluded: labels.append(1) else: labels.append(0) ret = find_consecutive_pairs(sampled_points) if ret: modify_labels = modify_labels_by_pairs(labels, ret) return labels def modify_labels_by_pairs(labels, pairs): """ 根据成对的索引对，将label列表中相应位置修改为1 参数: labels: 原始标签列表 pairs: 包含起始和结束索引的元组列表 返回: 修改后的标签列表 """ # 创建副本以避免修改原始列表 modified_labels = labels.copy() for start, end in pairs: # 确保索引在有效范围内 if start < len(modified_labels) and end < len(modified_labels): for i in range(start, end + 1): modified_labels[i] = 1 else: print(f"警告: 索引对 ({start}, {end}) 超出列表范围") return modified_labels def find_consecutive_pairs(arr): # 找到列表中连续相同元素的成对索引 if not arr: return [] result = [] start = 0 for i in range(1, len(arr)): if arr[i] != arr[i - 1]: if i - start > 1: # 至少有两个相同的元素 result.append((start, i - 1)) start = i # 检查最后一个序列 if len(arr) - start > 1: result.append((start, len(arr) - 1)) return result def visualize_results(stone_data: Dict[str, Any], centroid: Tuple[float, float], full_centroid: Tuple[float, float], equidistant_points: List[List[float]], equidistant_labels: List[int], equiangular_points: List[List[float]], equiangular_labels: List[int], full_equiangular_points: List[List[float]], output_path: str, stone_id: str): """可视化原图、等间距点和等角度点""" visible_edges = stone_data['visible_edges'] occluded_edges = stone_data['occluded_edges'] unoccluded_edges = stone_data['unoccluded_edges'] plt.rcParams['font.sans-serif'] = ['Times New Roman'] # 创建图形 fig, axes = plt.subplots(2, 2, figsize=(15, 12)) fig.suptitle(f'Stone {stone_id} Analysis', fontsize=16) # 1. 原图可视化 ax1 = axes[0, 0] x_visible = [p[0] for p in visible_edges] y_visible = [p[1] for p in visible_edges] # 绘制封闭多边形 ax1.plot(x_visible + [x_visible[0]], y_visible + [y_visible[0]], 'k-', linewidth=1, label='Boundary') # 标记遮挡和非遮挡边缘 if occluded_edges: x_occ = [p[0] for p in occluded_edges] y_occ = [p[1] for p in occluded_edges] ax1.scatter(x_occ, y_occ, c='red', s=10, label='Occluded', alpha=0.7) if unoccluded_edges: x_unocc = [p[0] for p in unoccluded_edges] y_unocc = [p[1] for p in unoccluded_edges] ax1.scatter(x_unocc, y_unocc, c='green', s=10, label='Unoccluded', alpha=0.7) # 绘制y轴正向 arrow_length = max(max(x_visible) - min(x_visible), max(y_visible) - min(y_visible)) * 0.1 ax1.arrow(centroid[0], centroid[1], 0, arrow_length, head_width=arrow_length * 0.2, head_length=arrow_length * 0.2, fc='green', ec='green', label='Y-axis') # 标记形心 ax1.scatter(centroid[0], centroid[1], c='red', s=50, marker='o', label='Visible centroid') ax1.set_title('Original Visible Edges') ax1.legend() ax1.set_aspect('equal') # 2. 等间距采样可视化 ax2 = axes[0, 1] ax2.plot(x_visible + [x_visible[0]], y_visible + [y_visible[0]], 'k-', linewidth=1, alpha=0.5) # 根据标签着色 for i, point in enumerate(equidistant_points): color = 'red' if equidistant_labels[i] == 1 else 'green' ax2.scatter(point[0], point[1], c=color, s=20) ax2.scatter(centroid[0], centroid[1], c='red', s=50, marker='o') # ax2.arrow(centroid[0], centroid[1], 0, arrow_length, # head_width=arrow_length * 0.2, head_length=arrow_length * 0.2, fc='red', ec='red') ax2.set_title(f'Equidistant Sampling ({len(equidistant_points)} points)') ax2.set_aspect('equal') # 3. 等角度采样可视化 ax3 = axes[1, 0] ax3.plot(x_visible + [x_visible[0]], y_visible + [y_visible[0]], 'k-', linewidth=1, alpha=0.5) # 原始轮廓 for i, point in enumerate(full_equiangular_points): # color = 'red' if equiangular_labels[i] == 1 else 'green' ax3.scatter(point[0], point[1], c='blue', s=20) # 根据标签着色 for i, point in enumerate(equiangular_points): # print(f'i={i},label={equiangular_labels[i]}') color = 'red' if equiangular_labels[i] == 1 else 'green' ax3.scatter(point[0], point[1], c=color, s=20) full_contour_ray = [] # 绘制角度射线 for i, point in enumerate(full_equiangular_points): angle = 2 * math.pi * i / len(full_equiangular_points) ray_length = math.sqrt((point[0] - full_centroid[0]) ** 2 + (point[1] - full_centroid[1]) ** 2) full_contour_ray.append(ray_length) dx = math.sin(angle) * ray_length dy = math.cos(angle) * ray_length color = 'brown' if equiangular_labels[i] == 1 else 'gray' ax3.plot([full_centroid[0], full_centroid[0] + dx], [full_centroid[1], full_centroid[1] + dy], color, linewidth=0.5, alpha=0.5) ax3.arrow(centroid[0], centroid[1], # 起始点 full_centroid[0] - centroid[0], full_centroid[1] - centroid[1], # 箭头方向向量 head_width=0.2*abs(full_centroid[1] - centroid[1]), head_length=0.3*abs(full_centroid[1] - centroid[1]), # 箭头头部参数 fc='green', ec='green', # 填充颜色和边缘颜色 length_includes_head=True, # 长度包含箭头头部 width=0.05) ax3.scatter(centroid[0], centroid[1], c='red', s=50, marker='o', label='Visible centroid') ax3.scatter(full_centroid[0], full_centroid[1], c='green', s=50, marker='o', label='Original centroid') ax3.legend() ax3.set_title(f'Equiangular Sampling ({len(equiangular_points)} points)') ax3.set_aspect('equal') # 4. 两种采样方法对比 ax4 = axes[1, 1] ax4.plot(x_visible + [x_visible[0]], y_visible + [y_visible[0]], 'k-', linewidth=1, alpha=0.3) # 绘制等间距点 for i, point in enumerate(equidistant_points): # color = 'red' if equidistant_labels[i] == 1 else 'green' # ax4.scatter(point[0], point[1], c='fuchsia', s=15, marker='o', alpha=0.7, label='Equidistant' if i == 0 else "") ax4.scatter(point[0], point[1], c='darkorange', s=15, marker='o', alpha=0.7, label='Equidistant' if i == 0 else "") # 绘制等角度点 for i, point in enumerate(equiangular_points): # color = 'red' if equiangular_labels[i] == 1 else 'green' # ax4.scatter(point[0], point[1], c='lime', s=15, marker='^', alpha=0.7, label='Equiangular' if i == 0 else "") ax4.scatter(point[0], point[1], c='black', s=15, marker='^', alpha=0.7, label='Equiangular' if i == 0 else "") # ax4.scatter(centroid[0], centroid[1], c='blue', s=50, marker='x', label='Centroid') ax4.set_title('Comparison: Equidistant vs Equiangular ') ax4.legend() ax4.set_aspect('equal') plt.tight_layout() # plt.show() # 保存图像 save_path = os.path.join(output_path, 'PAS visualization') if os.path.exists(save_path): pass else: os.makedirs(save_path) # 可用于创建多级目录 output_file = os.path.join(save_path, f'{stone_id}_analysis.png') plt.savefig(output_file, dpi=300, bbox_inches='tight') plt.close() print(f"可视化结果已保存: {output_file}") return output_file, full_contour_ray def process_stone(stone_data: Dict[str, Any], output_path: str, num_spacing_points: int = 100, num_angle_points: int = 36) -> Dict[str, Any]: """处理单个石头数据""" result = {} # 获取数据 visible_edges = stone_data['visible_edges'] occluded_edges = stone_data['occluded_edges'] unoccluded_edges = stone_data['unoccluded_edges'] full_edges = stone_data['full_edge'] stone_id = stone_data.get('stone_id', 'unknown') # print(f"正在处理石头: {stone_id}") # print(f" - 可见边缘点数: {len(visible_edges)}") # print(f" - 遮挡边缘点数: {len(occluded_edges)}") # print(f" - 未遮挡边缘点数: {len(unoccluded_edges)}") # 计算形心 full_centroid = calculate_polygon_centroid(full_edges) visible_centroid = calculate_polygon_centroid(visible_edges) # print(f" - 形心坐标: ({visible_centroid[0]:.2f}, {visible_centroid[1]:.2f})") # 等间距采样 - 使用样条插值方法 equidistant_points = interpolate_points_with_spline(visible_edges, num_spacing_points) equidistant_labels = assign_labels(equidistant_points, occluded_edges, unoccluded_edges, visible_edges) # print(f" - 等间距采样点: {len(equidistant_points)}") # 对可见边 等角度采样 - 使用原始版本 visible_equiangular_points_full_center, visible_angles = sample_by_angle_original(visible_edges, full_centroid, num_angle_points) visible_equiangular_labels_full_center= assign_labels(visible_equiangular_points_full_center, occluded_edges, unoccluded_edges, visible_edges) # print(f" - 等角度采样点: {len(visible_equiangular_points_full_center)}") # 对完整边 等角度采样 - 使用原始版本 full_equiangular_points_full_center, full_angles = sample_by_angle_original(full_edges, full_centroid, num_angle_points) # # 可视化 # viz_file, full_edges_ray = visualize_results(stone_data, visible_centroid, full_centroid, equidistant_points, # equidistant_labels, visible_equiangular_points_full_center, # visible_equiangular_labels_full_center, full_equiangular_points_full_center, output_path, stone_id) full_edges_ray = [] # 绘制角度射线 for i, point in enumerate(full_equiangular_points_full_center): # angle = 2 * math.pi * i / len(full_equiangular_points_full_center) ray_length = math.sqrt((point[0] - full_centroid[0]) ** 2 + (point[1] - full_centroid[1]) ** 2) full_edges_ray.append(ray_length) # 构建结果 result['stone_id'] = stone_id result['visible_centroid'] = [float(visible_centroid[0]), float(visible_centroid[1])] # 转换为可序列化格式 result['full_centroid'] = [float(full_centroid[0]), float(full_centroid[1])] # 转换为可序列化格式 result['centroid_diff'] = [float(full_centroid[0]) - float(visible_centroid[0]), float(full_centroid[1]) - float(visible_centroid[1])] result['equidistant_sampling'] = { 'points': [[float(p[0]), float(p[1])] for p in equidistant_points], 'labels': equidistant_labels } result['equiangular_sampling'] = { 'visible_points': [[float(p[0]), float(p[1])] for p in visible_equiangular_points_full_center], 'full_points': [[float(p[0]), float(p[1])] for p in full_equiangular_points_full_center], 'labels': visible_equiangular_labels_full_center, 'angles': visible_angles, 'full_edges_ray': full_edges_ray } return result def process_all_stones(stones_data: Dict[str, Any], output_path: str, num_spacing_points: int = 100, num_angle_points: int = 36) -> Dict[str, Any]: """处理所有石头数据""" results = {} # 确保输出目录存在 os.makedirs(output_path, exist_ok=True) for stone_id, stone_data in stones_data.items(): if stone_id == 'Total_unoccluded_edges': continue # print(f'正在处理{stone_id}') stone_data['stone_id'] = stone_id # 添加stone_id到数据中 results[stone_id] = process_stone(stone_data, output_path, num_spacing_points, num_angle_points) return results def save_results(results: Dict[str, Any], output_path: str): """保存处理结果到JSON文件""" output_file = os.path.join(output_path, 'sampling_results.json') with open(output_file, 'w') as f: # 转换numpy数组为Python原生类型 def convert_to_serializable(obj): if isinstance(obj, (np.int32, np.int64)): return int(obj) elif isinstance(obj, (np.float32, np.float64)): return float(obj) elif isinstance(obj, np.ndarray): return obj.tolist() elif isinstance(obj, dict): return {k: convert_to_serializable(v) for k, v in obj.items()} elif isinstance(obj, list): return [convert_to_serializable(item) for item in obj] else: return obj serializable_results = convert_to_serializable(results) json.dump(serializable_results, f, indent=2) print(f"处理结果已保存: {output_file}") return output_file def sample_PAS(input_file: str, output_path: str, num_spacing_points: int = 100, num_angle_points: int = 36): """主函数：通过参数传递处理数据 Args: input_file: 输入JSON数据文件路径 output_path: 输出目录路径 num_spacing_points: 等间距采样点数 (默认: 100) num_angle_points: 等角度采样点数 (默认: 36) """ # 检查输入文件是否存在 if not os.path.exists(input_file): print(f"错误: 文件 '{input_file}' 不存在") return None # 确保输出目录存在 os.makedirs(output_path, exist_ok=True) # 加载数据 try: with open(input_file, 'r') as f: data = json.load(f) except Exception as e: print(f"加载JSON文件时出错: {e}") return None # 检查数据结构 if 'stones' not in data: print("错误: JSON文件中未找到 'stones' 键") return None stones_data = data['stones'] print(f"找到 {len(stones_data)} 个石头需要处理") print(f"输入文件: {input_file}") print(f"输出目录: {output_path}") print(f"等间距采样点数: {num_spacing_points}") print(f"等角度采样点数: {num_angle_points}") print("开始处理...") # 处理所有石头 results = process_all_stones(stones_data, output_path, num_spacing_points, num_angle_points) # 保存结果 save_results(results, output_path) print("所有处理完成！") return results