parent
8501981f09
commit
2a284aacd0
|
@ -156,10 +156,12 @@ class YOLOLayer(nn.Module):
|
||||||
io = p.clone() # inference output
|
io = p.clone() # inference output
|
||||||
io[..., 0:2] = torch.sigmoid(io[..., 0:2]) + self.grid_xy # xy
|
io[..., 0:2] = torch.sigmoid(io[..., 0: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[..., :4] *= self.stride
|
||||||
|
|
||||||
# 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:] = torch.sigmoid(io[..., 4:]) # p_conf, p_cls
|
io[..., 4:] = torch.sigmoid(io[..., 4:]) # p_conf, p_cls
|
||||||
# io[..., 5:] = F.softmax(io[..., 5:], dim=4) # p_cls
|
# io[..., 5:] = F.softmax(io[..., 5:], dim=4) # p_cls
|
||||||
io[..., :4] *= self.stride
|
|
||||||
if self.nc == 1:
|
if self.nc == 1:
|
||||||
io[..., 5] = 1 # single-class model https://github.com/ultralytics/yolov3/issues/235
|
io[..., 5] = 1 # single-class model https://github.com/ultralytics/yolov3/issues/235
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue