Bienqq
/
car-detection-bayes
2
1
Fork 0
Code Issues Pull Requests Releases Wiki Activity

updates

This commit is contained in:
Glenn Jocher 2019-12-03 17:22:58 -08:00
parent cae901c2da
commit b9fa92d3f7
1 changed files with 7 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
models.py
View File

@ -185,18 +185,19 @@ class YOLOLayer(nn.Module):
elif ONNX_EXPORT:
# Constants CAN NOT BE BROADCAST, ensure correct shape!
ngu = self.ng.repeat((1, self.na * self.nx * self.ny, 1))
grid_xy = self.grid_xy.repeat((1, self.na, 1, 1, 1)).view((1, -1, 2))
anchor_wh = self.anchor_wh.repeat((1, 1, self.nx, self.ny, 1)).view((1, -1, 2)) / ngu
m = self.na * self.nx * self.ny
ngu = self.ng.repeat((1, m, 1))
grid_xy = self.grid_xy.repeat((1, self.na, 1, 1, 1)).view(1, m, 2)
anchor_wh = self.anchor_wh.repeat((1, 1, self.nx, self.ny, 1)).view(1, m, 2) / ngu
p = p.view(-1, 5 + self.nc)
p = p.view(m, 5 + self.nc)
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:85], 1) * p_conf # SSD-like conf
return torch.cat((xy / ngu[0], wh, p_conf, p_cls), 1).t()
# p = p.view(1, -1, 5 + self.nc)
# p = p.view(1, m, 5 + self.nc)
# xy = torch.sigmoid(p[..., 0:2]) + grid_xy # x, y
# wh = torch.exp(p[..., 2:4]) * anchor_wh # width, height
# p_conf = torch.sigmoid(p[..., 4:5]) # Conf
@ -278,7 +279,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[:4].t() # ONNX scores, boxes
return output[5:5 + 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