U
    Ph                     @  s  d dl mZ d dlZd dlmZmZmZmZ d dlm	Z	 d dl
Zd dlZd dlmZmZmZ d dlmZ d dlmZmZ d dlmZmZmZ d d	lmZmZ d d
lmZmZm Z m!Z!m"Z" d dl#m$Z$ e"dde!\Z%Z&e"ddd\Z'Z&e$( Z)G dd deZ*G dd deeZ+G dd deZ,G dd deZ-G dd deeZ.G dd deZ/G dd deZ0G dd  d eZ1G d!d" d"eZ2G d#d$ d$eZ3G d%d& d&eZ4dS )'    )annotationsN)CallableHashableIterableSequence)Any)IndexSelectionKeysCollectionNdarrayOrTensor)GaussianFilter)ResizeSpatialCrop)MapTransformRandomizable	Transform)generate_spatial_bounding_boxis_positive)InterpolateModeensure_tupleensure_tuple_repmin_versionoptional_import)PostFixzskimage.measurez0.14.2zscipy.ndimage.morphologydistance_transform_cdt)namec                   @  s:   e Zd ZdZddddddZdd	 Zd
ddddZdS )FindAllValidSlicesdz
    Find/List all valid slices in the label.
    Label is assumed to be a 4D Volume with shape CDHW, where C=1.

    Args:
        label: key to the label source.
        sids: key to store slices indices having valid label map.
    labelsidsstrr   r   c                 C  s   || _ || _d S Nr   )selfr   r    r"   S/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/apps/deepgrow/transforms.py__init__.   s    zFindAllValidSlicesd.__init__c                 C  sB   g }t |jd D ]$}t|d | dkr|| qt|S )N   r   )rangeshapenpsumappendasarray)r!   r   r   sidr"   r"   r#   _apply2   s
    zFindAllValidSlicesd._applyr   dictdatareturnc                 C  s   t |}t|| j tjr(|| j  n|| j }|jd dkrRtd|j dt|jdkrrtd|j d| 	|}|d k	rt|r||| j
< |S )Nr   r%   z5Only supports single channel labels, got label shape !   z5Only supports label with shape CDHW, got label shape )r.   
isinstancer   torchTensornumpyr'   
ValueErrorlenr-   r   )r!   r0   dr   r   r"   r"   r#   __call__9   s    *

zFindAllValidSlicesd.__call__N)r   r   )__name__
__module____qualname____doc__r$   r-   r;   r"   r"   r"   r#   r   $   s   	r   c                   @  s@   e Zd ZdZddddddd	d
dZdd Zdd Zdd ZdS )AddInitialSeedPointda  
    Add random guidance as initial seed point for a given label.

    Note that the label is of size (C, D, H, W) or (C, H, W)

    The guidance is of size (2, N, # of dims) where N is number of guidance added.
    # of dims = 4 when C, D, H, W; # of dims = 3 when (C, H, W)

    Args:
        label: label source.
        guidance: key to store guidance.
        sids: key that represents list of valid slice indices for the given label.
        sid: key that represents the slice to add initial seed point.  If not present, random sid will be chosen.
        connected_regions: maximum connected regions to use for adding initial points.
    r   guidancer   r,      r   int)r   rA   r   r,   connected_regionsc                 C  s(   || _ || _|| _d | _|| _|| _d S r    )r   sids_keysid_keyr,   rA   rD   )r!   r   rA   r   r,   rD   r"   r"   r#   r$   Y   s    zAddInitialSeedPointd.__init__c                 C  sT   | | jd }| | jd }|d k	rF|d ks4||krJ| jj|dd}nd }|| _d S )NF)replace)getrF   rE   Rchoicer,   )r!   r0   r,   r   r"   r"   r#   	randomizeh   s    zAddInitialSeedPointd.randomizec              	   C  s  t |jdkrdnd}dg|d  }|}|d k	rN|dkrNd}|d | tj }|dktj}|dkrztj|tddn|}t	|dkrt
