REVA-QCAV/utils/data_vis.py
milesial 4c0f0a7a7b Global cleanup, better logging and CLI
Former-commit-id: ff1ac0936c118d129bc8a8014958948d3b3883be
2019-10-26 23:17:48 +02:00

18 lines
508 B
Python

import matplotlib.pyplot as plt
def plot_img_and_mask(img, mask):
classes = mask.shape[2] 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[:, :, i])
else:
ax[1].set_title(f'Output mask')
ax[1].imshow(mask)
plt.xticks([]), plt.yticks([])
plt.show()