1. classification should be sigmoid 2. NMS for each class

Author: ersheng-ai

.gitignore

@@ -16,4 +16,5 @@ runs
 log
 
 *.jpg
+*.json
 data/outcome

run_coco.py

@@ -0,0 +1,23 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+from tool.torch_utils import *
+from models import *
+
+import sys
+import cv2
+from tool.utils import load_class_names, plot_boxes_cv2
+
+
+n_classes = 80
+cocoImageListFileName = "/home/erics/MS_COCO/val2017.txt"
+cocoClassIDFileName = "/home/erics/yolo_cpp_standalone/data/categories.txt"
+cocoClassNamesFileName = "/home/erics/yolo_cpp_standalone/data/coco.names"
+
+model = Yolov4(yolov4conv137weight=None, n_classes=n_classes, inference=True)
+
+pretrained_dict = torch.load(weightfile, map_location=torch.device('cuda'))
+model.load_state_dict(pretrained_dict)
+
+
+boxes = do_detect(model, sized, 0.4, 0.6, use_cuda)

tool/utils.py

@@ -182,6 +182,8 @@ def post_processing(img, conf_thresh, nms_thresh, output):
         box_array = box_array.cpu().detach().numpy()
         confs = confs.cpu().detach().numpy()
 
+    num_classes = confs.shape[2]
+
     # [batch, num, 4]
     box_array = box_array[:, :, 0]
 
@@ -199,16 +201,24 @@ def post_processing(img, conf_thresh, nms_thresh, output):
         l_max_conf = max_conf[i, argwhere]
         l_max_id = max_id[i, argwhere]
 
-        keep = nms_cpu(l_box_array, l_max_conf, nms_thresh)
-        
         bboxes = []
-        if (keep.size > 0):
-            l_box_array = l_box_array[keep, :]
-            l_max_conf = l_max_conf[keep]
-            l_max_id = l_max_id[keep]
-
-            for j in range(l_box_array.shape[0]):
-                bboxes.append([l_box_array[j, 0], l_box_array[j, 1], l_box_array[j, 2], l_box_array[j, 3], l_max_conf[j], l_max_conf[j], l_max_id[j]])
+        # nms for each class
+        for j in range(num_classes):
+
+            cls_argwhere = l_max_id == j
+            ll_box_array = l_box_array[cls_argwhere, :]
+            ll_max_conf = l_max_conf[cls_argwhere]
+            ll_max_id = l_max_id[cls_argwhere]
+
+            keep = nms_cpu(ll_box_array, ll_max_conf, nms_thresh)
+            
+            if (keep.size > 0):
+                ll_box_array = ll_box_array[keep, :]
+                ll_max_conf = ll_max_conf[keep]
+                ll_max_id = ll_max_id[keep]
+
+                for k in range(ll_box_array.shape[0]):
+                    bboxes.append([ll_box_array[k, 0], ll_box_array[k, 1], ll_box_array[k, 2], ll_box_array[k, 3], ll_max_conf[k], ll_max_conf[k], ll_max_id[k]])
 
         bboxes_batch.append(bboxes)
 

tool/yolo_layer.py

@@ -3,7 +3,7 @@
 from tool.torch_utils import *
 
 
-def yolo_forward_alternative(output, conf_thresh, num_classes, anchors, num_anchors, only_objectness=1,
+def yolo_forward_alternative(output, conf_thresh, num_classes, anchors, num_anchors, scale_x_y, only_objectness=1,
                               validation=False):
     # Output would be invalid if it does not satisfy this assert
     # assert (output.size(1) == (5 + num_classes) * num_anchors)
@@ -80,7 +80,7 @@ def yolo_forward_alternative(output, conf_thresh, num_classes, anchors, num_anch
     bxy = torch.sigmoid(bxy)
     bwh = torch.exp(bwh)
     det_confs = torch.sigmoid(det_confs)
-    cls_confs = torch.nn.Softmax(dim=2)(cls_confs)
+    cls_confs = torch.sigmoid(cls_confs)
 
     # Shape: [batch, num_anchors, 2, H * W]
     bxy = bxy.view(batch, num_anchors, 2, H * W)
@@ -165,7 +165,7 @@ def yolo_forward(output, conf_thresh, num_classes, anchors, num_anchors, scale_x
     bxy = torch.sigmoid(bxy) * scale_x_y - 0.5 * (scale_x_y - 1)
     bwh = torch.exp(bwh)
     det_confs = torch.sigmoid(det_confs)
-    cls_confs = torch.nn.Softmax(dim=2)(cls_confs)
+    cls_confs = torch.sigmoid(cls_confs)
 
     # Prepare C-x, C-y, P-w, P-h (None of them are torch related)
     grid_x = np.expand_dims(np.expand_dims(np.expand_dims(np.linspace(0, W - 1, W), axis=0).repeat(H, 0), axis=0), axis=0)