Add code/cube3d/training/sort_mappings.py

code/cube3d/training/sort_mappings.py

@@ -0,0 +1,75 @@
+import json
+import os
+
+def process_mappings(original_mapping_path, inverse_mapping_path, frequency_data_path, output_dir=None):
+    """
+    处理正向和反向映射关系，根据使用频率重新排序并生成新的ID
+    
+    参数:
+        original_mapping_path: 原始映射关系JSON文件路径
+        inverse_mapping_path: 反向映射关系JSON文件路径
+        frequency_data_path: 包含使用频率数据的JSON文件路径
+        output_dir: 输出新映射文件的目录，默认为原文件所在目录
+    """
+    # 读取原正向映射关系
+    with open(original_mapping_path, 'r', encoding='utf-8') as f:
+        original_mapping = json.load(f)
+    
+    # 读取原反向映射关系
+    with open(inverse_mapping_path, 'r', encoding='utf-8') as f:
+        inverse_mapping = json.load(f)
+    
+    # 读取频率数据
+    with open(frequency_data_path, 'r', encoding='utf-8') as f:
+        frequency_data = json.load(f)
+    
+    # 创建part_name到usage_count的映射
+    part_frequency = {item['part_name']: item['usage_count'] for item in frequency_data}
+    
+    # 提取所有part_name，确保不遗漏任何数据
+    all_parts = set(original_mapping.values())
+    all_parts.update(inverse_mapping.keys())
+    all_parts.update(part_frequency.keys())
+    
+    # 按照使用频率排序，频率高的在前
+    # 对于没有频率数据的part，使用0作为默认值
+    sorted_parts = sorted(all_parts, key=lambda x: part_frequency.get(x, 0), reverse=True)
+    
+    # 生成新的正向映射关系，ID从0开始
+    new_forward_mapping = {str(i): part for i, part in enumerate(sorted_parts)}
+    
+    # 生成新的反向映射关系 (part_name -> ID)
+    #new_inverse_mapping = {part: str(i) for i, part in enumerate(sorted_parts)}
+    new_inverse_mapping = {part: int(i) for i, part in enumerate(sorted_parts)}
+    
+    # 确定输出目录
+    if output_dir is None:
+        output_dir = os.path.dirname(original_mapping_path)
+    
+    # 构建输出文件路径
+    forward_output = os.path.join(output_dir, "new_forward_mapping.json")
+    inverse_output = os.path.join(output_dir, "new_inverse_mapping.json")
+    
+    # 保存新的正向映射
+    with open(forward_output, 'w', encoding='utf-8') as f:
+        json.dump(new_forward_mapping, f, ensure_ascii=False, indent=4)
+    
+    # 保存新的反向映射
+    with open(inverse_output, 'w', encoding='utf-8') as f:
+        json.dump(new_inverse_mapping, f, ensure_ascii=False, indent=4)
+    
+    print(f"新的正向映射已生成: {forward_output}")
+    print(f"新的反向映射已生成: {inverse_output}")
+    print(f"共处理 {len(new_forward_mapping)} 个part")
+    
+    return new_forward_mapping, new_inverse_mapping
+
+if __name__ == "__main__":
+    # 示例用法，根据实际情况修改文件路径
+    original_mapping_file = "/public/home/wangshuo/gap/assembly/data/car_1k/subset_self/label_mapping.json"  # 原正向映射文件
+    inverse_mapping_file = "/public/home/wangshuo/gap/assembly/data/car_1k/subset_self/label_inverse_mapping.json"   # 原反向映射文件
+    frequency_data_file = "/public/home/wangshuo/gap/assembly/data/car_1k/subset_self/label_frequency.json"     # 频率数据文件
+    
+    # 处理映射关系
+    process_mappings(original_mapping_file, inverse_mapping_file, frequency_data_file)
+