Add Yolo V2 support

tool/darknet2pytorch.py

@@ -1,7 +1,6 @@
 import torch.nn as nn
 import torch.nn.functional as F
 import numpy as np
-from tool.region_loss import RegionLoss
 from tool.yolo_layer import YoloLayer
 from tool.config import *
 from tool.torch_utils import *
@@ -208,12 +207,8 @@ def forward(self, x):
                 x = x1 * x2
                 outputs[ind] = x
             elif block['type'] == 'region':
-                continue
-                if self.loss:
-                    self.loss = self.loss + self.models[ind](x)
-                else:
-                    self.loss = self.models[ind](x)
-                outputs[ind] = None
+                boxes = self.models[ind](x)
+                out_boxes.append(boxes)
             elif block['type'] == 'yolo':
                 # if self.training:
                 #     pass
@@ -392,19 +387,19 @@ def create_network(self, blocks):
                 out_strides.append(prev_stride)
                 models.append(model)
             elif block['type'] == 'region':
-                loss = RegionLoss()
+                region = YoloLayer()
                 anchors = block['anchors'].split(',')
-                loss.anchors = [float(i) for i in anchors]
-                loss.num_classes = int(block['classes'])
-                loss.num_anchors = int(block['num'])
-                loss.anchor_step = len(loss.anchors) // loss.num_anchors
-                loss.object_scale = float(block['object_scale'])
-                loss.noobject_scale = float(block['noobject_scale'])
-                loss.class_scale = float(block['class_scale'])
-                loss.coord_scale = float(block['coord_scale'])
+                region.anchors = [float(i) for i in anchors]
+                region.num_classes = int(block['classes'])
+                region.num_anchors = int(block['num'])
+                region.anchor_step = len(region.anchors) // region.num_anchors
+                region.scale_x_y = 1.0 # thre is not such value in region config
+                region.anchor_mask = [int(i) for i in range(len(anchors) // 2)] # region has no anchor masks
+                region.stride = 1 # not implemented for region
+                region.multiply_confs = False # do not multiply detection and class confidence
                 out_filters.append(prev_filters)
                 out_strides.append(prev_stride)
-                models.append(loss)
+                models.append(region)
             elif block['type'] == 'yolo':
                 yolo_layer = YoloLayer()
                 anchors = block['anchors'].split(',')

tool/utils.py

@@ -137,8 +137,8 @@ def get_color(c, x, max_val):
             t_size = cv2.getTextSize(msg, 0, 0.7, thickness=bbox_thick // 2)[0]
             c1, c2 = (x1,y1), (x2, y2)
             c3 = (c1[0] + t_size[0], c1[1] - t_size[1] - 3)
-            cv2.rectangle(img, (x1,y1), (np.float32(c3[0]), np.float32(c3[1])), rgb, -1)
-            img = cv2.putText(img, msg, (c1[0], np.float32(c1[1] - 2)), cv2.FONT_HERSHEY_SIMPLEX,0.7, (0,0,0), bbox_thick//2,lineType=cv2.LINE_AA)
+            cv2.rectangle(img, (x1,y1), (int(c3[0]), int(c3[1])), rgb, -1)
+            img = cv2.putText(img, msg, (c1[0], int(c1[1] - 2)), cv2.FONT_HERSHEY_SIMPLEX,0.7, (0,0,0), bbox_thick//2,lineType=cv2.LINE_AA)
 
         img = cv2.rectangle(img, (x1, y1), (x2, y2), rgb, bbox_thick)
     if savename:

tool/yolo_layer.py

@@ -146,7 +146,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, multiply_confs=True):
     # Output would be invalid if it does not satisfy this assert
     # assert (output.size(1) == (5 + num_classes) * num_anchors)
 
@@ -280,7 +280,10 @@ def yolo_forward_dynamic(output, conf_thresh, num_classes, anchors, num_anchors,
     # det_confs: [batch, num_anchors * H * W]
 
     det_confs = det_confs.view(output.size(0), num_anchors * output.size(2) * output.size(3), 1)
-    confs = cls_confs * det_confs
+    if multiply_confs:
+        confs = cls_confs * det_confs
+    else:
+        confs = det_confs
 
     # boxes: [batch, num_anchors * H * W, 1, 4]
     # confs: [batch, num_anchors * H * W, num_classes]
@@ -292,7 +295,7 @@ class YoloLayer(nn.Module):
     model_out: while inference,is post-processing inside or outside the model
         true:outside
     '''
-    def __init__(self, anchor_mask=[], num_classes=0, anchors=[], num_anchors=1, stride=32, model_out=False):
+    def __init__(self, anchor_mask=[], num_classes=0, anchors=[], num_anchors=1, stride=32, model_out=False, multiply_confs=True):
         super(YoloLayer, self).__init__()
         self.anchor_mask = anchor_mask
         self.num_classes = num_classes
@@ -307,6 +310,7 @@ def __init__(self, anchor_mask=[], num_classes=0, anchors=[], num_anchors=1, str
         self.stride = stride
         self.seen = 0
         self.scale_x_y = 1
+        self.multiply_confs = multiply_confs
 
         self.model_out = model_out
 
@@ -318,5 +322,4 @@ def forward(self, output, target=None):
             masked_anchors += self.anchors[m * 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, multiply_confs=self.multiply_confs)