From c98128fe71a8676037a0605ab389c7473c743d07 Mon Sep 17 00:00:00 2001 From: KristenKehrer <34010022+KristenKehrer@users.noreply.github.com> Date: Sun, 2 Oct 2022 18:25:10 -0400 Subject: [PATCH] Allow list for Comet artifact class 'names' field (#9654) * Update __init__.py In the Comet logger, when I run train.py, it wants to download the data artifact. It was requiring me to format the 'names' field in the data artifact metadata as a dictionary, so I've changed this so that it also accepts a list. Signed-off-by: KristenKehrer <34010022+KristenKehrer@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update utils/loggers/comet/__init__.py Co-authored-by: Dhruv Nair Signed-off-by: KristenKehrer <34010022+KristenKehrer@users.noreply.github.com> Signed-off-by: KristenKehrer <34010022+KristenKehrer@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Dhruv Nair Co-authored-by: Glenn Jocher --- utils/loggers/comet/__init__.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/utils/loggers/comet/__init__.py b/utils/loggers/comet/__init__.py index ba5cecc..b0318f8 100644 --- a/utils/loggers/comet/__init__.py +++ b/utils/loggers/comet/__init__.py @@ -353,7 +353,14 @@ class CometLogger: metadata = logged_artifact.metadata data_dict = metadata.copy() data_dict["path"] = artifact_save_dir - data_dict["names"] = {int(k): v for k, v in metadata.get("names").items()} + + metadata_names = metadata.get("names") + if type(metadata_names) == dict: + data_dict["names"] = {int(k): v for k, v in metadata.get("names").items()} + elif type(metadata_names) == list: + data_dict["names"] = {int(k): v for k, v in zip(range(len(metadata_names)), metadata_names)} + else: + raise "Invalid 'names' field in dataset yaml file. Please use a list or dictionary" data_dict = self.update_data_paths(data_dict) return data_dict