base_predictor.py

# Implementation of this model is borrowed and modified
# (from torch to paddle) from here:
# https://github.com/MIC-DKFZ/nnUNet

# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
from functools import partial
from typing import Tuple, List, Union
from scipy.ndimage.filters import gaussian_filter

from .utils import no_op, pad_nd_image

import paddle
from paddle.amp import auto_cast
from tools.preprocess_utils import GenericPreprocessor, PreprocessorFor2D
from nnunet.transforms import default_2D_augmentation_params, default_3D_augmentation_params


class BasePredictor:
    """
    Basic predictor for nnunet.
    """

    def __init__(self):
        # The follow 4 attributes must be set in subclass.
        self.input_shape_must_be_divisible_by = None
        self.threeD = None
        self.num_classes = None
        self.inference_apply_nonlin = None

        self._gaussian_3d = self._patch_size_for_gaussian_3d = None
        self._gaussian_2d = self._patch_size_for_gaussian_2d = None

    def __call__(self, *args, **kwargs):
        raise NotImplementedError

    def predict_3D(self,
                   x: np.ndarray,
                   do_mirroring: bool,
                   mirror_axes: Tuple[int, ...]=(0, 1, 2),
                   use_sliding_window: bool=False,
                   step_size: float=0.5,
                   patch_size: Tuple[int, ...]=None,
                   regions_class_order: Tuple[int, ...]=None,
                   use_gaussian: bool=False,
                   pad_border_mode: str="constant",
                   pad_kwargs: dict=None,
                   verbose: bool=True,
                   mixed_precision: bool=True) -> Tuple[np.ndarray, np.ndarray]:
        """
        x: Input data. Must be a nd.ndarray of shape (c, x, y, z).
        do_mirroring: If True, use test time data augmentation in the form of mirroring.
        mirror_axes: Determines which axes to use for mirroing. Per default, mirroring is done along all three axes. Default: (0, 1, 2).
        use_sliding_window: if True, run sliding window prediction. Heavily recommended! Default: True.
        step_size: When running sliding window prediction, the step size determines the distance between adjacent
        predictions. The smaller the step size, the denser the predictions (and the longer it takes!). Step size is given
        as a fraction of the patch_size. 0.5 is the default and means that wen advance by patch_size * 0.5 between
        predictions. step_size cannot be larger than 1! Default: 0.5.
        patch_size: The patch size that was used for training the network. Do not use different patch sizes here,
        this will either crash or give potentially less accurate segmentations. Default: None.
        regions_class_order: Return region class by given order. Default: None.
        use_gaussian: (Only applies to sliding window prediction) If True, uses a Gaussian importance weighting
         to weigh predictions closer to the center of the current patch higher than those at the borders. The reason
         behind this is that the segmentation accuracy decreases towards the borders. Default (and recommended): True
        pad_border_mode: The type of padding image border. Default: 'constant'.
        pad_kwargs: Parameters for padding image border. Default: NOne.
        verbose: Whether print log. Default: True.
        mixed_precision: Whether use amp in inference. Default: True.
        """
        assert step_size <= 1, 'step_size must be smaller than 1. Otherwise there will be a gap between consecutive predictions'

        if verbose:
            print("debug: mirroring", do_mirroring, "mirror_axes", mirror_axes)

        if pad_kwargs is None:
            pad_kwargs = {'constant_values': 0}

        assert len(x.shape) == 4, "data must have shape (c,x,y,z)"

        if mixed_precision:
            context = auto_cast
        else:
            context = no_op

        with context():
            with paddle.no_grad():
                if self.threeD:
                    if use_sliding_window:
                        res = self._internal_predict_3D_3Dconv_tiled(
                            x,
                            step_size,
                            do_mirroring,
                            mirror_axes,
                            patch_size,
                            regions_class_order,
                            use_gaussian,
                            pad_border_mode,
                            pad_kwargs=pad_kwargs,
                            verbose=verbose)
                    else:
                        res = self._internal_predict_3D_3Dconv(
                            x,
                            patch_size,
                            do_mirroring,
                            mirror_axes,
                            regions_class_order,
                            pad_border_mode,
                            pad_kwargs=pad_kwargs,
                            verbose=verbose)
                else:
                    if use_sliding_window:
                        res = self._internal_predict_3D_2Dconv_tiled(
                            x, patch_size, do_mirroring, mirror_axes, step_size,
                            regions_class_order, use_gaussian, pad_border_mode,
                            pad_kwargs, False)
                    else:
                        res = self._internal_predict_3D_2Dconv(
                            x, patch_size, do_mirroring, mirror_axes,
                            regions_class_order, pad_border_mode, pad_kwargs,
                            False)
        return res

    @staticmethod
    def _get_gaussian(patch_size, sigma_scale=1. / 8) -> np.ndarray:
        tmp = np.zeros(patch_size)
        center_coords = [i // 2 for i in patch_size]
        sigmas = [i * sigma_scale for i in patch_size]
        tmp[tuple(center_coords)] = 1
        gaussian_importance_map = gaussian_filter(
            tmp, sigmas, 0, mode='constant', cval=0)
        gaussian_importance_map = gaussian_importance_map / np.max(
            gaussian_importance_map) * 1
        gaussian_importance_map = gaussian_importance_map.astype(np.float32)

        gaussian_importance_map[gaussian_importance_map == 0] = np.min(
            gaussian_importance_map[gaussian_importance_map != 0])

        return gaussian_importance_map

    @staticmethod
    def _compute_steps_for_sliding_window(patch_size: Tuple[int, ...],
                                          image_size: Tuple[int, ...],
                                          step_size: float) -> List[List[int]]:
        assert [i >= j for i, j in zip(image_size, patch_size)
                ], "image size must be as large or larger than patch_size"
        assert 0 < step_size <= 1, 'step_size must be larger than 0 and smaller or equal to 1'

        target_step_sizes_in_voxels = [i * step_size for i in patch_size]
        num_steps = [
            int(np.ceil((i - k) / j)) + 1
            for i, j, k in zip(image_size, target_step_sizes_in_voxels,
                               patch_size)
        ]

        steps = []
        for dim in range(len(patch_size)):
            max_step_value = image_size[dim] - patch_size[dim]
            if num_steps[dim] > 1:
                actual_step_size = max_step_value / (num_steps[dim] - 1)
            else:
                actual_step_size = 99999999999

            steps_here = [
                int(np.round(actual_step_size * i))
                for i in range(num_steps[dim])
            ]
            steps.append(steps_here)
        return steps

    def _internal_predict_3D_3Dconv_tiled(
            self,
            x: np.ndarray,
            step_size: float,
            do_mirroring: bool,
            mirror_axes: tuple,
            patch_size: tuple,
            regions_class_order: tuple,
            use_gaussian: bool,
            pad_border_mode: str,
            pad_kwargs: dict,
            verbose: bool) -> Tuple[np.ndarray, np.ndarray]:
        assert len(x.shape) == 4, "x must be (c, x, y, z)"
        assert patch_size is not None, "patch_size cannot be None for tiled prediction"

        data, slicer = pad_nd_image(x, patch_size, pad_border_mode, pad_kwargs,
                                    True, None)
        data_shape = data.shape

        steps = self._compute_steps_for_sliding_window(
            patch_size, data_shape[1:], step_size)
        num_tiles = len(steps[0]) * len(steps[1]) * len(steps[2])

        if verbose:
            print("step_size:", step_size)
            print("do mirror:", do_mirroring)
            print("data shape:", data_shape)
            print("patch size:", patch_size)
            print("steps (x, y, and z):", steps)
            print("number of tiles:", num_tiles)

        if use_gaussian and num_tiles > 1:
            if self._gaussian_3d is None or not all([
                    i == j
                    for i, j in zip(patch_size,
                                    self._patch_size_for_gaussian_3d)
            ]):
                gaussian_importance_map = self._get_gaussian(
                    patch_size, sigma_scale=1. / 8)

                self._gaussian_3d = gaussian_importance_map
                self._patch_size_for_gaussian_3d = patch_size
                if verbose:
                    print("computing Gaussian done")
            else:
                if verbose:
                    print("using precomputed Gaussian")
                gaussian_importance_map = self._gaussian_3d

            gaussian_importance_map = paddle.to_tensor(gaussian_importance_map)
        else:
            gaussian_importance_map = None

        if use_gaussian and num_tiles > 1:
            add_for_nb_of_preds = self._gaussian_3d
        else:
            add_for_nb_of_preds = np.ones(patch_size, dtype=np.float32)
        aggregated_results = np.zeros(
            [self.num_classes] + list(data.shape[1:]), dtype=np.float32)
        aggregated_nb_of_predictions = np.zeros(
            [self.num_classes] + list(data.shape[1:]), dtype=np.float32)

        for x in steps[0]:
            lb_x = x
            ub_x = x + patch_size[0]
            for y in steps[1]:
                lb_y = y
                ub_y = y + patch_size[1]
                for z in steps[2]:
                    lb_z = z
                    ub_z = z + patch_size[2]

                    predicted_patch = self._internal_maybe_mirror_and_pred_3D(
                        data[None, :, lb_x:ub_x, lb_y:ub_y, lb_z:ub_z],
                        mirror_axes, do_mirroring, gaussian_importance_map)[0]

                    predicted_patch = predicted_patch.numpy()
                    aggregated_results[:, lb_x:ub_x, lb_y:ub_y, lb_z:
                                       ub_z] += predicted_patch
                    aggregated_nb_of_predictions[:, lb_x:ub_x, lb_y:ub_y, lb_z:
                                                 ub_z] += add_for_nb_of_preds

        slicer = tuple([
            slice(0, aggregated_results.shape[i])
            for i in range(len(aggregated_results.shape) - (len(slicer) - 1))
        ] + slicer[1:])
        aggregated_results = aggregated_results[slicer]
        aggregated_nb_of_predictions = aggregated_nb_of_predictions[slicer]
        aggregated_results /= aggregated_nb_of_predictions

        del aggregated_nb_of_predictions
        if regions_class_order is None:
            predicted_segmentation = aggregated_results.argmax(0)
        else:
            class_probabilities_here = aggregated_results
            predicted_segmentation = np.zeros(
                class_probabilities_here.shape[1:], dtype=np.float32)
            for i, c in enumerate(regions_class_order):
                predicted_segmentation[class_probabilities_here[i] > 0.5] = c

        if verbose:
            print("prediction done")
        return predicted_segmentation, aggregated_results

    def _internal_predict_2D_2Dconv(
            self,
            x: np.ndarray,
            min_size: Tuple[int, int],
            do_mirroring: bool,
            mirror_axes: tuple=(0, 1, 2),
            regions_class_order: tuple=None,
            pad_border_mode: str="constant",
            pad_kwargs: dict=None,
            verbose: bool=True) -> Tuple[np.ndarray, np.ndarray]:
        assert len(x.shape) == 3, "x must be (c, x, y), but got {}.".format(
            x.shape)

        assert self.input_shape_must_be_divisible_by is not None, 'input_shape_must_be_divisible_by must be set but got None'
        if verbose:
            print("do mirror:", do_mirroring)

        data, slicer = pad_nd_image(x, min_size, pad_border_mode, pad_kwargs,
                                    True, self.input_shape_must_be_divisible_by)
        predicted_probabilities = self._internal_maybe_mirror_and_pred_2D(
            data[None], mirror_axes, do_mirroring, None)[0]

        slicer = tuple([
            slice(0, predicted_probabilities.shape[i])
            for i in range(
                len(predicted_probabilities.shape) - (len(slicer) - 1))
        ] + slicer[1:])
        predicted_probabilities = predicted_probabilities[slicer]

        if regions_class_order is None:
            predicted_segmentation = predicted_probabilities.argmax(0)
            predicted_segmentation = predicted_segmentation.numpy()
            predicted_probabilities = predicted_probabilities.numpy()
        else:
            predicted_probabilities = predicted_probabilities.numpy()
            predicted_segmentation = np.zeros(
                predicted_probabilities.shape[1:], dtype=np.float32)
            for i, c in enumerate(regions_class_order):
                predicted_segmentation[predicted_probabilities[i] > 0.5] = c

        return predicted_segmentation, predicted_probabilities

    def _internal_predict_3D_3Dconv(
            self,
            x: np.ndarray,
            min_size: Tuple[int, ...],
            do_mirroring: bool,
            mirror_axes: tuple=(0, 1, 2),
            regions_class_order: tuple=None,
            pad_border_mode: str="constant",
            pad_kwargs: dict=None,
            verbose: bool=True) -> Tuple[np.ndarray, np.ndarray]:
        assert len(x.shape) == 4, "x must be (c, x, y, z), but got {}.".format(
            x.shape)

        assert self.input_shape_must_be_divisible_by is not None, 'input_shape_must_be_divisible_by must be set but got None'
        if verbose:
            print("do mirror:", do_mirroring)

        data, slicer = pad_nd_image(x, min_size, pad_border_mode, pad_kwargs,
                                    True, self.input_shape_must_be_divisible_by)
        predicted_probabilities = self._internal_maybe_mirror_and_pred_3D(
            data[None], mirror_axes, do_mirroring, None)[0]

        slicer = tuple([
            slice(0, predicted_probabilities.shape[i])
            for i in range(
                len(predicted_probabilities.shape) - (len(slicer) - 1))
        ] + slicer[1:])
        predicted_probabilities = predicted_probabilities[slicer]

        if regions_class_order is None:
            predicted_segmentation = predicted_probabilities.argmax(0)
            predicted_segmentation = predicted_segmentation.numpy()
            predicted_probabilities = predicted_probabilities.numpy()
        else:
            predicted_probabilities = predicted_probabilities.numpy()
            predicted_segmentation = np.zeros(
                predicted_probabilities.shape[1:], dtype=np.float32)
            for i, c in enumerate(regions_class_order):
                predicted_segmentation[predicted_probabilities[i] > 0.5] = c

        return predicted_segmentation, predicted_probabilities

    def _internal_maybe_mirror_and_pred_3D(self,
                                           x: Union[np.ndarray, paddle.Tensor],
                                           mirror_axes: tuple,
                                           do_mirroring: bool=True,
                                           mult: np.ndarray or
                                           paddle.Tensor=None) -> paddle.Tensor:
        assert len(
            x.shape) == 5, 'x must be (b, c, x, y, z), but got {}.'.format(
                x.shape)

        x = paddle.to_tensor(x).astype('float32')
        result = paddle.zeros(
            [1, self.num_classes] + list(x.shape[2:]), dtype='float32')

        if mult is not None:
            mult = paddle.to_tensor(mult).astype('float32')

        if do_mirroring:
            mirror_idx = 8
            num_results = 2**len(mirror_axes)
        else:
            mirror_idx = 1
            num_results = 1

        for m in range(mirror_idx):
            if m == 0:
                pred = self.inference_apply_nonlin(self(x))
                result += 1 / num_results * pred

            if m == 1 and (2 in mirror_axes):
                pred = self.inference_apply_nonlin(self(paddle.flip(x, (4, ))))
                result += 1 / num_results * paddle.flip(pred, (4, ))

            if m == 2 and (1 in mirror_axes):
                pred = self.inference_apply_nonlin(self(paddle.flip(x, (3, ))))
                result += 1 / num_results * paddle.flip(pred, (3, ))

            if m == 3 and (2 in mirror_axes) and (1 in mirror_axes):
                pred = self.inference_apply_nonlin(self(paddle.flip(x, (4, 3))))
                result += 1 / num_results * paddle.flip(pred, (4, 3))

            if m == 4 and (0 in mirror_axes):
                pred = self.inference_apply_nonlin(self(paddle.flip(x, (2, ))))
                result += 1 / num_results * paddle.flip(pred, (2, ))

            if m == 5 and (0 in mirror_axes) and (2 in mirror_axes):
                pred = self.inference_apply_nonlin(self(paddle.flip(x, (4, 2))))
                result += 1 / num_results * paddle.flip(pred, (4, 2))

            if m == 6 and (0 in mirror_axes) and (1 in mirror_axes):
                pred = self.inference_apply_nonlin(self(paddle.flip(x, (3, 2))))
                result += 1 / num_results * paddle.flip(pred, (3, 2))

            if m == 7 and (0 in mirror_axes) and (1 in mirror_axes) and (
                    2 in mirror_axes):
                pred = self.inference_apply_nonlin(
                    self(paddle.flip(x, (4, 3, 2))))
                result += 1 / num_results * paddle.flip(pred, (4, 3, 2))

        if mult is not None:
            result[:, :] *= mult
        return result

    def _internal_maybe_mirror_and_pred_2D(self,
                                           x: Union[np.ndarray, paddle.Tensor],
                                           mirror_axes: tuple,
                                           do_mirroring: bool=True,
                                           mult: np.ndarray or
                                           paddle.Tensor=None) -> paddle.Tensor:
        assert len(x.shape) == 4, 'x must be (b, c, x, y), but got {}.'.format(
            x.shape)

        x = paddle.to_tensor(x).astype('float32')
        result = paddle.zeros(
            [x.shape[0], self.num_classes] + list(x.shape[2:]), dtype='float32')

        if mult is not None:
            mult = paddle.to_tensor(mult).astype('float32')

        if do_mirroring:
            mirror_idx = 4
            num_results = 2**len(mirror_axes)
        else:
            mirror_idx = 1
            num_results = 1

        for m in range(mirror_idx):
            if m == 0:
                pred = self.inference_apply_nonlin(self(x))
                result += 1 / num_results * pred
            if m == 1 and (1 in mirror_axes):
                pred = self.inference_apply_nonlin(self(paddle.flip(x, (3, ))))
                result += 1 / num_results * paddle.flip(pred, (3, ))

            if m == 2 and (0 in mirror_axes):
                pred = self.inference_apply_nonlin(self(paddle.flip(x, (2, ))))
                result += 1 / num_results * paddle.flip(pred, (2, ))

            if m == 3 and (0 in mirror_axes) and (1 in mirror_axes):
                pred = self.inference_apply_nonlin(self(paddle.flip(x, (3, 2))))
                result += 1 / num_results * paddle.flip(pred, (3, 2))
        if mult is not None:
            result[:, :] *= mult
        return result

    def _internal_predict_2D_2Dconv_tiled(
            self,
            x: np.ndarray,
            step_size: float,
            do_mirroring: bool,
            mirror_axes: tuple,
            patch_size: tuple,
            regions_class_order: tuple,
            use_gaussian: bool,
            pad_border_mode: str,
            pad_kwargs: dict,
            verbose: bool) -> Tuple[np.ndarray, np.ndarray]:
        assert len(x.shape) == 3, "x must be (c, x, y), but got {}.".format(
            x.shape)
        assert patch_size is not None, "patch_size cannot be None for tiled prediction."

        data, slicer = pad_nd_image(x, patch_size, pad_border_mode, pad_kwargs,
                                    True, None)
        data_shape = data.shape

        steps = self._compute_steps_for_sliding_window(
            patch_size, data_shape[1:], step_size)
        num_tiles = len(steps[0]) * len(steps[1])

        if verbose:
            print("step_size:", step_size)
            print("do mirror:", do_mirroring)
            print("data shape:", data_shape)
            print("patch size:", patch_size)
            print("steps (x, y, and z):", steps)
            print("number of tiles:", num_tiles)

        if use_gaussian and num_tiles > 1:
            if self._gaussian_2d is None or not all([
                    i == j
                    for i, j in zip(patch_size,
                                    self._patch_size_for_gaussian_2d)
            ]):
                if verbose:
                    print('computing Gaussian')
                gaussian_importance_map = self._get_gaussian(
                    patch_size, sigma_scale=1. / 8)
                self._gaussian_2d = gaussian_importance_map
                self._patch_size_for_gaussian_2d = patch_size
            else:
                if verbose:
                    print("using precomputed Gaussian")
                gaussian_importance_map = self._gaussian_2d

            gaussian_importance_map = paddle.to_tensor(gaussian_importance_map)
        else:
            gaussian_importance_map = None

        if use_gaussian and num_tiles > 1:
            add_for_nb_of_preds = self._gaussian_2d
        else:
            add_for_nb_of_preds = np.ones(patch_size, dtype=np.float32)
        aggregated_results = np.zeros(
            [self.num_classes] + list(data.shape[1:]), dtype=np.float32)
        aggregated_nb_of_predictions = np.zeros(
            [self.num_classes] + list(data.shape[1:]), dtype=np.float32)

        for x in steps[0]:
            lb_x = x
            ub_x = x + patch_size[0]
            for y in steps[1]:
                lb_y = y
                ub_y = y + patch_size[1]

                predicted_patch = self._internal_maybe_mirror_and_pred_2D(
                    data[None, :, lb_x:ub_x, lb_y:ub_y], mirror_axes,
                    do_mirroring, gaussian_importance_map)[0]

                predicted_patch = predicted_patch.numpy()
                aggregated_results[:, lb_x:ub_x, lb_y:ub_y] += predicted_patch
                aggregated_nb_of_predictions[:, lb_x:ub_x, lb_y:
                                             ub_y] += add_for_nb_of_preds

        slicer = tuple([
            slice(0, aggregated_results.shape[i])
            for i in range(len(aggregated_results.shape) - (len(slicer) - 1))
        ] + slicer[1:])
        aggregated_results = aggregated_results[slicer]
        aggregated_nb_of_predictions = aggregated_nb_of_predictions[slicer]

        class_probabilities = aggregated_results / aggregated_nb_of_predictions

        if regions_class_order is None:
            predicted_segmentation = class_probabilities.argmax(0)
        else:
            class_probabilities_here = class_probabilities
            predicted_segmentation = np.zeros(
                class_probabilities_here.shape[1:], dtype=np.float32)
            for i, c in enumerate(regions_class_order):
                predicted_segmentation[class_probabilities_here[i] > 0.5] = c

        if verbose:
            print("prediction done")
        return predicted_segmentation, class_probabilities

    def _internal_predict_3D_2Dconv(
            self,
            x: np.ndarray,
            min_size: Tuple[int, int],
            do_mirroring: bool,
            mirror_axes: tuple=(0, 1),
            regions_class_order: tuple=None,
            pad_border_mode: str="constant",
            pad_kwargs: dict=None,
            verbose: bool=True) -> Tuple[np.ndarray, np.ndarray]:
        assert len(x.shape) == 4, "x must be (c, x, y, z), but got {}.".format(
            x.shape)
        predicted_segmentation = []
        softmax_pred = []
        for s in range(x.shape[1]):
            pred_seg, softmax_pres = self._internal_predict_2D_2Dconv(
                x[:, s], min_size, do_mirroring, mirror_axes,
                regions_class_order, pad_border_mode, pad_kwargs, verbose)
            predicted_segmentation.append(pred_seg[None])
            softmax_pred.append(softmax_pres[None])
        predicted_segmentation = np.vstack(predicted_segmentation)
        softmax_pred = np.vstack(softmax_pred).transpose((1, 0, 2, 3))
        return predicted_segmentation, softmax_pred

    def _internal_predict_3D_2Dconv_tiled(
            self,
            x: np.ndarray,
            patch_size: Tuple[int, int],
            do_mirroring: bool,
            mirror_axes: tuple=(0, 1),
            step_size: float=0.5,
            regions_class_order: tuple=None,
            use_gaussian: bool=False,
            pad_border_mode: str="edge",
            pad_kwargs: dict=None,
            verbose: bool=True) -> Tuple[np.ndarray, np.ndarray]:
        assert len(x.shape) == 4, "x must be (c, x, y, z), but got {}.".format(
            x.shape)
        predicted_segmentation = []
        softmax_pred = []
        for s in range(x.shape[1]):
            pred_seg, softmax_pres = self._internal_predict_2D_2Dconv_tiled(
                x[:, s], step_size, do_mirroring, mirror_axes, patch_size,
                regions_class_order, use_gaussian, pad_border_mode, pad_kwargs,
                verbose)
            predicted_segmentation.append(pred_seg[None])
            softmax_pred.append(softmax_pres[None])

        predicted_segmentation = np.vstack(predicted_segmentation)
        softmax_pred = np.vstack(softmax_pred).transpose((1, 0, 2, 3))
        return predicted_segmentation, softmax_pred


class DynamicPredictor(BasePredictor):
    def __init__(self, model):
        super().__init__()
        assert hasattr(
            model, 'net_num_pool_op_kernel_sizes'
        ), "BasePredictor only used for nnunet predict, but not found net_num_pool_op_kernel_sizes in your model."
        self.input_shape_must_be_divisible_by = np.prod(
            model.net_num_pool_op_kernel_sizes, 0, dtype=np.int64)
        self.threeD = model.threeD
        self.num_classes = model.num_classes
        self.inference_apply_nonlin = partial(
            paddle.nn.functional.softmax, axis=1)
        self.model = model

    def __call__(self, x):
        x = self.model(x)[0]
        return x


class MultiFoldsPredictor(BasePredictor):
    def __init__(self, model, param_paths):
        super().__init__()
        assert hasattr(
            model, 'net_num_pool_op_kernel_sizes'
        ), "BasePredictor only used for nnunet predict, but not found net_num_pool_op_kernel_sizes in your model."
        self.input_shape_must_be_divisible_by = np.prod(
            model.net_num_pool_op_kernel_sizes, 0, dtype=np.int64)
        self.threeD = model.threeD
        self.num_classes = model.num_classes
        self.inference_apply_nonlin = partial(
            paddle.nn.functional.softmax, axis=1)
        self.model = model
        self.param_list = [
            paddle.load(param_path) for param_path in param_paths
        ]
        self.plans = model.plans
        self.stage = model.stage

        if self.threeD:
            self.data_aug_params = default_3D_augmentation_params
        else:
            self.data_aug_params = default_2D_augmentation_params

        self.intensity_properties = self.plans['dataset_properties'][
            'intensityproperties']
        self.normalization_schemes = self.plans['normalization_schemes']
        self.use_mask_for_norm = self.plans['use_mask_for_norm']
        if self.plans.get('transpose_forward') is None or self.plans.get(
                'transpose_backward') is None:
            print(
                "WARNING! You seem to have data that was preprocessed with a previous version of nnU-Net. "
                "You should rerun preprocessing. We will proceed and assume that both transpose_foward "
                "and transpose_backward are [0, 1, 2]. If that is not correct then weird things will happen!"
            )
            self.plans['transpose_forward'] = [0, 1, 2]
            self.plans['transpose_backward'] = [0, 1, 2]
        self.transpose_forward = self.plans['transpose_forward']
        self.transpose_backward = self.plans['transpose_backward']

    def __call__(self, x):
        x = self.model(x)[0]
        return x

    def preprocess_patient(self, input_files):
        if self.threeD:
            preprocessor_class = GenericPreprocessor
        else:
            preprocessor_class = PreprocessorFor2D

        preprocessor = preprocessor_class(
            self.normalization_schemes, self.use_mask_for_norm,
            self.transpose_forward, self.intensity_properties)
        d, s, properties = preprocessor.preprocess_test_case(
            input_files,
            self.plans['plans_per_stage'][self.stage]['current_spacing'])
        return d, s, properties

    def predict_preprocessed_data_return_seg_and_softmax(
            self,
            data: np.ndarray,
            do_mirroring: bool=True,
            mirror_axes: Tuple[int]=None,
            use_sliding_window: bool=True,
            step_size: float=0.5,
            use_gaussian: bool=True,
            pad_border_mode: str='constant',
            pad_kwargs: dict=None,
            verbose: bool=True,
            mixed_precision=True):
        if pad_border_mode == 'constant' and pad_kwargs is None:
            pad_kwargs = {'constant_values': 0}
        if do_mirroring and mirror_axes is None:
            mirror_axes = self.data_aug_params['mirror_axes']
        if do_mirroring:
            assert self.data_aug_params["do_mirror"], "Cannot do mirroring as test time augmentation when training " \
                                                      "was done without mirroring"
        self.model.eval()
        res = self.predict_3D(
            x=data,
            do_mirroring=do_mirroring,
            mirror_axes=mirror_axes,
            use_sliding_window=use_sliding_window,
            step_size=step_size,
            patch_size=self.model.patch_size,
            regions_class_order=None,
            use_gaussian=use_gaussian,
            pad_border_mode=pad_border_mode,
            pad_kwargs=pad_kwargs,
            verbose=verbose,
            mixed_precision=mixed_precision)
        return res

    def multi_folds_predict_preprocessed_data_return_seg_and_softmax(
            self,
            data: np.ndarray,
            do_mirroring: bool=True,
            mirror_axes: Tuple[int]=None,
            use_sliding_window: bool=True,
            step_size: float=0.5,
            use_gaussian: bool=True,
            pad_border_mode: str='constant',
            pad_kwargs: dict=None,
            verbose: bool=True,
            mixed_precision=True):
        softmax_res = None
        for params in self.param_list:
            self.model.set_state_dict(params)
            x = self.predict_preprocessed_data_return_seg_and_softmax(
                data=data,
                do_mirroring=do_mirroring,
                mirror_axes=mirror_axes,
                use_sliding_window=use_sliding_window,
                step_size=step_size,
                use_gaussian=use_gaussian,
                pad_border_mode=pad_border_mode,
                pad_kwargs=pad_kwargs,
                verbose=verbose,
                mixed_precision=mixed_precision)[1]
            if softmax_res is None:
                softmax_res = x
            else:
                softmax_res += x
        return softmax_res / len(self.param_list)