| 70 | return self.outputs[0]["shape"], self.outputs[0]["dtype"] |
| 71 | |
| 72 | def infer(self, batch, top=1): |
| 73 | # Process I/O and execute the network |
| 74 | input = {self.inputs[0]["name"]: tf.convert_to_tensor(batch)} |
| 75 | output = self.pred_fn(**input) |
| 76 | output = output[self.outputs[0]["name"]].numpy() |
| 77 | |
| 78 | # Read and process the results |
| 79 | classes = np.argmax(output, axis=1) |
| 80 | scores = np.max(output, axis=1) |
| 81 | top = max(top, output.shape[1]) |
| 82 | top_classes = np.flip(np.argsort(output, axis=1), axis=1)[:, 0:top] |
| 83 | top_scores = np.flip(np.sort(output, axis=1), axis=1)[:, 0:top] |
| 84 | |
| 85 | return classes, scores, [top_classes, top_scores] |
| 86 | |
| 87 | |
| 88 | def main(args): |