Update val.py

mrljwlm · Jan 16, 2022 · bcffd1f · bcffd1f
1 parent 1f4a99a
commit bcffd1f
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/val.py b/val.py
@@ -158,7 +158,7 @@ def run(data,
     iouv = torch.linspace(0.5, 0.95, 10).to(device)  # iou vector for [email protected]:0.95
     niou = iouv.numel()
 
-    names = {k: v for k, v in enumerate(model.names if hasattr(model, 'names') else model.module.names)}
+    names = model.names if hasattr(model, 'names') else model.module.names
     # Dataloader
     if not training:
         model.warmup(imgsz=(1, 3, imgsz, imgsz), half=half)  # warmup
@@ -169,7 +169,7 @@ def run(data,
 
     seen = 0
     confusion_matrix = ConfusionMatrix(nc=nc)
-    # names = {k: v for k, v in enumerate(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)}
     class_map = coco80_to_coco91_class() if is_coco else list(range(1000))
     s = ('%20s' + '%11s' * 6) % ('Class', 'Images', 'Labels', 'P', 'R', '[email protected]', '[email protected]:.95')
     dt, p, r, f1, mp, mr, map50, map = [0.0, 0.0, 0.0], 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0