You got some explaining to do!

pytorch_yolo/darknet2pytorch.py

@@ -126,10 +126,11 @@ def forward(self, x):
 
 # support route shortcut and reorg
 class Darknet(nn.Module):
-    def __init__(self, cfgfile, inference=False, attack_mode=False):
+    def __init__(self, cfgfile, inference=False, attack_mode=False, explain_mode=False):
         super(Darknet, self).__init__()
         self.inference = inference
         self.attack_mode = attack_mode
+        self.explain_mode = explain_mode
         self.training = not self.inference
 
         self.blocks = parse_cfg(cfgfile)
@@ -233,7 +234,7 @@ def forward(self, x):
         if self.attack_mode or self.training:
             return out_boxes
         else:
-            return get_region_boxes(out_boxes)
+            return get_region_boxes(out_boxes,self.explain_mode)
 
     def print_network(self):
         print_cfg(self.blocks)
@@ -421,6 +422,7 @@ def create_network(self, blocks):
                 yolo_layer.stride = prev_stride
                 yolo_layer.scale_x_y = float(block['scale_x_y'])
                 yolo_layer.attack_mode = self.attack_mode
+                yolo_layer.explain_mode = self.explain_mode
                 # yolo_layer.object_scale = float(block['object_scale'])
                 # yolo_layer.noobject_scale = float(block['noobject_scale'])
                 # yolo_layer.class_scale = float(block['class_scale'])

pytorch_yolo/torch_utils.py

@@ -45,23 +45,29 @@ def bbox_ious(boxes1, boxes2, x1y1x2y2=True):
     return carea / uarea
 
 
-def get_region_boxes(boxes_and_confs):
+def get_region_boxes(boxes_and_confs, explain_mode=False):
 
     # print('Getting boxes from boxes and confs ...')
 
     boxes_list = []
     confs_list = []
+    obj_list = []
 
     for item in boxes_and_confs:
         boxes_list.append(item[0])
         confs_list.append(item[1])
+        if explain_mode:
+            obj_list.append(item[2])
 
     # boxes: [batch, num1 + num2 + num3, 1, 4]
     # confs: [batch, num1 + num2 + num3, num_classes]
     boxes = torch.cat(boxes_list, dim=1)
     confs = torch.cat(confs_list, dim=1)
-
-    return [boxes, confs]
+    if explain_mode:
+        objs = torch.cat(obj_list, dim=1)
+        return [boxes, confs, objs]
+    else:
+        return [boxes, confs]
 
 
 def convert2cpu(gpu_matrix):
@@ -103,3 +109,83 @@ def do_detect(model, img, conf_thresh, nms_thresh, use_cuda=1):
 
         return utils.post_processing(img, conf_thresh, nms_thresh, output)
 
+def post_processing_torch(conf_thresh, nms_thresh, output):
+    # [num, 1, 4]
+    box_array = output[0][0].cpu()
+    # [num, num_classes]
+    confs = output[1][0].cpu()
+
+    num_classes = confs.shape[1]
+
+    # [batch, num, 4]
+    box_array = box_array[:, 0]
+
+    # [num, num_classes] --> [num]
+    max_conf,max_id = torch.max(confs, axis=1)
+
+    argwhere = max_conf > conf_thresh
+    l_box_array = box_array[argwhere, :]
+    l_max_conf = max_conf[argwhere]
+    l_max_id = max_id[argwhere]
+    l_prob_array = confs[argwhere, :]
+
+    bboxes = []
+    probs = []
+    # 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]
+        ll_prob_array = l_prob_array[cls_argwhere, :]
+
+        keep = nms_torch(ll_box_array, ll_max_conf, nms_thresh)
+
+        if torch.numel(keep):
+            ll_box_array = ll_box_array[keep, :]
+            ll_max_conf = ll_max_conf[keep]
+            ll_max_id = ll_max_id[keep]
+            ll_prob_array = ll_prob_array[keep, :]
+
+            for k in range(ll_box_array.shape[0]):
+                bboxes.append(ll_box_array[k, :])
+                probs.append(ll_prob_array[k, :])
+
+    return torch.vstack(bboxes), torch.vstack(probs)
+
+def nms_torch(boxes, confs, nms_thresh=0.5, min_mode=False):
+    # print(boxes.shape)
+    x1 = boxes[:, 0]
+    y1 = boxes[:, 1]
+    x2 = boxes[:, 2]
+    y2 = boxes[:, 3]
+
+    areas = (x2 - x1) * (y2 - y1)
+    order = torch.argsort(confs,descending=True)
+
+    keep = []
+    while torch.numel(order) > 0:
+        idx_self = order[0]
+        idx_other = order[1:]
+
+        keep.append(idx_self)
+
+        xx1 = torch.maximum(x1[idx_self], x1[idx_other])
+        yy1 = torch.maximum(y1[idx_self], y1[idx_other])
+        xx2 = torch.minimum(x2[idx_self], x2[idx_other])
+        yy2 = torch.minimum(y2[idx_self], y2[idx_other])
+
+        w = torch.maximum(torch.tensor(0.0), xx2 - xx1)
+        h = torch.maximum(torch.tensor(0.0), yy2 - yy1)
+        inter = w * h
+
+        if min_mode:
+            over = inter / torch.minimum(areas[order[0]], areas[order[1:]])
+        else:
+            over = inter / (areas[order[0]] + areas[order[1:]] - inter)
+
+        inds = torch.where(over <= nms_thresh)[0]
+        order = order[inds + 1]
+
+    return torch.tensor(keep)

pytorch_yolo/yolo_layer.py

@@ -155,7 +155,7 @@ def yolo_forward(output, conf_thresh, num_classes, anchors, num_anchors, scale_x
 
 
 def yolo_forward_dynamic(output, conf_thresh, num_classes, anchors, num_anchors, scale_x_y, only_objectness=1,
-                         validation=False):
+                         validation=False, explain_mode=False):
     # Output would be invalid if it does not satisfy this assert
     # assert (output.size(1) == (5 + num_classes) * num_anchors)
 
@@ -308,7 +308,10 @@ def yolo_forward_dynamic(output, conf_thresh, num_classes, anchors, num_anchors,
     # boxes: [batch, num_anchors * H * W, 1, 4]
     # confs: [batch, num_anchors * H * W, num_classes]
 
-    return boxes, confs
+    if explain_mode:
+        return boxes, cls_confs, det_confs
+    else:
+        return  boxes, confs
 
 
 class YoloLayer(nn.Module):
@@ -344,4 +347,4 @@ def forward(self, output, target=None):
                                            self.anchor_step:(m + 1) * self.anchor_step]
         masked_anchors = [anchor / self.stride for anchor in masked_anchors]
 
-        return yolo_forward_dynamic(output, self.thresh, self.num_classes, masked_anchors, len(self.anchor_mask), scale_x_y=self.scale_x_y)
+        return yolo_forward_dynamic(output, self.thresh, self.num_classes, masked_anchors, len(self.anchor_mask), scale_x_y=self.scale_x_y, explain_mode=self.explain_mode)