An intuitive function to plot the images.

make_plot_one_hot.py

def make_plot(tfdata, take_batch=1, title=True, figsize=(20, 20)):
    '''ref: https://gist.github.com/innat/4dc4080cfdf5cf20ef0fc93d3623ca9b
    '''
    font = {
        "family": "serif",
        "color": "darkred",
        "weight": "normal",
        "size": 10,
    }

    for images, labels in tfdata.take(take_batch):
        plt.figure(figsize=figsize, layout="tight")
        xy = int(np.ceil(images.shape[0] * 0.5))
        
        for i in range(images.shape[0]):
            plt.subplot(xy, xy, i + 1)
            plt.imshow(tf.cast(images[i], dtype=tf.uint8))
            if title:
                plt.title(tcls_names[tf.argmax(labels[i], axis=-1)], fontdict=font)
            plt.axis("off")

    plt.show()