dg }td|dkrdn| jd D ]}|dkr||ktj}t|dkr|| qt| }	t|	d	 }
t| dkd }| jj|d|
| t|
|  d
}|	| }tt||j  d }|d |d< |dks|dkr|| q||d ||d |d g qt||gt | gS )N      r%   r         ?)
backgroundzNot a valid Label      ?sizep)r9   r'   r(   newaxisastypefloat32measurer   rC   maxAssertionErrorr&   rD   r)   r*   r   flattenexpwhererI   rJ   r+   unravel_index	transposetolist)r!   r   r,   
dimensionsdefault_guidancedimsblobs_labelspos_guidanceridxdistanceprobabilityidxseeddstgr"   r"   r#   r-   r   s8      
$  zAddInitialSeedPointd._applyc                 C  sD   t |}| | t| || j | jjtdd	 || j
< |S )NFcopy)r.   rK   jsondumpsr-   r   r,   rW   rC   ra   rA   )r!   r0   r:   r"   r"   r#   r;      s    
.zAddInitialSeedPointd.__call__N)r   rA   r   r,   rB   )r<   r=   r>   r?   r$   rK   r-   r;   r"   r"   r"   r#   r@   H   s        
&r@   c                   @  s>   e Zd ZdZddddddd	d
Zdd Zdd Zdd ZdS )AddGuidanceSignaldaW  
    Add Guidance signal for input image.

    Based on the "guidance" points, apply gaussian to them and add them as new channel for input image.

    Args:
        image: key to the image source.
        guidance: key to store guidance.
        sigma: standard deviation for Gaussian kernel.
        number_intensity_ch: channel index.

    imagerA   rM   r%   r   rC   rs   rA   sigmanumber_intensity_chc                 C  s   || _ || _|| _|| _d S r    rt   )r!   rs   rA   ru   rv   r"   r"   r#   r$      s    zAddGuidanceSignald.__init__c                 C  sN  t |jdkrdnd}t|tjr*| n|}t|trBt|n|}|dkr~tj	t ||jd |jd |jd ftj
d}n&tj	t ||jd |jd ftj
d}|j}t|D ]\}}|D ]}tt|dk rq|dkrVtdtt|d |d d }	tdtt|d |d d }
tdtt|d |d d }d	|||	|
|f< qtdtt|d |d d }	tdtt|d |d d }
d	|||	|
f< qt|| dkrt|| }tt |j| jd
}||dd}|dd}|   ||< || t||  t|| t||   ||< q|S )NrL   rM   rU   rN   dtyper   r%   rQ   )ru   )r9   r'   r4   r(   ndarrayra   r   rp   loadszerosrX   	enumerateanyr+   rZ   minrC   r5   tensorr   ru   	unsqueezesqueezedetachcpur7   )r!   rs   rA   rb   signalsshapeig_ipointp1p2p3signal_tensorpt_gaussianr"   r"   r#   _get_signal   s6    0&
     8zAddGuidanceSignald._get_signalc                 C  sP   |  ||}t|tjr(|   }|dd| j df }tj	||gddS )Nr   .axis)
r   r4   r5   r6   r   r   r7   rv   r(   concatenate)r!   rs   rA   r   r"   r"   r#   r-      s
    zAddGuidanceSignald._applyc                 C  s2   t |}|| j }|| j }| |||| j< |S r    )r.   rs   rA   r-   )r!   r0   r:   rs   rA   r"   r"   r#   r;      s
    

zAddGuidanceSignald.__call__N)rs   rA   rM   r%   )r<   r=   r>   r?   r$   r   r-   r;   r"   r"   r"   r#   rr      s
   "	rr   c                   @  s@   e Zd ZdZdddddddZed	d
 Zdd Zdd ZdS )FindDiscrepancyRegionsda  
    Find discrepancy between prediction and actual during click interactions during training.

    Args:
        label: key to label source.
        pred: key to prediction source.
        discrepancy: key to store discrepancies found between label and prediction.

    r   preddiscrepancyr   r   r   r   c                 C  s   || _ || _|| _d S r    r   )r!   r   r   r   r"   r"   r#   r$      s    z FindDiscrepancyRegionsd.__init__c                 C  sP   | dk tj} |dk tj}| | }|dk tj}|dk  tj}||gS )NrO   r   )rW   r(   rX   )r   r   	disparitypos_disparityneg_disparityr"   r"   r#   r      s    z!FindDiscrepancyRegionsd.disparityc                 C  s   |  ||S r    )r   )r!   r   r   r"   r"   r#   r-     s    zFindDiscrepancyRegionsd._applyc                 C  s2   t |}|| j }|| j }| |||| j< |S r    )r.   r   r   r-   r   )r!   r0   r:   r   r   r"   r"   r#   r;     s
    

