stack86/CodeFormer

1import argparse2import os3from os import path as osp4 5# from basicsr.utils.download_util import download_file_from_google_drive6import gdown7 8 9def download_pretrained_models(method, file_ids):10    save_path_root = f'./weights/{method}'11    os.makedirs(save_path_root, exist_ok=True)12 13    for file_name, file_id in file_ids.items():14        file_url = 'https://drive.google.com/uc?id='+file_id15        save_path = osp.abspath(osp.join(save_path_root, file_name))16        if osp.exists(save_path):17            user_response = input(f'{file_name} already exist. Do you want to cover it? Y/N\n')18            if user_response.lower() == 'y':19                print(f'Covering {file_name} to {save_path}')20                gdown.download(file_url, save_path, quiet=False)21                # download_file_from_google_drive(file_id, save_path)22            elif user_response.lower() == 'n':23                print(f'Skipping {file_name}')24            else:25                raise ValueError('Wrong input. Only accepts Y/N.')26        else:27            print(f'Downloading {file_name} to {save_path}')28            gdown.download(file_url, save_path, quiet=False)29            # download_file_from_google_drive(file_id, save_path)30 31if __name__ == '__main__':32    parser = argparse.ArgumentParser()33 34    parser.add_argument(35        'method',36        type=str,37        help=("Options: 'CodeFormer' 'facelib'. Set to 'all' to download all the models."))38    args = parser.parse_args()39 40    # file name: file id41    # 'dlib': {42    #     'mmod_human_face_detector-4cb19393.dat': '1qD-OqY8M6j4PWUP_FtqfwUPFPRMu6ubX',43    #     'shape_predictor_5_face_landmarks-c4b1e980.dat': '1vF3WBUApw4662v9Pw6wke3uk1qxnmLdg',44    #     'shape_predictor_68_face_landmarks-fbdc2cb8.dat': '1tJyIVdCHaU6IDMDx86BZCxLGZfsWB8yq'45    # }46    file_ids = {47        'CodeFormer': {48            'codeformer.pth': '1v_E_vZvP-dQPF55Kc5SRCjaKTQXDz-JB'49        },50        'facelib': {51            'yolov5l-face.pth': '131578zMA6B2x8VQHyHfa6GEPtulMCNzV',52            'parsing_parsenet.pth': '16pkohyZZ8ViHGBk3QtVqxLZKzdo466bK'53        }54    }55 56    if args.method == 'all':57        for method in file_ids.keys():58            download_pretrained_models(method, file_ids[method])59    else:60        download_pretrained_models(args.method, file_ids[args.method])