import os
import pymeshlab
import gc
# 设置输入和输出文件夹
input_folder = r"D:\\project\\3dmodels\\modelnet\\test\\WatermarkedModels"


# 批量处理文件并转换格式
def convert_format(input_folder):
    skipped_files = []  # 用来存储跳过的文件

    # 只处理输入目录下的文件，不递归子目录
    for file in os.listdir(input_folder):
        input_file = os.path.join(input_folder, file)

        # 检查是否是文件
        if os.path.isfile(input_file):
            # 获取文件扩展名
            ext = file.split('.')[-1]

            # 打印处理的文件名和扩展名
            print(f"Processing: {input_file} (Extension: {ext})")

            # 创建新的保存目录 "typeTransed"（如果不存在）
            output_dir = os.path.join(input_folder, 'typeTransed')
            if not os.path.exists(output_dir):
                os.makedirs(output_dir)

            # 根据文件扩展名决定转换方式
            if file.endswith('.obj') or file.endswith('.ply') or file.endswith('.stl'):
                ms = pymeshlab.MeshSet()
                try:
                    ms.load_new_mesh(input_file)

                    # obj -> ply 和 obj -> stl
                    if file.endswith('.obj'):
                        ply_output = os.path.join(output_dir, file.replace('.obj', f'obj.ply'))
                        stl_output = os.path.join(output_dir, file.replace('.obj', f'obj.stl'))
                        print(f"Converting .obj to .ply and .stl: {ply_output}, {stl_output}")
                        ms.save_current_mesh(ply_output)
                        ms.save_current_mesh(stl_output)
                    
                    # ply -> obj 和 ply -> stl
                    elif file.endswith('.ply'):
                        obj_output = os.path.join(output_dir, file.replace('.ply', f'ply.obj'))
                        stl_output = os.path.join(output_dir, file.replace('.ply', f'ply.stl'))
                        print(f"Converting .ply to .obj and .stl: {obj_output}, {stl_output}")
                        ms.save_current_mesh(obj_output)
                        ms.save_current_mesh(stl_output)
                    
                    # stl -> obj 和 stl -> ply
                    elif file.endswith('.stl'):
                        obj_output = os.path.join(output_dir, file.replace('.stl', f'stl.obj'))
                        ply_output = os.path.join(output_dir, file.replace('.stl', f'stl.ply'))
                        print(f"Converting .stl to .obj and .ply: {obj_output}, {ply_output}")
                        ms.save_current_mesh(obj_output)
                        ms.save_current_mesh(ply_output)

                except MemoryError:
                    skipped_files.append(input_file)  # 将出错的文件记录下来
                    continue
                finally:
                    # 每处理完一个文件，手动回收内存
                    gc.collect()

    # 打印所有跳过的文件
    if skipped_files:
        print("\nThe following files were skipped due to memory errors:")
        for file in skipped_files:
            print(file)



# 批量处理文件（不做格式转换）
def process_mesh(input_folder):
    skipped_files = []  # 用来存储跳过的文件

    # 只处理输入目录下的文件，不递归子目录
    for file in os.listdir(input_folder):
        input_file = os.path.join(input_folder, file)

        # 检查是否是文件
        if os.path.isfile(input_file):
            # 获取文件扩展名
            ext = file.split('.')[-1]

            # 打印处理的文件名和扩展名
            print(f"Processing: {input_file} (Extension: {ext})")

            # 根据文件扩展名决定是否进行处理
            if file.endswith('.obj') or file.endswith('.ply') or file.endswith('.stl'):
                ms = pymeshlab.MeshSet()
                try:
                    ms.load_new_mesh(input_file)

                    # 对模型进行等比缩放、旋转和平移
                    # 等比缩放2倍
                    ms.apply_filter('compute_matrix_from_scaling_or_normalization', axisx=2.0, axisy=2.0, axisz=2.0, uniformflag=True, freeze=True)
                    # 旋转30度（绕Z轴旋转）
                    ms.apply_filter('compute_matrix_from_rotation', rotaxis='Z axis', rotcenter='origin', angle=30.0, freeze=True)
                    # 平移(10, 20, 30)
                    ms.apply_filter('compute_matrix_from_translation', traslmethod='XYZ translation', axisx=10.0, axisy=20.0, axisz=30.0, freeze=True)

                    # 保存处理后的模型（可选：保留原文件扩展名）
                    output_file = os.path.join(input_folder, "processed", file)  # 可以创建一个 "processed" 子目录来存放处理后的文件
                    os.makedirs(os.path.dirname(output_file), exist_ok=True)
                    ms.save_current_mesh(output_file)
                    print(f"Saved processed file: {output_file}")

                except MemoryError:
                    skipped_files.append(input_file)  # 将出错的文件记录下来
                    continue
                finally:
                    # 每处理完一个文件，手动回收内存
                    gc.collect()

    # 打印所有跳过的文件
    if skipped_files:
        print("\nThe following files were skipped due to memory errors:")
        for file in skipped_files:
            print(file)


# 执行攻击
convert_format(input_folder) #格式转换
process_mesh(input_folder) #几何攻击

print("批量攻击完成！")