U {PhI@sddlmZddlmZddlmZmZddlmZddl m Z ddl m Z m Z e deje d\ZZervdd lmZne deje d \ZZGd d d Zd S)) annotations) TYPE_CHECKING) IgniteInfoKeysCollection)IterationEvents) Decollated) min_versionoptional_importz ignite.engineEvents)Enginer c @sLeZdZdZddddddddd d d Zd d dddZd d dddZdS)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_COMPLETEDTNFstrboolzKeysCollection | None)eventdetachdecollate_batch batch_keysdecollate_output output_keysallow_missing_keyscCsR|}|dkrtd||_|r0t|||dnd|_|rHt|||dnd|_dS)N)r ITERATION_COMPLETEDz;event should be `MODEL_COMPLETED` or `ITERATION_COMPLETED`.)keysrr)upper ValueErrorrrbatch_transformoutput_transform)selfrrrrrrrrS/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/handlers/decollate_batch.py__init__3s zDecollateBatch.__init__r None)enginereturncCs,|jdkr|tj|n|tj|dS)g Args: engine: Ignite Engine, it can be a trainer, validator or evaluator. r N)radd_event_handlerrr r rrr"rrrattachNs zDecollateBatch.attachcCs`|jdk r.t|jjttfr.||jj|j_|jdk r\t|jjttfr\||jj|j_dS)r$N)r isinstancestatebatchlistdictroutputr&rrr__call__XszDecollateBatch.__call__)r TTNTNF)__name__ __module__ __qualname____doc__r r'r.rrrrr s r N) __future__rtypingr monai.configrrmonai.engines.utilsrmonai.transformsr monai.utilsrr OPT_IMPORT_VERSIONr _ ignite.enginer r rrrr s