(graph, num_classes, num_anchors)
| 29 | |
| 30 | |
| 31 | def build_decoder(graph, num_classes, num_anchors): |
| 32 | # Read the anchors into a (4, 1917) tensor. |
| 33 | anchors = get_anchors(graph) |
| 34 | |
| 35 | # MLMultiArray inputs of neural networks must have 1 or 3 dimensions. |
| 36 | # We only have 2, so add an unused dimension of size one at the back. |
| 37 | input_features = [ |
| 38 | ("scores", datatypes.Array(num_classes + 1, num_anchors, 1)), |
| 39 | ("boxes", datatypes.Array(4, num_anchors, 1)) |
| 40 | ] |
| 41 | |
| 42 | # The outputs of the decoder model should match the inputs of the next |
| 43 | # model in the pipeline, NonMaximumSuppression. This expects the number |
| 44 | # of bounding boxes in the first dimension. |
| 45 | output_features = [ |
| 46 | ("raw_confidence", datatypes.Array(num_anchors, num_classes)), |
| 47 | ("raw_coordinates", datatypes.Array(num_anchors, 4)) |
| 48 | ] |
| 49 | |
| 50 | builder = neural_network.NeuralNetworkBuilder(input_features, output_features) |
| 51 | |
| 52 | # (num_classes+1, num_anchors, 1) --> (1, num_anchors, num_classes+1) |
| 53 | builder.add_permute( |
| 54 | name="permute_scores", |
| 55 | dim=(0, 3, 2, 1), |
| 56 | input_name="scores", |
| 57 | output_name="permute_scores_output") |
| 58 | |
| 59 | # Strip off the "unknown" class (at index 0). |
| 60 | builder.add_slice( |
| 61 | name="slice_scores", |
| 62 | input_name="permute_scores_output", |
| 63 | output_name="raw_confidence", |
| 64 | axis="width", |
| 65 | start_index=1, |
| 66 | end_index=num_classes + 1) |
| 67 | |
| 68 | # Grab the y, x coordinates (channels 0-1). |
| 69 | builder.add_slice( |
| 70 | name="slice_yx", |
| 71 | input_name="boxes", |
| 72 | output_name="slice_yx_output", |
| 73 | axis="channel", |
| 74 | start_index=0, |
| 75 | end_index=2) |
| 76 | |
| 77 | # boxes_yx / 10 |
| 78 | builder.add_elementwise( |
| 79 | name="scale_yx", |
| 80 | input_names="slice_yx_output", |
| 81 | output_name="scale_yx_output", |
| 82 | mode="MULTIPLY", |
| 83 | alpha=0.1) |
| 84 | |
| 85 | # Split the anchors into two (2, 1917, 1) arrays. |
| 86 | anchors_yx = np.expand_dims(anchors[:2, :], axis=-1) |
| 87 | anchors_hw = np.expand_dims(anchors[2:, :], axis=-1) |
| 88 |
no test coverage detected