Source code for monai.apps.deepgrow.interaction

# Copyright 2020 - 2021 MONAI Consortium
# 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
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Callable, Dict, Sequence, Union

import torch

from monai.engines import SupervisedEvaluator, SupervisedTrainer
from monai.engines.workflow import Events
from monai.transforms import Compose
from monai.utils.enums import CommonKeys

[docs]class Interaction: """ Ignite handler used to introduce interactions (simulation of clicks) for Deepgrow Training/Evaluation. This implementation is based on: Sakinis et al., Interactive segmentation of medical images through fully convolutional neural networks. (2019) Args: transforms: execute additional transformation during every iteration (before train). Typically, several Tensor based transforms composed by `Compose`. max_interactions: maximum number of interactions per iteration train: training or evaluation key_probability: field name to fill probability for every interaction """ def __init__( self, transforms: Union[Sequence[Callable], Callable], max_interactions: int, train: bool, key_probability: str = "probability", ) -> None: if not isinstance(transforms, Compose): transforms = Compose(transforms) self.transforms = transforms self.max_interactions = max_interactions self.train = train self.key_probability = key_probability def attach(self, engine: Union[SupervisedTrainer, SupervisedEvaluator]) -> None: if not engine.has_event_handler(self, Events.ITERATION_STARTED): engine.add_event_handler(Events.ITERATION_STARTED, self) def __call__(self, engine: Union[SupervisedTrainer, SupervisedEvaluator], batchdata: Dict[str, torch.Tensor]): if batchdata is None: raise ValueError("Must provide batch data for current iteration.") for j in range(self.max_interactions): inputs, _ = engine.prepare_batch(batchdata) inputs = with torch.no_grad(): if engine.amp: with torch.cuda.amp.autocast(): predictions = engine.inferer(inputs, else: predictions = engine.inferer(inputs, batchdata.update({CommonKeys.PRED: predictions}) batchdata[self.key_probability] = torch.as_tensor( ([1.0 - ((1.0 / self.max_interactions) * j)] if self.train else [1.0]) * len(inputs) ) batchdata = self.transforms(batchdata) return engine._iteration(engine, batchdata)