z FindDiscrepancyRegionsd.__call__N)r   r   r   )	r<   r=   r>   r?   r$   staticmethodr   r-   r;   r"   r"   r"   r#   r      s   

	r   c                   @  sN   e Zd ZdZdddddddZdd
dZdd Zdd Zdd Zdd Z	d	S )AddRandomGuidanceda  
    Add random guidance based on discrepancies that were found between label and prediction.
    input shape is as below:
    Guidance is of shape (2, N, # of dim)
    Discrepancy is of shape (2, C, D, H, W) or (2, C, H, W)
    Probability is of shape (1)

    Args:
        guidance: key to guidance source.
        discrepancy: key that represents discrepancies found between label and prediction.
        probability: key that represents click/interaction probability.

    rA   r   ri   r   )rA   r   ri   c                 C  s   || _ || _|| _d | _d S r    )rA   r   ri   _will_interact)r!   rA   r   ri   r"   r"   r#   r$     s    zAddRandomGuidanced.__init__Nc                 C  s,   || j  }| jjddg|d| gd| _d S )NTFrQ   )rT   )ri   rI   rJ   r   )r!   r0   ri   r"   r"   r#   rK   "  s    
zAddRandomGuidanced.randomizec                 C  s   t | }t|d }t| dkd }t|dkdkr| jj|d|| t||  d}|| }tt	||j
  d }|d |d< |S d S )NrQ   r   r%   rR   )r   r\   r(   r]   r^   r)   rI   rJ   r+   r_   r'   r`   ra   )r!   r   rh   ri   rj   rk   rl   rm   r"   r"   r#   find_guidance&  s    $ z AddRandomGuidanced.find_guidancec                 C  sx   |sdS |d }|d }t |dk}t |dk}t |t |k}|r^|r^| |d fS |st|rtd | |fS dS )N)NNr   r%   )r(   r)   r   )r!   r   Zwill_interact	pos_discrZ	neg_discrZcan_be_positiveZcan_be_negativeZcorrect_posr"   r"   r#   add_guidance4  s    zAddRandomGuidanced.add_guidancec                 C  s   t |tjr| n|}t |tr,t|n|}| || j\}}|rl|d 	| |d 	dgt
|  |r|d 	dgt
|  |d 	| ttj|td S )Nr   r%   rN   rx   )r4   r(   rz   ra   r   rp   r{   r   r   r*   r9   rq   r+   rC   )r!   rA   r   posnegr"   r"   r#   r-   F  s    zAddRandomGuidanced._applyc                 C  s<   t |}|| j }|| j }| | | |||| j< |S r    )r.   rA   r   rK   r-   )r!   r0   r:   rA   r   r"   r"   r#   r;   S  s    


zAddRandomGuidanced.__call__)rA   r   ri   )N)
r<   r=   r>   r?   r$   rK   r   r   r-   r;   r"   r"   r"   r#   r     s   
r   c                      sb   e Zd ZdZeddddeddddd	fd
ddddddddddddddd fddZdd Z  ZS )SpatialCropForegroundda  
    Crop only the foreground object of the expected images.

    Difference VS :py:class:`monai.transforms.CropForegroundd`:

      1. If the bounding box is smaller than spatial size in all dimensions then this transform will crop the
         object using box's center and spatial_size.

      2. This transform will set "start_coord_key", "end_coord_key", "original_shape_key" and "cropped_shape_key"
         in data[{key}_{meta_key_postfix}]

    The typical usage is to help training and evaluation if the valid part is small in the whole medical image.
    The valid part can be determined by any field in the data with `source_key`, for example:

    - Select values > 0 in image field as the foreground and crop on all fields specified by `keys`.
    - Select label = 3 in label field as the foreground to crop on all fields specified by `keys`.
    - Select label > 0 in the third channel of a One-Hot label field as the foreground to crop all `keys` fields.

    Users can define arbitrary function to select expected foreground from the whole source image or specified
    channels. And it can also add margin to every dim of the bounding box of foreground object.

    Args:
        keys: keys of the corresponding items to be transformed.
            See also: :py:class:`monai.transforms.MapTransform`
        source_key: data source to generate the bounding box of foreground, can be image or label, etc.
        spatial_size: minimal spatial size of the image patch e.g. [128, 128, 128] to fit in.
        select_fn: function to select expected foreground, default is to select values > 0.
        channel_indices: if defined, select foreground only on the specified channels
            of image. if None, select foreground on the whole image.
        margin: add margin value to spatial dims of the bounding box, if only 1 value provided, use it for all dims.
        allow_smaller: when computing box size with `margin`, whether allow the image size to be smaller
            than box size, default to `True`. if the margined size is bigger than image size, will pad with
            specified `mode`.
        meta_keys: explicitly indicate the key of the corresponding metadata dictionary.
            for example, for data with key `image`, the metadata by default is in `image_meta_dict`.
            the metadata is a dictionary object which contains: filename, original_shape, etc.
            it can be a sequence of string, map to the `keys`.
            if None, will try to construct meta_keys by `key_{meta_key_postfix}`.
        meta_key_postfix: if meta_keys is None, use `{key}_{meta_key_postfix}` to fetch/store the metadata according
            to the key data, default is `meta_dict`, the metadata is a dictionary object.
            For example, to handle key `image`,  read/write affine matrices from the
            metadata `image_meta_dict` dictionary's `affine` field.
        start_coord_key: key to record the start coordinate of spatial bounding box for foreground.
        end_coord_key: key to record the end coordinate of spatial bounding box for foreground.
        original_shape_key: key to record original shape for foreground.
        cropped_shape_key: key to record cropped shape for foreground.
        allow_missing_keys: don't raise exception if key is missing.
    Nr   Tforeground_start_coordforeground_end_coordforeground_original_shapeforeground_cropped_shapeFr	   r   zSequence[int] | np.ndarrayr   zIndexSelection | NonerC   boolKeysCollection | NoneNone)keys
