From 0ccbef2f1225da08ba955f8db5c2070958c64214 Mon Sep 17 00:00:00 2001 From: Ayush Chaurasia Date: Tue, 6 Apr 2021 20:27:13 +0530 Subject: [PATCH] Add support for list-of-directory data format for wandb (#2719) --- utils/wandb_logging/wandb_utils.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/utils/wandb_logging/wandb_utils.py b/utils/wandb_logging/wandb_utils.py index 86038e199dc8..d407e6cd54fb 100644 --- a/utils/wandb_logging/wandb_utils.py +++ b/utils/wandb_logging/wandb_utils.py @@ -57,14 +57,14 @@ def process_wandb_config_ddp_mode(opt): with open(opt.data) as f: data_dict = yaml.load(f, Loader=yaml.SafeLoader) # data dict train_dir, val_dir = None, None - if data_dict['train'].startswith(WANDB_ARTIFACT_PREFIX): + if isinstance(data_dict['train'], str) and data_dict['train'].startswith(WANDB_ARTIFACT_PREFIX): api = wandb.Api() train_artifact = api.artifact(remove_prefix(data_dict['train']) + ':' + opt.artifact_alias) train_dir = train_artifact.download() train_path = Path(train_dir) / 'data/images/' data_dict['train'] = str(train_path) - if data_dict['val'].startswith(WANDB_ARTIFACT_PREFIX): + if isinstance(data_dict['val'], str) and data_dict['val'].startswith(WANDB_ARTIFACT_PREFIX): api = wandb.Api() val_artifact = api.artifact(remove_prefix(data_dict['val']) + ':' + opt.artifact_alias) val_dir = val_artifact.download() @@ -158,7 +158,7 @@ def setup_training(self, opt, data_dict): return data_dict def download_dataset_artifact(self, path, alias): - if path and path.startswith(WANDB_ARTIFACT_PREFIX): + if isinstance(path, str) and path.startswith(WANDB_ARTIFACT_PREFIX): dataset_artifact = wandb.use_artifact(remove_prefix(path, WANDB_ARTIFACT_PREFIX) + ":" + alias) assert dataset_artifact is not None, "'Error: W&B dataset artifact doesn\'t exist'" datadir = dataset_artifact.download() @@ -229,7 +229,9 @@ def map_val_table_path(self): def create_dataset_table(self, dataset, class_to_id, name='dataset'): # TODO: Explore multiprocessing to slpit this loop parallely| This is essential for speeding up the the logging artifact = wandb.Artifact(name=name, type="dataset") - for img_file in tqdm([dataset.path]) if Path(dataset.path).is_dir() else tqdm(dataset.img_files): + img_files = tqdm([dataset.path]) if isinstance(dataset.path, str) and Path(dataset.path).is_dir() else None + img_files = tqdm(dataset.img_files) if not img_files else img_files + for img_file in img_files: if Path(img_file).is_dir(): artifact.add_dir(img_file, name='data/images') labels_path = 'labels'.join(dataset.path.rsplit('images', 1))