U Phj@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@seZdZdZddZddZddZdd Zd d d d d dddZd!ddZ ddd ddddZ d"dddddZ e ddZ dS)# AlgoEnsemblez, The base class of Ensemble methods cCsg|_d|_g|_g|_dS)Nmean)algosmode infer_files algo_ensembleselfr+Z/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/apps/auto3dseg/ensemble_builder.py__init__6szAlgoEnsemble.__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_algocCs|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_ensembletestingstrz str | listNone)datarootdata_list_or_pathdata_keyreturncCsg|_t|tr||_ndt|trtt|}||krLt||d|d\|_}q|t|dr`|jdkr|t d|dnt ddS) 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' isinstancelistr6rload_config_filerhasattrrAloggerinfo ValueError)r*r8r9r:r=_r+r+r,set_infer_filesVs    zAlgoEnsemble.set_infer_filesFcstdd|Dr dd|D}|jdkrLt|}tttj|dddS|jd krfd d|D}rvt|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+cpurMr+r+r, {sz.AlgoEnsemble.ensemble_pred..r$rTdimkeepdimsigmoidvotecsg|]}t|dddqS)rTrT)rrMrWr+r,rSs) num_classesN) anyr&rrr torchTensorrshape)r*predsrWprobclassesr+rYr, ensemble_predms     zAlgoEnsemble.ensemble_preddict)algo_spec_paramparam algo_namer;cCs@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*rdrerfZ_param_to_override_paramkvr+r+r,_apply_algo_specific_params  z'AlgoEnsemble._apply_algo_specific_paramN dict | NonerC) pred_paramr;c Cs|dkr int|}|j}d|kr,|d}d|krF|d}||}d|krj|d}t|ddgd|_|dd }d |krt|d }|d i}g} tr|r|d d d krt t |ddnt |D]\} } g} |j D]H} t | t j}| t j}||||}|j| g|d}| |d qd |krz|j| |d}Wn.tk rv|jdd| D|d}YnX||}t|drd|jkr|jd}n tdd}ntd|j| |d}| |q| 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$rX supportedrWFimage_save_funcalgo_spec_paramsrArzEnsembling (rank 0)...)desc) predict_filespredict_paramsrYcSsg|]}|dqSrQ)to)rNrIr+r+r,rSsz)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 rr0ALGOrmpredictappendrb BaseExceptionrErxkeysr )r*rorefilesslicesr&rWZ img_saverrsoutputsrIfiler_r2Zinfer_algo_nameZinfer_instancerjpredZensemble_predsresr+r+r,__call__sX           zAlgoEnsemble.__call__cOstdSrK)NotImplementedError)r*argskwargsr+r+r, collect_algosszAlgoEnsemble.collect_algos)r5)F)N)__name__ __module__ __qualname____doc__r-r/r3r4rJrbrmrrrr+r+r+r,r#1s   Pr#csBeZdZdZdddfdd ZddZddd d d 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). intn_bestcst||_dSrK)superr-r)r*r __class__r+r,r-s zAlgoEnsembleBestN.__init__cCst|jtjg}t|S)z' Sort the best_metrics )rr%rSCOREnpargsorttolist)r*scoresr+r+r, sort_scoreszAlgoEnsembleBestN.sort_scorer<r7)rr;csdkr|j|tkrNtdtddtdtfddtD}t|dd }t|j|_|D]}|t|jkr|j |qd 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)rNirrranksr+r,rSsz3AlgoEnsembleBestN.collect_algos..T)reverseN) rrrr r~sortedrr%r(rz)r*rindicesidxr+rr,rs $  zAlgoEnsembleBestN.collect_algos)r)r<)rrrrr-rr __classcell__r+r+rr,rsrcs6eZdZdZd ddfdd Zddd 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 rrn_foldcst||_dSrK)rr-r)r*rrr+r,r-#s zAlgoEnsembleBestByFold.__init__r7r;c Csg|_t|jD]}d}d}|jD]~}|tjdd}z t|}Wn4tk rz}ztd|d|W5d}~XYnX||kr"|tj |kr"|}|tj }q"|j |qdS)zX Rank the algos by finding the best model in each cross-validation fold gNrIzmodel identifier z is not number.) r(rangerr%rr0splitrrHrr)r*f_idx best_scoreZ best_modelr2r1Zalgo_iderrr+r+r,r's  $ z$AlgoEnsembleBestByFold.collect_algos)r)rrrrr-rrr+r+rr,rsrc@sTeZdZdZddddddZddd d d d d dZdddd dddZddZdS)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() NzSequence[dict[str, Any]]z str | None)historydata_src_cfg_namec Csg|_|tdd|_|dk r:tjt|r:|j||D]z}|tj }|tj }| }|j }tj |dd}tj|st|j dtj|st|d||||q>dS)NF)globalsscriptszinfer.pyz+ is not a directory. Please check the path.z% is not found. Please check the path.)r.r data_src_cfgospathexistsr6 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__r6r z float | Noner7)r1rrr;cCs6|dkrtdtj|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)rHrr0rrr.r)r*r1rrr2r+r+r,ris zAlgoEnsembleBuilder.add_infererr#r)ensemblerrr;cOs:||j|j||||jd|jd||_dS)zj Set the ensemble method. Args: ensemble: the AlgoEnsemble to build. r8r=N)r/r.rrJrr)r*rrrr+r+r,set_ensemble_methodys  z'AlgoEnsembleBuilder.set_ensemble_methodcCs|jS)zGet the ensemble)rr)r+r+r, get_ensemblesz AlgoEnsembleBuilder.get_ensemble)N)N)rrrrr-rrrr+r+r+r,r=s rc @seZdZdZd#dddddd d d d d Zd$dd d dddZddZd d dddZd%dd dddZddZ d&dd dddZ d 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_dirrrTr6rboolrr7)rwork_dirnum_foldensemble_method_namemgpurr;cKs||_||_||_||_||_t||_d|_d|_d ddt t j Dt j ttjddtjddtjd d d |_dS) Nrr,cSsg|] }t|qSr+)r6)rNxr+r+r,rSsz+EnsembleRunner.__init__.. NUM_NODESMN_START_METHODbcprun CMD_PREFIX)CUDA_VISIBLE_DEVICES n_devicesrrr)rrrrrrrrA world_sizerrr\cuda device_countrrenvironr}device_setting)r*rrrrrrr+r+r,r-s    zEnsembleRunner.__init__)rrr;cKsft|ddgd|_|jdkr6|dd}t|d|_n,|jdkrPt|jd|_ntd|jd d S) 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. rrrprrrzEnsemble method z is not implemented.N)rrrzrensemble_methodrrr)r*rrrr+r+r,rs    z"EnsembleRunner.set_ensemble_methodc Ks|dd}|dkr6tj|jd}td|dtj|sbtj|ddtd|d t |j }| d d }d ||d d|dd|dddd|d||ddd }t t|\}}|r||n*t|D] }||kr||||iq|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 predictionsr8rroutput_postfixr output_dtypez $np.uint8resampleF data_root_dirseparate_folder) _target_rrrr print_logsavepath_in_metadictrr)rzrrrrrFrGrmakedirsrrDrr}rrrirC) r*rrZ input_yamlr save_imageZare_all_args_save_image extra_argsrr+r+r,'_pop_kwargs_to_get_image_save_transforms4            z6EnsembleRunner._pop_kwargs_to_get_image_save_transform)rr;cKs2tt|\}}|r |j|nt|ddS)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)rrrrirH)r*rare_all_args_presentrr+r+r,set_image_save_transform s z'EnsembleRunner.set_image_save_transform)rr;cCs |dkrtd|||_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)rHr)r*rr+r+r, set_num_foldszEnsembleRunner.set_num_foldcCs|jr4tjdddt|_t|_tt |j |j d|j |j f|j|jf|j}t|jdd}dd|D}|rtd d d|Dd d d|D}t|d krtd|jdt||j}| |j||_|jj}t||jkrBt|d kr tddS|jt|kr<||jgng}nt|d|jdd|j}||j_|j|jd<||jd<td|jD]}t|tjq|d}td|d|j|jd|jrt dS)Nncclzenv://)backend init_method)rF) only_trainedcSsg|]}|tjs|qSr+r IS_TRAINEDrNhr+r+r,rS6s z+EnsembleRunner.ensemble..zEnsembling step will skip cSsg|]}|tjqSr+)rr0rr+r+r,rS9szJ untrained algos.Generally it means these algos did not complete training.cSsg|]}|tjr|qSr+rrr+r+r,rS<s rz&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_divisiblerArrz4Auto3Dseg picked the following networks to ensemble:rz7Auto3Dseg ensemble prediction outputs will be saved in .)ro)!rdistinit_process_groupget_world_sizerget_rankrArF addFilterrrrrrrrr rwarningrrHrrrrZ ensemblerr'rGrr4rr0destroy_process_group)r*rrZhistory_untrainedbuilderr'r2rr+r+r,r(s^         "   zEnsembleRunner.ensembleNrn)rr;cCs>|dk r2|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)rrirr6r _create_cmd)r*rr+r+r,run`s  zEnsembleRunner.runrc Csxt|jddkr:t|jddkr:td|dSd|jd|jd|jd|jd }|j rt |j t r|j D]\}}|d |d |7}q|t j}t|jd |d <t|jddkr8|jd dkrt|jd dtd|jddtd||jdd}t||jd|jddn s:            ?,"P