U |Ph@sddlmZddlZddlZddlmZddlZddlm Z m Z ddl m Z ddl mZmZmZddlmZede jed \ZZerdd lmZnede jed \ZZGd d d ZdS)) annotationsN) TYPE_CHECKING) DtypeLike IgniteInfo) FolderLayout) ProbMapKeys min_versionoptional_import) CommonKeysz ignite.engineEvents)Enginer c@sveZdZdZdddejdfdddddd d d d Zd d dddZd d dddZdd dddZ d d dddZ dS)ProbMapProducera; Event handler triggered on completing every iteration to calculate and save the probability map. This handler use metadata from MetaTensor to create the probability map. This can be simply achieved by using `monai.data.SlidingPatchWSIDataset` or `monai.data.MaskedPatchWSIDataset` as the dataset. z./predNstrrz str | NoneNone) output_diroutput_postfixprob_keydtypenamereturncCsZt||ddddd|_t||_||_||_||_i|_i|_ d|_ d|_ t |_dS)a Args: output_dir: output directory to save probability maps. output_postfix: a string appended to all output file names. prob_key: the key associated to the probability output of the model dtype: the data type in which the probability map is stored. Default np.float64. name: identifier of logging.logger to use, defaulting to `engine.logger`. z.npyFTr)rpostfix extensionparentmakedirs data_root_dirrN)r folder_layoutlogging getLoggerlogger_namerrprob_mapcounternum_done_images num_images threadingLocklock)selfrrrrrr*T/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/handlers/probability_maps.py__init__(s" zProbMapProducer.__init__r )enginercCs|jjj}t||_|D]:}|tj}|tj|j|<t j |tj |j d|j |<q|jdkrf|j|_||tjs|tj|||jtjs|tj|jdS)zg Args: engine: Ignite Engine, it can be a trainer, validator or evaluator. )rN) data_loaderdataset image_datalenr%rNAMECOUNTr#npzerosSIZErr"r!r has_event_handlerr ITERATION_COMPLETEDadd_event_handlerfinalize COMPLETED)r)r-r0samplerr*r*r+attachLs    zProbMapProducer.attachc Cst|jjtrt|jjts$td|jjtjjt j }|jjtjjt j }|jj|j }t |||D]X\}}}||j|t|<|j0|j|d8<|j|dkr||W5QRXqjdS)z This method assumes self.batch_transform will extract metadata from the input batch. Args: engine: Ignite Engine, it can be a trainer, validator or evaluator. z@engine.state.batch and engine.state.output must be dictionaries.rN) isinstancestatebatchdictoutput ValueErrorr IMAGEmetarr2LOCATIONrzipr"tupler(r# save_prob_map)r)r-nameslocsprobsrlocprobr*r*r+__call__bszProbMapProducer.__call__)rrc Csd|j|}t||j||jd7_|jd|d|jd|jd|j|=|j |=dS)z This method save the probability map for an image, when its inference is finished, and delete that probability map from memory. Args: name: the name of image to be saved. r>zInference of 'z ' is done [/z]!N) rfilenamer4saver"r$r infor%r#)r)r file_pathr*r*r+rJus  $zProbMapProducer.save_prob_mapcCs"|jd|jd|jddS)NzProbability map is created for rQz images!)r rTr$r%)r)r-r*r*r+r:szProbMapProducer.finalize) __name__ __module__ __qualname____doc__r4float64r,r=rPrJr:r*r*r*r+r s $r ) __future__rrr&typingrnumpyr4 monai.configrrmonai.data.folder_layoutr monai.utilsrrr monai.utils.enumsr OPT_IMPORT_VERSIONr _ ignite.enginer r r*r*r*r+ s