This commit is contained in:
Glenn Jocher 2019-03-04 15:52:18 +01:00
parent 65ed0d1122
commit 8069fc452a
2 changed files with 5 additions and 5 deletions

View File

@ -150,7 +150,7 @@ class YOLOLayer(nn.Module):
p_conf = p[..., 4] # Conf p_conf = p[..., 4] # Conf
p_cls = p[..., 5:] # Class p_cls = p[..., 5:] # Class
txy, twh, mask, tcls = build_targets(targets, self.anchor_vec, self.nA, self.nC, nG) txy, twh, mask, tcls = build_targets(targets, self.anchor_vec, self.nA, self.nC, nG.cpu())
tcls = tcls[mask] tcls = tcls[mask]
if xy.is_cuda: if xy.is_cuda:

View File

@ -231,7 +231,7 @@ def bbox_iou(box1, box2, x1y1x2y2=True):
return inter_area / (b1_area + b2_area - inter_area + 1e-16) return inter_area / (b1_area + b2_area - inter_area + 1e-16)
def build_targets(target, anchor_wh, nA, nC, nG): def build_targets(target, anchor_vec, nA, nC, nG):
""" """
returns nT, nCorrect, tx, ty, tw, th, tconf, tcls returns nT, nCorrect, tx, ty, tw, th, tconf, tcls
""" """
@ -255,7 +255,7 @@ def build_targets(target, anchor_wh, nA, nC, nG):
# iou of targets-anchors (using wh only) # iou of targets-anchors (using wh only)
box1 = gwh box1 = gwh
box2 = anchor_wh.unsqueeze(1) box2 = anchor_vec.unsqueeze(1)
inter_area = torch.min(box1, box2).prod(2) inter_area = torch.min(box1, box2).prod(2)
iou = inter_area / (box1.prod(1) + box2.prod(2) - inter_area + 1e-16) iou = inter_area / (box1.prod(1) + box2.prod(2) - inter_area + 1e-16)
@ -290,8 +290,8 @@ def build_targets(target, anchor_wh, nA, nC, nG):
txy[b, a, gj, gi] = gxy - gxy.floor() txy[b, a, gj, gi] = gxy - gxy.floor()
# Width and height # Width and height
twh[b, a, gj, gi] = torch.log(gwh / anchor_wh[a]) # yolo method twh[b, a, gj, gi] = torch.log(gwh / anchor_vec[a]) # yolo method
# twh[b, a, gj, gi] = torch.sqrt(gwh / anchor_wh[a]) / 2 # power method # twh[b, a, gj, gi] = torch.sqrt(gwh / anchor_vec[a]) / 2 # power method
# One-hot encoding of label # One-hot encoding of label
tcls[b, a, gj, gi, tc] = 1 tcls[b, a, gj, gi, tc] = 1