mirror of
https://github.com/Laurent2916/REVA-QCAV.git
synced 2024-09-19 19:45:28 +00:00
eb95dab01e
fix visualization glitch, when running predict.py Former-commit-id: 291e41e01bbcd783c3b7ac4b9e59d9d8427e0bf5
18 lines
531 B
Python
18 lines
531 B
Python
import matplotlib.pyplot as plt
|
|
|
|
|
|
def plot_img_and_mask(img, mask):
|
|
classes = mask.shape[0] if len(mask.shape) > 2 else 1
|
|
fig, ax = plt.subplots(1, classes + 1)
|
|
ax[0].set_title('Input image')
|
|
ax[0].imshow(img)
|
|
if classes > 1:
|
|
for i in range(classes):
|
|
ax[i + 1].set_title(f'Output mask (class {i + 1})')
|
|
ax[i + 1].imshow(mask[1, :, :])
|
|
else:
|
|
ax[1].set_title(f'Output mask')
|
|
ax[1].imshow(mask)
|
|
plt.xticks([]), plt.yticks([])
|
|
plt.show()
|