Spaces:
Paused
Paused
| 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) | |