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
|
p_boxes = None
|
||||||
if batch_report:
|
if batch_report:
|
||||||
# Predictd boxes: add offset and scale with anchors (in grid space, i.e. 0-13)
|
# Predictd boxes: add offset and scale with anchors (in grid space, i.e. 0-13)
|
||||||
gx = self.grid_x[:, :, :nG, :nG]
|
gx = x.data + self.grid_x[:, :, :nG, :nG]
|
||||||
gy = self.grid_y[:, :, :nG, :nG]
|
gy = y.data + self.grid_y[:, :, :nG, :nG]
|
||||||
p_boxes = torch.stack((x.data + gx - width / 2,
|
p_boxes = torch.stack((gx - width / 2,
|
||||||
y.data + gy - height / 2,
|
gy - height / 2,
|
||||||
x.data + gx + width / 2,
|
gx + width / 2,
|
||||||
y.data + gy + height / 2), 4) # x1y1x2y2
|
gy + height / 2), 4) # x1y1x2y2
|
||||||
|
|
||||||
tx, ty, tw, th, mask, tcls, TP, FP, FN, TC = \
|
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)
|
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)
|
pipeline = Pipeline(input_features, output_features)
|
||||||
|
|
||||||
# Add 3rd dimension of size 1 (apparently not needed, produces error on compile)
|
# Add 3rd dimension of size 1 (apparently not needed, produces error on compile)
|
||||||
# ssd_output = coreml_model._spec.description.output
|
ssd_output = coreml_model._spec.description.output
|
||||||
# ssd_output[0].type.multiArrayType.shape[:] = [num_classes, num_anchors, 1]
|
ssd_output[0].type.multiArrayType.shape[:] = [num_classes, num_anchors, 1]
|
||||||
# ssd_output[1].type.multiArrayType.shape[:] = [4, num_anchors, 1]
|
ssd_output[1].type.multiArrayType.shape[:] = [4, num_anchors, 1]
|
||||||
|
|
||||||
# And now we can add the three models, in order:
|
# And now we can add the three models, in order:
|
||||||
pipeline.add_model(coreml_model)
|
pipeline.add_model(coreml_model)
|
||||||
|
|
Loading…
Reference in New Issue