updates
This commit is contained in:
parent
cc018c73ad
commit
36a06a1e90
12
models.py
12
models.py
|
@ -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)
|
||||
|
|
|
@ -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)
|
||||
|
|
Loading…
Reference in New Issue