import os

# 指定目录路径
directory_path = "/home/anmin/code/SLfM/slfm/checkpoints/HM3D/wsSL"  # 请替换为实际目录的路径

'''
zip(['vision','audio', 'pretext', 'optimizer'], [net_audio, net_pretext, optimizer])
args.exp, f'{name}_latest.pth.tar')
'''

def rename(o,t):
    audio_best_path = os.path.join(dir_path, o+"_best.pth.tar")
    new_path = os.path.join(dir_path, t+"_best.pth.tar")
    os.rename(audio_best_path, new_path)
    audio_best_path = os.path.join(dir_path, o+"_latest.pth.tar")
    new_path = os.path.join(dir_path, t+"_latest.pth.tar")
    os.rename(audio_best_path, new_path)
    print(f'Renamed {audio_best_path} to {new_path}')
            
flag = 0
# 获取目录下所有的二级目录
for root, dirs, files in os.walk(directory_path):
    for dir in dirs:
        # 获取二级目录的完整路径
        dir_path = os.path.join(root, dir)

        # 检查是否存在 "audio_best.pth.tar" 文件
        audio_best_path = os.path.join(dir_path, "pretext_best.pth.tar")
        if os.path.exists(audio_best_path):
            # 将 "audio_best.pth.tar" 重命名为 "pretext_best.pth.tar"
            rename('pretext','ws')
            # rename('audio','pretext')
            # rename('vision','audio')    
            
            
            flag += 1
            
print(f'Total {flag} files renamed.')
