o  iN3@s"dZddlmZddlZddlZddlZddlmZddlm Z ddl m Z ddl Z ddl mmZddlmZmZddlmZdd lmZdd lmZd d lmZd d lmZgdZd@ddZdAddZ dBdCddZ!dBdDdd Z" "dEdFd,d-Z#  ! "  ! . / ! dGdHd8d9Z$dIdJd>d?Z%dS)Kzr Utilities for accessing Nvidia MMARs See Also: - https://docs.nvidia.com/clara/clara-train-sdk/pt/mmar.html ) annotationsN)Mapping)Path)Any)download_and_extractlogger)PathLike)copy_model_state)optional_import) MODEL_DESC)RemoteMMARKeys)get_model_spec download_mmarload_from_mmaridx int | strreturn dict | AnycCsZt|tr t|St|tr+|}tD]}t|tj|kr*|Sq|S)zlget model specification by `idx`. `idx` could be index of the constant tuple of dict or the actual model ID.) isinstanceintr strstriplowergetKeysID)rkeycandrX/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/apps/mmars/mmars.pyr)s   r2cCs d}ddddggdgddd }td d |d dg}||d <||d <||d<t|}|d|}tddd\}} | rH||} | ntdt| j} i} | dD]*} | dD]#}|d}d |d i| |<|dD]}|ddkr|d| |d<qpq^qX| S)Nz         rMr" model_namerversion model_prefixcCsd||d|dS)Nz%https://api.ngc.nvidia.com/v2/models/z /versions/z/zipr)rNrOrPrrr _get_ngc_urlbsrQcCsd||S)Nz&https://ngc.nvidia.com/catalog/models/r)rNrPrrr _get_ngc_doc_urlfsrRTitem str | Mappingmmar_dirPathLike | Noneprogressboolapirrc Cs|stddd\}}|rt|d}ntdt|}|rtt|tr,|tj|n|}t |dkr>td|dg} | D]6\} } |d krP| d nt |} t | | } || d }t | || d d | d |dddd|d| |qD| std|d| dSt|tst|}|tjd} |dkrt |} |tjd | }||}|tjpt |tj| dd}t |||d|tj||tj|tj|tjd|d|S)a' Download and extract Medical Model Archive (MMAR) from Nvidia Clara Train. See Also: - https://docs.nvidia.com/clara/ - Nvidia NGC Registry CLI - https://docs.nvidia.com/clara/clara-train-sdk/pt/mmar.html Args: item: the corresponding model item from `MODEL_DESC`. Or when api is True, the substring to query NGC's model name field. mmar_dir: target directory to store the MMAR, default is `mmars` subfolder under `torch.hub get_dir()`. progress: whether to display a progress bar. api: whether to query NGC and download via api version: which version of MMAR to download. -1 means the latest from ngc. Examples:: >>> from monai.apps import download_mmar >>> download_mmar("clara_pt_prostate_mri_segmentation_1", mmar_dir=".") >>> download_mmar("prostate_mri_segmentation", mmar_dir=".", api=True) Returns: The local directory of the downloaded model. If api is True, a list of local directories of downloaded models. z torch.hubget_dirr4mmarszTmmar_dir=None, but no suitable default directory computed. Upgrade Pytorch to 1.6+ ?rz&api query returns no item for pattern z. Please change or shorten it.rSr7r*_z.zipNmd5zipF)rAfilepath output_dirhash_val hash_type file_typehas_baserXz'api query download no item for pattern r z nvidia/med/)rOrP.)r rr<rMrrrrNAMElenitemsrrQrappendrVERSIONZURLZ FILE_TYPEZHASH_VALZ HASH_TYPE)rTrVrXrZrOr[Zhas_homeZ _mmar_dirrIZmodel_dir_listkvver download_url model_dirZmodel_fullnameZ model_urlrrr rjs`$       rFrKMapping | str | int map_location Any | None pretrained weights_only model_key model_filerc Cs |rtjt|trt|tjn|i}t|tst|}t|||||d} | dur1tj dd} | | tj | } t d| tjd| d| jdrf|sWtd |r^td tjj| |d Stj| |d d } |rv| || Stt|  di|id} | rt| ts| | tjtj dd}t| }t|}Wdn1swYt|}t||id} | stt| |id} | rt| tstd| tjd| tj d| dr| d}tj|}n5| dr| dddd\}}t||d\}}|std| ddd| dntd| dt d|| dd}|rF|d(i|}t d|n|}|rpt|| || d d \}}}|r`|rpt d!t!|d"t!|dt d#| tj"pt#|tjd$d%}t d&|d'|S))ak Download and extract Medical Model Archive (MMAR) model weights from Nvidia Clara Train. Args: item: the corresponding model item from `MODEL_DESC`. mmar_dir: : target directory to store the MMAR, default is mmars subfolder under `torch.hub get_dir()`. progress: whether to display a progress bar when downloading the content. version: version number of the MMAR. Set it to `-1` to use `item[Keys.VERSION]`. map_location: pytorch API parameter for `torch.load` or `torch.jit.load`. pretrained: whether to load the pretrained weights after initializing a network module. weights_only: whether to load only the weights instead of initializing the network module and assign weights. model_key: a key to search in the model file or config file for the model dictionary. Currently this function assumes that the model dictionary has `{"[name|path]": "test.module", "args": {'kw': 'test'}}`. api: whether to query NGC API to get model infomation. model_file: the relative path to the model file within an MMAR. Examples:: >>> from monai.apps import load_from_mmar >>> unet_model = load_from_mmar("clara_pt_prostate_mri_segmentation_1", mmar_dir=".", map_location="cpu") >>> print(unet_model) See Also: https://docs.nvidia.com/clara/ )rTrVrXrOrZNmodelszmodel.ptz *** "z" available at rfz.tsz4Loading a ScriptModule, 'pretrained' option ignored.z6Loading a ScriptModule, 'weights_only' option ignored.)rrT)rrruZ train_confrdefaultconfigzconfig_train.jsonz4Could not load model config dictionary from config: z, or from model file: r*pathr )moduler*z!Could not load MMAR model config r"z*, Please make sure MMAR's sub-folders in 'z[' is on the PYTHONPATH.See also: https://docs.nvidia.com/clara/clara-train-sdk/pt/byom.htmlzCould not load model config z *** Model: argsz*** Model params: )inplacez&*** Loading model state -- unchanged: z , changed: z ---z nvidia:med:)rPz#For more information, please visit  r)$rrgrrrrrosr|joinrZ MODEL_FILErinfor*endswithwarningswarntorchjitload_get_valr8Z CONFIG_FILEopenr9r< monai_nets__dict__rsplitr r warningrhZDOCrR)rTrVrXrOrrrtrurvrZrwrpZ _model_filerIZ model_configZ json_pathf conf_dictrNZ model_clsZ model_moduleZhas_clsZ model_kwargsZ model_instr]changedZ unchangedZdoc_urlrrr rs|%$                 r input_dictrrrzcCsN||vr||S|D]}||}t|tr$t||dd}|dur$|Sq |S)z Search for the item with `key` in `config_dict`. Returns: the first occurrence of `key` in a breadth first search. Nry)rrr)rrrzZsub_dictvalZ found_valrrr r-s r)rrrr)rr!)r")rNrrOrrPrrr)rNrrPrrr)NTTrS) rTrUrVrWrXrYrZrYrOrrr) NTrSNTFrKTN)rTrqrVrWrXrYrOrrrrsrtrYrurYrvrrZrYrwrWrr)rKN)rrrrrzrsrrs)&__doc__ __future__rr9rrcollections.abcrpathlibrtypingrrmonai.networks.netsnetworksnetsrmonai.apps.utilsrrmonai.config.type_definitionsrmonai.networks.utilsr monai.utils.moduler model_descr r r__all__rrMrQrRrrrrrrr sD           -  T q