From f930bc17ad7217e63a964a3113867cc8322a5e97 Mon Sep 17 00:00:00 2001 From: mckay Date: Wed, 4 Dec 2024 21:25:03 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E7=BB=9F=E8=AE=A1brep=E6=95=B0?= =?UTF-8?q?=E6=8D=AE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- brep2sdf/scripts/__init__.py | 0 brep2sdf/scripts/brep_stats.py | 194 +++++++++++++++++++++++++++++++++ 2 files changed, 194 insertions(+) create mode 100644 brep2sdf/scripts/__init__.py create mode 100644 brep2sdf/scripts/brep_stats.py diff --git a/brep2sdf/scripts/__init__.py b/brep2sdf/scripts/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/brep2sdf/scripts/brep_stats.py b/brep2sdf/scripts/brep_stats.py new file mode 100644 index 0000000..e2e7bd3 --- /dev/null +++ b/brep2sdf/scripts/brep_stats.py @@ -0,0 +1,194 @@ +import os +import numpy as np +from collections import defaultdict +from tqdm import tqdm +import pickle +from brep2sdf.scripts.process_brep import parse_solid +from brep2sdf.utils.logger import logger + +def analyze_brep_stats(step_path_list): + """分析BRep文件的统计信息 + + Args: + step_path_list: STEP文件路径列表 + + Returns: + stats_dict: 包含统计信息的字典 + """ + stats = defaultdict(list) + failed_files = [] + + for step_path in tqdm(step_path_list, desc="分析BRep文件"): + try: + # 解析STEP文件 + data = parse_solid(step_path) + if data is None: + failed_files.append(os.path.basename(step_path)) + continue + + # 收集统计信息 + stats['num_faces'].append(len(data['surf_wcs'])) + stats['num_edges'].append(len(data['edge_wcs'])) + stats['num_vertices'].append(len(data['corner_unique'])) + + # 计算每个面的点数 + face_points = [len(face) for face in data['surf_wcs']] + stats['points_per_face'].extend(face_points) + + # 计算每条边的采样点数 + edge_points = [len(edge) for edge in data['edge_wcs']] + stats['points_per_edge'].extend(edge_points) + + except Exception as e: + logger.error(f"处理文件 {step_path} 时出错: {str(e)}") + failed_files.append(os.path.basename(step_path)) + continue + + # 计算最终统计结果 + stats_dict = { + 'total_files': len(step_path_list), + 'successful_files': len(step_path_list) - len(failed_files), + 'failed_files': failed_files, + 'success_rate': (len(step_path_list) - len(failed_files)) / len(step_path_list) * 100, + + 'face_stats': { + 'min': np.min(stats['num_faces']), + 'max': np.max(stats['num_faces']), + 'mean': np.mean(stats['num_faces']), + 'median': np.median(stats['num_faces']), + 'std': np.std(stats['num_faces']) + }, + + 'edge_stats': { + 'min': np.min(stats['num_edges']), + 'max': np.max(stats['num_edges']), + 'mean': np.mean(stats['num_edges']), + 'median': np.median(stats['num_edges']), + 'std': np.std(stats['num_edges']) + }, + + 'vertex_stats': { + 'min': np.min(stats['num_vertices']), + 'max': np.max(stats['num_vertices']), + 'mean': np.mean(stats['num_vertices']), + 'median': np.median(stats['num_vertices']), + 'std': np.std(stats['num_vertices']) + }, + + 'face_points_stats': { + 'min': np.min(stats['points_per_face']), + 'max': np.max(stats['points_per_face']), + 'mean': np.mean(stats['points_per_face']), + 'median': np.median(stats['points_per_face']), + 'std': np.std(stats['points_per_face']) + }, + + 'edge_points_stats': { + 'min': np.min(stats['points_per_edge']), + 'max': np.max(stats['points_per_edge']), + 'mean': np.mean(stats['points_per_edge']), + 'median': np.median(stats['points_per_edge']), + 'std': np.std(stats['points_per_edge']) + } + } + + return stats_dict + +def print_stats(stats_dict): + """打印统计信息""" + logger.info("\n=== BRep文件统计信息 ===") + + # 基本信息 + logger.info(f"\n处理文件总数: {stats_dict['total_files']}") + logger.info(f"成功处理文件数: {stats_dict['successful_files']}") + logger.info(f"失败文件数: {len(stats_dict['failed_files'])}") + logger.info(f"成功率: {stats_dict['success_rate']:.2f}%") + + # 几何元素统计 + for name, stats in [ + ('面', stats_dict['face_stats']), + ('边', stats_dict['edge_stats']), + ('顶点', stats_dict['vertex_stats']) + ]: + logger.info(f"\n{name}统计:") + logger.info(f" 最小值: {stats['min']}") + logger.info(f" 最大值: {stats['max']}") + logger.info(f" 平均值: {stats['mean']:.2f}") + logger.info(f" 中位数: {stats['median']}") + logger.info(f" 标准差: {stats['std']:.2f}") + + # 点云统计 + for name, stats in [ + ('面点云', stats_dict['face_points_stats']), + ('边点云', stats_dict['edge_points_stats']) + ]: + logger.info(f"\n{name}统计:") + logger.info(f" 最小点数: {stats['min']}") + logger.info(f" 最大点数: {stats['max']}") + logger.info(f" 平均点数: {stats['mean']:.2f}") + logger.info(f" 中位数点数: {stats['median']}") + logger.info(f" 点数标准差: {stats['std']:.2f}") + + # 失败文件列表 + if stats_dict['failed_files']: + logger.info("\n失败文件列表:") + for file in stats_dict['failed_files']: + logger.info(f" - {file}") + +def get_step_files(root_dir): + """递归获取所有STEP文件路径 + + Args: + root_dir: 数据集根目录 + + Returns: + dict: 按数据集分割的文件路径字典 + """ + step_files_dict = {'train': [], 'val': [], 'test': []} + + for split in ['train', 'val', 'test']: + split_dir = os.path.join(root_dir, split) + if not os.path.exists(split_dir): + logger.warning(f"目录不存在: {split_dir}") + continue + + files = [os.path.join(split_dir, f) + for f in os.listdir(split_dir) + if f.endswith('.step')] + step_files_dict[split] = files + logger.info(f"找到 {split} 集 STEP 文件: {len(files)} 个") + + return step_files_dict + +def main(): + """主函数""" + # 数据集根目录 + root_dir = '/mnt/mynewdisk/dataset/furniture/step/furniture_dataset_step/' + if not os.path.exists(root_dir): + logger.error(f"数据集根目录不存在: {root_dir}") + return + + # 获取所有数据集的STEP文件 + step_files_dict = get_step_files(root_dir) + + # 统计每个数据集的信息 + for split, files in step_files_dict.items(): + if not files: + logger.warning(f"{split} 集没有找到STEP文件") + continue + + logger.info(f"\n=== {split.upper()} 集统计 ===") + stats = analyze_brep_stats(files) + print_stats(stats) + + # 保存统计结果 + save_dir = '/home/wch/brep2sdf/test_data/stats' + os.makedirs(save_dir, exist_ok=True) + save_path = os.path.join(save_dir, f'{split}_stats.pkl') + with open(save_path, 'wb') as f: + pickle.dump(stats, f) + logger.info(f"统计结果已保存到: {save_path}") + +if __name__ == '__main__': + main() +