Function map_classes_to_indices

monai/transforms/utils.py:482–544

Filter out indices of every class of the input label data, return the indices after fattening. It can handle both One-Hot format label and Argmax format label, must provide `num_classes` for Argmax label. For example: ``label = np.array([[[0, 1, 2], [2, 0, 1], [1, 2, 0]]])`` an

(
    label: NdarrayOrTensor,
    num_classes: int | None = None,
    image: NdarrayOrTensor | None = None,
    image_threshold: float = 0.0,
    max_samples_per_class: int | None = None,
)

def map_classes_to_indices(
    label: NdarrayOrTensor,
    num_classes: int | None = None,
    image: NdarrayOrTensor | None = None,
    image_threshold: float = 0.0,
    max_samples_per_class: int | None = None,
) -> list[NdarrayOrTensor]:
    """
    Filter out indices of every class of the input label data, return the indices after fattening.
    It can handle both One-Hot format label and Argmax format label, must provide `num_classes` for
    Argmax label.

    For example:
    ``label = np.array([[[0, 1, 2], [2, 0, 1], [1, 2, 0]]])`` and `num_classes=3`, will return a list
    which contains the indices of the 3 classes:
    ``[np.array([0, 4, 8]), np.array([1, 5, 6]), np.array([2, 3, 7])]``

    Args:
        label: use the label data to get the indices of every class.
        num_classes: number of classes for argmax label, not necessary for One-Hot label.
        image: if image is not None, only return the indices of every class that are within the valid
            region of the image (``image > image_threshold``).
        image_threshold: if enabled `image`, use ``image > image_threshold`` to
            determine the valid image content area and select class indices only in this area.
        max_samples_per_class: maximum length of indices in each class to reduce memory consumption.
            Default is None, no subsampling.

    """
    check_non_lazy_pending_ops(label, name="map_classes_to_indices")
    img_flat: NdarrayOrTensor | None = None
    if image is not None:
        check_non_lazy_pending_ops(image, name="map_classes_to_indices")
        img_flat = ravel((image > image_threshold).any(0))  # type: ignore

    # assuming the first dimension is channel
    channels = len(label)

    num_classes_: int = channels
    if channels == 1:
        if num_classes is None:
            raise ValueError("channels==1 indicates not using One-Hot format label, must provide ``num_classes``.")
        num_classes_ = num_classes

    indices: list[NdarrayOrTensor] = []
    for c in range(num_classes_):
        if channels > 1:
            label_flat = ravel(convert_data_type(label[c], dtype=bool)[0])
        else:
            label_flat = ravel(label == c)
        if img_flat is not None:
            label_flat = img_flat & label_flat
        # no need to save the indices in GPU, otherwise, still need to move to CPU at runtime when crop by indices
        output_type = torch.Tensor if isinstance(label, monai.data.MetaTensor) else None
        cls_indices: NdarrayOrTensor = convert_data_type(
            nonzero(label_flat), output_type=output_type, device=torch.device("cpu")
        )[0]
        if max_samples_per_class and len(cls_indices) > max_samples_per_class and len(cls_indices) > 1:
            sample_id = np.round(np.linspace(0, len(cls_indices) - 1, max_samples_per_class)).astype(int)