Spaces:
Paused
Paused
| import os | |
| import numpy as np | |
| from process_single_ldr import process_ldr_data, process_ldr_flatten, process_ldr_flatten_bottom | |
| import re | |
| # 处理所有ldr文件 | |
| def process_all_ldr_in_folder(folder_path, output_folder): | |
| all_data = [] | |
| all_label_inverse_mappings = [] # 存储所有文件的label_inverse_mapping | |
| filenames = [] # 存储所有文件的文件名 | |
| line_count = 0 | |
| ldr_files = [f for f in os.listdir(input_folder) if | |
| os.path.isfile(os.path.join(input_folder, f)) and | |
| re.match(r'modified_car_\d+_车底_rot\.ldr', f)] | |
| ldr_files.sort(key=lambda x: int(re.search(r'(\d+)', x).group(1))) | |
| #import ipdb; ipdb.set_trace() | |
| # 遍历文件夹中的所有文件夹 | |
| #for file in os.listdir(folder_path): | |
| for idx, file in enumerate(ldr_files, 1): | |
| # if file == 'modified_car_81_rot.ldr' or file == 'modified_car_1320_rot.ldr' or file == 'modified_car_1688_rot.ldr' : | |
| # continue | |
| if file.endswith('.ldr') and file.startswith('modified'): # 只处理ldr文件 | |
| file_path = os.path.join(folder_path, file) | |
| with open(file_path, 'r') as f: | |
| lines = f.readlines() | |
| #if len(lines)>110: | |
| #print(file, len(lines)) | |
| # #break | |
| line_count += len(lines) | |
| #data, _ = process_ldr_data(lines) | |
| data, _ = process_ldr_flatten_bottom(lines) | |
| if data.shape[0]>100: | |
| # import ipdb; ipdb.set_trace() | |
| print(file, idx) | |
| print(file, data.shape[0]) | |
| #sort | |
| sort_cols = data[:, [-4, -5, -3]] | |
| sort_idx = np.lexsort((sort_cols[:, 2], sort_cols[:, 1], sort_cols[:, 0])) | |
| data = data[sort_idx] | |
| all_data.append(data) | |
| filenames.append(file) | |
| #print(f"Processed {file}") | |
| print(f"Total lines processed: {line_count}") #1263 | |
| #import ipdb; ipdb.set_trace() | |
| output_file = os.path.join(output_folder, "bottom_train.npz") | |
| np.savez_compressed(output_file, | |
| data=np.array(all_data, dtype=object), # 将所有数据以列表的形式存储 | |
| filenames=filenames) # 也可以保存文件名，方便后续查找 | |
| print(f"All LDR data have been processed and saved to {output_file}") | |
| # 主程序 | |
| if __name__ == "__main__": | |
| input_folder = '/public/home/wangshuo/gap/assembly/data/car_1k/subset_bottom_300/ldr_rot_expand_trans' # 输入文件夹路径 | |
| output_folder = '/public/home/wangshuo/gap/assembly/data/car_1k/subset_bottom_300' # 输出文件夹路径 | |
| os.makedirs(output_folder, exist_ok=True) # 确保输出文件夹存在 | |
| # 处理所有ldr文件并保存为npz文件 | |
| process_all_ldr_in_folder(input_folder, output_folder) | |