<commit_before>import glob
import json
import warnings
import copy
import os
from collections import OrderedDict
from functools import partial
from tqdm import tqdm
import colorspacious
import numpy as np
from skimage import draw, img_as_float
from skimage.io import imread, imsave
from skimage.transform import resize
from skimage.color import gray2rgb
from skimage.util import crop
from scipy.spatial import cKDTree
from scipy.cluster import vq


options = {'imread': {},
           'colorspace': {"name": "J'a'b'",
                          "ciecam02_space": colorspacious.CIECAM02Space.sRGB,
                          "luoetal2006_space": colorspacious.CAM02UCS}}


def set_options(imread=None, colorspace=None):
    """
    Set global options

    Parameters
    ----------
    imread : dict
        keyword arguments passed through to every call to ``imread``
        e.g., ``{'plugin': 'matplotlib'}``
    colorspace : string or dict
        colorspace used for color comparisions; see colorspacious documentation
        for details
    """
    global options
    if imread is not None:
        options['imread'].update(imread)
    if colorspace is not None:
        options['colorspace'] = colorspace


def basic_mosaic(image, pool, grid_dims, *, mask=None, depth=1):
    """
    Make a mosaic in one step with some basic settings.

    See documentation (or the source code of this function) for more
    powerful features and customization.

    Parameters
    ----------
    image : array
    pool : dict-like
        output from :func:`make_pool`; or any mapping of
        arguments for opening an image file to a vector characterizing it:
        e.g., ``{(filename,): [1, 2, 3]}``
    grid_dims : tuple
        Number of tiles along height, width.
    mask : array or None
        must have same shape as ``image``
    depth : int, optional
        Each tile can be subdividing this many times in regions of high
        contrast or along mask edges (if applicable). Default is 0.

    Returns
    -------
    mosaic : array

    Example
    -------
    Before making the mosaic, you need a collection of images to use as tiles.
    A collection of analyzed images is a "pool". Analyzing the images takes
    much more time that making the mosaic, so it is a separate step.
    >>> pool = make_pool('directory_of_images/*.jpg')

    Load an image to turn into mosaic.
    >>> from skimage.io import imread, imsave
    >>> my_image = imread('my_image.jpg')

    Make the mosaic and save it.
    >>> mosaic = basic_mosaic(my_image, (15, 15))
    >>> imsave('my_mosaic.jpg', mosaic)
    """
    image = img_as_float(image)
    image = rescale_commensurate(image, grid_dims, depth)
    if mask is not None:
        mask = rescale_commensurate(mask)
    tiles = partition(image, grid_dims=grid_dims, mask=mask, depth=depth)
    matcher = SimpleMatcher(pool)
    converted = colorspacious.cspace_convert(image, "sRGB1",
                                             options['colorspace'])
    tile_colors = [dominant_color(sample_pixels(converted[tile], 1000))
                   for tile in tqdm(tiles, desc='analyzing tiles')]
    matches = []
    for tile_color in tqdm(tile_colors, desc='matching'):
        matches.append(matcher.match(tile_color))
    canvas = np.ones_like(image)  # white canvas same shape as input image
    return draw_mosaic(canvas, tiles, matches)


def rescale_commensurate(image, grid_dims, depth=0):
    """
    For given grid dimensions and grid subdivision depth, scale image.

    The image is rescaled so that its shape can be evenly split into tiles.
    If necessary, one dimension is cropped to fit.

    Parameters
    ----------
    image : array
    grid_dims : tuple
        Number of tiles along height, width.
    depth : int, optional
        Each tile can be subdivided this many times. Default is 0.

    Returns
    -------
    rescaled_image : array
    """
    factor = np.array(grid_dims) * 2**depth
    new_shape = [int(factor[i] * np.ceil(image.shape[i] / factor[i]))
                 for i in [0, 1]]
    return crop_to_fit(image, new_shape)


