U
    PhF'                     @  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ddddddd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labelFz
list[dict]strintzSequence[float] | floatzPathLike | NoneboolzTransform | None)datalist
output_dir	dimensionpixdim	image_key	label_keybase_dirlimitrelative_path
transformsreturnc
                 C  s  |dkrt dt| s t d|	dkr4t|||n|	}	g }
t| D ]f\}}|r`||kr` q|| }||d}|rtj||}|rtj||nd}tj|}|rtj|nd}t	
d| d|r|nd  |	||||i}|| }||}t	
d|j d|dkrdn|j  t|dd	}|dk	rJt|dd	}t	
d
|j d|dkrfdn|j  |dkrt|||||d}nt|||||d}|
| qD|
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"    r8   P/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/apps/deepgrow/dataset.pycreate_dataset   sT    9
(
(
r:   c              	   C  sl   |d k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;   r>   r8   r8   r9   r&      s     
r&   c              
   C  s  g }d}d}d}t |jd D ]}	||	df }
|d k	rD||	df nd }|d k	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 kr|	d|r|
|tj dn|i qt| }||dk }t|t|}|D ]}| 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 qq|dkrtd| d td| |j||d k	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: {})ranger/   r0   sumr)   r*   r+   makedirssaveappendreplacepathsepuniqueflattenmaxr%   r   astypefloat32r-   warningr.   format)r    r!   r"   r#   r   	data_listimage_countlabel_countunique_labels_countsidr   r   image_file_prefix
image_fileunique_labelsr5   label_file_prefix
label_file
curr_labelr8   r8   r9   r2      sd    


r2   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 kr|d	|r|
|t j d
n|
i nt	|
 }||dk }t|t|}|D ]}|	 dt|d}t j|d|}|d7 }||ktj}t jt j|ddd t|| |d7 }||rB|
|t j d
n|
|r\||t j d
n|t|d q|dkrtd| d td| |j||d k	r|jnd || |S )Nr   rA   rB   rC   rD   TrE   rG   r   rH   rI   rJ   rK   rL   rN   rO   rP   rQ   )r)   r*   r+   rT   r0   rU   rV   rW   rX   rY   rZ   r[   r%   r   r\   r]   r-   r^   r.   r_   r/   )r    r!   r"   r#   r   r`   ra   rb   rc   re   rf   rg   r5   rh   ri   rj   r8   r8   r9   r3      sT    &

r3   )r   r   Nr   FN)
__future__r   r-   r)   collections.abcr   numpyr0   monai.configr   monai.transformsr   r   r   r   r	   r
   r   monai.utilsr   r:   r&   r2   r3   r8   r8   r8   r9   <module>   s"   $      $nC