brep2sdf/data/scripts/gen_name_list.py

import os

'''
# 统计合法的数据
# 1。 有obj 
    在 ../gt_mesh_ 下面有 {id}.obj
# 2. 有点云
    在 ../gt_point 下面有 {id}_50k.xyz
# 3. 有二面角 
    在 ../gt_point 下面有 {id}.ptangle

保存id到 ../name_list.txt
'''

def get_valid_ids(check_pt_exists=False):
    # 定义路径
    mesh_dir = "../gt_mesh"
    point_dir = "../gt_point"
    pt_dir = "../output_data"
    
    # 获取所有可能的ID
    obj_files = [f for f in os.listdir(mesh_dir) if f.endswith('.obj')]
    potential_ids = [f[:-4] for f in obj_files]  # 去掉.obj后缀
    
    valid_ids = []
    
    for id in potential_ids:
        # 检查三个条件
        has_obj = os.path.exists(os.path.join(mesh_dir, f"{id}.obj"))
        has_xyz = os.path.exists(os.path.join(point_dir, f"{id}_50k.xyz"))
        has_ptangle = os.path.exists(os.path.join(point_dir, f"{id}.ptangle"))
        has_pt = not check_pt_exists or os.path.exists(os.path.join(pt_dir, f"{id}.pt"))
        if has_obj and has_xyz and has_ptangle and has_pt:
            valid_ids.append(id)
    
    return valid_ids

if __name__ == "__main__":
    valid_ids = get_valid_ids(check_pt_exists=True)
    print(f"找到 {len(valid_ids)} 个有效ID:")
    
    # 保存到文件
    with open("../name_list.txt", "w") as f:
        for id in valid_ids:
            print(id)  # 控制台输出
            f.write(f"{id}\n")  # 写入文件
    
    print("结果已保存到 ../name_list.txt")