def sample_pixels(image, size, replace=True):
    """
    Randomly sample pixels from an image.

    This is a wrapper around ``np.random.choice`` (which only works directly
    on 1-dimensional arrays).

    Parameters
    ----------
    image : array
    size : int
        number of pixels to sample
    replace : boolean, optional
        whether to sample with or without replacement; default True
    """
    num_pixels = np.product(image.shape[:-1])
    pixels = image.reshape(num_pixels, image.shape[-1])
    random_indexes = np.random.choice(num_pixels, size=size, replace=replace)
    return pixels[random_indexes]


def dominant_color(pixels, n_clusters=5):
    """
    Cluster colors and identify the "central" color of the largest cluster.

    Parameters
    ----------
    pixels : array
        List of pixels. The second axis is expected to be the color axis.
    n_clusters : int, optional
        number of clusters; default 5

    Returns
    -------
    dominant_color : array
    """
    colors, dist = vq.kmeans(pixels, n_clusters)
    vecs, dist = vq.vq(pixels, colors)
    counts, bins = np.histogram(vecs, len(colors))
    return colors[counts.argmax()]


def make_pool(glob_string, *, pool=None, skip_read_failures=True,
              analyzer=dominant_color):
    """
    Analyze a collection of images.

    For each file:
    1. Read image.
    2. Convert to perceptually-uniform color space.
    3. Characterize the colors in the image as a vector.

    A progress bar is displayed and then hidden after completion.

    Parameters
    ----------
    glob_string : string
        a filepath with optional wildcards, like `'*.jpg'`
    pool : dict-like, optional
        dict-like data structure to hold results; if None, dict is used
    skip_read_failures: bool, optional
        If True (default), convert any exceptions that occur while reading a
        file into warnings and continue.
    analyzer : callable, optional
        Function with signature: ``f(img) -> arr`` where ``arr`` is a vector.
        The default analyzer is :func:`dominant_color`.

    Returns
    -------
    cache : dict-like
        mapping arguments for opening file to analyzer's result, e.g.:
        ``{(filename,): [1, 2, 3]}``
    """
    if pool is None:
        pool = {}
    filenames = glob.glob(glob_string)
    for filename in tqdm(filenames, desc='analyzing pool'):
        try:
            raw_image = imread(filename, **options['imread'])
        except Exception as err:
            if skip_read_failures:
                warnings.warn("Skipping {}; raised exception:\n    {}"
                              "".format(filename, err))
                continue
            raise
        image = standardize_image(raw_image)
        # Convert color to perceptually-uniform color space.
        sample = sample_pixels(image, 1000)
        converted_sample = colorspacious.cspace_convert(sample, "sRGB1",
                                                        options['colorspace'])
        vector = analyzer(converted_sample)
        pool[(filename,)] = vector
    return pool


def standardize_image(image):
    """
    Ensure that image is float 0-1 RGB with no alpha.

    Parameters
    ----------
    image : array

    Returns
    -------
    image : array
        may or may not be a copy of the original
    """
    image = img_as_float(image)  # ensure float scaled 0-1
    # If there is no color axis, create one.
    if image.ndim == 2:
        image = gray2rgb(image)
    # Assume last axis is color axis. If alpha channel exists, drop it.
    if image.shape[-1] == 4:
        image = image[:, :, :-1]
    return image


class SimpleMatcher:
    """
    This simple matcher returns the closest match.

    It maintains an internal tree representation of the pool for fast lookups.

    Parameters
    ----------
    pool : dict
    """
    def __init__(self, pool):
        self._pool = OrderedDict(pool)
        self._args = list(self._pool.keys())
        data = np.array([vector for vector in self._pool.values()])
        self._tree = cKDTree(data)

    def match(self, vector):
        """
        Return the key of the pool image that is "nearest" (in color space).

        Parameters
        ----------
        vector : array
            characterizing the color to be matched

        Returns
        -------
        args : tuple
            arguments that specify how to open the image
        """
        distance, index = self._tree.query(vector, k=1)
        return self._args[index]


