# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/01p_utils.visualization.ipynb.

# %% auto 0
__all__ = ['show_image', 'get_image_grid', 'show_grid', 'normalize_axes', 'plot_grid', 'image_grid', 'grid_from_sequence',
           'image_size', 'image_sizes_hist']

# %% ../../nbs/01p_utils.visualization.ipynb 3
from fastcore.all import *
import random
import polvo as pv
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import ImageGrid

# %% ../../nbs/01p_utils.visualization.ipynb 4
@delegates(plt.subplots)
def show_image(image, ax=None, cmap=None, show:bool=False, **kwargs):
    if ax is None: fig, ax = plt.subplots(**kwargs)
    ax.imshow(image, cmap=cmap)
    ax.set_axis_off()
    if show: plt.show()
    return ax

# %% ../../nbs/01p_utils.visualization.ipynb 8
@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/01p_utils.visualization.ipynb 9
def show_grid(grid, shows, show=True):
    for ax, show in zip(grid, shows): show(ax=ax)
    if show: plt.show()

# %% ../../nbs/01p_utils.visualization.ipynb 10
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/01p_utils.visualization.ipynb 11
@delegates(get_image_grid, but=['nitems'])
def plot_grid(shows, show=True, xmax=None, ymax=None, **kwargs):
    "Quickly plot a grid of images."
    fig, grid = get_image_grid(nitems=len(shows), **kwargs)
    if xmax or ymax: normalize_axes(grid, xmax, ymax)
    show_grid(grid, shows, show=show)
    return fig, grid

# %% ../../nbs/01p_utils.visualization.ipynb 12
@delegates(plot_grid, but=['shows'])
def image_grid(images, *args, **kwargs):
    shows = [partial(show_image, o) for o in images]
    return plot_grid(shows, *args, **kwargs)

# %% ../../nbs/01p_utils.visualization.ipynb 14
@delegates(image_grid)
def grid_from_sequence(sequence, get_image, nitems=9, idxs=None, **kwargs):
    "Generates a grid of images from a subset of items from the sequence."
    idxs = idxs or random.sample(range(0, len(sequence)), nitems)
    images = []
    for idx in idxs:
        item = sequence[idx]
        images.append(get_image(item))
    
    return image_grid(images, **kwargs)

# %% ../../nbs/01p_utils.visualization.ipynb 16
def image_size(image_file):
    with pv.open_image(image_file) as image:
        return image.size

# %% ../../nbs/01p_utils.visualization.ipynb 17
def image_sizes_hist(
    image_files # Sequence of image filepaths.
):
    "Plot a histogram of widths and heights."
    sizes = parallel(image_size, image_files, progress=pv.pbar)
    widths, heights = zip(*sizes)
    plt.hist(widths, label='width')
    plt.hist(heights, label='heights')
    plt.legend()
