U Ph] @sfddlmZddlmZmZddlZddlmZmZddl m Z ddl m Z dgZ GdddeZdS) ) annotations)AnySequenceN) PrepareBatchPrepareBatchExtraInput) ensure_tuple)HoVerNetBranchPrepareBatchHoVerNetc@s8eZdZdZdddddZdd d d d d dddZdS)r a Customized prepare batch callable for trainers or evaluators which support label to be a dictionary. Extra items are specified by the `extra_keys` parameter and are extracted from the input dictionary (ie. the batch). This assumes label is a dictionary. Args: extra_keys: If a sequence of strings is provided, values from the input dictionary are extracted from those keys and passed to the network as extra positional arguments. z Sequence[str]None) extra_keysreturncCs4tt|dkr&tdtt|t||_dS)Nz(length of `extra_keys` should be 2, get )lenr ValueErrorr prepare_batch)selfr rW/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/apps/pathology/engines/utils.py__init__$szPrepareBatchHoVerNet.__init__NFzdict[str, torch.Tensor]zstr | torch.device | Noneboolrz7tuple[torch.Tensor, dict[HoVerNetBranch, torch.Tensor]]) batchdatadevice non_blockingkwargsr c Ks@|j|||f|\}}}}tj|tj|dtj|di} || fS)z Args `batchdata`, `device`, `non_blocking` refer to the ignite API: https://pytorch.org/ignite/v0.4.8/generated/ignite.engine.create_supervised_trainer.html. `kwargs` supports other args for `Tensor.to()` API. r)rrNPNCHV) rrrrrimage_labelZ extra_label_labelrrr__call__)s zPrepareBatchHoVerNet.__call__)NF)__name__ __module__ __qualname____doc__rr"rrrrr s  ) __future__rtypingrrtorchZ monai.enginesrr monai.utilsrmonai.utils.enumsr__all__r rrrr s