def draw_mosaic(image, tiles, matches, scale=1, resized_copy_cache=None):
    """
    Assemble the mosaic, the final result.

    Parameters
    ----------
    image : array
        the "canvas" on which to draw the tiles, modified in place
    tiles : list
        list of pairs of slice objects
    matches : list
        for each tile in ``tiles``, a tuple of arguments for opening the
        matching image file
    scale : int, optional
        Scale up tiles for higher resolution image; default is 1.
        Any not-integer input will be cast to int.
    resized_copy_cache : dict or None, optional
        cache of images from the pool, sized to fit tiles
        entries look like: ``(pool_key, (height, width))``

    Returns
    -------
    image : array

    Example
    -------
    Basic usage:
    >>> draw_mosaic(image, tiles, matches)

    Cache the resized pool images to speed up repeated drawings:
    >>> cache = {}  # any mutable mapping
    >>> draw_mosiac(image, tiles, matches, resized_copy_cache=cache)

    The above populated ``cache`` with every resized pool image used in a tile.
    Now, ``draw_mosaic``will check the cache before loading the pool image and
    resizing it, which is the most expensive step.
    >>> draw_mosiac(image, tiles, matches, resized_copy_cache=cache)

    """
    scale = int(scale)
    if resized_copy_cache is None:
        resized_copy_cache = {}
    for tile, match_args in zip(tqdm(tiles, desc='drawing mosaic'), matches):
        if scale != 1:
            tile = tuple(slice(scale * s.start, scale * s.stop)
                         for s in tile)
        target_shape = _tile_shape(tile)
        cache_key = (match_args, target_shape)
        try:
            sized_match_image = resized_copy_cache[cache_key]
        except KeyError:
            target_shape = _tile_shape(tile)
            raw_match_image = imread(*match_args, **options['imread'])
            match_image = standardize_image(raw_match_image)
            sized_match_image = crop_to_fit(match_image, target_shape)
            resized_copy_cache[cache_key] = sized_match_image
        image[tile] = sized_match_image
    return image


