per-class mAP report

models.py

@@ -99,6 +99,8 @@ class YOLOLayer(nn.Module):
         self.scaled_anchors = torch.FloatTensor([(a_w / stride, a_h / stride) for a_w, a_h in anchors])
         self.anchor_w = self.scaled_anchors[:, 0:1].view((1, nA, 1, 1))
         self.anchor_h = self.scaled_anchors[:, 1:2].view((1, nA, 1, 1))
+        self.weights = class_weights()
+
 
     def forward(self, p, targets=None, requestPrecision=False):
         FT = torch.cuda.FloatTensor if p.is_cuda else torch.FloatTensor
@@ -110,6 +112,7 @@ class YOLOLayer(nn.Module):
         if p.is_cuda and not self.grid_x.is_cuda:
             self.grid_x, self.grid_y = self.grid_x.cuda(), self.grid_y.cuda()
             self.anchor_w, self.anchor_h = self.anchor_w.cuda(), self.anchor_h.cuda()
+            self.weights = self.weights.cuda()
 
         # p.view(12, 255, 13, 13) -- > (12, 3, 13, 13, 80)  # (bs, anchors, grid, grid, classes + xywh)
         p = p.view(bs, self.nA, self.bbox_attrs, nG, nG).permute(0, 1, 3, 4, 2).contiguous()  # prediction

test.py

@@ -1,4 +1,5 @@
 import argparse
+
 from models import *
 from utils.datasets import *
 from utils.utils import *
@@ -48,9 +49,11 @@ dataloader = load_images_and_labels(test_path, batch_size=opt.batch_size, img_si
 
 print('Compute mAP...')
 
+nC = 80  # number of classes
 correct = 0
 targets = None
 outputs, mAPs, TP, confidence, pred_class, target_class = [], [], [], [], [], []
+AP_accum, AP_accum_count = np.zeros(nC), np.zeros(nC)
 for batch_i, (imgs, targets) in enumerate(dataloader):
     imgs = imgs.to(device)
 
@@ -105,7 +108,11 @@ for batch_i, (imgs, targets) in enumerate(dataloader):
                     correct.append(0)
 
         # Compute Average Precision (AP) per class
-        AP = ap_per_class(tp=correct, conf=detections[:, 4], pred_cls=detections[:, 6], target_cls=target_cls)
+        AP, AP_class = ap_per_class(tp=correct, conf=detections[:, 4], pred_cls=detections[:, 6], target_cls=target_cls)
+
+        # Accumulate AP per class
+        AP_accum_count += np.bincount(AP_class, minlength=nC)
+        AP_accum += np.bincount(AP_class, minlength=nC, weights=AP)
 
         # Compute mean AP for this image
         mAP = AP.mean()
@@ -116,4 +123,10 @@ for batch_i, (imgs, targets) in enumerate(dataloader):
         # Print image mAP and running mean mAP
         print('+ Sample [%d/%d] AP: %.4f (%.4f)' % (len(mAPs), len(dataloader) * opt.batch_size, mAP, np.mean(mAPs)))
 
+# Print mAP per class
+classes = load_classes(opt.class_path)  # Extracts class labels from file
+for i, c in enumerate(classes):
+    print('%15s: %-.4f' % (c, AP_accum[i] / AP_accum_count[i]))
+
+# Print mAP
 print('Mean Average Precision: %.4f' % np.mean(mAPs))

utils/utils.py

@@ -130,7 +130,7 @@ def ap_per_class(tp, conf, pred_cls, target_cls):
             # AP from recall-precision curve
             ap.append(compute_ap(recall, precision))
 
-    return np.array(ap)
+    return np.array(ap), unique_classes.astype('int32')
 
 
 def compute_ap(recall, precision):