utils.py 文件源码

def mask_grid(pos, grid, radius, one_is_free=True):
    """Mask a grid.

    Parameters
    ----------
    pos : tuple
        the center of the circle (row, col)
        e.g. (2,3) = center at 3rd row, 4th column
    grid : numpy.ndarray
        should be a 2d matrix
        e.g. 8x8, 16x16, 28x28, 40x40
    imsize : tuple
        the grid size
    radius : int
        the length of the radius
    one_is_free : bool
        if True, then 1 is freezone, 0 is block
        if False, then 0 is freezone, 1 is block

    Returns
    -------
    masked_grid : numpy.ndarray
        the masked grid
    """
    mask = np.zeros_like(grid)
    new_grid = grid.copy()
    rr, cc = draw.circle(pos[0], pos[1], radius=radius,
                         shape=mask.shape)
    mask[rr, cc] = 1
    if one_is_free:
        return new_grid*mask, (rr, cc)
    else:
        masked_img = np.ones_like(new_grid)
        masked_img[rr, cc] = new_grid[rr, cc]
        return masked_img