From 5c5ce02fc1e8dde6d91c717cc52974bb9b1d11cf Mon Sep 17 00:00:00 2001 From: Lorenzo Date: Fri, 14 May 2021 10:17:39 +0200 Subject: [PATCH] create torch_dir if doesn't exists --- monoloco/predict.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/monoloco/predict.py b/monoloco/predict.py index e157027..66eabcf 100644 --- a/monoloco/predict.py +++ b/monoloco/predict.py @@ -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)