From 1e1cffae8b3726101dcbb1e563edae44d78b5d46 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 23 Dec 2019 23:34:30 -0800 Subject: [PATCH] updates --- models.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/models.py b/models.py index 48815452..c3bf5a18 100755 --- a/models.py +++ b/models.py @@ -182,11 +182,10 @@ class YOLOLayer(nn.Module): anchor_wh = self.anchor_wh.repeat((1, 1, self.nx, self.ny, 1)).view(1, m, 2) / ngu p = p.view(m, self.no) - xy = torch.sigmoid(p[..., 0:2]) + grid_xy[0] # x, y - wh = torch.exp(p[..., 2:4]) * anchor_wh[0] # width, height - p_conf = torch.sigmoid(p[:, 4:5]) # Conf - p_cls = F.softmax(p[:, 5:self.no], 1) * p_conf # SSD-like conf - return torch.cat((xy / ngu[0], wh, p_conf, p_cls), 1).t() + xy = torch.sigmoid(p[:, 0:2]) + grid_xy[0] # x, y + wh = torch.exp(p[:, 2:4]) * anchor_wh[0] # width, height + p_cls = F.softmax(p[:, 5:self.no], 1) * torch.sigmoid(p[:, 4:5]) # SSD-like conf + return torch.cat((xy / ngu[0], wh, p_cls), 1).t() # p = p.view(1, m, self.no) # xy = torch.sigmoid(p[..., 0:2]) + grid_xy # x, y @@ -203,7 +202,7 @@ class YOLOLayer(nn.Module): else: # inference # s = 1.5 # scale_xy (pxy = pxy * s - (s - 1) / 2) io = p.clone() # inference output - io[..., 0:2] = torch.sigmoid(io[..., 0:2]) + self.grid_xy # xy + io[..., :2] = torch.sigmoid(io[..., :2]) + self.grid_xy # xy io[..., 2:4] = torch.exp(io[..., 2:4]) * self.anchor_wh # wh yolo method # io[..., 2:4] = ((torch.sigmoid(io[..., 2:4]) * 2) ** 3) * self.anchor_wh # wh power method io[..., :4] *= self.stride @@ -270,7 +269,7 @@ class Darknet(nn.Module): elif ONNX_EXPORT: output = torch.cat(output, 1) # cat 3 layers 85 x (507, 2028, 8112) to 85 x 10647 nc = self.module_list[self.yolo_layers[0]].nc # number of classes - return output[5:5 + nc].t(), output[0:4].t() # ONNX scores, boxes + return output[4:4 + nc].t(), output[0:4].t() # ONNX scores, boxes else: io, p = list(zip(*output)) # inference output, training output return torch.cat(io, 1), p