source_keyspatial_size	select_fnchannel_indicesmarginallow_smaller	meta_keysmeta_key_postfixstart_coord_keyend_coord_keyoriginal_shape_keycropped_shape_keyallow_missing_keysr1   c                   s   t  || || _t|| _|| _|| _|| _|| _|d krNt	d t
| jnt|| _t
| jt
| jkrttdt	|	t
| j| _|
| _|| _|| _|| _d S Nz.meta_keys should have the same length as keys.)superr$   r   listr   r   r   r   r   r   r9   r   r   r   r8   r   r   r   r   r   )r!   r   r   r   r   r   r   r   r   r   r   r   r   r   r   	__class__r"   r#   r$     s    
"zSpatialCropForegroundd.__init__c                 C  sH  t |}t|| j | j| j| j| j\}}ttj	||gddj
tdd}tt||j
tdd}tt|| jrt|| jd}tdd |jD }tdd |jD }nt||d	}| || j| jD ]p\}}	}
|	p| d
|
 }	|||	 | j< |||	 | j< || j||	 | j< ||| }|j||	 | j< |||< q|S )Nr   r   Frn   
roi_centerroi_sizec                 S  s   g | ]
}|j qS r"   start.0sr"   r"   r#   
<listcomp>  s     z3SpatialCropForegroundd.__call__.<locals>.<listcomp>c                 S  s   g | ]
}|j qS r"   stopr   r"   r"   r#   r     s     	roi_startroi_end_)r.   r   r   r   r   r   r   r   r(   meanrW   rC   subtractalllessr   r   arraysliceskey_iteratorr   r   r   r   r'   r   r   )r!   r0   r:   	box_startbox_endcentercurrent_sizecropperkeymeta_keyr   rs   r"   r"   r#   r;     s0         
zSpatialCropForegroundd.__call__)	r<   r=   r>   r?   r   DEFAULT_POST_FIXr$   r;   __classcell__r"   r"   r   r#   r   ]  s   60"r   c                   @  sT   e Zd ZdZdddddddd	ef	d
d
d
d
dddd
dd
d
ddZdd Zdd Zd	S )AddGuidanceFromPointsda  
    Add guidance based on user clicks.

    We assume the input is loaded by LoadImaged and has the shape of (H, W, D) originally.
    Clicks always specify the coordinates in (H, W, D)

    If depth_first is True:

        Input is now of shape (D, H, W), will return guidance that specifies the coordinates in (D, H, W)

    else:

        Input is now of shape (H, W, D), will return guidance that specifies the coordinates in (H, W, D)

    Args:
        ref_image: key to reference image to fetch current and original image details.
        guidance: output key to store guidance.
        foreground: key that represents user foreground (+ve) clicks.
        background: key that represents user background (-ve) clicks.
        axis: axis that represents slices in 3D volume. (axis to Depth)
        depth_first: if depth (slices) is positioned at first dimension.
        spatial_dims: dimensions based on model used for deepgrow (2D vs 3D).
        slice_key: key that represents applicable slice to add guidance.
        meta_keys: explicitly indicate the key of the metadata dictionary of `ref_image`.
            for example, for data with key `image`, the metadata by default is in `image_meta_dict`.
            the metadata is a dictionary object which contains: filename, original_shape, etc.
            if None, will try to construct meta_keys by `{ref_image}_{meta_key_postfix}`.
        meta_key_postfix: if meta_key is None, use `{ref_image}_{meta_key_postfix}` to fetch the metadata according
            to the key data, default is `meta_dict`, the metadata is a dictionary object.
            For example, to handle key `image`,  read/write affine matrices from the
            metadata `image_meta_dict` dictionary's `affine` field.

    rA   
