Random batch subimages numpy

random_subimages.py

def random_subimages_idx(n, d1, d2, h, w):
    hs = np.random.randint(d1-h, size=n)
    ws = np.random.randint(d2-w, size=n)
    grid = np.meshgrid(np.arange(h), np.arange(w), indexing='ij')
    i = np.tile(np.arange(n)[:, None, None], (h, w))
    xs = grid[0]+hs[:, None, None]
    ys = grid[1]+ws[:, None, None]
    return i, xs, ys

def random_subimages(images, h, w):
    return images[random_subimages_idx(*images.shape, h, w)]