Add code/cube3d/training/ldr_coordinate_analyzer.py

code/cube3d/training/ldr_coordinate_analyzer.py

@@ -0,0 +1,194 @@
+import os
+import re
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.ticker import MaxNLocator
+import pandas as pd
+from glob import glob
+
+# 设置中文字体
+plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
+plt.rcParams["axes.unicode_minus"] = False  # 正确显示负号
+
+def parse_ldr_file(file_path):
+    """解析LDR文件，提取零件的XYZ坐标（修复：用空格分割，兼容不同格式）"""
+    coordinates = []
+    
+    with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
+        for line_num, line in enumerate(f, 1):
+            line = line.strip()
+            # 只处理以1开头的零件行，且排除空行
+            if line.startswith('1') and len(line.split()) >= 5:
+                try:
+                    # 核心修复：用split()按任意空格分割，取第3、4、5个元素（索引2、3、4）
+                    parts = line.split()
+                    x = float(parts[2])  # 原正则的group(3)对应split后的索引2
+                    y = float(parts[3])  # 原正则的group(4)对应split后的索引3
+                    z = float(parts[4])  # 原正则的group(5)对应split后的索引4
+                    coordinates.append((x, y, z))
+                except (ValueError, IndexError) as e:
+                    # 详细打印错误行，方便排查格式问题
+                    print(f"警告: 文件 {os.path.basename(file_path)} 第 {line_num} 行解析失败 - {e}")
+                    print(f"  错误行内容: {line}")
+    
+    return np.array(coordinates) if coordinates else None
+
+def translate_coordinates(coords):
+    """将坐标平移到原点附近，平移量为max_x/2, max_y/2, max_z/2"""
+    if coords is None or len(coords) == 0:
+        return None, (0, 0, 0)
+    
+    # 计算各维度最大值（确保坐标≥0的前提下）
+    max_x = np.max(coords[:, 0])
+    max_y = np.max(coords[:, 1])
+    max_z = np.max(coords[:, 2])
+    
+    # 计算平移量（向原点方向移动一半最大值）
+    tx = max_x / 2
+    ty = max_y / 2
+    tz = max_z / 2
+    
+    # 执行平移
+    translated = coords - np.array([tx, ty, tz])
+    
+    return translated, (tx, ty, tz)
+
+def analyze_ldr_folder(folder_path):
+    """分析文件夹中所有LDR文件的坐标分布"""
+    all_coords = []
+    # 递归获取所有LDR文件（包括子文件夹）
+    ldr_files = glob(os.path.join(folder_path, '**', '*.ldr'), recursive=True)
+    
+    if not ldr_files:
+        print(f"错误：在 {folder_path} 中未找到任何LDR文件")
+        return None
+    
+    print(f"找到 {len(ldr_files)} 个LDR文件，开始分析...")
+    
+    for file_idx, file_path in enumerate(ldr_files, 1):
+        # 解析单个文件的坐标
+        coords = parse_ldr_file(file_path)
+        if coords is None or len(coords) == 0:
+            print(f"文件 {file_idx}/{len(ldr_files)}: {os.path.basename(file_path)} 中未找到有效坐标")
+            continue
+        
+        # 平移坐标到原点附近
+        translated_coords, translation = translate_coordinates(coords)
+        
+        # 收集所有平移后的坐标
+        all_coords.append(translated_coords)
+        
+        # 每处理10个文件或最后一个文件时，打印进度
+        if file_idx % 10 == 0 or file_idx == len(ldr_files):
+            total_parts = sum(len(c) for c in all_coords)
+            print(f"已处理 {file_idx}/{len(ldr_files)} 个文件，累计零件数: {total_parts}")
+    
+    if not all_coords:
+        print("未找到任何有效坐标数据，终止分析")
+        return None
+    
+    # 合并所有文件的坐标（按行堆叠）
+    all_coords = np.vstack(all_coords)
+    print(f"\n分析完成！共处理 {len(all_coords)} 个零件坐标")
+    
+    return all_coords
+
+def plot_coordinate_distributions(coords, output_dir):
+    """绘制XYZ三个维度的坐标分布直方图"""
+    if coords is None or len(coords) == 0:
+        print("没有可绘制的坐标数据")
+        return
+    
+    # 创建输出目录（不存在则自动创建）
+    os.makedirs(output_dir, exist_ok=True)
+    
+    # 提取XYZ三个维度的坐标
+    x_coords = coords[:, 0]
+    y_coords = coords[:, 1]
+    z_coords = coords[:, 2]
+    
+    # 计算各维度的统计信息（最小值、最大值、平均值、标准差）
+    stats = {
+        'X轴': {
+            '最小值': np.min(x_coords),
+            '最大值': np.max(x_coords),
+            '平均值': np.mean(x_coords),
+            '标准差': np.std(x_coords)
+        },
+        'Y轴': {
+            '最小值': np.min(y_coords),
+            '最大值': np.max(y_coords),
+            '平均值': np.mean(y_coords),
+            '标准差': np.std(y_coords)
+        },
+        'Z轴': {
+            '最小值': np.min(z_coords),
+            '最大值': np.max(z_coords),
+            '平均值': np.mean(z_coords),
+            '标准差': np.std(z_coords)
+        }
+    }
+    
+    # 打印统计信息（控制台直观查看）
+    print("\n=== 坐标分布统计信息 ===")
+    for axis, info in stats.items():
+        print(f"{axis}:")
+        for key, val in info.items():
+            print(f"  {key}: {val:.2f}")
+    
+    # 保存统计信息到CSV文件（便于后续分析）
+    stats_df = pd.DataFrame(stats).T
+    stats_csv_path = os.path.join(output_dir, 'coordinate_stats.csv')
+    stats_df.to_csv(stats_csv_path, encoding='utf-8-sig')
+    print(f"\n统计信息已保存到: {stats_csv_path}")
+    
+    # 绘制XYZ三个维度的分布直方图（1行3列布局）
+    fig, axes = plt.subplots(1, 3, figsize=(18, 6))
+    bins = 50  # 直方图分箱数（可根据数据量调整，50适合中等数据量）
+    
+    # X轴坐标分布（红色）
+    axes[0].hist(x_coords, bins=bins, alpha=0.7, color='#FF6B6B')
+    axes[0].set_title(f'X轴坐标分布\n（均值: {stats["X轴"]["平均值"]:.2f}, 标准差: {stats["X轴"]["标准差"]:.2f}）', pad=15)
+    axes[0].set_xlabel('坐标值', fontsize=12)
+    axes[0].set_ylabel('频数', fontsize=12)
+    axes[0].grid(True, alpha=0.3)
+    axes[0].xaxis.set_major_locator(MaxNLocator(6))  # 限制X轴刻度数量，避免拥挤
+    
+    # Y轴坐标分布（绿色）
+    axes[1].hist(y_coords, bins=bins, alpha=0.7, color='#4ECDC4')
+    axes[1].set_title(f'Y轴坐标分布\n（均值: {stats["Y轴"]["平均值"]:.2f}, 标准差: {stats["Y轴"]["标准差"]:.2f}）', pad=15)
+    axes[1].set_xlabel('坐标值', fontsize=12)
+    axes[1].set_ylabel('频数', fontsize=12)
+    axes[1].grid(True, alpha=0.3)
+    axes[1].xaxis.set_major_locator(MaxNLocator(6))
+    
+    # Z轴坐标分布（蓝色）
+    axes[2].hist(z_coords, bins=bins, alpha=0.7, color='#45B7D1')
+    axes[2].set_title(f'Z轴坐标分布\n（均值: {stats["Z轴"]["平均值"]:.2f}, 标准差: {stats["Z轴"]["标准差"]:.2f}）', pad=15)
+    axes[2].set_xlabel('坐标值', fontsize=12)
+    axes[2].set_ylabel('频数', fontsize=12)
+    axes[2].grid(True, alpha=0.3)
+    axes[2].xaxis.set_major_locator(MaxNLocator(6))
+    
+    # 调整子图间距，避免标题/标签重叠
+    plt.tight_layout()
+    
+    # 保存图片（高分辨率300dpi，适合后续使用）
+    plot_path = os.path.join(output_dir, 'coordinate_distributions.png')
+    plt.savefig(plot_path, dpi=300, bbox_inches='tight')
+    print(f"坐标分布图已保存到: {plot_path}")
+    
+    # 显示图片（如果运行环境支持GUI）
+    plt.show()
+
+if __name__ == "__main__":
+    # 配置路径（请根据实际情况修改）
+    LDR_FOLDER = "/public/home/wangshuo/gap/assembly/data/car_1k/subset_self/ldr_l30_rotrans_expand_wom"  # LDR文件根目录
+    OUTPUT_DIR = "./ldr_coordinate_analysis"  # 结果输出目录（统计CSV+分布图）
+    
+    # 1. 分析所有LDR文件的坐标
+    all_translated_coords = analyze_ldr_folder(LDR_FOLDER)
+    
+    # 2. 绘制并保存坐标分布图
+    if all_translated_coords is not None:
+        plot_coordinate_distributions(all_translated_coords, OUTPUT_DIR)