Spaces:
Paused
Paused
File size: 4,303 Bytes
fa0e257 |
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 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 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)
|