diff --git a/tool/darknet2pytorch.py b/tool/darknet2pytorch.py index adf9d8f1..c0da249b 100644 --- a/tool/darknet2pytorch.py +++ b/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(',') diff --git a/tool/utils.py b/tool/utils.py index a42e6264..94e38d06 100644 --- a/tool/utils.py +++ b/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: diff --git a/tool/yolo_layer.py b/tool/yolo_layer.py index c3c904a5..fffc4505 100644 --- a/tool/yolo_layer.py +++ b/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)