Weights & Biases (W&B) Feature Addition (#1235)

* Add wandb metric logging and bounding box debugging

* Improve formatting, readability

* Remove mutliple path for init, improve formatting

* Add wandb params

* Remove typecasting in bbox coordinates and reformat

* Cleanup

* add wandb to requirements.txt

* minor updates to test.py

* general reorg

* reduce --log-imgs to 10

* clean wandb import

* reverse wandb import assert

* add except AssertionError to try import

* move wandb init to all global ranks

* replace print() with logger.info()

* replace print() with logger.info()

* move wandb.init() bug fix

* project PosixPath to basename bug fix

Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>

test.py

@@ -33,7 +33,9 @@ def test(data,
          save_dir=Path(''),  # for saving images
          save_txt=False,  # for auto-labelling
          save_conf=False,
-         plots=True):
+         plots=True,
+         log_imgs=0):  # number of logged images
+
     # Initialize/load model and set device
     training = model is not None
     if training:  # called by train.py
@@ -77,6 +79,13 @@ def test(data,
     iouv = torch.linspace(0.5, 0.95, 10).to(device)  # iou vector for mAP@0.5:0.95
     niou = iouv.numel()
 
+    # Logging
+    log_imgs = min(log_imgs, 100)  # ceil
+    try:
+        import wandb  # Weights & Biases
+    except ImportError:
+        log_imgs = 0
+
     # Dataloader
     if not training:
         img = torch.zeros((1, 3, imgsz, imgsz), device=device)  # init img
@@ -91,7 +100,7 @@ def test(data,
     s = ('%20s' + '%12s' * 6) % ('Class', 'Images', 'Targets', 'P', 'R', 'mAP@.5', 'mAP@.5:.95')
     p, r, f1, mp, mr, map50, map, t0, t1 = 0., 0., 0., 0., 0., 0., 0., 0., 0.
     loss = torch.zeros(3, device=device)
-    jdict, stats, ap, ap_class = [], [], [], []
+    jdict, stats, ap, ap_class, wandb_images = [], [], [], [], []
     for batch_i, (img, targets, paths, shapes) in enumerate(tqdm(dataloader, desc=s)):
         img = img.to(device, non_blocking=True)
         img = img.half() if half else img.float()  # uint8 to fp16/32
@@ -139,6 +148,14 @@ def test(data,
                     with open(str(out / Path(paths[si]).stem) + '.txt', 'a') as f:
                         f.write(('%g ' * len(line) + '\n') % line)
 
+            # W&B logging
+            if len(wandb_images) < log_imgs:
+                bbox_data = [{"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
+                              "class_id": int(cls),
+                              "scores": {"class_score": conf},
+                              "domain": "pixel"} for *xyxy, conf, cls in pred.clone().tolist()]
+                wandb_images.append(wandb.Image(img[si], boxes={"predictions": {"box_data": bbox_data}}))
+
             # Clip boxes to image bounds
             clip_coords(pred, (height, width))
 
@@ -196,6 +213,10 @@ def test(data,
             f = save_dir / f'test_batch{batch_i}_pred.jpg'
             plot_images(img, output_to_target(output, width, height), paths, str(f), names)  # predictions
 
+    # W&B logging
+    if wandb_images:
+        wandb.log({"outputs": wandb_images})
+
     # Compute statistics
     stats = [np.concatenate(x, 0) for x in zip(*stats)]  # to numpy
     if len(stats) and stats[0].any():