stack86/CodeFormer

1import argparse2import os3from os import path as osp4 5from basicsr.utils.download_util import load_file_from_url6 7 8def download_pretrained_models(method, file_urls):9    save_path_root = f'./weights/{method}'10    os.makedirs(save_path_root, exist_ok=True)11 12    for file_name, file_url in file_urls.items():13        save_path = load_file_from_url(url=file_url, model_dir=save_path_root, progress=True, file_name=file_name)14 15 16if __name__ == '__main__':17    parser = argparse.ArgumentParser()18 19    parser.add_argument(20        'method',21        type=str,22        help=("Options: 'CodeFormer' 'facelib'. Set to 'all' to download all the models."))23    args = parser.parse_args()24 25    file_urls = {26        'CodeFormer': {27            'codeformer.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth'28        },29        'facelib': {30            # 'yolov5l-face.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/yolov5l-face.pth',31            'detection_Resnet50_Final.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/detection_Resnet50_Final.pth',32            'parsing_parsenet.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/parsing_parsenet.pth'33        }34    }35 36    if args.method == 'all':37        for method in file_urls.keys():38            download_pretrained_models(method, file_urls[method])39    else:40        download_pretrained_models(args.method, file_urls[args.method])