foregroundrP   r   TrM   sliceNr   rC   r   
str | None)
	ref_imagerA   r   rP   r   depth_firstspatial_dims	slice_keyr   r   c                 C  s@   || _ || _|| _|| _|| _|| _|| _|| _|	| _|
| _	d S r    )
r   rA   r   rP   r   r   rb   r   r   r   )r!   r   rA   r   rP   r   r   r   r   r   r   r"   r"   r#   r$     s    zAddGuidanceFromPointsd.__init__c                   sp  g  }}| j dkrt|}|| ttt|d d | jf } d krX|d nt fdd|D }	t|rt|}|t	|d d | jf |	k | d d dd f 
t }t|rt|}|t	|d d | jf |	k | d d dd f 
t }|||	g}
nPt|r@t||j
tdd }t|rdt||j
tdd }||g}
|
S )NrM   r   c                 3  s   | ]}| kr|V  qd S r    r"   )r   x	slice_numr"   r#   	<genexpr>  s      z0AddGuidanceFromPointsd._apply.<locals>.<genexpr>r%   Frn   )rb   r   extendr(   uniquer   r   nextr9   r^   rW   rC   ra   multiply)r!   
pos_clicks
neg_clicksfactorr   r   r   pointsr   	slice_idxrA   r"   r   r#   r-   	  s&    
"&
>

>

zAddGuidanceFromPointsd._applyc                 C  s.  t |}| jp| j d| j }||kr8td| dd|| krLtd|| d }t|| j j}| jr| jdkrtdt	
|d}t	|| }g }| j| jfD ]X}|| }	tt	j|	td	}	| jrtt|	D ]}
tt	
|	|
 d|	|
