Update utils.py

fix visualization glitch, when running predict.py


Former-commit-id: 291e41e01bbcd783c3b7ac4b9e59d9d8427e0bf5
This commit is contained in:
pleaseRedo 2022-04-01 01:07:13 +08:00
parent 2ca43802cc
commit eb95dab01e

View file

@ -9,7 +9,7 @@ def plot_img_and_mask(img, mask):
if classes > 1: if classes > 1:
for i in range(classes): for i in range(classes):
ax[i + 1].set_title(f'Output mask (class {i + 1})') ax[i + 1].set_title(f'Output mask (class {i + 1})')
ax[i + 1].imshow(mask[:, :, i]) ax[i + 1].imshow(mask[1, :, :])
else: else:
ax[1].set_title(f'Output mask') ax[1].set_title(f'Output mask')
ax[1].imshow(mask) ax[1].imshow(mask)