diff --git a/test.py b/test.py
index e274990..a270fbc 100644
--- a/test.py
+++ b/test.py
@@ -95,7 +95,7 @@ def test(data,
                                        hyp=None, augment=False, cache=False, pad=0.5, rect=True)[0]
 
     seen = 0
-    names = model.names if hasattr(model, 'names') else model.module.names
+    names = {k: v for k, v in enumerate(model.names if hasattr(model, 'names') else model.module.names)}
     coco91class = coco80_to_coco91_class()
     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.
@@ -150,11 +150,13 @@ def test(data,
 
             # 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}}))
+                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.clone().tolist()]
+                boxes = {"predictions": {"box_data": box_data, "class_labels": names}}
+                wandb_images.append(wandb.Image(img[si], boxes=boxes))
 
             # Clip boxes to image bounds
             clip_coords(pred, (height, width))
diff --git a/train.py b/train.py
index 86d2b9d..2df54ca 100644
--- a/train.py
+++ b/train.py
@@ -121,7 +121,7 @@ def train(hyp, opt, device, tb_writer=None, wandb=None):
     # Logging
     if wandb and wandb.run is None:
         id = ckpt.get('wandb_id') if 'ckpt' in locals() else None
-        wandb_run = wandb.init(config=opt, resume="allow", project=os.path.basename(log_dir), id=id)
+        wandb_run = wandb.init(config=opt, resume="allow", project="YOLOv5", name=os.path.basename(log_dir), id=id)
 
     # Resume
     start_epoch, best_fitness = 0, 0.0