Add code/cube3d/training/batch_process_ldrs.py

code/cube3d/training/batch_process_ldrs.py

@@ -0,0 +1,71 @@
+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)