File size: 2,598 Bytes
2dc379c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 | import os
import warnings
pretrained_sevirlr_vae_name = "pretrained_sevirlr_vae_8x8x64_v1.pt"
pretrained_sevirlr_earthformerunet_name = "pretrained_sevirlr_earthformerunet_v1.pt"
pretrained_sevirlr_alignment_name = "pretrained_sevirlr_alignment_avg_x_cuboid_v1.pt"
pretrained_i3d_400_name = "pretrained_i3d_400.pt"
pretrained_i3d_600_name = "pretrained_i3d_600.pt"
# file_id_dict = {
# pretrained_sevirlr_vae_name: "10OicEQuOPzSKDp5WYF3zDHsL-COywe98",
# pretrained_sevirlr_earthformerunet_name: "1cVB0Sm2V4OMTLxNNEXlb2__ONqSAjUDJ",
# pretrained_sevirlr_alignment_name: "1CzrzNVDVTyc8ivnWqEOiJnrN6wraY6fy",
# pretrained_i3d_400_name: "1UwfpQEIi1jJqGmTmRWcIClEpXF3jKwOq",
# pretrained_i3d_600_name: "1wiueaTRfi0DqJUUaIF3l_ze8OMNF5NIQ",
# }
file_id_dict = {
pretrained_sevirlr_vae_name: "EZ-BfMnX-dhFrccXtLuCB78BPs3xBc-ke1BG4_9Jf7UblQ?e=bMplKs",
pretrained_sevirlr_earthformerunet_name: "EUp432GSXplCuhN7-WUF6YsBK4WrWiL8A9RZCA1Pf9D0Ag?e=eY6hcQ",
pretrained_sevirlr_alignment_name: "EWHH7y6w9D5Dg01YNX99IFQBA3tCR2a7s7z7Xv0BiNLV7Q?e=Hd9jCH",
pretrained_i3d_400_name: "ESSxcaYvlrlAvXnsJoQ8P-kBusWJiM1D8pOu7wcNEVmzcw?e=HqWz8F",
pretrained_i3d_600_name: "EU6tZPSExoZIgoAwu5hTkjoBBJu1RBFepjFsP68Msb6JFA?e=31aIfa",
}
def download_pretrained_weights(ckpt_name, save_dir=None, exist_ok=False):
r"""
Download pretrained weights from Google Drive.
Parameters
----------
ckpt_name: str
save_dir: str
exist_ok: bool
"""
ckpt_path = os.path.join(save_dir, ckpt_name)
if os.path.exists(ckpt_path) and not exist_ok:
warnings.warn(f"Checkpoint file {ckpt_path} already exists!")
else:
os.makedirs(save_dir, exist_ok=True)
file_id = file_id_dict[ckpt_name]
# os.system(f"wget --load-cookies /tmp/cookies.txt "
# f"\"https://docs.google.com/uc?export=download&confirm="
# f"$(wget --quiet --save-cookies /tmp/cookies.txt --keep-session-cookies --no-check-certificate "
# f"'https://docs.google.com/uc?export=download&id={file_id}'"
# f" -O- | sed -rn 's/.*confirm=([0-9A-Za-z_]+).*/\1\n/p')"
# f"&id={file_id}\" -O {ckpt_path} && rm -rf /tmp/cookies.txt")
os.system(f"wget https://hkustconnect-my.sharepoint.com/:u:/g/personal/zgaoag_connect_ust_hk/{file_id}"
f"\&download=1 -O {ckpt_path}")
if __name__ == "__main__":
download_pretrained_weights("pretrained_i3d_400.pt",save_dir = "/cm/archive/vinhbk1/prediff_log_temp/pretrained_weights/metrics",exist_ok=True) |