select_random_crop_centers

select_random_crop_centers.py

import skimage.measure, skimage.filters.rank
import numpy as np

def select_random_crop_centers(mask, n, w, b):
    """
    mask: binary image indicating which pixels are "good"
    n: number of crops to generate
    w: size (width and height) of crops
    b: block reduction size (larger numbers trade off speed for accuracy)

    Returns an array of (y, x) points representing the center of each crop.
    Subtract and add w/2 to each center to get the corners of the crops.

    Note that w must be even, and w / b must be an even integer.
    """
    assert w % 2 == 0, "w must be even"
    assert w % b == 0, "b must divide w"
    assert w / b % 2 == 0, "w / b must be even"
    ws = round(w / b) + 1
    mask_reduced = skimage.measure.block_reduce(mask, (b, b), np.min)
    mask_eroded = skimage.filters.rank.minimum(
        mask_reduced, np.ones((ws, ws), bool)
    )
    mask_eroded = mask_eroded.astype(bool)
    cmin = w // 2
    cmax = np.array(mask.shape) - w // 2
    centers = np.zeros((0,2), int)
    while len(centers) < n:
        y = np.random.randint(cmin, cmax[0], size=n)
        x = np.random.randint(cmin, cmax[1], size=n)
        ys = y // b
        xs = x // b
        keep = mask_eroded[ys, xs]
        centers = np.vstack([centers, np.vstack([y[keep], x[keep]]).T])
    return centers[:n]