o  ix @snddlmZddlmZddlmZddlZddlmZm Z ddl m Z ddl m Z dgZGd ddeZdS) ) annotations)Sequence)AnyN) PrepareBatchPrepareBatchExtraInput) ensure_tuple)HoVerNetBranchPrepareBatchHoVerNetc@s*eZdZdZdddZ  ddddZdS)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. extra_keys Sequence[str]returnNonecCs4tt|dkrtdtt|t||_dS)Nz(length of `extra_keys` should be 2, get )lenr ValueErrorr prepare_batch)selfr rd/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/apps/pathology/engines/utils.py__init__%szPrepareBatchHoVerNet.__init__NF batchdatadict[str, torch.Tensor]devicestr | torch.device | None non_blockingboolkwargsr7tuple[torch.Tensor, dict[HoVerNetBranch, torch.Tensor]]c KsD|j|||fi|\}}}}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__)r r r r )NF) rrrrrrrrr r)__name__ __module__ __qualname____doc__rr&rrrrr s  ) __future__rcollections.abcrtypingrtorchZ monai.enginesrr monai.utilsrmonai.utils.enumsr__all__r rrrrs