def download_pretrained_models(method, file_ids):
save_path_root = f'./experiments/pretrained_models/{method}'
os.makedirs(save_path_root, exist_ok=True)
for file_name, file_id in file_ids.items():
save_path = osp.abspath(osp.join(save_path_root, file_name))
if osp.exists(save_path):
user_response = input(f'{file_name} already exist. Do you want to cover it? Y/N\n')
if user_response.lower() == 'y':
print(f'Covering {file_name} to {save_path}')
download_file_from_google_drive(file_id, save_path)
elif user_response.lower() == 'n':
print(f'Skipping {file_name}')
else:
raise ValueError('Wrong input. Only accepts Y/N.')
else:
print(f'Downloading {file_name} to {save_path}')
download_file_from_google_drive(file_id, save_path)
def download_dataset(dataset, file_ids):
save_path_root = './datasets/'
os.makedirs(save_path_root, exist_ok=True)
for file_name, file_id in file_ids.items():
save_path = osp.abspath(osp.join(save_path_root, file_name))
if osp.exists(save_path):
user_response = input(
f'{file_name} already exist. Do you want to cover it? Y/N\n')
if user_response.lower() == 'y':
print(f'Covering {file_name} to {save_path}')
download_file_from_google_drive(file_id, save_path)
elif user_response.lower() == 'n':
print(f'Skipping {file_name}')
else:
raise ValueError('Wrong input. Only accpets Y/N.')
else:
print(f'Downloading {file_name} to {save_path}')
download_file_from_google_drive(file_id, save_path)
# unzip
if save_path.endswith('.zip'):
extracted_path = save_path.replace('.zip', '')
print(f'Extract {save_path} to {extracted_path}')
import zipfile
with zipfile.ZipFile(save_path, 'r') as zip_ref:
zip_ref.extractall(extracted_path)
file_name = file_name.replace('.zip', '')
subfolder = osp.join(extracted_path, file_name)
if osp.isdir(subfolder):
print(f'Move {subfolder} to {extracted_path}')
import shutil
for path in glob.glob(osp.join(subfolder, '*')):
shutil.move(path, extracted_path)
shutil.rmtree(subfolder)
import argparse
from basicsr.utils.download_util import download_file_from_google_drive
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--id', type=str, help='File id')
parser.add_argument('--output', type=str, help='Save path')
args = parser.parse_args()
download_file_from_google_drive(args.id, args.save_path)