Update utils.py

fix visualization glitch, when running predict.py Former-commit-id: 291e41e01bbcd783c3b7ac4b9e59d9d8427e0bf5
2022-04-01 01:07:13 +08:00 · 2022-04-01 01:07:13 +08:00 · eb95dab01e
parent 2ca43802cc
commit eb95dab01e
1 changed files with 1 additions and 1 deletions
--- 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)