# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_common.ipynb.

# %% auto 0
__all__ = ['get_image_grid', 'show_grid', 'normalize_axes', 'image_grid', 'image_size', 'image_sizes_hist', 'one_batch_with_idxs',
           'show_tfms']

# %% ../nbs/00_common.ipynb 3
#TODO: optimize imports
from fastai.vision.all import *
from mpl_toolkits.axes_grid1 import ImageGrid

# %% ../nbs/00_common.ipynb 12
@delegates(ImageGrid)
def get_image_grid(
    nitems=None,
    ncols=3,
    nrows=None,
    pad=0.0,
    figsize=None,
    **kwargs
):
    "Creates axes optimized for displaying images."
    nrows = nrows or int(np.ceil(nitems / ncols))
    figsize = figsize or (14, 4 * nrows)

    fig = plt.figure(figsize=figsize)
    grid = ImageGrid(fig, 111, (nrows, ncols), axes_pad=pad, **kwargs)
    
    return fig, grid

# %% ../nbs/00_common.ipynb 13
def show_grid(grid, shows, show=True):
    for ax, show in zip(grid, shows): show(ctx=ax)
    if show: plt.show()

# %% ../nbs/00_common.ipynb 14
def normalize_axes(grid, xmax, ymax):
    "Expand all axes to have the same dimensions."
    for ax in grid:
        ax.set_xlim(right=xmax)
        ax.set_ylim(bottom=ymax)

# %% ../nbs/00_common.ipynb 15
@delegates(get_image_grid, but=['nitems'])
def image_grid(shows, show=True, xmax=None, ymax=None, **kwargs):
    "Quickly plot a grid of images."
    fig, grid = get_image_grid(nitems=len(shows), **kwargs)
    show_grid(grid, shows, show=show)
    normalize_axes(grid, xmax, ymax)
    return fig, grid

# %% ../nbs/00_common.ipynb 17
def image_size(image_file):
    with Image.open(str(image_file)) as image:
        return image.size

# %% ../nbs/00_common.ipynb 18
def image_sizes_hist(
    image_files # Sequence of image filepaths.
):
    "Plot a histogram of widths and heights."
    sizes = parallel(image_size, image_files, progress=progress_bar)
    widths, heights = zip(*sizes)
    plt.hist(widths, label='width')
    plt.hist(heights, label='heights')
    plt.legend()

# %% ../nbs/00_common.ipynb 21
def one_batch_with_idxs(
    dl, # A `DataLoader` instance. Commonly `dls.train` or `dls.valid`
    unique_idx=None, # If specified, show only images with index `unique_idx`.
):
    "Show original and transformed version of images."
    old_get_idxs = dl.get_idxs

    idxs = dl.get_idxs() if unique_idx is None else [unique_idx] * len(dl.items)
    dl.get_idxs = lambda: idxs
    x, y = dl.one_batch()

    dl.get_idxs = old_get_idxs
    
    return (x, y), idxs

# %% ../nbs/00_common.ipynb 22
@delegates(image_grid)
def show_tfms(
    dl, # A `DataLoader` instance. Commonly `dls.train` or `dls.valid`
    unique_idx=None, # If specified, show only images with index `unique_idx`.
    max_n=5, # Maximum number of samples.
    pad=0.2,
    ncols=2,
    **kwargs,
):
    "Show original and transformed version of images."
    (xs, ys), idxs = one_batch_with_idxs(dl, unique_idx=unique_idx)
    
    shows = []
    xmax, ymax = 0, 0 # Has to be passed to show grid so images are not cut when displayed
    for i, x, idx in zip(range(max_n), xs, idxs):
        orig = dl.dataset[idx]
        image = orig[0]
        xmax = max(xmax, x.shape[2], image.shape[1])
        ymax = max(ymax, x.shape[1], image.shape[0])
        
        def _s(ctx, orig=orig, idx=idx): return dl.dataset.show(orig, ctx=show_title(idx, ctx=ctx))
        shows.append(_s)

    fig, grid = get_image_grid(nrows=len(shows), ncols=ncols, pad=pad, **kwargs)
    
    dl.show_batch((xs, ys), ctxs=grid[1::2], show=True)
    show_grid(grid[::2], shows, show=False)
    normalize_axes(grid, xmax=xmax, ymax=ymax)
    
    return fig, grid