def _subdivide(tile):
    "Create four tiles from the four quadrants of the input tile."
    tile_dims = [(s.stop - s.start) // 2 for s in tile]
    subtiles = []
    for y in (0, 1):
        for x in (0, 1):
            subtile = (slice(tile[0].start + y * tile_dims[0],
                             tile[0].start + (1 + y) * tile_dims[0]),
                       slice(tile[1].start + x * tile_dims[1],
                             tile[1].start + (1 + x) * tile_dims[1]))
            subtiles.append(subtile)
    return subtiles


def partition(image, grid_dims, mask=None, depth=0, hdr=80):
    """
    Parition the target image into tiles.

    Optionally, subdivide tiles that contain high contrast, creating a grid
    with tiles of varied size.

    Parameters
    ----------
    grid_dims : int or tuple
        number of (largest) tiles along each dimension
    mask : array, optional
        Tiles that straddle a mask edge will be subdivided, creating a smooth
        edge.
    depth : int, optional
        Default is 0. Maximum times a tile can be subdivided.
    hdr : float
        "High Dynamic Range" --- level of contrast that trigger subdivision

    Returns
    -------
    tiles : list
        list of pairs of slice objects
    """
    # Validate inputs.
    if isinstance(grid_dims, int):
        tile_dims, = image.ndims * (grid_dims,)
    for i in (0, 1):
        image_dim = image.shape[i]
        grid_dim = grid_dims[i]
        if image_dim % grid_dim*2**depth != 0:
            raise ValueError("Image dimensions must be evenly divisible by "
                             "dimensions of the (subdivided) grid. "
                             "Dimension {image_dim} is not "
                             "evenly divisible by {grid_dim}*2**{depth} "
                             "".format(image_dim=image_dim, grid_dim=grid_dim,
                                       depth=depth))

    # Partition into equal-sized tiles (Python slice objects).
    tile_height = image.shape[0] // grid_dims[0]
    tile_width = image.shape[1] // grid_dims[1]
    tiles = []
    total = np.product(grid_dims)
    with tqdm(total=total, desc='partitioning depth 0') as pbar:
        for y in range(grid_dims[0]):
            for x in range(grid_dims[1]):
                tile = (slice(y * tile_height, (1 + y) * tile_height),
                        slice(x * tile_width, (1 + x) * tile_width))
                tiles.append(tile)
                pbar.update()

    # Discard any tiles that reside fully outside the mask.
    if mask is not None:
        tiles = [tile for tile in tiles if np.any(mask[tile])]

    # If depth > 0, subdivide any tiles that straddle a mask edge or that
    # contain an image with high contrast.
    for d in range(depth):
        new_tiles = []
        for tile in tqdm(tiles, desc='partitioning depth %d' % d):
            if ((mask is not None) and
                    np.any(mask[tile]) and np.any(~mask[tile])):
                # This tile straddles a mask edge.
                subtiles = _subdivide(tile)
                # Discard subtiles that reside fully outside the mask.
                subtiles = [tile for tile in subtiles if np.any(mask[tile])]
                new_tiles.extend(subtiles)
            elif False:  # TODO hdr check
                new_tiles.extend(_subdivide(tile))
            else:
                new_tiles.append(tile)
        tiles = new_tiles
    return tiles


def color_hist(image, bins=256):
    """
    Compute the distribution of each color channel.

    Parameters
    ----------
    image : array
        The last axis is expected to be the color axis.
    bins : int, optional
        default is 256

    Returns
    -------
    counts : array
        dimensions ``(num_chanels, bins)``

    Note
    ----
    The image is expected to be float with domain [0, 1]. Histrogram bins
    span full domain, divided equally into the number of bins specified.
    """
    image = np.asarray(image)
    num_pixels = np.product(image.shape[:-1])
    pixels = image.reshape(num_pixels, image.shape[-1])
    func = partial(np.histogram, bins=bins, range=(0, 1), density=True)
    counts, bin_edges = np.apply_along_axis(func, 0, pixels)
    # no need to return bin_edges because np.histogram options are locked down
    return np.asarray([channel for channel in counts])


def palette_map(from_palette, to_palette):
    """
    Build a function that maps from one color palette into another.

    Parameters
    ----------
    from_palette : array
        histogram of colors on [0, 1]
    from_palette : array
        histogram of colors on [0, 1]

    Returns
    -------
    f : function
    """
    old_N = from_palette.shape[-1]
    new_N = to_palette.shape[-1]
    # cumulative distribution functions
    old_cdf = np.cumsum(from_palette) / old_N
    new_cdf = np.cumsum(to_palette) / new_N
    bins = np.linspace(0, 1, old_N)

    def f(image):
        """
        Rescale colors in image from the old color palette to the new one.
        """
        image = np.asarray(image)
        num_pixels = np.product(image.shape[:-1])
        pixels = image.reshape(num_pixels, image.shape[-1])

        # Identify a color bin for each pixels.
        old_x = np.apply_along_axis(partial(np.searchsorted, bins), 0, pixels)

        # Where in the original gamut did this color fall?
        old_y = old_cdf[old_x]

        # Find that same position in the new gamut. What is the color?
        new_x = np.apply_along_axis(
            partial(np.searchsorted, new_cdf), 0, old_y)
        return new_x.reshape(image.shape) / new_N

    return f

<commit_msg>ENH:  Generalize palette stuff to spaces not on [0, 1]."<commit_after>import glob
import json
import warnings
import copy
import os
from collections import OrderedDict
from functools import partial
from tqdm import tqdm
import colorspacious
import numpy as np
from skimage import draw, img_as_float
from skimage.io import imread, imsave
from skimage.transform import resize
from skimage.color import gray2rgb
from skimage.util import crop
from scipy.spatial import cKDTree
from scipy.cluster import vq


options = {'imread': {},
           'colorspace': {"name": "J'a'b'",
                          "ciecam02_space": colorspacious.CIECAM02Space.sRGB,
                          "luoetal2006_space": colorspacious.CAM02UCS}}


def set_options(imread=None, colorspace=None):
    """
    Set global options

    Parameters
    ----------
    imread : dict
        keyword arguments passed through to every call to ``imread``
        e.g., ``{'plugin': 'matplotlib'}``
    colorspace : string or dict
        colorspace used for color comparisions; see colorspacious documentation
        for details
    """
    global options
    if imread is not None:
        options['imread'].update(imread)
    if colorspace is not None:
        options['colorspace'] = colorspace


def basic_mosaic(image, pool, grid_dims, *, mask=None, depth=1):
    """
    Make a mosaic in one step with some basic settings.

    See documentation (or the source code of this function) for more
    powerful features and customization.

    Parameters
    ----------
    image : array
    pool : dict-like
        output from :func:`make_pool`; or any mapping of
        arguments for opening an image file to a vector characterizing it:
        e.g., ``{(filename,): [1, 2, 3]}``
    grid_dims : tuple
        Number of tiles along height, width.
    mask : array or None
        must have same shape as ``image``
    depth : int, optional
        Each tile can be subdividing this many times in regions of high
        contrast or along mask edges (if applicable). Default is 0.

    Returns
    -------
    mosaic : array

    Example
    -------
    Before making the mosaic, you need a collection of images to use as tiles.
    A collection of analyzed images is a "pool". Analyzing the images takes
    much more time that making the mosaic, so it is a separate step.
    >>> pool = make_pool('directory_of_images/*.jpg')

    Load an image to turn into mosaic.
    >>> from skimage.io import imread, imsave
    >>> my_image = imread('my_image.jpg')

    Make the mosaic and save it.
    >>> mosaic = basic_mosaic(my_image, (15, 15))
    >>> imsave('my_mosaic.jpg', mosaic)
    """
    image = img_as_float(image)
    image = rescale_commensurate(image, grid_dims, depth)
    if mask is not None:
        mask = rescale_commensurate(mask)
    tiles = partition(image, grid_dims=grid_dims, mask=mask, depth=depth)
    matcher = SimpleMatcher(pool)
    converted = colorspacious.cspace_convert(image, "sRGB1",
                                             options['colorspace'])
    tile_colors = [dominant_color(sample_pixels(converted[tile], 1000))
                   for tile in tqdm(tiles, desc='analyzing tiles')]
    matches = []
    for tile_color in tqdm(tile_colors, desc='matching'):
        matches.append(matcher.match(tile_color))
    canvas = np.ones_like(image)  # white canvas same shape as input image
    return draw_mosaic(canvas, tiles, matches)


def rescale_commensurate(image, grid_dims, depth=0):
    """
    For given grid dimensions and grid subdivision depth, scale image.

    The image is rescaled so that its shape can be evenly split into tiles.
    If necessary, one dimension is cropped to fit.

    Parameters
    ----------
    image : array
    grid_dims : tuple
        Number of tiles along height, width.
    depth : int, optional
        Each tile can be subdivided this many times. Default is 0.

    Returns
    -------
    rescaled_image : array
    """
    factor = np.array(grid_dims) * 2**depth
    new_shape = [int(factor[i] * np.ceil(image.shape[i] / factor[i]))
                 for i in [0, 1]]
    return crop_to_fit(image, new_shape)


def sample_pixels(image, size, replace=True):
    """
    Randomly sample pixels from an image.

    This is a wrapper around ``np.random.choice`` (which only works directly
    on 1-dimensional arrays).

    Parameters
    ----------
    image : array
    size : int
        number of pixels to sample
    replace : boolean, optional
        whether to sample with or without replacement; default True
    """
    num_pixels = np.product(image.shape[:-1])
    pixels = image.reshape(num_pixels, image.shape[-1])
    random_indexes = np.random.choice(num_pixels, size=size, replace=replace)
    return pixels[random_indexes]


def dominant_color(pixels, n_clusters=5):
    """
    Cluster colors and identify the "central" color of the largest cluster.

    Parameters
    ----------
    pixels : array
        List of pixels. The second axis is expected to be the color axis.
    n_clusters : int, optional
        number of clusters; default 5

    Returns
    -------
    dominant_color : array
    """
    colors, dist = vq.kmeans(pixels, n_clusters)
    vecs, dist = vq.vq(pixels, colors)
    counts, bins = np.histogram(vecs, len(colors))
    return colors[counts.argmax()]


def make_pool(glob_string, *, pool=None, skip_read_failures=True,
              analyzer=dominant_color):
    """
    Analyze a collection of images.

    For each file:
    1. Read image.
    2. Convert to perceptually-uniform color space.
    3. Characterize the colors in the image as a vector.

    A progress bar is displayed and then hidden after completion.

    Parameters
    ----------
    glob_string : string
        a filepath with optional wildcards, like `'*.jpg'`
    pool : dict-like, optional
        dict-like data structure to hold results; if None, dict is used
    skip_read_failures: bool, optional
        If True (default), convert any exceptions that occur while reading a
        file into warnings and continue.
    analyzer : callable, optional
        Function with signature: ``f(img) -> arr`` where ``arr`` is a vector.
        The default analyzer is :func:`dominant_color`.

    Returns
    -------
    cache : dict-like
        mapping arguments for opening file to analyzer's result, e.g.:
        ``{(filename,): [1, 2, 3]}``
    """
    if pool is None:
        pool = {}
    filenames = glob.glob(glob_string)
    for filename in tqdm(filenames, desc='analyzing pool'):
        try:
            raw_image = imread(filename, **options['imread'])
        except Exception as err:
            if skip_read_failures:
                warnings.warn("Skipping {}; raised exception:\n    {}"
                              "".format(filename, err))
                continue
            raise
        image = standardize_image(raw_image)
        # Convert color to perceptually-uniform color space.
        sample = sample_pixels(image, 1000)
        converted_sample = colorspacious.cspace_convert(sample, "sRGB1",
                                                        options['colorspace'])
        vector = analyzer(converted_sample)
        pool[(filename,)] = vector
    return pool


def standardize_image(image):
    """
    Ensure that image is float 0-1 RGB with no alpha.

    Parameters
    ----------
    image : array

    Returns
    -------
    image : array
        may or may not be a copy of the original
    """
    image = img_as_float(image)  # ensure float scaled 0-1
    # If there is no color axis, create one.
    if image.ndim == 2:
        image = gray2rgb(image)
    # Assume last axis is color axis. If alpha channel exists, drop it.
    if image.shape[-1] == 4:
        image = image[:, :, :-1]
    return image


class SimpleMatcher:
    """
    This simple matcher returns the closest match.

    It maintains an internal tree representation of the pool for fast lookups.

    Parameters
    ----------
    pool : dict
    """
    def __init__(self, pool):
        self._pool = OrderedDict(pool)
        self._args = list(self._pool.keys())
        data = np.array([vector for vector in self._pool.values()])
        self._tree = cKDTree(data)

    def match(self, vector):
        """
        Return the key of the pool image that is "nearest" (in color space).

        Parameters
        ----------
        vector : array
            characterizing the color to be matched

        Returns
        -------
        args : tuple
            arguments that specify how to open the image
        """
        distance, index = self._tree.query(vector, k=1)
        return self._args[index]


def draw_mosaic(image, tiles, matches, scale=1, resized_copy_cache=None):
    """
    Assemble the mosaic, the final result.

    Parameters
    ----------
    image : array
        the "canvas" on which to draw the tiles, modified in place
    tiles : list
        list of pairs of slice objects
    matches : list
        for each tile in ``tiles``, a tuple of arguments for opening the
        matching image file
    scale : int, optional
        Scale up tiles for higher resolution image; default is 1.
        Any not-integer input will be cast to int.
    resized_copy_cache : dict or None, optional
        cache of images from the pool, sized to fit tiles
        entries look like: ``(pool_key, (height, width))``

    Returns
    -------
    image : array

    Example
    -------
    Basic usage:
    >>> draw_mosaic(image, tiles, matches)

    Cache the resized pool images to speed up repeated drawings:
    >>> cache = {}  # any mutable mapping
    >>> draw_mosiac(image, tiles, matches, resized_copy_cache=cache)

    The above populated ``cache`` with every resized pool image used in a tile.
    Now, ``draw_mosaic``will check the cache before loading the pool image and
    resizing it, which is the most expensive step.
    >>> draw_mosiac(image, tiles, matches, resized_copy_cache=cache)

    """
    scale = int(scale)
    if resized_copy_cache is None:
        resized_copy_cache = {}
    for tile, match_args in zip(tqdm(tiles, desc='drawing mosaic'), matches):
        if scale != 1:
            tile = tuple(slice(scale * s.start, scale * s.stop)
                         for s in tile)
        target_shape = _tile_shape(tile)
        cache_key = (match_args, target_shape)
        try:
            sized_match_image = resized_copy_cache[cache_key]
        except KeyError:
            target_shape = _tile_shape(tile)
            raw_match_image = imread(*match_args, **options['imread'])
            match_image = standardize_image(raw_match_image)
            sized_match_image = crop_to_fit(match_image, target_shape)
            resized_copy_cache[cache_key] = sized_match_image
        image[tile] = sized_match_image
    return image


def _subdivide(tile):
    "Create four tiles from the four quadrants of the input tile."
    tile_dims = [(s.stop - s.start) // 2 for s in tile]
    subtiles = []
    for y in (0, 1):
        for x in (0, 1):
            subtile = (slice(tile[0].start + y * tile_dims[0],
                             tile[0].start + (1 + y) * tile_dims[0]),
                       slice(tile[1].start + x * tile_dims[1],
                             tile[1].start + (1 + x) * tile_dims[1]))
            subtiles.append(subtile)
    return subtiles


def partition(image, grid_dims, mask=None, depth=0, hdr=80):
    """
    Parition the target image into tiles.

    Optionally, subdivide tiles that contain high contrast, creating a grid
    with tiles of varied size.

    Parameters
    ----------
    grid_dims : int or tuple
        number of (largest) tiles along each dimension
    mask : array, optional
        Tiles that straddle a mask edge will be subdivided, creating a smooth
        edge.
    depth : int, optional
        Default is 0. Maximum times a tile can be subdivided.
    hdr : float
        "High Dynamic Range" --- level of contrast that trigger subdivision

    Returns
    -------
    tiles : list
        list of pairs of slice objects
    """
    # Validate inputs.
    if isinstance(grid_dims, int):
        tile_dims, = image.ndims * (grid_dims,)
    for i in (0, 1):
        image_dim = image.shape[i]
        grid_dim = grid_dims[i]
        if image_dim % grid_dim*2**depth != 0:
            raise ValueError("Image dimensions must be evenly divisible by "
                             "dimensions of the (subdivided) grid. "
                             "Dimension {image_dim} is not "
                             "evenly divisible by {grid_dim}*2**{depth} "
                             "".format(image_dim=image_dim, grid_dim=grid_dim,
                                       depth=depth))

    # Partition into equal-sized tiles (Python slice objects).
    tile_height = image.shape[0] // grid_dims[0]
    tile_width = image.shape[1] // grid_dims[1]
    tiles = []
    total = np.product(grid_dims)
    with tqdm(total=total, desc='partitioning depth 0') as pbar:
        for y in range(grid_dims[0]):
            for x in range(grid_dims[1]):
                tile = (slice(y * tile_height, (1 + y) * tile_height),
                        slice(x * tile_width, (1 + x) * tile_width))
                tiles.append(tile)
                pbar.update()

    # Discard any tiles that reside fully outside the mask.
    if mask is not None:
        tiles = [tile for tile in tiles if np.any(mask[tile])]

    # If depth > 0, subdivide any tiles that straddle a mask edge or that
    # contain an image with high contrast.
    for d in range(depth):
        new_tiles = []
        for tile in tqdm(tiles, desc='partitioning depth %d' % d):
            if ((mask is not None) and
                    np.any(mask[tile]) and np.any(~mask[tile])):
                # This tile straddles a mask edge.
                subtiles = _subdivide(tile)
                # Discard subtiles that reside fully outside the mask.
                subtiles = [tile for tile in subtiles if np.any(mask[tile])]
                new_tiles.extend(subtiles)
            elif False:  # TODO hdr check
                new_tiles.extend(_subdivide(tile))
            else:
                new_tiles.append(tile)
        tiles = new_tiles
    return tiles


def color_hist(image, bins=256, **kwargs):
    """
    Compute the distribution of each color channel.

    This wraps ``numpy.histogram``, merely adding data munging relevant to
    image array with color channels. See numpy documentation for details on the
    meaning of the parameters.

    Parameters
    ----------
    image : array
        The last axis is expected to be the color axis.
    bins : int or list
        default 256; passed through to ``numpy.historgram``
    density : bool
        True by default; passed through to ``numpy.histrogram``.
    kwargs :
        passed through to ``numpy.histogram``

    Returns
    -------
    counts, bins : array, array
        Each array has dimensions ``(num_chanels, bins)``.

    Note
    ----
    The image is expected to be float with domain [0, 1]. Histrogram bins
    span full domain, divided equally into the number of bins specified.
    """
    image = np.asarray(image)
    num_pixels = np.product(image.shape[:-1])
    pixels = image.reshape(num_pixels, image.shape[-1])
    func = partial(np.histogram, bins=bins, density=True, **kwargs)
    counts, bins = np.apply_along_axis(func, 0, pixels)
    return np.asarray(list(counts)), np.asarray(list(bins))


def palette_map(from_palette, to_palette):
    """
    Build a function that maps from one color palette into another.

    Parameters
    ----------
    from_palette : tuple
        histogram arrays ``(count, bins)``
    from_palette : array
        histogram arrays ``(count, bins)``

    Returns
    -------
    f : function
    """
    old_counts, old_bins = from_palette
    new_counts, new_bins = to_palette
    # take bins from just one color
    old_bins = old_bins[0]
    new_bins = new_bins[0]
    # cumulative distribution functions
    old_cdf = np.cumsum(np.asarray(list(old_counts)))
    new_cdf = np.cumsum(np.asarray(list(new_counts)))

    def f(image):
        """
        Rescale colors in image from the old color palette to the new one.
        """
        image = np.asarray(image)
        num_pixels = np.product(image.shape[:-1])
        pixels = image.reshape(num_pixels, image.shape[-1])

        # Identify a color bin for each pixels.
        old_x = np.apply_along_axis(
            partial(np.searchsorted, old_bins), 0, pixels)

        # Where in the original gamut did this color fall?
        old_y = old_cdf[old_x]

        # Find that same position in the new gamut. What is the color?
        new_x = np.apply_along_axis(
            partial(np.searchsorted, new_cdf), 0, old_y)
        return new_bins[new_x].reshape(image.shape)

    return f


{'fraction_non_alphanumeric': 0.098941941671912, 'fraction_numerical': 0.006893864460630039, 'mean_word_length': 2.2077684837996348, 'pattern_counts': {'":': 6, '<': 3, '<?xml version=': 0, '>': 61, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3071099', 'n_tokens_mistral': 9863, 'n_tokens_neox': 9485, 'n_words': 3247}