This commit is contained in:
Glenn Jocher 2018-12-31 12:31:38 +01:00
parent cc018c73ad
commit 36a06a1e90
2 changed files with 9 additions and 9 deletions

View File

@ -169,12 +169,12 @@ class YOLOLayer(nn.Module):
p_boxes = None
if batch_report:
# Predictd boxes: add offset and scale with anchors (in grid space, i.e. 0-13)
gx = self.grid_x[:, :, :nG, :nG]
gy = self.grid_y[:, :, :nG, :nG]
p_boxes = torch.stack((x.data + gx - width / 2,
y.data + gy - height / 2,
x.data + gx + width / 2,
y.data + gy + height / 2), 4) # x1y1x2y2
gx = x.data + self.grid_x[:, :, :nG, :nG]
gy = y.data + self.grid_y[:, :, :nG, :nG]
p_boxes = torch.stack((gx - width / 2,
gy - height / 2,
gx + width / 2,
gy + height / 2), 4) # x1y1x2y2
tx, ty, tw, th, mask, tcls, TP, FP, FN, TC = \
build_targets(p_boxes, p_conf, p_cls, targets, self.scaled_anchors, self.nA, self.nC, nG, batch_report)

View File

@ -105,9 +105,9 @@ def main():
pipeline = Pipeline(input_features, output_features)
# Add 3rd dimension of size 1 (apparently not needed, produces error on compile)
# ssd_output = coreml_model._spec.description.output
# ssd_output[0].type.multiArrayType.shape[:] = [num_classes, num_anchors, 1]
# ssd_output[1].type.multiArrayType.shape[:] = [4, num_anchors, 1]
ssd_output = coreml_model._spec.description.output
ssd_output[0].type.multiArrayType.shape[:] = [num_classes, num_anchors, 1]
ssd_output[1].type.multiArrayType.shape[:] = [4, num_anchors, 1]
# And now we can add the three models, in order:
pipeline.add_model(coreml_model)