< q||	 q| |d |d ||| j|| j< |S )
Nr   zMissing meta_dict z	 in data!spatial_shapez%Missing "spatial_shape" in meta_dict!r   z-Depth first means the depth axis should be 0.r%   rx   )r.   r   r   r   RuntimeErrorr   r'   r   r   r(   rollr   r   rP   rC   r&   r9   r*   r-   rH   r   rA   )r!   r0   r:   meta_dict_keyoriginal_shapecurrent_shaper   Zfg_bg_clicksr   clicksr   r"   r"   r#   r;   #  s.    
&zAddGuidanceFromPointsd.__call__)r<   r=   r>   r?   r   r$   r-   r;   r"   r"   r"   r#   r     s   %"r   c                      sf   e Zd ZdZddedddddfd	d
dddd
d
d
d
d
ddd fddZdd ZdddddZ  ZS )SpatialCropGuidanceda  
    Crop image based on guidance with minimal spatial size.

    - If the bounding box is smaller than spatial size in all dimensions then this transform will crop the
      object using box's center and spatial_size.

    - This transform will set "start_coord_key", "end_coord_key", "original_shape_key" and "cropped_shape_key"
      in data[{key}_{meta_key_postfix}]

    Input data is of shape (C, spatial_1, [spatial_2, ...])

    Args:
        keys: keys of the corresponding items to be transformed.
        guidance: key to the guidance. It is used to generate the bounding box of foreground
        spatial_size: minimal spatial size of the image patch e.g. [128, 128, 128] to fit in.
        margin: add margin value to spatial dims of the bounding box, if only 1 value provided, use it for all dims.
        meta_keys: explicitly indicate the key of the corresponding metadata dictionary.
            for example, for data with key `image`, the metadata by default is in `image_meta_dict`.
            the metadata is a dictionary object which contains: filename, original_shape, etc.
            it can be a sequence of string, map to the `keys`.
            if None, will try to construct meta_keys by `key_{meta_key_postfix}`.
        meta_key_postfix: if meta_keys is None, use `key_{postfix}` to fetch the metadata according
            to the key data, default is `meta_dict`, the metadata is a dictionary object.
            For example, to handle key `image`,  read/write affine matrices from the
            metadata `image_meta_dict` dictionary's `affine` field.
        start_coord_key: key to record the start coordinate of spatial bounding box for foreground.
        end_coord_key: key to record the end coordinate of spatial bounding box for foreground.
        original_shape_key: key to record original shape for foreground.
        cropped_shape_key: key to record cropped shape for foreground.
        allow_missing_keys: don't raise exception if key is missing.
       Nr   r   r   r   Fr	   r   zIterable[int]rC   r   r   r   )r   rA   r   r   r   r   r   r   r   r   r   r1   c                   s   t  || || _t|| _|| _|d kr<td t| jnt	|| _
t| jt| j
krbtdt|t| j| _|| _|| _|	| _|
| _d S r   )r   r$   rA   r   r   r   r   r9   r   r   r   r8   r   r   r   r   r   )r!   r   rA   r   r   r   r   r   r   r   r   r   r   r"   r#   r$   b  s    
"zSpatialCropGuidanced.__init__c                 C  s   t |}t| j|}|D ]}|dk rtdqdg| }dg| }t|D ]V}|d|f }	tt|	||  d}
t|| t|	||  d }|
| ||< ||< qJ||fS )Nr   z+margin value should not be negative number..r%   )r9   r   r   r8   r&   rZ   r   )r!   r   	img_shapendimr   mr   r   didtmin_dmax_dr"   r"   r#   bounding_box~  s    


z!SpatialCropGuidanced.bounding_boxr   r.   r/   c                 C  sj  t |}| |}|dkr|S || j }|| jdd  }| t|d |d  |\}}ttj||gddj	t
dd}| j}	tt||j	t
dd}
|	t|
 d  }	t|	t|
k rt|
t|	 }t|dd|  |	 }	tt|
|	r&t|dkr|	d d |d< t||	d	}nt||d
}tdd |jD }tdd |jD }| || j| jD ]\}}}t|| jdd  |std|p| d| }||| | j< ||| | j< || j|| | j< ||| }|j|| | j< |||< qp|d |d  }}t|r6t|| ng }t|rTt|| ng }||g|| j< |S )Nr"   r%   r   r   Frn   rL   rM   r   r   c                 S  s   g | ]
}|j qS r"   r   r   r"   r"   r#   r     s     z1SpatialCropGuidanced.__call__.<locals>.<listcomp>c                 S  s   g | ]
}|j qS r"   r   r   r"   r"   r#   r     s     z>All the image specified in keys should have same spatial shaper   )r.   	first_keyrA   r'   r  r(   r   r   r   rW   rC   r   r   r9   r   r   r   r   r   r   r   array_equalr   r   r   r   r   ra   )r!   r0   r:   r  rA   original_spatial_shaper   r   r   r   box_sizediffr   r   r   r   rs   r   r   r   r   r"   r"   r#   r;     sH    

