o  i@sddlmZddlmZmZddlZddlZddlm Z m Z ddl m Z m Z ddlmZddlmZddlmZGd d d ZdS) ) annotations)CallableSequenceN)decollate_batchlist_data_collate)SupervisedEvaluatorSupervisedTrainer)IterationEvents)Compose) CommonKeysc@s,eZdZdZ   ddddZdddZdS) InteractionaC Ignite process_function used to introduce interactions (simulation of clicks) for DeepEdit Training/Evaluation. More details about this can be found at: Diaz-Pinto et al., MONAI Label: A framework for AI-assisted Interactive Labeling of 3D Medical Images. (2022) https://arxiv.org/abs/2203.12362 Args: deepgrow_probability: probability of simulating clicks in an iteration transforms: execute additional transformation during every iteration (before train). Typically, several Tensor based transforms composed by `Compose`. train: True for training mode or False for evaluation mode click_probability_key: key to click/interaction probability label_names: Dict of label names max_interactions: maximum number of interactions per iteration N probabilitydeepgrow_probabilityfloat transformsSequence[Callable] | Callabletrainbool label_namesNone | dict[str, int]click_probability_keystrmax_interactionsintreturnNonecCs:||_t|ts t|n||_||_||_||_||_dS)N)r isinstancer rrrrr)selfrrrrrrra/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/apps/deepedit/interaction.py__init__-s   zInteraction.__init__engine'SupervisedTrainer | SupervisedEvaluator batchdatadict[str, torch.Tensor]dictc Cs|durtdtjjddg|jd|jgdrt|jD]}||\}}||j j }| t j |jt/|jr]td|||j}Wdn1sWwYn|||j}Wdn1snwY|tj|it|dd}tt|D]}|jrdd|j|nd|||j<|||||<qt|}| t jqn&t|dd}tdt|d tjD]}|d tj|d 9<qt|}||j _ |!||S) Nz.Must provide batch data for current iteration.TFr)pcuda)detachg?r)" ValueErrornprandomchoicerranger prepare_batchtostatedevice fire_eventr INNER_ITERATION_STARTEDnetworkevaltorchno_gradampautocastinfererupdater PREDrlenrrrrINNER_ITERATION_COMPLETEDIMAGEbatch _iteration) rr"r$jinputs_ predictionsZbatchdata_listirrr __call__=s@         zInteraction.__call__)Nr r)rrrrrrrrrrrrrr)r"r#r$r%rr&)__name__ __module__ __qualname____doc__r!rHrrrr r s r ) __future__rcollections.abcrrnumpyr+r7 monai.datarr monai.enginesrrmonai.engines.utilsr monai.transformsr monai.utils.enumsr r rrrr s