Spaces:
Paused
Paused
| import os | |
| import json | |
| import re | |
| # 提取文件名中起始的连续数字（核心逻辑） | |
| def extract_starting_numbers(filename): | |
| # 匹配字符串开头的连续数字 | |
| match = re.match(r'^\d+', filename) | |
| if match: | |
| return match.group() # 返回起始数字字符串（如"4592"） | |
| return filename # 无起始数字则返回原文件名（小写处理后） | |
| # 处理单个ldr文件数据 | |
| def process_ldr_data(lines, label_mapping, label_inverse_mapping, label_counter): | |
| all_labels = [] | |
| for line in lines: | |
| if line.startswith('1'): # 只处理零件数据行 | |
| parts = line.split() | |
| if len(parts) < 15: | |
| print(f"Skipping invalid line: {line.strip()}") | |
| continue | |
| # 1. 统一转为小写 | |
| filename = parts[14].lower() | |
| # 2. 提取起始连续数字作为标识 | |
| part_identifier = extract_starting_numbers(filename) | |
| print(f"Processing file: {filename} → Identifier: {part_identifier}") | |
| if part_identifier not in label_mapping: | |
| label_mapping[part_identifier] = label_counter | |
| label_inverse_mapping[label_counter] = part_identifier | |
| label_counter += 1 | |
| all_labels.append(label_mapping[part_identifier]) | |
| return label_mapping, label_inverse_mapping, label_counter | |
| # 处理一个文件夹中的所有ldr文件 | |
| def process_all_ldr_in_folder(folder_path): | |
| overall_label_mapping = {} | |
| overall_label_inverse_mapping = {} | |
| label_counter = 0 | |
| for root, dirs, files in os.walk(folder_path): | |
| for file in files: | |
| if file.endswith('.ldr'): | |
| file_path = os.path.join(root, file) | |
| with open(file_path, 'r') as f: | |
| lines = f.readlines() | |
| overall_label_mapping, overall_label_inverse_mapping, label_counter = process_ldr_data( | |
| lines, overall_label_mapping, overall_label_inverse_mapping, label_counter) | |
| return overall_label_mapping, overall_label_inverse_mapping | |
| # 保存mapping到文件 | |
| def save_mappings(label_mapping, label_inverse_mapping, output_dir): | |
| os.makedirs(output_dir, exist_ok=True) | |
| with open(os.path.join(output_dir, 'label_mapping.json'), 'w') as f: | |
| json.dump(label_mapping, f, indent=4) | |
| with open(os.path.join(output_dir, 'label_inverse_mapping.json'), 'w') as f: | |
| json.dump(label_inverse_mapping, f, indent=4) | |
| 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' | |
| label_mapping, label_inverse_mapping = process_all_ldr_in_folder(input_folder) | |
| save_mappings(label_mapping, label_inverse_mapping, output_folder) | |
| print(f"Label mappings have been saved to {output_folder}") | |