Function predict_with_inferer

monai/apps/detection/utils/predict_utils.py:92–141

Predict network dict output with an inferer. Compared with directly output network(images), it enables a sliding window inferer that can be used to handle large inputs. Args: images: input of the network, Tensor sized (B, C, H, W) or (B, C, H, W, D) network: a network

(
    images: Tensor, network: nn.Module, keys: list[str], inferer: SlidingWindowInferer | None = None
)

def predict_with_inferer(
    images: Tensor, network: nn.Module, keys: list[str], inferer: SlidingWindowInferer | None = None
) -> dict[str, list[Tensor]]:
    """
    Predict network dict output with an inferer. Compared with directly output network(images),
    it enables a sliding window inferer that can be used to handle large inputs.

    Args:
        images: input of the network, Tensor sized (B, C, H, W) or  (B, C, H, W, D)
        network: a network that takes an image Tensor sized (B, C, H, W) or (B, C, H, W, D) as input
            and outputs a dictionary Dict[str, List[Tensor]] or Dict[str, Tensor].
        keys: the keys in the output dict, should be network output keys or a subset of them.
        inferer: a SlidingWindowInferer to handle large inputs.

    Return:
        The predicted head_output from network, a Dict[str, List[Tensor]]

    Example:
        .. code-block:: python

            # define a naive network
            import torch
            import monai
            class NaiveNet(torch.nn.Module):
                def __init__(self, ):
                    super().__init__()

                def forward(self, images: torch.Tensor):
                    return {"cls": torch.randn(images.shape), "box_reg": [torch.randn(images.shape)]}

            # create a predictor
            network = NaiveNet()
            inferer = monai.inferers.SlidingWindowInferer(
                roi_size = (128, 128, 128),
                overlap = 0.25,
                cache_roi_weight_map = True,
            )
            network_output_keys=["cls", "box_reg"]
            images = torch.randn((2, 3, 512, 512, 512))  # a large input
            head_outputs = predict_with_inferer(images, network, network_output_keys, inferer)

    """
    if inferer is None:
        raise ValueError("Please set inferer as a monai.inferers.inferer.SlidingWindowInferer(*)")
    head_outputs_sequence = inferer(images, _network_sequence_output, network, keys=keys)
    num_output_levels: int = len(head_outputs_sequence) // len(keys)
    head_outputs = {}
    for i, k in enumerate(keys):
        head_outputs[k] = list(head_outputs_sequence[num_output_levels * i : num_output_levels * (i + 1)])
    return head_outputs