o  iI@sddlmZddlmZddlmZddlmZddlm Z ddl m Z m Z m Z e de je d\ZZer;dd lmZn e de je d \ZZGd d d Zd S)) annotations) TYPE_CHECKING)KeysCollection)IterationEvents) Decollated) IgniteInfo min_versionoptional_importz ignite.engineEvents)Enginer c@s>eZdZdZ       ddddZdddZdddZdS)DecollateBatcha Ignite handler to execute the `decollate batch` logic for `engine.state.batch` and `engine.state.output`. Typical usage is to set `decollate=False` in the engine and execute some postprocessing logic first then decollate the batch, otherwise, engine will decollate batch before the postprocessing. Args: event: expected EVENT to attach the handler, should be "MODEL_COMPLETED" or "ITERATION_COMPLETED". default to "MODEL_COMPLETED". detach: whether to detach the tensors. scalars tensors will be detached into number types instead of torch tensors. decollate_batch: whether to decollate `engine.state.batch` of ignite engine. batch_keys: if `decollate_batch=True`, specify the keys of the corresponding items to decollate in `engine.state.batch`, note that it will delete other keys not specified. if None, will decollate all the keys. it replicates the scalar values to every item of the decollated list. decollate_output: whether to decollate `engine.state.output` of ignite engine. output_keys: if `decollate_output=True`, specify the keys of the corresponding items to decollate in `engine.state.output`, note that it will delete other keys not specified. if None, will decollate all the keys. it replicates the scalar values to every item of the decollated list. allow_missing_keys: don't raise exception if key is missing. MODEL_COMPLETEDTNFeventstrdetachbooldecollate_batch batch_keysKeysCollection | Nonedecollate_output output_keysallow_missing_keyscCsX|}|dvr td||_|rt|||dnd|_|r't|||d|_dSd|_dS)N)r ITERATION_COMPLETEDz;event should be `MODEL_COMPLETED` or `ITERATION_COMPLETED`.)keysrr)upper ValueErrorrrbatch_transformoutput_transform)selfrrrrrrrr`/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/handlers/decollate_batch.py__init__3s zDecollateBatch.__init__enginer returnNonecCs.|jdkr|tj|dS|tj|dS)g Args: engine: Ignite Engine, it can be a trainer, validator or evaluator. r N)radd_event_handlerrr r rrr"rrr attachNs zDecollateBatch.attachcCsh|jdurt|jjttfr||jj|j_|jdur0t|jjttfr2||jj|j_dSdSdS)r%N)r isinstancestatebatchlistdictroutputr'rrr __call__Xs zDecollateBatch.__call__)r TTNTNF)rrrrrrrrrrrrrr)r"r r#r$)__name__ __module__ __qualname____doc__r!r(r/rrrr r s  r N) __future__rtypingr monai.configrmonai.engines.utilsrmonai.transformsr monai.utilsrrr OPT_IMPORT_VERSIONr _ ignite.enginer r rrrr s