Repository: Project-MONAI/MONAI
File: monai/transforms/utility/dictionary.py (Python)

Code before:
    """

    backend = ClassesToIndices.backend

    def __init__(
        self,
        keys: KeysCollection,
        indices_postfix: str = "_cls_indices",
        num_classes: int | None = None,
        image_key: str | None = None,
        image_threshold: float = 0.0,
        output_shape: Sequence[int] | None = None,
        max_indices_per_class: int | None = None,
        allow_missing_keys: bool = False,
    ) -> None:
        super().__init__(keys, allow_missing_keys)
        self.indices_postfix = indices_postfix
        self.image_key = image_key
        self.max_indices_per_class = max_indices_per_class
        self.converter = ClassesToIndices(num_classes, image_threshold, output_shape)

    def __call__(self, data: Mapping[Hashable, Any]):
        d = dict(data)
        image = d[self.image_key] if self.image_key else None
        for key in self.key_iterator(d):
            out = self.converter(d[key], image)

            if self.max_indices_per_class is not None:
                for c in range(len(out)):
                    if len(out[c]) > self.max_indices_per_class:
                        ind = np.random.choice(len(out[c]), size=self.max_indices_per_class, replace=False)
                        out[c] = out[c][ind]

            d[str(key) + self.indices_postfix] = out

        return d


class ConvertToMultiChannelBasedOnBratsClassesd(MapTransform):
    """
    Dictionary-based wrapper of :py:class:`monai.transforms.ConvertToMultiChannelBasedOnBratsClasses`.
    Convert labels to multi channels based on brats18 classes:
    label 1 is the necrotic and non-enhancing tumor core
    label 2 is the peritumoral edema
    label 4 is the GD-enhancing tumor
    The possible classes are TC (Tumor core), WT (Whole tumor)
    and ET (Enhancing tumor).
    """

    backend = ConvertToMultiChannelBasedOnBratsClasses.backend


Review comment (suggestion):
if it's feasible in general, perhaps we consider subsampling when generating the indices?

https://github.com/Project-MONAI/MONAI/blob/9e33ff25f34ebd11e026c49634743d9389564765/monai/transforms/utils.py#L332-L341


Code after:

    """

    backend = ClassesToIndices.backend

    def __init__(
        self,
        keys: KeysCollection,
        indices_postfix: str = "_cls_indices",
        num_classes: int | None = None,
        image_key: str | None = None,
        image_threshold: float = 0.0,
        output_shape: Sequence[int] | None = None,
        max_samples_per_class: int | None = None,
        allow_missing_keys: bool = False,
    ) -> None:
        super().__init__(keys, allow_missing_keys)
        self.indices_postfix = indices_postfix
        self.image_key = image_key
        self.converter = ClassesToIndices(num_classes, image_threshold, output_shape, max_samples_per_class)

    def __call__(self, data: Mapping[Hashable, Any]):
        d = dict(data)
        image = d[self.image_key] if self.image_key else None
        for key in self.key_iterator(d):
            d[str(key) + self.indices_postfix] = self.converter(d[key], image)

        return d


class ConvertToMultiChannelBasedOnBratsClassesd(MapTransform):
    """
    Dictionary-based wrapper of :py:class:`monai.transforms.ConvertToMultiChannelBasedOnBratsClasses`.
    Convert labels to multi channels based on brats18 classes:
    label 1 is the necrotic and non-enhancing tumor core
    label 2 is the peritumoral edema
    label 4 is the GD-enhancing tumor
    The possible classes are TC (Tumor core), WT (Whole tumor)
    and ET (Enhancing tumor).
    """

    backend = ConvertToMultiChannelBasedOnBratsClasses.backend

    def __init__(self, keys: KeysCollection, allow_missing_keys: bool = False):
        super().__init__(keys, allow_missing_keys)
        self.converter = ConvertToMultiChannelBasedOnBratsClasses()

    def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> dict[Hashable, NdarrayOrTensor]:
        d = dict(data)
        for key in self.key_iterator(d):
            d[key] = self.converter(d[key])