o
    ii                     @  s   d dl mZ d dlZd dlZd dlZd dlZd dlmZ d dl	m
Z
mZmZ eddd\ZZed\ZZejjeZg dZG d	d
 d
e
Zd ddZd!ddZd"ddZdS )#    )annotationsN)ConfigParser)StrEnumensure_tupleoptional_importtqdm)namenibabel)analyze_datacreate_new_data_copycreate_new_dataset_json
NNUNETModec                   @  s   e Zd ZdZdZdZdZdS )r   2d
3d_fullresZ	3d_lowresZ3d_cascade_fullresN)__name__
__module____qualname__N_2DN_3D_FULLRESN_3D_LOWRESN_3D_CASCADE_FULLRES r   r   Y/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/apps/nnunet/utils.pyr      s
    r   datalist_jsondictdata_dirstrreturntuple[int, int]c                 C  s   t jjddddtj|| d d d }| dkr"| d nd}t	d|  d}t
t| d D ]!}t jjddddtj|| d | d	 }t|t| }q6t	d
|  ||fS )z
    Analyze (training) data

    Args:
        datalist_json: original data list .json (required by most monai tutorials).
        data_dir: raw data directory.
    T
image_onlyensure_channel_firstsimple_keystrainingr   image      znum_input_channels: labelznum_foreground_classes: )monai
transforms	LoadImageospathjoindimsizeloggerinforangelenmaxint)r   r   imgnum_input_channelsnum_foreground_classes_isegr   r   r   r
   &   s   r
   test_keyr7   r5   output_datafolderNonec                 C  s~  d}dg | g i}t td| gddgddgD ]%\}}}	|du r!qtd| d	 tr6ttt|| ntt|| D ]}
t|| |
 t	rQ|| |
 d
 n|| |
 }d| }|d7 }t
jjddddtj||}|jd }| }t|D ](}t||df |}dt|d dd  }t|tj|||| d  q}t|| |
 t	rd|| |
 v rt
jjddddtj||| |
 d }|jd }| tj}|jdkr|jd dkr|d n|}tt||tj||	|d  t|| |
 t	rt|| |
 }||d< || | n|| || |
 |d tj|tj|dddddd q>qdS )a  
    Create and organize a new copy of data to meet the requirements of nnU-Net V2

    Args:
        test_key: key for test data in the data list .json.
        datalist_json: original data list .json (required by most monai tutorials).
        data_dir: raw data directory.
        num_input_channels: number of input (image) channels.
        output_datafolder: output folder.
    r   r#   imagesTrimagesTslabelsTrlabelsTsNzconverting data section: z...r$   Zcase_r&   Tr   original_affine._i'  .nii.gzr'   r%   )r   .new_name)r$   rF   zdatalist.jsonjsonFconfigfilepathfmt	sort_keysindentensure_ascii) listzipr0   r1   has_tqdmr   r2   r3   
isinstancer   r(   r)   r*   r+   r,   r-   metanumpynibNifti1Imager   saveastypenpuint8ndimshapecopydeepcopyappendr   export_config_file)r;   r   r   r7   r<   _indexZnew_datalist_json_keyZ_folderZ_label_folder_kZorig_img_nameZimg_namendaaffine_lZoutimgindex_valr   r   r   r   ?   sb   ,


""
$0r   modalityr8   num_training_dataoutput_filepathc                 C  s   i }t | } i |d< t|D ]}| | |d t|< qi |d< d|d d< t|D ]}|d |d d|d  < q)||d< d|d	< tj||d
dddd dS )a  
    Create a new copy of dataset .json to meet the requirements of nnU-Net V2

    Args:
        modality: image modality, could a string or a list of strings.
        num_foreground_classes: number of foreground classes.
        num_input_channels: number of input (image) channels.
        num_training_data: number of training data.
        output_filepath: output file path/name.
    channel_nameslabelsr   
backgroundr&   classnumTrainingrE   Zfile_endingrG   Tr%   FrH   N)r   r2   r   r   r`   )ri   r8   r7   rj   rk   Znew_json_data_jr   r   r   r      s(   
r   )r   r   r   r   r   r   )r;   r   r   r   r   r   r7   r5   r<   r   r   r=   )ri   r   r8   r5   r7   r5   rj   r5   rk   r   r   r=   )
__future__r   r]   r+   rT   rY   r(   monai.bundler   monai.utilsr   r   r   r   rQ   rU   rC   appsutils
get_loggerr   r0   __all__r   r
   r   r   r   r   r   r   <module>   s   

J