Helper for debugging FastAI image data bunches

visualize_data_bunch.py

from fastai.vision import *
import pandas as pd
import matplotlib

from matplotlib.gridspec import GridSpec

def visualize_data_bunch(db, figsize=(16,16), layoutRows=3, layoutCols=5):
    graphRows = 4
    layout = { "rows": layoutRows + graphRows, "cols": layoutCols }
    graphCols = int(layoutCols * 0.75) 

    number_of_random_transforms = len(list(filter(lambda t: isinstance(t, RandTransform), db.train_ds.tfms)))
    total_number_of_transforms = len(db.train_ds.tfms)

    labels = db.classes
    
    train_ds_counts = pd.value_counts(db.train_ds.y.items, sort=False)
    valid_ds_counts = pd.value_counts(db.valid_ds.y.items, sort=False)
    
    width = 0.35

    with plt.style.context('Solarize_Light2'):
        fig = plt.figure(figsize=figsize, constrained_layout=True)
        gs = GridSpec(layout["rows"], layout["cols"], figure=fig)
        
        ax1 = plt.subplot(gs.new_subplotspec((0, 0), colspan=graphCols, rowspan=graphRows ))
        ax2 = plt.subplot(gs.new_subplotspec((0, graphCols), colspan=layoutCols - graphCols, rowspan=graphRows))
        
        ax1.barh(labels, train_ds_counts.values, width, label='Training')
        ax1.barh(labels, valid_ds_counts.values, width, label='Validation')

        ax1.set_xlabel('Number of items')
        ax1.set_ylabel('Classes')
        ax1.set_title('Training and Validation datasets')
        ax1.legend()

        ax2.axis('off')
        ax2.text(0, 0.9, "Training dataset: %d items" % (len(db.train_ds)))
        ax2.text(0, 0.8, "Validation dataset: %d items" % (len(db.valid_ds)))
        ax2.text(0, 0.7, "Batch size: %d items" % (db.batch_size))
        ax2.text(0, 0.6, "Image transformations: %d (%d random)" % (total_number_of_transforms, number_of_random_transforms))
        
        ds_item_index = random.randint(0, len(db.train_ds) - 1)
        
        for imageRow in range(graphRows, layout["rows"]):
            for imageCol in range(0, layout["cols"]):
                ax = plt.subplot(gs.new_subplotspec((imageRow, imageCol), colspan=1))
                x,y = db.train_ds[ds_item_index]
                x.show(ax, y=y)
        
    plt.show()