" zSpatialCropGuidanced.__call__)	r<   r=   r>   r?   r   r$   r  r;   r   r"   r"   r   r#   r   A  s   %*r   c                   @  s@   e Zd ZdZdedfddddddddd	Zd
ddddZdS )ResizeGuidanceda  
    Resize the guidance based on cropped vs resized image.

    This transform assumes that the images have been cropped and resized. And the shape after cropped is store inside
    the meta dict of ref image.

    Args:
        guidance: key to guidance
        ref_image: key to reference image to fetch current and original image details
        meta_keys: explicitly indicate the key of the metadata dictionary of `ref_image`.
            for example, for data with key `image`, the metadata by default is in `image_meta_dict`.
            the metadata is a dictionary object which contains: filename, original_shape, etc.
            if None, will try to construct meta_keys by `{ref_image}_{meta_key_postfix}`.
        meta_key_postfix: if meta_key is None, use `{ref_image}_{meta_key_postfix}` to fetch the metadata according
            to the key data, default is `meta_dict`, the metadata is a dictionary object.
            For example, to handle key `image`,  read/write affine matrices from the
            metadata `image_meta_dict` dictionary's `affine` field.
        cropped_shape_key: key that records cropped shape for foreground.
    Nr   r   r   r   )rA   r   r   r   r   r1   c                 C  s"   || _ || _|| _|| _|| _d S r    )rA   r   r   r   r   )r!   rA   r   r   r   r   r"   r"   r#   r$     s
    zResizeGuidanced.__init__r   r.   r/   c                 C  s   t |}|| j }|| jp*| j d| j  }|| j jdd  }|| j dd  }t||}|d |d  }}	t	|rt
