From d9b742f61a65999d78dacc9c555ed2956bcbe6e5 Mon Sep 17 00:00:00 2001 From: Jim Martens Date: Tue, 2 Jul 2019 16:00:58 +0200 Subject: [PATCH] Improve argmax usage Signed-off-by: Jim Martens --- src/twomartens/masterthesis/cli.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/twomartens/masterthesis/cli.py b/src/twomartens/masterthesis/cli.py index ca1bea9..dda77da 100644 --- a/src/twomartens/masterthesis/cli.py +++ b/src/twomartens/masterthesis/cli.py @@ -121,7 +121,7 @@ def _ssd_train(args: argparse.Namespace) -> None: ymin = instance[-11] * image_size xmax = instance[-10] * image_size ymax = instance[-9] * image_size - class_id = np.argmax(instance[:-12]) + class_id = np.argmax(instance[:-12], axis=0) color = colors[class_id] label = f"{classes_to_names[class_id]}" current_axis.add_patch(