From eb95dab01e70d4042b7ee46f7b9ba7ef82a20823 Mon Sep 17 00:00:00 2001 From: pleaseRedo Date: Fri, 1 Apr 2022 01:07:13 +0800 Subject: [PATCH] Update utils.py fix visualization glitch, when running predict.py Former-commit-id: 291e41e01bbcd783c3b7ac4b9e59d9d8427e0bf5 --- utils/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/utils/utils.py b/utils/utils.py index 1d48738..c1eb578 100644 --- a/utils/utils.py +++ b/utils/utils.py @@ -9,7 +9,7 @@ def plot_img_and_mask(img, mask): if classes > 1: for i in range(classes): ax[i + 1].set_title(f'Output mask (class {i + 1})') - ax[i + 1].imshow(mask[:, :, i]) + ax[i + 1].imshow(mask[1, :, :]) else: ax[1].set_title(f'Output mask') ax[1].imshow(mask)