Add code/cube3d/training/rotate_bacthLdr.py

code/cube3d/training/rotate_bacthLdr.py

@@ -0,0 +1,106 @@
+import os
+import numpy as np
+import pandas as pd
+from math import cos, sin, radians
+# 从原有文件导入所需函数
+from  rotateLdr import (checkEncoding, readTextFile, 
+                          apply_rotation, apply_rotation_matrix,
+                          getModulePart, process_ldr)
+
+def create_rotation_matrix(y_rotation_code):
+    """
+    根据旋转编码创建Y轴旋转矩阵
+    0: 不旋转, 1: 绕Y轴90度, 2: 绕Y轴180度, 3: 绕Y轴270度
+    """
+    rotation_angles = {0: 0, 1: -90, 2: -180, 3: -270}
+    y_angle = rotation_angles.get(y_rotation_code, 0)
+    y_rad = radians(y_angle)
+    
+    # 绕Y轴旋转矩阵
+    rot_mat = [
+        [cos(y_rad), 0, sin(y_rad)],
+        [0, 1, 0],
+        [-sin(y_rad), 0, cos(y_rad)]
+    ]
+    return rot_mat, y_angle
+
+def extract_prefix(filename):
+    """从文件名中提取前缀（去除扩展名和可能的_车底）"""
+    # 先去除文件扩展名
+    prefix = os.path.splitext(str(filename))[0]
+    # 去除可能存在的_车底后缀
+    if prefix.endswith("_车底"):
+        prefix = prefix[:-3]  # "_车底"是3个字符，截取前面部分
+    return prefix
+
+def batch_process_ldr(input_dir, output_dir, excel_path):
+    """
+    批量处理LDR文件，通过前缀匹配Excel中的记录
+    """
+    # 创建输出目录
+    os.makedirs(output_dir, exist_ok=True)
+    
+    # 读取Excel文件
+    try:
+        df = pd.read_excel(excel_path, engine='openpyxl')
+        # 确保旋转编码列是整数类型
+        rotation_code_col = df.columns[1]  # 第二列是旋转编码
+        df[rotation_code_col] = df[rotation_code_col].astype(int, errors='ignore')
+        
+        # 提取文件名前缀并添加到DataFrame
+        df['prefix'] = df['Name'].apply(extract_prefix)
+        print(f"成功读取Excel文件，共 {len(df)} 条记录")
+    except Exception as e:
+        print(f"读取Excel文件出错: {e}")
+        return
+    
+    # 获取所有LDR文件
+    ldr_files = [f for f in os.listdir(input_dir) if f.lower().endswith('.ldr')]
+    print(f"在 {input_dir} 中找到 {len(ldr_files)} 个LDR文件")
+    
+    # 处理每个LDR文件
+    processed_count = 0
+    for ldr_file in ldr_files:
+        # 提取LDR文件的前缀和扩展名
+        file_base, file_ext = os.path.splitext(ldr_file)
+        # 构建带_rot后缀的输出文件名 (例如 car_29.ldr → car_29_rot.ldr)
+        output_filename = f"{file_base}_rot{file_ext}"
+        # 提取LDR文件的前缀
+        ldr_prefix = extract_prefix(ldr_file)
+        
+        # 查找匹配的前缀
+        matched_rows = df[df['prefix'] == ldr_prefix]
+        if matched_rows.empty:
+            print(f"警告: 在Excel中未找到与 {ldr_file} (前缀: {ldr_prefix}) 匹配的记录，跳过")
+            continue
+        
+        # 获取旋转编码
+        try:
+            rotation_code = int(matched_rows.iloc[0][rotation_code_col])
+            if rotation_code not in [0, 1, 2, 3]:
+                print(f"警告: {ldr_file} 的旋转编码 {rotation_code} 无效，使用默认值0")
+                rotation_code = 0
+        except (ValueError, IndexError, KeyError) as e:
+            print(f"警告: {ldr_file} 的旋转编码获取失败: {e}，使用默认值0")
+            rotation_code = 0
+        
+        # 创建旋转矩阵
+        rot_mat, y_angle = create_rotation_matrix(rotation_code)
+        print(f"处理文件: {ldr_file}, 旋转编码: {rotation_code} (绕Y轴旋转 {y_angle} 度)")
+        
+        # 处理文件
+        input_path = os.path.join(input_dir, ldr_file)
+        output_path = os.path.join(output_dir, output_filename)
+        
+        if process_ldr(input_path, output_path, rot_mat):
+            processed_count += 1
+    
+    print(f"批量处理完成，成功处理 {processed_count}/{len(ldr_files)} 个文件")
+
+
+if __name__ == "__main__":
+    INPUT_DIR = "/public/home/wangshuo/gap/assembly/data/car_1k/subset_bottom_300/ldr_bottom_gather_test/"  # 包含多个LDR文件的目录
+    OUTPUT_DIR = "/public/home/wangshuo/gap/assembly/data/car_1k/subset_bottom_300/ldr_rot_test/"  # 输出目录
+    EXCEL_PATH = "/public/home/wangshuo/gap/assembly/data/car_1k/subset_bottom_300/Pose_test.xlsx"  # 包含旋转角度的Excel文件
+    batch_process_ldr(INPUT_DIR, OUTPUT_DIR, EXCEL_PATH)
+