U
    Ph;3                     @  sX  d Z ddlmZ ddlZddlZddlZddlmZ ddlm	Z	 ddl
mZ ddlZddlm  mZ ddlmZmZ ddlmZ dd	lmZ dd
lmZ ddlmZ ddlmZ dddgZdddddZd4ddZ d5dddddddZ!d6ddddddZ"d7d"d#d$d$d%d&d'd(dZ#d8d+d#d$d%d,d$d$dd$d#d-d.d/dZ$d9d0dd,d,d1d2d3Z%dS ):zr
Utilities for accessing Nvidia MMARs

See Also:
    - https://docs.nvidia.com/clara/clara-train-sdk/pt/mmar.html
    )annotationsN)Mapping)Path)Any)download_and_extractlogger)PathLike)copy_model_state)optional_import   )
MODEL_DESC)RemoteMMARKeysget_model_specdownload_mmarload_from_mmarz	int | strz
dict | Any)idxreturnc                 C  sZ   t | trt|  S t | trV|   }tD ](}t|tj  |kr,|  S q,| S )zlget model specification by `idx`. `idx` could be index of the constant tuple of dict or the actual model ID.)	
isinstanceintr   strstriplowergetKeysID)r   keycand r   K/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/apps/mmars/mmars.pyr   )   s    


2   c                 C  s   d}ddddgdddd	d
gdddd	ddddddg
dd}t d	d|  ddg}||d< ||d< ||d< t|}| d| }tddd\}}	|	r||}
|
  ntdt|
j}i }|d D ]T}|d D ]F}|d
 }d	|d	 i||< |d D ] }|d d kr|d! || d"< qqq|S )#Nz<https://api.ngc.nvidia.com/v2/search/catalog/resources/MODEL scoreDESC)fieldvaluealldescriptionZdisplayNamenameZ
resourceIdZisPublic
attributesZguestAccessZorgNameZteamNameZdateModifiedlabelsr   )queryZorderByZqueryFieldsfieldspage*r,   ZpageSizefiltersz?q=requestsr   r'   z6NGC API requires requests package.  Please install it.results	resourcesr   ZlatestVersionIdStrr$   latest)dictjsondumpsr
   raise_for_status
ValueErrorloadstext)patternZ
page_index	page_sizeurl
query_dictfilterZ	query_strfull_urlrequests_gethas_requestsresp
model_list
model_dictresultmodelZcurrent_res_id	attributer   r   r   _get_all_ngc_models5   sL    


rI   r    r   )
model_nameversionmodel_prefixr   c                 C  s   d| |  d| dS )Nz%https://api.ngc.nvidia.com/v2/models/z
/versions/z/zipr   )rJ   rK   rL   r   r   r   _get_ngc_urlb   s    rM   )rJ   rL   r   c                 C  s   d| |  S )Nz&https://ngc.nvidia.com/catalog/models/r   )rJ   rL   r   r   r   _get_ngc_doc_urlf   s    rN   Tzstr | MappingzPathLike | Noneboolr   r   )itemmmar_dirprogressapirK   r   c                 C  s  |s0t ddd\}}|r(t| d }ntdt|}|rtt| trZ| tj|  n|  }t	|dkr~td|  dg }	|
 D ]l\}
}|d	kr|d
 nt|}t|
|}||d  }t|||d  d| d |dddd|d |	| q|	std|  d|	d S t| ts*t| } | tjd}|dkrJt|}| tj  d| }|| }| tjpt| tj |dd}t||| d| tj   || tj | tj | tj d|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.hubget_dirr0   mmarszTmmar_dir=None, but no suitable default directory computed. Upgrade Pytorch to 1.6+ ?r   z&api query returns no item for pattern z.  Please change or shorten it.rO   r3   r'   _z.zipNmd5zipF)r=   filepath
output_dirhash_val	hash_type	file_typehas_baserS   z'api query download no item for pattern r   znvidia/med/)rK   rL   .)r
   r   r8   rI   r   r   r   r   NAMElenitemsr   rM   r   appendr   VERSIONURLZ	FILE_TYPEZHASH_VALZ	HASH_TYPE)rQ   rR   rS   rT   rK   rU   Zhas_homeZ	_mmar_dirrE   Zmodel_dir_listkvverdownload_url	model_dirZmodel_fullnameZ	model_urlr   r   r   r   j   s`    $


"
FrG   zMapping | str | intz
Any | Noner   )rQ   rR   rS   rK   map_location
pretrainedweights_only	model_keyrT   
model_filer   c
              	   C  s  |r(t jt| tr t| t j n|  i} t| ts:t| } t| ||||d}
|	dkrbtj	dd}	|
| 
t j|	 }td| 
t j d|
 d |jdr|std	 |rtd
 tjj||dS tj||d}|r|
||S tt|
di |i d}|rt|tsf|
| 
t jtj	dd }t|}t|}W 5 Q R X t|}t||i d}|s~tt||i d}|rt|tstd| 
t j d| 
t j d|
dr|d }tj| }nj|
dr0|
dddd\}}t||d\}}|s@td|
dd d|
 dntd| dtd|  |
dd}|r~|f |}td|  n| }|rt||
||dd\}}}|r|rt d t!| d!t!| d td" | 
t j"pt#| t j d#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/
    )rQ   rR   rS   rK   rT   Nmodelszmodel.ptz
*** "z" available at r`   z.tsz4Loading a ScriptModule, 'pretrained' option ignored.z6Loading a ScriptModule, 'weights_only' option ignored.)rl   Z
train_confr   defaultconfigzconfig_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.htmlzCould not load model config z*** Model: argsz*** Model params: T)inplacez&*** Loading model state -- unchanged: z, changed: z
---znvidia:med:)rL   z#For more information, please visit 
)$r   ra   r   r   r   r   r   osru   joinr   Z
MODEL_FILEr   infor'   endswithwarningswarntorchjitload_get_valr4   ZCONFIG_FILEopenr5   r8   
monai_nets__dict__rsplitr
   r	   warningrb   ZDOCrN   )rQ   rR   rS   rK   rl   rm   rn   ro   rT   rp   rk   Z_model_filerE   Zmodel_configZ	json_pathf	conf_dictrJ   Z	model_clsZmodel_moduleZhas_clsZmodel_kwargsZ
model_instrW   changed	unchangedZdoc_urlr   r   r   r      sp    %$
 


 
 
 r   )
input_dictr   rs   r   c                 C  sN   || kr| | S | D ]4}| | }t |trt||dd}|dk	r|  S q|S )z
    Search for the item with `key` in `config_dict`.
    Returns: the first occurrence of `key` in a breadth first search.
    Nrr   )r   r   r   )r   r   rs   Zsub_dictvalZ	found_valr   r   r   r   -  s    

r   )r   r   )r    )r    )NTTrO   )	NTrO   NTFrG   TN)rG   N)&__doc__
__future__r   r5   rz   r~   collections.abcr   pathlibr   typingr   r   monai.networks.netsnetworksnetsr   monai.apps.utilsr   r   monai.config.type_definitionsr   monai.networks.utilsr	   monai.utils.moduler
   
model_descr   r   r   __all__r   rI   rM   rN   r   r   r   r   r   r   r   <module>   sH   

-       T         $q