Function weighted_patch_samples

monai/transforms/utils.py:547–593

Computes `n_samples` of random patch sampling locations, given the sampling weight map `w` and patch `spatial_size`. Args: spatial_size: length of each spatial dimension of the patch. w: weight map, the weights must be non-negative. each element denotes a sampling weight of

(
    spatial_size: int | Sequence[int],
    w: NdarrayOrTensor,
    n_samples: int = 1,
    r_state: np.random.RandomState | None = None,
)

def weighted_patch_samples(
    spatial_size: int | Sequence[int],
    w: NdarrayOrTensor,
    n_samples: int = 1,
    r_state: np.random.RandomState | None = None,
) -> list:
    """
    Computes `n_samples` of random patch sampling locations, given the sampling weight map `w` and patch `spatial_size`.

    Args:
        spatial_size: length of each spatial dimension of the patch.
        w: weight map, the weights must be non-negative. each element denotes a sampling weight of the spatial location.
            0 indicates no sampling.
            The weight map shape is assumed ``(spatial_dim_0, spatial_dim_1, ..., spatial_dim_n)``.
        n_samples: number of patch samples
        r_state: a random state container

    Returns:
        a list of `n_samples` N-D integers representing the spatial sampling location of patches.

    """
    check_non_lazy_pending_ops(w, name="weighted_patch_samples")
    if w is None:
        raise ValueError("w must be an ND array, got None.")
    if r_state is None:
        r_state = np.random.RandomState()
    img_size = np.asarray(w.shape, dtype=int)
    win_size = np.asarray(fall_back_tuple(spatial_size, img_size), dtype=int)

    s = tuple(slice(w // 2, m - w + w // 2) if m > w else slice(m // 2, m // 2 + 1) for w, m in zip(win_size, img_size))
    v = w[s]  # weight map in the 'valid' mode
    v_size = v.shape
    v = ravel(v)  # always copy
    if (v < 0).any():
        v -= v.min()  # shifting to non-negative
    v = cumsum(v)
    if not v[-1] or not isfinite(v[-1]) or v[-1] < 0:  # uniform sampling
        idx = r_state.randint(0, len(v), size=n_samples)
    else:
        r_samples = r_state.random(n_samples)
        r, *_ = convert_to_dst_type(r_samples, v, dtype=r_samples.dtype)
        idx = searchsorted(v, r * v[-1], right=True)  # type: ignore
    idx, *_ = convert_to_dst_type(idx, v, dtype=torch.int)  # type: ignore
    # compensate 'valid' mode
    diff = np.minimum(win_size, img_size) // 2
    diff, *_ = convert_to_dst_type(diff, v)  # type: ignore
    return [unravel_index(i, v_size) + diff for i in idx]


def correct_crop_centers(