Fixed wandb logger KeyError (ultralytics#6637)

fhkiel-mlaip · Feb 14, 2022 · c679476 · c679476
1 parent ce6e5b3
commit c679476
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/utils/loggers/wandb/wandb_utils.py b/utils/loggers/wandb/wandb_utils.py
@@ -487,7 +487,7 @@ def val_one_image(self, pred, predn, path, names, im):
             if self.current_epoch % self.bbox_interval == 0:
                 box_data = [{"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
                              "class_id": int(cls),
-                             "box_caption": f"{names[cls]} {conf:.3f}",
+                             "box_caption": f"{names[int(cls)]} {conf:.3f}",
                              "scores": {"class_score": conf},
                              "domain": "pixel"} for *xyxy, conf, cls in pred.tolist()]
                 boxes = {"predictions": {"box_data": box_data, "class_labels": names}}  # inference-space