Spaces:
Paused
Paused
File size: 2,959 Bytes
5ec45ba |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
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}")
|