Reservoir sampling over a collection of images

[starcalibre]

Code by @jni Submitted by @starcalibre

Reservoir sampling over a collection of images for purposes of image illumination correction.

def _reservoir_sampled_image(ims_iter, random_state=None):
    """Return an image where each pixel is sampled from a list of images.
    The idea is to get a sample of image intensity throughout a collection
    of images, to know what the "standard range" is for this type of image.
    The implementation uses a "reservoir" image to sample while remaining
    space-efficient, and only needs to hold about four images at one time
    (the reservoir, the current sample, a random image for sampling, and
    a thresholded version of the random image).
    Parameters
    ----------
    ims_iter : iterable of arrays
        An iterable over numpy arrays (representing images).
    random_state : None, int, or numpy RandomState instance, optional
        An optional random number generator or seed from which to draw
        samples.
    Returns
    -------
    sampled : array, same shape as input
        The sampled "image".
    Examples
    --------
    >>> ims = iter(np.arange(27).reshape((3, 3, 3)))
    >>> _reservoir_sampled_image(ims, 0)
    array([[ 0,  1,  2],
           [ 3, 13, 23],
           [24, 25,  8]])
    """
    random = normalise_random_state(random_state)
    ims_iter = iter(ims_iter)  # ensure iterator and not e.g. list
    sampled = next(ims_iter)
    for k, im in enumerate(ims_iter, start=2):
        to_replace = random.rand(*im.shape) < (1 / k)
        sampled[to_replace] = im[to_replace]
    return sampled