Improved W&B integration (#2125)

test.py

@@ -35,8 +35,9 @@ def test(data,
          save_hybrid=False,  # for hybrid auto-labelling
          save_conf=False,  # save auto-label confidences
          plots=True,
-         log_imgs=0,  # number of logged images
-         compute_loss=None):
+         wandb_logger=None,
+         compute_loss=None,
+         is_coco=False):
     # Initialize/load model and set device
     training = model is not None
     if training:  # called by train.py
@@ -66,21 +67,19 @@ def test(data,
 
     # Configure
     model.eval()
-    is_coco = data.endswith('coco.yaml')  # is COCO dataset
+    if isinstance(data, str):
+        is_coco = data.endswith('coco.yaml')
         with open(data) as f:
-        data = yaml.load(f, Loader=yaml.SafeLoader)  # model dict
+            data = yaml.load(f, Loader=yaml.SafeLoader)
     check_dataset(data)  # check
     nc = 1 if single_cls else int(data['nc'])  # number of classes
     iouv = torch.linspace(0.5, 0.95, 10).to(device)  # iou vector for mAP@0.5:0.95
     niou = iouv.numel()
 
     # Logging
-    log_imgs, wandb = min(log_imgs, 100), None  # ceil
-    try:
-        import wandb  # Weights & Biases
-    except ImportError:
     log_imgs = 0
-
+    if wandb_logger and wandb_logger.wandb:
+        log_imgs = min(wandb_logger.log_imgs, 100)
     # Dataloader
     if not training:
         if device.type != 'cpu':
@@ -147,15 +146,17 @@ def test(data,
                     with open(save_dir / 'labels' / (path.stem + '.txt'), 'a') as f:
                         f.write(('%g ' * len(line)).rstrip() % line + '\n')
 
-            # W&B logging
-            if plots and len(wandb_images) < log_imgs:
+            # W&B logging - Media Panel Plots
+            if len(wandb_images) < log_imgs and wandb_logger.current_epoch > 0:  # Check for test operation
+                if wandb_logger.current_epoch % wandb_logger.bbox_interval == 0:
                     box_data = [{"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
                                  "class_id": int(cls),
                                  "box_caption": "%s %.3f" % (names[cls], conf),
                                  "scores": {"class_score": conf},
                                  "domain": "pixel"} for *xyxy, conf, cls in pred.tolist()]
                     boxes = {"predictions": {"box_data": box_data, "class_labels": names}}  # inference-space
-                wandb_images.append(wandb.Image(img[si], boxes=boxes, caption=path.name))
+                    wandb_images.append(wandb_logger.wandb.Image(img[si], boxes=boxes, caption=path.name))
+                wandb_logger.log_training_progress(predn, path, names)  # logs dsviz tables
 
             # Append to pycocotools JSON dictionary
             if save_json:
@@ -239,9 +240,11 @@ def test(data,
     # Plots
     if plots:
         confusion_matrix.plot(save_dir=save_dir, names=list(names.values()))
-        if wandb and wandb.run:
-            val_batches = [wandb.Image(str(f), caption=f.name) for f in sorted(save_dir.glob('test*.jpg'))]
-            wandb.log({"Images": wandb_images, "Validation": val_batches}, commit=False)
+        if wandb_logger and wandb_logger.wandb:
+            val_batches = [wandb_logger.wandb.Image(str(f), caption=f.name) for f in sorted(save_dir.glob('test*.jpg'))]
+            wandb_logger.log({"Validation": val_batches})
+    if wandb_images:
+        wandb_logger.log({"Bounding Box Debugger/Images": wandb_images})
 
     # Save JSON
     if save_json and len(jdict):

utils/wandb_logging/log_dataset.py

@@ -12,20 +12,7 @@ WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'
 def create_dataset_artifact(opt):
     with open(opt.data) as f:
         data = yaml.load(f, Loader=yaml.SafeLoader)  # data dict
-    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')
-    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:
-        yaml.dump(data, f)
-    print("New Config file => ", path)
+    logger = WandbLogger(opt, '', None, data, job_type='Dataset Creation')
 
 
 if __name__ == '__main__':
@@ -33,7 +20,6 @@ if __name__ == '__main__':
     parser.add_argument('--data', type=str, default='data/coco128.yaml', help='data.yaml path')
     parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset')
     parser.add_argument('--project', type=str, default='YOLOv5', help='name of W&B Project')
-    parser.add_argument('--overwrite_config', action='store_true', help='overwrite data.yaml')
     opt = parser.parse_args()
 
     create_dataset_artifact(opt)