o
    iF'                     @  s   d dl mZ d dlZd dlZd dlmZ d dlZd dlm	Z	 d dl
mZmZmZmZmZmZmZ d dlmZ 				 			d$d%ddZdd Zd d! Zd"d# ZdS )&    )annotationsN)Sequence)PathLike)ComposeEnsureChannelFirstd
LoadImagedOrientationdSpacingdSqueezeDimd	Transform)GridSampleModeimagelabelFdatalist
list[dict]
output_dirstr	dimensionintpixdimSequence[float] | float	image_key	label_keybase_dirPathLike | Nonelimitrelative_pathbool
transformsTransform | Nonereturnc
                 C  s  |dvrt dt| st d|	du rt|||n|	}	g }
t| D ]\}}|r/||kr/ |
S || }||d}|rMtj||}|rKtj||nd}tj|}|r[tj|nd}t	
d| d|rg|nd  |	||||i}|| }||}t	
d|j d|du rdn|j  t|dd	}|durt|dd	}t	
d
|j d|du rdn|j  |dkrt|||||d}n	t|||||d}|
| q"|
S )a8  
    Utility to pre-process and create dataset list for Deepgrow training over on existing one.
    The input data list is normally a list of images and labels (3D volume) that needs pre-processing
    for Deepgrow training pipeline.

    Args:
        datalist: A list of data dictionary. Each entry should at least contain 'image_key': <image filename>.
            For example, typical input data can be a list of dictionaries::

                [{'image': <image filename>, 'label': <label filename>}]

        output_dir: target directory to store the training data for Deepgrow Training
        pixdim: output voxel spacing.
        dimension: dimension for Deepgrow training.  It can be 2 or 3.
        image_key: image key in input datalist. Defaults to 'image'.
        label_key: label key in input datalist. Defaults to 'label'.
        base_dir: base directory in case related path is used for the keys in datalist.  Defaults to None.
        limit: limit number of inputs for pre-processing.  Defaults to 0 (no limit).
        relative_path: output keys values should be based on relative path.  Defaults to False.
        transforms: explicit transforms to execute operations on input data.

    Raises:
        ValueError: When ``dimension`` is not one of [2, 3]
        ValueError: When ``datalist`` is Empty

    Returns:
        A new datalist that contains path to the images/labels after pre-processing.

    Example::

        datalist = create_dataset(
            datalist=[{'image': 'img1.nii', 'label': 'label1.nii'}],
            base_dir=None,
            output_dir=output_2d,
            dimension=2,
            image_key='image',
            label_key='label',
            pixdim=(1.0, 1.0),
            limit=0,
            relative_path=True
        )

        print(datalist[0]["image"], datalist[0]["label"])
    )      zEDimension can be only 2 or 3 as Deepgrow supports only 2D/3D TrainingzInput datalist is emptyNzImage: z	; Label: zImage (transform): r   zImage (final): r!   )vol_idx	vol_image	vol_labeldataset_dirr   )
ValueErrorlen_default_transforms	enumerategetospathjoinabspathlogginginfoshapenpmoveaxis_save_data_2d_save_data_3dextend)r   r   r   r   r   r   r   r   r   r   Znew_datalistidxitemr   r   datar%   r&    r<   ]/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/apps/deepgrow/dataset.pycreate_dataset   sV   9(
&&r>   c              	   C  sl   |d u r| gn| |g}t |dkrtjtjgntjg}tt|dt|dt|ddt|||dt	|dgS )Nr!   )keysRAS)r?   axcodes)r?   r   mode)
r)   r   BILINEARNEARESTr   r   r   r   r	   r
   )r   r   r   r?   rB   r<   r<   r=   r*      s    
r*   c              
   C  s
  g }d}d}d}t |jd D ]}	||	df }
|d ur!||	df nd }|d ur/t|dkr/qd| dd|	d}tj|d|}|d7 }tjtj|dd	d
 t||
 |d7 }|d u rr|	d|rm|
|tj dn|i qt| }||dk }t|t|}|D ]W}| dt|d}tj|d|}|d7 }tjtj|dd	d
 ||ktj}t|| |d7 }|	|r|
|tj dn||r|
|tj dn|t|d qq|dkrtd| d td| |j||d ur|jnd || |S )Nr   .vol_idx_0>4dZ_slice_z0>3dimages.npyTexist_ok   r    _region_0>2dlabelsr   r   region   Unique labels - exceeds 20. Please check if this is correct.E{} => Image Shape: {} => {}; Label Shape: {} => {}; Unique Labels: {})ranger3   r4   sumr-   r.   r/   makedirssaveappendreplacepathsepuniqueflattenmaxr)   r   astypefloat32r1   warningr2   format)r$   r%   r&   r'   r   	data_listimage_countlabel_countunique_labels_countsidr   r   image_file_prefix
image_fileunique_labelsr9   label_file_prefix
label_file
curr_labelr<   r<   r=   r6      sf   
r6   c              
   C  s  g }d}d}d}d| d}	t j|d|	}
|
d7 }
t jt j|ddd t|
| |d7 }|d u rG|d	|rB|
|t j d
n|
i nnt	|
 }||dk }t|t|}|D ]W}|	 dt|d}t j|d|}|d7 }||ktj}t jt j|ddd t|| |d7 }||r|
|t j d
n|
|r||t j d
n|t|d q]|dkrtd| d td| |j||d ur|jnd || |S )Nr   rE   rF   rG   rH   TrI   rK   r   rL   rM   rN   rO   rP   rR   rS   rT   rU   )r-   r.   r/   rX   r4   rY   rZ   r[   r\   r]   r^   r_   r)   r   r`   ra   r1   rb   r2   rc   r3   )r$   r%   r&   r'   r   rd   re   rf   rg   ri   rj   rk   r9   rl   rm   rn   r<   r<   r=   r7      sT   &
r7   )r   r   Nr   FN)r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r    r   )
__future__r   r1   r-   collections.abcr   numpyr4   monai.configr   monai.transformsr   r   r   r   r	   r
   r   monai.utilsr   r>   r*   r6   r7   r<   r<   r<   r=   <module>   s$   $nC