code/cube3d/training/ldr_coordinate_analyzer.py

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)