o  ia@sddlmZddlZddlZddlmZddlmZmZm Z ddl Z ddl Z ddl mZddlmZddlmZedd d \ZZedd d \ZZgd Z       d>d?d#d$ZGd%d&d&e jjZd@dAd+d,ZdBdCd2d3Z  4dDdEd:d;ZdBdCd)r?r@rArBrCr)rprr&r)__name__ __module__ __qualname____doc__rTr __classcell__r<r<rnr=rsFrr>rArBrCc Cs:ddlm}|dddtdddddd}t|||}|S)a Initializes and returns a `nnUNetMONAIModelWrapper` containing the corresponding `nnUNetPredictor`. The model folder should contain the following files, created during training: - dataset.json: from the nnUNet results folder - plans.json: from the nnUNet results folder - nnunet_checkpoint.pth: The nnUNet checkpoint file, containing the nnUNet training configuration - model.pt: The checkpoint file containing the model weights. The returned wrapper object can be used for inference with MONAI framework: Example:: from monai.bundle.nnunet import get_nnunet_monai_predictor model_folder = 'path/to/monai_bundle/model' model_name = 'model.pt' wrapper = get_nnunet_monai_predictor(model_folder, model_name) # Perform inference input_data = ... output = wrapper(input_data) Parameters ---------- model_folder : Union[str, Path] The folder where the model is stored. model_name : str, optional The name of the model file, by default "model.pt". Returns ------- ModelnnUNetWrapper A wrapper object that contains the nnUNetPredictor and the loaded model. r)nnUNetPredictorg?TFr)Ztile_step_sizeZ use_gaussianZ use_mirroringr#verboseZverbose_preprocessingZ allow_tqdm)Z(nnunetv2.inference.predict_from_raw_datarr1r#r)rArCrr?wrapperr<r<r=r)s &  r nnunet_configdictbundle_root_folderr.NonecCsd}d}d}d|vr|d}d|vr|d}d|vr|d}ddlm}||d }ttjd ||d |d |}tjt|d |d dd} tjt|d |ddd} i} | d| d<| d| d<| d| d<t| t|ddt|dd |j dddi} | d| d<t| t|dd |di} | d| d<t| t|dd |dtj tj |ddst t|dt|ddtj tj |ddst t|dt|dddSdS)a Convert nnUNet model checkpoints and configuration to MONAI bundle format. Parameters ---------- nnunet_config : dict Configuration dictionary for nnUNet, containing keys such as 'dataset_name_or_id', 'nnunet_configuration', 'nnunet_trainer', and 'nnunet_plans'. bundle_root_folder : str Root folder where the MONAI bundle will be saved. fold : int, optional Fold number of the nnUNet model to be converted, by default 0. Returns ------- None rrZ 3d_fullresr: nnunet_plansnnunet_configurationrmaybe_convert_to_dataset_namernnUNet_results__fold_checkpoint_final.pthTr*checkpoint_best.pthrMrLrKmodelsrHparentsexist_okr,r>z best_model.ptrGrFN)-nnunetv2.utilities.dataset_name_id_conversionrrosenvironrXr1r5savemkdirpathexistsr shutilcopy)rrrr:rrr dataset_namennunet_model_folderZnnunet_checkpoint_finalZnnunet_checkpoint_bestnnunet_checkpointmonai_last_checkpointmonai_best_checkpointr<r<r=r_sT           rmodel plans_file dataset_file model_ckptmodel_key_in_ckptUnion[torch.nn.Module, Any]c Csddlm}ddlm}ddlm}ddlm}||} ||} || } | |} || | | } | | }d}|| j | j | j | |j d|d}|durK|Stj|dd }||||S) a- Load and initialize a nnUNet network based on nnUNet plans and configuration. Parameters ---------- plans_file : str Path to the JSON file containing the nnUNet plans. dataset_file : str Path to the JSON file containing the dataset information. configuration : str The configuration name to be used from the plans. model_ckpt : Optional[str], optional Path to the model checkpoint file. If None, the network is returned without loading weights (default is None). model_key_in_ckpt : str, optional The key in the checkpoint file that contains the model state dictionary (default is "model"). Returns ------- network : torch.nn.Module The initialized neural network, with weights loaded if `model_ckpt` is provided. r)r )get_network_from_plansrOrDT)Z allow_initdeep_supervisionNr*)Z4batchgenerators.utilities.file_and_folder_operationsr Z)nnunetv2.utilities.get_network_from_plansrr]rPrUrEr[rcr`rarbrdr1r5r8)rrrrrr rrPrErjrgrerfrmrhrRr6r;r<r<r=rs2        rcCsRddlm}d}d}d|vr|d}d|vr|d}ddlm}ddlm} d9d:dd}ttjd||d|d|d} ttjd||d} t| d|j d d dt j |dd d} |t|dd|d d d!} g} | D]}| t |td"td# q}| | d$}t j |d%|d&|d#d d}|t|dd|d'd d!}g}|D]}| t|td(td# q||d$}t j |d%|d)|d#d d}|d*| d*<|| d+<|d+D] }|d+|| d+|<q|| d,<|| d-<d| d.<d | d/<t | t| d|d0|| d+<|d*| d*<|d+D] }|d+|| d+|<q4t | t| d|d1tjtj| d2sft|d3| tjtj| d4s{t|d5| tjtj| d6st| d7| tjtj| d8st|d| d Sd S);av Convert a MONAI bundle to nnU-Net format. Parameters ---------- nnunet_config : dict Configuration dictionary for nnU-Net. Expected keys are: - "dataset_name_or_id": str, name or ID of the dataset. - "nnunet_trainer": str, optional, name of the nnU-Net trainer (default is "nnUNetTrainer"). - "nnunet_plans": str, optional, name of the nnU-Net plans (default is "nnUNetPlans"). bundle_root_folder : str Path to the root folder of the MONAI bundle. fold : int, optional Fold number for cross-validation (default is 0). Returns ------- None r)odictrrr:r) nnUNetLoggerrNTfolderrBprefixr%suffixsortrr& list[str]cs,fddt|D}|r||S)NcsBg|]}|rdus|jrdus|jr|jqS)N)rYr startswithendswith).0irrr<r= szDconvert_monai_bundle_to_nnunet..subfiles..)riterdirr)rrrrresr<rr=subfiless  z0convert_monai_bundle_to_nnunet..subfilesrrrZ __3d_fullresZnnUNet_preprocessedrrz/models/nnunet_checkpoint.pthr*rZcheckpoint_epoch)rrzcheckpoint_epoch=z.ptz /models/fold_z/checkpoint_epoch=Zcheckpoint_key_metriczcheckpoint_key_metric=z/checkpoint_key_metric=optimizer_stater, current_epochloggingZ _best_emaZgrad_scaler_staterrrFz/models/dataset.jsonrGz/models/plans.jsonzdataset_fingerprint.jsonz/dataset_fingerprint.jsonrH)NNT) rrBrr%rr%rrr&r)rZ'nnunetv2.training.logging.nnunet_loggerrrrrrrrXrr1r5rZr.lenrrZget_checkpointrrrr rr)rrrrr:rrrrrZnnunet_preprocess_model_folderrZlatest_checkpointsepochsZlatest_checkpointZ final_epochrZbest_checkpointsZ key_metricsZbest_checkpointZbest_key_metricrkeyr<r<r=rs      ""      r)rrFFFFrN)rrrrrrrrrrrrr rr!rr"rr#rr$r%r&rr)rArBrCrr&r)r)rrrrrr.r&r)Nr) rrrrrrrr%rrr&r) __future__rrrpathlibrtypingrrrryrr1torch.backendsrmonai.data.meta_tensorr monai.utilsr r _r __all__r nnModulerrrrrr<r<r<r=s<      h  6N =