charles/monoloco
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

create torch_dir if doesn't exists

Browse Source
This commit is contained in:
Lorenzo 2021-05-14 10:17:39 +02:00
parent e34f68f5a4
commit 5c5ce02fc1
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
monoloco/predict.py
View File

@ -53,13 +53,15 @@ def get_torch_checkpoints_dir():
def download_checkpoints(args):
torch_dir = get_torch_checkpoints_dir()
os.makedirs(torch_dir, exist_ok=True)
if args.checkpoint is None:
os.makedirs(torch_dir, exist_ok=True)
pifpaf_model = os.path.join(torch_dir, 'shufflenetv2k30-201104-224654-cocokp-d75ed641.pkl')
else:
pifpaf_model = args.checkpoint
dic_models = {'keypoints': pifpaf_model}
if not os.path.exists(pifpaf_model):
assert DOWNLOAD is not None, "pip install gdown to download pifpaf model, or pass it as --checkpoint"
assert DOWNLOAD is not None, "install gdown to download pifpaf model, or pass it as --checkpoint"
LOG.info('Downloading OpenPifPaf model in %s', torch_dir)
DOWNLOAD(OPENPIFPAF_MODEL, pifpaf_model, quiet=False)
@ -83,6 +85,7 @@ def download_checkpoints(args):
model = os.path.join(torch_dir, name)
dic_models[args.mode] = model
if not os.path.exists(model):
os.makedirs(torch_dir, exist_ok=True)
assert DOWNLOAD is not None, "pip install gdown to download monoloco model, or pass it as --model"
LOG.info('Downloading model in %s', torch_dir)
DOWNLOAD(path, model, quiet=False)