Class SplitDimd

monai/transforms/utility/dictionary.py:341–404

class SplitDimd(MapTransform, MultiSampleTrait):
    backend = SplitDim.backend

    def __init__(
        self,
        keys: KeysCollection,
        output_postfixes: Sequence[str] | None = None,
        dim: int = 0,
        keepdim: bool = True,
        update_meta: bool = True,
        list_output: bool = False,
        allow_missing_keys: bool = False,
    ) -> None:
        """
        Args:
            keys: keys of the corresponding items to be transformed.
                See also: :py:class:`monai.transforms.compose.MapTransform`
            output_postfixes: the postfixes to construct keys to store split data.
                for example: if the key of input data is `pred` and split 2 classes, the output
                data keys will be: pred_(output_postfixes[0]), pred_(output_postfixes[1])
                if None, using the index number: `pred_0`, `pred_1`, ... `pred_N`.
            dim: which dimension of input image is the channel, default to 0.
            keepdim: if `True`, output will have singleton in the split dimension. If `False`, this
                dimension will be squeezed.
            update_meta: if `True`, copy `[key]_meta_dict` for each output and update affine to
                reflect the cropped image
            list_output: it `True`, the output will be a list of dictionaries with the same keys as original.
            allow_missing_keys: don't raise exception if key is missing.
        """
        super().__init__(keys, allow_missing_keys)
        self.output_postfixes = output_postfixes
        self.splitter = SplitDim(dim, keepdim, update_meta)
        self.list_output = list_output
        if self.list_output is None and self.output_postfixes is not None:
            raise ValueError("`output_postfixes` should not be provided when `list_output` is `True`.")

    def __call__(
        self, data: Mapping[Hashable, torch.Tensor]
    ) -> dict[Hashable, torch.Tensor] | list[dict[Hashable, torch.Tensor]]:
        d = dict(data)
        all_keys = list(set(self.key_iterator(d)))

        if self.list_output:
            output = []
            results = [self.splitter(d[key]) for key in all_keys]
            for row in zip(*results):
                new_dict = dict(zip(all_keys, row))
                # fill in the extra keys with unmodified data
                for k in set(d.keys()).difference(set(all_keys)):
                    new_dict[k] = deepcopy(d[k])
                output.append(new_dict)
            return output

        for key in all_keys:
            rets = self.splitter(d[key])
            postfixes: Sequence = list(range(len(rets))) if self.output_postfixes is None else self.output_postfixes
            if len(postfixes) != len(rets):
                raise ValueError(f"count of splits must match output_postfixes, {len(postfixes)} != {len(rets)}.")