o  ij@sddlmZddlZddlmZmZddlmZmZddl m Z ddl m Z m Z ddlmZddlZddlZddlmZddlmZdd lmZmZdd lmZdd lmZdd lm Z m!Z!m"Z"m#Z#m$Z$dd l%m&Z&ddl'm(Z(ddl)m*Z*m+Z+m,Z,ddl-m.Z.ddl/m0Z0ddl1m2Z2m3Z3ddl4m5Z5m6Z6e6ddd\Z7Z8ee9dZ:GdddeZ;Gddde;ZGdd d Z?dS)!) annotationsN)ABCabstractmethod)MappingSequence)deepcopy)Anycast)warn) BundleAlgo)get_name_from_algo_idimport_bundle_algo_history) get_logger)concat_val_to_np)_prepare_cmd_bcprun_prepare_cmd_torchrun_run_cmd_bcprun_run_cmd_torchrun datafold_read) ConfigParser)partition_dataset) MeanEnsemble SaveImage VoteEnsemble) RankFilter)AlgoKeys) check_kwargs_exist_in_class_init prob2class)look_up_optionoptional_importtqdm)name) module_namec@sheZdZdZddZddZddZdd Zd%d&ddZd'ddZ d(ddZ d)d*d!d"Z e d#d$Z dS)+ AlgoEnsemblez, The base class of Ensemble methods cCsg|_d|_g|_g|_dS)Nmean)algosmode infer_files algo_ensembleselfr+g/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/apps/auto3dseg/ensemble_builder.py__init__6s zAlgoEnsemble.__init__cCst||_dS)z0 Register model in the ensemble N)rr%)r* infer_algosr+r+r, set_algos<szAlgoEnsemble.set_algoscCs&|jD] }||tjkr|SqdS)zn Get a model by identifier. Args: identifier: the name of the bundleAlgo N)r%rID)r* identifieralgor+r+r,get_algoBs zAlgoEnsemble.get_algocC|jS)z Get the algo ensemble after ranking or a empty list if ranking was not started. Returns: A list of Algo )r(r)r+r+r,get_algo_ensembleMszAlgoEnsemble.get_algo_ensembletestingdatarootstrdata_list_or_path str | listdata_keyreturnNonecCsg|_t|tr ||_dSt|tr?t|}||vr(t||d|d\|_}dSt|dr2|jdkr=t d|ddSdSt d) z Set the files to perform model inference. Args: dataroot: the path of the files data_list_or_path: the data source file path )datalistbasedirfoldkeyrankrz#Datalist file has no testing key - z$. No data for inference is specifiedzUnsupported parameter typeN) r' isinstancelistr8rload_config_filerhasattrrCloggerinfo ValueError)r*r7r9r;r?_r+r+r,set_infer_filesVs    zAlgoEnsemble.set_infer_filesFcstdd|Drdd|D}|jdkr&t|}tttj|dddS|jd krGfd d|D}r;t|St|djdd |Sd S) a ensemble the results using either "mean" or "vote" method Args: preds: a list of probability prediction in Tensor-Like format. sigmoid: use the sigmoid function to threshold probability one-hot map, otherwise argmax is used. Defaults to False Returns: a tensor which is the ensembled prediction. css|]}|j VqdSN)is_cuda.0pr+r+r, zsz-AlgoEnsemble.ensemble_pred..cSsg|]}|qSr+cpurOr+r+r, {z.AlgoEnsemble.ensemble_pred..r$rTdimkeepdimsigmoidvotecsg|] }t|dddqS)rTrW)rrOrZr+r,rUs) num_classesN) anyr&rrr torchTensorrshape)r*predsrZprobclassesr+r\r, ensemble_predms     zAlgoEnsemble.ensemble_predalgo_spec_paramdictparam algo_namecCs@t|}t|}|D]\}}||kr||q |S)a Apply the model-specific params to the prediction params based on the name of the Algo. Args: algo_spec_param: a dict that has structure of {"": ""}. param: the prediction params to override. algo_name: name of the Algo Returns: param after being updated with the model-specific param )ritemslowerupdate)r*rfrhriZ_param_to_override_paramkvr+r+r,_apply_algo_specific_params  z'AlgoEnsemble._apply_algo_specific_paramN pred_param dict | NonerEc Cs|durint|}|j}d|vr|d}d|vr#|d}||}d|vr5|d}t|ddgd|_|dd }d |vrGt|d }|d i}g} trc|rc|d d d krct t |ddnt |D]~\} } g} |j D]$} t | t j}| t j}||||}|j| g|d}| |d qpd |vrz |j| |d}Wnty|jdd| D|d}Ynw||}t|drd|jvr|jd}ntdd}n td|j| |d}| |qg| S)aG Use the ensembled model to predict result. Args: pred_param: prediction parameter dictionary. The key has two groups: the first one will be consumed in this function, and the second group will be passed to the `InferClass` to override the parameters of the class functions. The first group contains: - ``"infer_files"``: file paths to the images to read in a list. - ``"files_slices"``: a value type of `slice`. The files_slices will slice the ``"infer_files"`` and only make prediction on the infer_files[file_slices]. - ``"mode"``: ensemble mode. Currently "mean" and "vote" (majority voting) schemes are supported. - ``"image_save_func"``: a dictionary used to instantiate the ``SaveImage`` transform. When specified, the ensemble prediction will save the prediction files, instead of keeping the files in the memory. Example: `{"_target_": "SaveImage", "output_dir": "./"}` - ``"sigmoid"``: use the sigmoid function (e.g. x > 0.5) to convert the prediction probability map to the label class prediction, otherwise argmax(x) is used. - ``"algo_spec_params"``: a dictionary to add pred_params that are specific to a model. The dict has a format of {"": ""}. The parameters in the second group is defined in the ``config`` of each Algo templates. Please check: https://github.com/Project-MONAI/research-contributions/tree/main/auto3dseg/algorithm_templates Returns: A list of tensors or file paths, depending on whether ``"image_save_func"`` is set. Nr'Z files_slicesr&r$r[ supportedrZFimage_save_funcalgo_spec_paramsrCrzEnsembling (rank 0)...)desc) predict_filespredict_paramsr\cSsg|]}|dqSrS)to)rPrKr+r+r,rUz)AlgoEnsemble.__call__..metasaved_tozImage save path not returned.z\Prediction returned in list instead of disk, provide image_save_func to avoid out of memory.)rr'poprr&rget_parsed_contenthas_tqdmget enumerater r(r rr0ALGOrppredictappendre BaseExceptionrGr|keysr )r*rqrhfilesslicesr&rZZ img_saverrvoutputsrKfilerbr2Zinfer_algo_nameZinfer_instancermpredZensemble_predsresr+r+r,__call__sZ           zAlgoEnsemble.__call__cOstrM)NotImplementedError)r*argskwargsr+r+r, collect_algosszAlgoEnsemble.collect_algos)r6)r7r8r9r:r;r8r<r=)F)rfrgrhrgrir8r<rgrM)rqrrr<rE)__name__ __module__ __qualname____doc__r-r/r3r5rLrerprrrr+r+r+r,r#1s   Pr#cs8eZdZdZddfdd ZddZddd d ZZS)AlgoEnsembleBestNz Ensemble method that select N model out of all using the models' best_metric scores Args: n_best: number of models to pick for ensemble (N). n_bestintct||_dSrM)superr-r)r*r __class__r+r,r-  zAlgoEnsembleBestN.__init__cCst|jtjg}t|S)z' Sort the best_metrics )rr%rSCOREnpargsorttolist)r*scoresr+r+r, sort_scoreszAlgoEnsembleBestN.sort_scorer>r<r=csdkr|j|tkr'tdtddtdtfddtD}t|dd }t|j|_|D]}|t|jkrP|j |qAd S) zQ Rank the algos by finding the top N (n_best) validation scores. rzFound z% available algos (pre-defined n_best=z). All z will be used.cs$g|]\}}|tkr|qSr+)len)rPirrranksr+r,rUs$z3AlgoEnsembleBestN.collect_algos..T)reverseN) rrrr rsortedrr%r(r~)r*rindicesidxr+rr,rs $   zAlgoEnsembleBestN.collect_algosr)rr)r>)rrr<r=)rrrrr-rr __classcell__r+r+rr,rs rcs.eZdZdZd d fdd Zd d d ZZS)AlgoEnsembleBestByFoldz Ensemble method that select the best models that are the tops in each fold. Args: n_fold: number of cross-validation folds used in training rn_foldrcrrM)rr-r)r*rrr+r,r-#rzAlgoEnsembleBestByFold.__init__r<r=c Csg|_t|jD]J}d}d}|jD]:}|tjdd}zt|}Wnty8}z td|d|d}~ww||krK|tj |krK|}|tj }q|j |qdS)zX Rank the algos by finding the best model in each cross-validation fold gNrKzmodel identifier z is not number.) r(rangerr%rr0splitrrJrr)r*f_idx best_scoreZ best_modelr2r1Zalgo_iderrr+r+r,r's$   z$AlgoEnsembleBestByFold.collect_algosr)rrr<r=)rrrrr-rrr+r+rr,rsrc@s:eZdZdZddddZddddZdddZddZdS) AlgoEnsembleBuildera Build ensemble workflow from configs and arguments. Args: history: a collection of trained bundleAlgo algorithms. data_src_cfg_name: filename of the data source. Examples: .. code-block:: python builder = AlgoEnsembleBuilder(history, data_src_cfg) builder.set_ensemble_method(BundleAlgoEnsembleBestN(3)) ensemble = builder.get_ensemble() NhistorySequence[dict[str, Any]]data_src_cfg_name str | Nonec Csg|_|tdd|_|durtjt|r|j||D]=}|tj }|tj }| }|j }tj |dd}tj|sHt|j dtj|sUt|d||||qdS)NF)globalsscriptszinfer.pyz+ is not a directory. Please check the path.z% is not found. Please check the path.)r.r data_src_cfgospathexistsr8 read_configrr0r get_score output_pathjoinisdirr isfile add_inferer) r*rr algo_dictr!gen_algo best_metric algo_pathZ infer_pathr+r+r,r-Os"      zAlgoEnsembleBuilder.__init__r1r8rr r float | Noner<r=cCs6|durtdtj|tj|tj|i}|j|dS)z Add model inferer to the builder. Args: identifier: name of the bundleAlgo. gen_algo: a trained BundleAlgo model object. best_metric: the best metric in validation of the trained model. Nz+Feature to re-validate is to be implemented)rJrr0rrr.r)r*r1rrr2r+r+r,ris zAlgoEnsembleBuilder.add_infererensembler#rrrcOs>||j|j|i|||jd|jd||_dS)zj Set the ensemble method. Args: ensemble: the AlgoEnsemble to build. r7r?N)r/r.rrLrr)r*rrrr+r+r,set_ensemble_methodys  z'AlgoEnsembleBuilder.set_ensemble_methodcCr4)zGet the ensemble)rr)r+r+r, get_ensemblesz AlgoEnsembleBuilder.get_ensemblerM)rrrr)r1r8rr rrr<r=)rr#rrrrr<r=)rrrrr-rrrr+r+r+r,r=s     rc@sleZdZdZ    d%d&ddZd'd(ddZddZd)ddZd*d+ddZddZ d,d-d!d"Z d.d#d$Z dS)/EnsembleRunnera The Runner for ensembler. It ensembles predictions and saves them to the disk with a support of using multi-GPU. Args: data_src_cfg_name: filename of the data source. work_dir: working directory to save the intermediate and final results. Default is `./work_dir`. num_fold: number of fold. Default is 5. ensemble_method_name: method to ensemble predictions from different model. Default is AlgoEnsembleBestByFold. Supported methods: ["AlgoEnsembleBestN", "AlgoEnsembleBestByFold"]. mgpu: if using multi-gpu. Default is True. kwargs: additional image writing, ensembling parameters and prediction parameters for the ensemble inference. - for image saving, please check the supported parameters in SaveImage transform. - for prediction parameters, please check the supported parameters in the ``AlgoEnsemble`` callables. - for ensemble parameters, please check the documentation of the selected AlgoEnsemble callable. Example: .. code-block:: python ensemble_runner = EnsembleRunner(data_src_cfg_name, work_dir, ensemble_method_name, mgpu=device_setting['n_devices']>1, **kwargs, **pred_params) ensemble_runner.run(device_setting) ./work_dirrrTrr8work_dirnum_foldrensemble_method_namemgpuboolrrr<r=cKs||_||_||_||_||_t||_d|_d|_d ddt t j Dt j ttjddtjddtjd d d |_dS) Nrr,cSsg|]}t|qSr+)r8)rPxr+r+r,rUrVz+EnsembleRunner.__init__.. NUM_NODESMN_START_METHODbcprun CMD_PREFIX)CUDA_VISIBLE_DEVICES n_devicesrrr)rrrrrrrrC world_sizerrr_cuda device_countrrenvironrdevice_setting)r*rrrrrrr+r+r,r-s     zEnsembleRunner.__init__cKsft|ddgd|_|jdkr|dd}t|d|_d S|jdkr*t|jd|_d Std|jd ) a Set the bundle ensemble method Args: ensemble_method_name: the name of the ensemble method. Only two methods are supported "AlgoEnsembleBestN" and "AlgoEnsembleBestByFold". kwargs: the keyword arguments used to define the ensemble method. Currently only ``n_best`` for ``AlgoEnsembleBestN`` is supported. rrrsr)r)rzEnsemble method z is not implemented.N)rrr~rensemble_methodrrr)r*rrrr+r+r,rs    z"EnsembleRunner.set_ensemble_methodc Ks|dd}|durtj|jd}td|dtj|s1tj|ddtd|d t |j }| d d }d ||d d|dd|dddd|d||ddd }t t|\}}|rm|||St|D]}||vr||||iqq|S)a Pop the kwargs used to define ImageSave class for the ensemble output. Args: kwargs: image writing parameters for the ensemble inference. The kwargs format follows SaveImage transform. For more information, check https://docs.monai.io/en/stable/transforms.html#saveimage . Returns: save_image: a dictionary that can be used to instantiate a SaveImage class in ConfigParser. output_dirNZensemble_outputz!The output_dir is not specified. z+ will be used to save ensemble predictions.T)exist_okz Directory z( is created to save ensemble predictionsr7rroutput_postfixr output_dtypez $np.uint8resampleF data_root_dirseparate_folder) _target_rrrr print_logsavepath_in_metadictrr)r~rrrrrHrIrmakedirsrrFrrrrrlrE) r*rrZ input_yamlr save_imageZare_all_args_save_image extra_argsrr+r+r,'_pop_kwargs_to_get_image_save_transforms8            z6EnsembleRunner._pop_kwargs_to_get_image_save_transformcKs0tt|\}}|r|j|dSt|d)a$ Set the ensemble output transform. Args: kwargs: image writing parameters for the ensemble inference. The kwargs format follows SaveImage transform. For more information, check https://docs.monai.io/en/stable/transforms.html#saveimage . z are not supported in monai.transforms.SaveImage,Check https://docs.monai.io/en/stable/transforms.html#saveimage for more information.N)rrrrlrJ)r*rare_all_args_presentrr+r+r,set_image_save_transform s  z'EnsembleRunner.set_image_save_transformcCs |dkr td|||_dS)z Set the number of cross validation folds for all algos. Args: num_fold: a positive integer to define the number of folds. rzEnum_fold is expected to be an integer greater than zero. Now it gets N)rJr)r*rr+r+r, set_num_folds zEnsembleRunner.set_num_foldcCs|jrtjdddt|_t|_tt |j |j d|j |j fi|j|jdi|j}t|jdd}dd|D}|rZtd d d|Dd d d|D}t|d kritd|jdt||j}| |j||_|jj}t||jkrt|d krtddS|jt|kr||jgng}n t|d|jdd|j}||j_|j|jd<||jd<td|jD] }t|tjq|d}td|d|j|jd|jrt dSdS)Nncclzenv://)backend init_method)rF) only_trainedcSsg|] }|tjs|qSr+r IS_TRAINEDrPhr+r+r,rU6z+EnsembleRunner.ensemble..zEnsembling step will skip cSsg|]}|tjqSr+)rr0rr+r+r,rU9r{zJ untrained algos.Generally it means these algos did not complete training.cSsg|] }|tjr|qSr+rrr+r+r,rU<rrz&Could not find the trained results in z8. Possibly the required training step was not completed.z=No testing files for inference is provided. Ensembler ending.)datashufflenum_partitionseven_divisiblerCruz4Auto3Dseg picked the following networks to ensemble:rz7Auto3Dseg ensemble prediction outputs will be saved in .)rqr+)!rdistinit_process_groupget_world_sizerget_rankrCrH addFilterrrrrrrrr rwarningrrJrrrrZ ensemblerr'rIrr5rr0destroy_process_group)r*rrZhistory_untrainedbuilderr'r2rr+r+r,r(sZ                zEnsembleRunner.ensembleNrrrcCs>|dur|j|tt|jdd|jd<|dS)a Load the run function in the training script of each model. Training parameter is predefined by the algo_config.yaml file, which is pre-filled by the fill_template_config function in the same instance. Args: device_setting: device related settings, should follow the device_setting in auto_runner.set_device_info. 'CUDA_VISIBLE_DEVICES' should be a string e.g. '0,1,2,3' Nrrr)rrlrr8r _create_cmd)r*rr+r+r,run`s  zEnsembleRunner.runc Csxt|jddkrt|jddkrtd|dSd|jd|jd|jd|jd }|j rMt |j t rM|j D]\}}|d |d |7}q>t j}t|jd |d <t|jddkr|jd dkrut|jd dtd|jddtd||jdd}t||jd|jdddStd|jddtd|}t|d|jd|dddS)NrrrzEnsembling using single GPU!zQmonai.apps.auto3dseg EnsembleRunner ensemble --data_src_cfg_name z --work_dir z --num_fold z( --ensemble_method_name z --mgpu Truez --=rrrzN is not supported yet. Try modify EnsembleRunner._create_cmd for your cluster.zEnsembling on z nodes!z-m r) cmd_prefix)nrQzEnsembling using z GPU!T)nnodesnproc_per_nodeenvcheck)rrrHrIrrrrrrrDrrjrrcopyr8rrrrr)r*Zbase_cmdrnro ps_environcmdr+r+r,rosB$   zEnsembleRunner._create_cmd)rrrT)rr8rr8rrrr8rrrrr<r=)r)rr8rrr<r=)rrr<r=r)rrr<r=rM)rrrr<r=r) rrrrr-rrrrrrrr+r+r+r,rs   0  8r)@ __future__rrabcrrcollections.abcrrrrtypingrr warningsr numpyrr_torch.distributed distributedr monai.apps.auto3dseg.bundle_genr monai.apps.auto3dseg.utilsr r monai.apps.utilsrmonai.auto3dsegrmonai.auto3dseg.utilsrrrrr monai.bundler monai.datarmonai.transformsrrr monai.utilsrmonai.utils.enumsrmonai.utils.miscrrmonai.utils.modulerrr rrrHr#rrrrr+r+r+r,s<            ?,"P