Ayush Chaurasia glenn-jocher commited on
Commit
96fcde4
·
unverified ·
1 Parent(s): 2062765

W&B feature improvements (#1258)

Browse files

* W&B feature improvements

This PR add:
* Class to id labels. Now, the caption of bounding boxes will display the class name and the class confidence score.
* The project name is set to "Yolov5" and the run name will be set to opt.logdir

* cleanup

* remove parenthesis on caption

Co-authored-by: Glenn Jocher <[email protected]>

Files changed (2) hide show
  1. test.py +8 -6
  2. train.py +1 -1
test.py CHANGED
@@ -95,7 +95,7 @@ def test(data,
95
  hyp=None, augment=False, cache=False, pad=0.5, rect=True)[0]
96
 
97
  seen = 0
98
- names = model.names if hasattr(model, 'names') else model.module.names
99
  coco91class = coco80_to_coco91_class()
100
  s = ('%20s' + '%12s' * 6) % ('Class', 'Images', 'Targets', 'P', 'R', '[email protected]', '[email protected]:.95')
101
  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,
150
 
151
  # W&B logging
152
  if len(wandb_images) < log_imgs:
153
- bbox_data = [{"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
154
- "class_id": int(cls),
155
- "scores": {"class_score": conf},
156
- "domain": "pixel"} for *xyxy, conf, cls in pred.clone().tolist()]
157
- wandb_images.append(wandb.Image(img[si], boxes={"predictions": {"box_data": bbox_data}}))
 
 
158
 
159
  # Clip boxes to image bounds
160
  clip_coords(pred, (height, width))
 
95
  hyp=None, augment=False, cache=False, pad=0.5, rect=True)[0]
96
 
97
  seen = 0
98
+ names = {k: v for k, v in enumerate(model.names if hasattr(model, 'names') else model.module.names)}
99
  coco91class = coco80_to_coco91_class()
100
  s = ('%20s' + '%12s' * 6) % ('Class', 'Images', 'Targets', 'P', 'R', '[email protected]', '[email protected]:.95')
101
  p, r, f1, mp, mr, map50, map, t0, t1 = 0., 0., 0., 0., 0., 0., 0., 0., 0.
 
150
 
151
  # W&B logging
152
  if len(wandb_images) < log_imgs:
153
+ box_data = [{"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
154
+ "class_id": int(cls),
155
+ "box_caption": "%s %.3f" % (names[cls], conf),
156
+ "scores": {"class_score": conf},
157
+ "domain": "pixel"} for *xyxy, conf, cls in pred.clone().tolist()]
158
+ boxes = {"predictions": {"box_data": box_data, "class_labels": names}}
159
+ wandb_images.append(wandb.Image(img[si], boxes=boxes))
160
 
161
  # Clip boxes to image bounds
162
  clip_coords(pred, (height, width))
train.py CHANGED
@@ -121,7 +121,7 @@ def train(hyp, opt, device, tb_writer=None, wandb=None):
121
  # Logging
122
  if wandb and wandb.run is None:
123
  id = ckpt.get('wandb_id') if 'ckpt' in locals() else None
124
- wandb_run = wandb.init(config=opt, resume="allow", project=os.path.basename(log_dir), id=id)
125
 
126
  # Resume
127
  start_epoch, best_fitness = 0, 0.0
 
121
  # Logging
122
  if wandb and wandb.run is None:
123
  id = ckpt.get('wandb_id') if 'ckpt' in locals() else None
124
+ wandb_run = wandb.init(config=opt, resume="allow", project="YOLOv5", name=os.path.basename(log_dir), id=id)
125
 
126
  # Resume
127
  start_epoch, best_fitness = 0, 0.0