o
    i                     @  s   d dl mZ d dlmZ d dlmZ d dlmZmZ d dl	m
Z
mZmZ ede
jed\ZZer7d dlmZ n
ede
jed	\ZZG d
d dZdS )    )annotations)Callable)TYPE_CHECKING)IterationEventsengine_apply_transform)
IgniteInfomin_versionoptional_importzignite.engineEvents)Enginer   c                   @  s0   e Zd ZdZddd	d
ZdddZdddZdS )PostProcessinga  
    Ignite handler to execute additional post processing after the post processing in engines.
    So users can insert other handlers between engine postprocessing and this post processing handler.
    If using components from `monai.transforms` as the `transform`, recommend to decollate `engine.state.batch`
    and `engine.state.batch` in the engine(set `decollate=True`) or in the `DecollateBatch` handler first.

    MODEL_COMPLETED	transformr   eventstrreturnNonec                 C  s(   || _ | }|dvrtd|| _dS )aQ  
        Args:
            transform: callable function to execute on the `engine.state.batch` and `engine.state.output`.
                can also be composed transforms.
            event: expected EVENT to attach the handler, should be "MODEL_COMPLETED" or "ITERATION_COMPLETED".
                default to "MODEL_COMPLETED".

        )r   ITERATION_COMPLETEDz;event should be `MODEL_COMPLETED` or `ITERATION_COMPLETED`.N)r   upper
ValueErrorr   )selfr   r    r   _/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/handlers/postprocessing.py__init__$   s
   	
zPostProcessing.__init__enginer   c                 C  s.   | j dkr|tj|  dS |tj|  dS )g
        Args:
            engine: Ignite Engine, it can be a trainer, validator or evaluator.
        r   N)r   add_event_handlerr   r   r
   r   )r   r   r   r   r   attach3   s   
zPostProcessing.attachc                 C  s   t |jjtrt |jjts"t|jj|jj| jd\|j_|j_dS tt|jj|jjD ]\}\}}t||| j\|jj|< |jj|< q-dS )r   )batchoutputr   N)	
isinstancestater   listr   r   r   	enumeratezip)r   r   ibor   r   r   __call__=   s   "$zPostProcessing.__call__N)r   )r   r   r   r   r   r   )r   r   r   r   )__name__
__module____qualname____doc__r   r   r(   r   r   r   r   r      s
    

r   N)
__future__r   collections.abcr   typingr   monai.engines.utilsr   r   monai.utilsr   r   r	   OPT_IMPORT_VERSIONr
   _ignite.enginer   r   r   r   r   r   <module>   s   