Add tqdm bar, support for .txt dataset format

train.py

@@ -470,7 +470,6 @@ def train(hyp, opt, device, tb_writer=None):
     parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
     parser.add_argument('--quad', action='store_true', help='quad dataloader')
     parser.add_argument('--linear-lr', action='store_true', help='linear LR')
-    # W&B logging Args:
     parser.add_argument('--log-imgs', type=int, default=16, help='number of images for W&B logging, max 100')
     parser.add_argument('--bbox_interval', type=int, default=-1, help='Set bounding-box image logging interval for W&B')
     parser.add_argument('--save_period', type=int, default=-1, help='Save model artifact after every "save_period" epoch')

utils/wandb_logging/log_dataset.py

@@ -15,12 +15,12 @@ def create_dataset_artifact(opt):
     logger = WandbLogger(opt, '', None, data, job_type='create_dataset')
     nc, names = (1, ['item']) if opt.single_cls else (int(data['nc']), data['names'])
     names = {k: v for k, v in enumerate(names)}  # to index dictionary
-    logger.log_dataset_artifact(LoadImagesAndLabels(data['train']), names, name='train')  # trainset
-    logger.log_dataset_artifact(LoadImagesAndLabels(data['val']), names, name='val')  # valset
-
-    # Update data.yaml with artifact links
-    data['train'] = WANDB_ARTIFACT_PREFIX + str(Path(opt.project) / 'train')
-    data['val'] = WANDB_ARTIFACT_PREFIX + str(Path(opt.project) / 'val')
+    logger.log_dataset_artifact(LoadImagesAndLabels(data['train']), names, name='train') if data.get('train') else None
+    logger.log_dataset_artifact(LoadImagesAndLabels(data['val']), names, name='val')  if data.get('val') else None
+    if data.get('train'):
+        data['train'] = WANDB_ARTIFACT_PREFIX + str(Path(opt.project) / 'train')
+    if data.get('val'):
+        data['val'] = WANDB_ARTIFACT_PREFIX + str(Path(opt.project) / 'val')
     path = opt.data if opt.overwrite_config else opt.data.replace('.', '_wandb.')  # updated data.yaml path
     data.pop('download', None)  # download via artifact instead of predefined field 'download:'
     with open(path, 'w') as f:

utils/wandb_logging/wandb_utils.py

@@ -4,10 +4,12 @@
 from datetime import datetime
 from pathlib import Path
 
+from tqdm import tqdm
 import torch
 
 sys.path.append(str(Path(__file__).parent.parent.parent))  # add utils/ to path
 from utils.general import colorstr, xywh2xyxy
+from utils.datasets import img2label_paths
 
 try:
     import wandb
@@ -64,16 +66,14 @@ def download_dataset_artifact(self, path, alias):
             assert dataset_artifact is not None, "'Error: W&B dataset artifact doesn\'t exist'"
             datadir = dataset_artifact.download()
             labels_zip = Path(datadir) / "data/labels.zip"
-            shutil.unpack_archive(labels_zip, Path(datadir) / 'data/labels', 'zip')
-            print("Downloaded dataset to : ", datadir)
+            shutil.unpack_archive(labels_zip, Path(datadir) / 'data/labels', 'zip') if labels_zip.exists() else None
             return datadir, dataset_artifact
         return None, None
 
     def download_model_artifact(self, name):
         model_artifact = wandb.use_artifact(name + ":latest")
         assert model_artifact is not None, 'Error: W&B model artifact doesn\'t exist'
         modeldir = model_artifact.download()
-        print("Downloaded model to : ", modeldir)
         return modeldir, model_artifact
 
     def log_model(self, path, opt, epoch):
@@ -92,16 +92,19 @@ def log_model(self, path, opt, epoch):
 
     def log_dataset_artifact(self, dataset, class_to_id, name='dataset'):
         artifact = wandb.Artifact(name=name, type="dataset")
-        image_path = dataset.path
-        artifact.add_dir(image_path, name='data/images')
+        for img_file in [dataset.path] if Path(dataset.path).is_dir() else dataset.img_files:
+            if Path(img_file).is_dir():
+                artifact.add_dir(img_file, name='data/images')
+            else:
+                artifact.add_file(img_file, name='data/images/'+Path(img_file).name)
+                label_file = Path(img2label_paths([img_file])[0])
+                artifact.add_file(str(label_file), name='data/labels/'+label_file.name) if label_file.exists() else None
         table = wandb.Table(columns=["id", "train_image", "Classes"])
         class_set = wandb.Classes([{'id': id, 'name': name} for id, name in class_to_id.items()])
-        for si, (img, labels, paths, shapes) in enumerate(dataset):
+        for si, (img, labels, paths, shapes) in enumerate(tqdm(dataset)):
             height, width = shapes[0]
-            labels[:, 2:] = (xywh2xyxy(labels[:, 2:].view(-1, 4)))
-            labels[:, 2:] *= torch.Tensor([width, height, width, height])
-            box_data = []
-            img_classes = {}
+            labels[:, 2:] = (xywh2xyxy(labels[:, 2:].view(-1, 4))) * torch.Tensor([width, height, width, height])
+            box_data, img_classes = [], {}
             for cls, *xyxy in labels[:, 1:].tolist():
                 cls = int(cls)
                 box_data.append({"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
@@ -113,13 +116,13 @@ def log_dataset_artifact(self, dataset, class_to_id, name='dataset'):
             boxes = {"ground_truth": {"box_data": box_data, "class_labels": class_to_id}}  # inference-space
             table.add_data(si, wandb.Image(paths, classes=class_set, boxes=boxes), json.dumps(img_classes))
         artifact.add(table, name)
-        labels_path = 'labels'.join(image_path.rsplit('images', 1))
-        zip_path = Path(labels_path).parent / (name + '_labels.zip')
-        if not zip_path.is_file():  # make_archive won't check if file exists
-            shutil.make_archive(zip_path.with_suffix(''), 'zip', labels_path)
-        artifact.add_file(str(zip_path), name='data/labels.zip')
+        if Path(dataset.path).is_dir():
+            labels_path = 'labels'.join(dataset.path.rsplit('images', 1))
+            zip_path = Path(labels_path).parent / (name + '_labels.zip')
+            if not zip_path.is_file():  # make_archive won't check if file exists
+                shutil.make_archive(zip_path.with_suffix(''), 'zip', labels_path)
+            artifact.add_file(str(zip_path), name='data/labels.zip')
         wandb.log_artifact(artifact)
-        print("Saving data to W&B...")
 
     def log(self, log_dict):
         if self.wandb_run: