updates
This commit is contained in:
parent
05a9a6205f
commit
1e1cffae8b
13
models.py
13
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
|
anchor_wh = self.anchor_wh.repeat((1, 1, self.nx, self.ny, 1)).view(1, m, 2) / ngu
|
||||||
|
|
||||||
p = p.view(m, self.no)
|
p = p.view(m, self.no)
|
||||||
xy = torch.sigmoid(p[..., 0:2]) + grid_xy[0] # x, y
|
xy = torch.sigmoid(p[:, 0:2]) + grid_xy[0] # x, y
|
||||||
wh = torch.exp(p[..., 2:4]) * anchor_wh[0] # width, height
|
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) * torch.sigmoid(p[:, 4:5]) # SSD-like conf
|
||||||
p_cls = F.softmax(p[:, 5:self.no], 1) * p_conf # SSD-like conf
|
return torch.cat((xy / ngu[0], wh, p_cls), 1).t()
|
||||||
return torch.cat((xy / ngu[0], wh, p_conf, p_cls), 1).t()
|
|
||||||
|
|
||||||
# p = p.view(1, m, self.no)
|
# p = p.view(1, m, self.no)
|
||||||
# xy = torch.sigmoid(p[..., 0:2]) + grid_xy # x, y
|
# xy = torch.sigmoid(p[..., 0:2]) + grid_xy # x, y
|
||||||
|
@ -203,7 +202,7 @@ class YOLOLayer(nn.Module):
|
||||||
else: # inference
|
else: # inference
|
||||||
# s = 1.5 # scale_xy (pxy = pxy * s - (s - 1) / 2)
|
# s = 1.5 # scale_xy (pxy = pxy * s - (s - 1) / 2)
|
||||||
io = p.clone() # inference output
|
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.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[..., 2:4] = ((torch.sigmoid(io[..., 2:4]) * 2) ** 3) * self.anchor_wh # wh power method
|
||||||
io[..., :4] *= self.stride
|
io[..., :4] *= self.stride
|
||||||
|
@ -270,7 +269,7 @@ class Darknet(nn.Module):
|
||||||
elif ONNX_EXPORT:
|
elif ONNX_EXPORT:
|
||||||
output = torch.cat(output, 1) # cat 3 layers 85 x (507, 2028, 8112) to 85 x 10647
|
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
|
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:
|
else:
|
||||||
io, p = list(zip(*output)) # inference output, training output
|
io, p = list(zip(*output)) # inference output, training output
|
||||||
return torch.cat(io, 1), p
|
return torch.cat(io, 1), p
|
||||||
|
|
Loading…
Reference in New Issue