||jtdd ng }
t	|	rt
|	|jtdd ng }|
|g|| j< |S )Nr   r%   r   Frn   )r.   rA   r   r   r   r'   r   r(   divider9   r   rW   rC   ra   )r!   r0   r:   rA   	meta_dictr   cropped_shaper   r   r   r   r   r"   r"   r#   r;     s    
&&zResizeGuidanced.__call__)r<   r=   r>   r?   r   r$   r;   r"   r"   r"   r#   r    s   r  c                      sf   e Zd ZdZdejddedddddf
dd	d
dddd	d	d	d	d	d
dd fddZdddddZ  Z	S )RestoreLabelda
  
    Restores label based on the ref image.

    The ref_image is assumed that it went through the following transforms:

        1. Fetch2DSliced (If 2D)
        2. Spacingd
        3. SpatialCropGuidanced
        4. Resized

    And its shape is assumed to be (C, D, H, W)

    This transform tries to undo these operation so that the result label can be overlapped with original volume.
    It does the following operation:

        1. Undo Resized
        2. Undo SpatialCropGuidanced
        3. Undo Spacingd
        4. Undo Fetch2DSliced

    The resulting label is of shape (D, H, W)

    Args:
        keys: keys of the corresponding items to be transformed.
        ref_image: reference image to fetch current and original image details
        slice_only: apply only to an applicable slice, in case of 2D model/prediction
        mode: {``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``, ``"mean"``,
            ``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``}
            One of the listed string values or a user supplied function for padding. Defaults to ``"constant"``.
            See also: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html
        align_corners: Geometrically, we consider the pixels of the input as squares rather than points.
            See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.grid_sample.html
            It also can be a sequence of bool, each element corresponds to a key in ``keys``.
        meta_keys: explicitly indicate the key of the corresponding metadata dictionary.
            for example, for data with key `image`, the metadata by default is in `image_meta_dict`.
            the metadata is a dictionary object which contains: filename, original_shape, etc.
            it can be a sequence of string, map to the `keys`.
            if None, will try to construct meta_keys by `key_{meta_key_postfix}`.
        meta_key_postfix: if meta_key is None, use `key_{meta_key_postfix} to fetch the metadata according
            to the key data, default is `meta_dict`, the metadata is a dictionary object.
            For example, to handle key `image`,  read/write affine matrices from the
            metadata `image_meta_dict` dictionary's `affine` field.
        start_coord_key: key that records the start coordinate of spatial bounding box for foreground.
        end_coord_key: key that records the end coordinate of spatial bounding box for foreground.
        original_shape_key: key that records original shape for foreground.
        cropped_shape_key: key that records cropped shape for foreground.
        allow_missing_keys: don't raise exception if key is missing.
    FNr   r   r   r   r	   r   r   z7Sequence[InterpolateMode | str] | InterpolateMode | strz#Sequence[bool | None] | bool | Noner   r   )r   r   
slice_onlymodealign_cornersr   r   r   r   r   r   r   r1   c                   s   t  || || _|| _t|t| j| _t|t| j| _|d krVtd t| jnt	|| _
t| jt| j
kr|td|| _|| _|	| _|
| _|| _d S r   )r   r$   r   r  r   r9   r   r  r  r   r   r8   r   r   r   r   r   )r!   r   r   r  r  r  r   r   r   r   r   r   r   r   r"   r#   r$   (  s    "zRestoreLabeld.__init__r   r.   r/   c              	   C  s*  t |}|| j d| j  }| || j| j| jD ]\}}}}|| }|j}	|| j }
t	
t	|	|
rt|
dd  |d}||||d}|| j }t	j|t	jd}|| j }|| j }tt|t|jdd  }ttd gdd t|d | |d | D  }|||< |jdd  }tt	|d d}|t| d  }t	
t	||rtt||d}||||d}|d	}|d ks| jrt|jd
kr|n|d }n"|d	 d }t	t|}|||< |||< |p| d| j }||}|d krt  }|||< ||d	< |d |d< q4|S )Nr   r%   )r   r  )r  r  rx   c                 S  s   g | ]\}}t ||qS r"   )r   )r   r   er"   r"   r#   r   [  s     z*RestoreLabeld.__call__.<locals>.<listcomp>r   r   rL   r   original_affineaffine)r.   r   r   r   r  r  r   r'   r   r(   r~   	not_equalr   r   r|   rX   r   r   r   r9   tupler   zipr   r   rH   r  )r!   r0   r:   r	  r   r  r  r   rs   r   r
  resizerr   resultr   r   r   r   r   r   r   r   final_resultmetar"   r"   r#   r;   E  sN    $



,


zRestoreLabeld.__call__)
r<   r=   r>   r?   r   NEARESTr   r$   r;   r   r"   r"   r   r#   r    s   5,r  c                      sL   e Zd ZdZdddedfdddd	dd
d fddZdd Zdd Z  ZS )Fetch2DSlicedaT  
    Fetch one slice in case of a 3D volume.

    The volume only contains spatial coordinates.

    Args:
        keys: keys of the corresponding items to be transformed.
        guidance: key that represents guidance.
        axis: axis that represents slice in 3D volume.
        meta_keys: explicitly indicate the key of the corresponding metadata dictionary.
            for example, for data with key `image`, the metadata by default is in `image_meta_dict`.
            the metadata is a dictionary object which contains: filename, original_shape, etc.
            it can be a sequence of string, map to the `keys`.
            if None, will try to construct meta_keys by `key_{meta_key_postfix}`.
        meta_key_postfix: use `key_{meta_key_postfix}` to fetch the metadata according to the key data,
            default is `meta_dict`, the metadata is a dictionary object.
            For example, to handle key `image`,  read/write affine matrices from the
            metadata `image_meta_dict` dictionary's `affine` field.
        allow_missing_keys: don't raise exception if key is missing.
    rA   r   NFr	   r   rC   r   r   )r   rA   r   r   r   r   c                   sn   t  || || _|| _|d kr2td t| jnt|| _t| jt| jkrXt	dt|t| j| _
d S r   )r   r$   rA   r   r   r9   r   r   r   r8   r   )r!   r   rA   r   r   r   r   r   r"   r#   r$     s    	"zFetch2DSliced.__init__c                 C  sX   |d }g }t |jD ],\}}|| jkr2||n|td| q|t| t|fS )NrM   r   )r}   r'   r   r*   r   r  )r!   rs   rA   r   rj   r   size_ir"   r"   r#   r-     s
    &zFetch2DSliced._applyc           	      C  s~   t |}|| j }t|dk r&td| || j| jD ]@\}}}| || |\}}|||< |||pr| d|  d< q8|S )NrL   z&Guidance does not container slice_idx!r   r   )r.   rA   r9   r   r   r   r   r-   )	r!   r0   r:   rA   r   r   r   Z	img_slicerj   r"   r"   r#   r;     s    
zFetch2DSliced.__call__)	r<   r=   r>   r?   r   r$   r-   r;   r   r"   r"   r   r#   r  ~  s   r  )5
__future__r   rp   collections.abcr   r   r   r   typingr   r7   r(   r5   monai.configr   r	   r
   monai.networks.layersr   monai.transformsr   r   monai.transforms.transformr   r   r   monai.transforms.utilsr   r   monai.utilsr   r   r   r   r   monai.utils.enumsr   rY   r   r   r  r   r   r@   rr   r   r   r   r   r   r  r  r  r"   r"   r"   r#   <module>   s8   $WH&Pqs 3 	