U
    Ph+                  	   @  s  d dl mZ d dl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 d dlmZ d dl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dddddddg	ZG dd deeZ G dd deeZ!G dd deeZ"e" Z#Z$e! Z%Z&e  Z'Z(dS )    )annotations)HashableMappingSequence)AnyN)KeysCollectionSequenceStr)NdarrayOrTensor)get_track_meta)RandSmoothDeformRandSmoothFieldAdjustContrastRandSmoothFieldAdjustIntensity)MapTransformRandomizableTransform)GridSampleModeGridSamplePadModeInterpolateModeconvert_to_tensorensure_tuple_repRandSmoothFieldAdjustContrastdRandSmoothFieldAdjustIntensitydRandSmoothDeformdRandSmoothFieldAdjustContrastDRandSmoothFieldAdjustIntensityDRandSmoothDeformD!RandSmoothFieldAdjustContrastDict"RandSmoothFieldAdjustIntensityDictRandSmoothDeformDictc                      s   e Zd ZdZejZdejddddfddddd	d
dddd	ddZdddd d fddZ	dddd fddZ
dddddZ  ZS ) r   a  
    Dictionary version of RandSmoothFieldAdjustContrast.

    The field is randomized once per invocation by default so the same field is applied to every selected key. The
    `mode` parameter specifying interpolation mode for the field can be a single value or a sequence of values with
    one for each key in `keys`.

    Args:
        keys: key names to apply the augment to
        spatial_size: size of input arrays, all arrays stated in `keys` must have same dimensions
        rand_size: size of the randomized field to start from
        pad: number of pixels/voxels along the edges of the field to pad with 0
        mode: interpolation mode to use when upsampling
        align_corners: if True align the corners when upsampling field
        prob: probability transform is applied
        gamma: (min, max) range for exponential field
        device: Pytorch device to define field on
    r   N皙?)g      ?g      @r   Sequence[int]intr   bool | NonefloatSequence[float] | floattorch.device | None	keysspatial_size	rand_sizepadmodealign_cornersprobgammadevicec
           
   
   C  sN   t | | t| | t|t| j| _t|||| jd |d||	d| _d S Nr         ?)r'   r(   r)   r*   r+   r,   r-   r.   )	r   __init__r   r   lenr&   r*   r   trans
selfr&   r'   r(   r)   r*   r+   r,   r-   r.    r6   ]/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/transforms/smooth_field/dictionary.pyr1   B   s    z'RandSmoothFieldAdjustContrastd.__init__
int | Nonenp.random.RandomState | Noneseedstatereturnc                   s    t  || | j|| | S Nsuperset_random_stater3   r5   r;   r<   	__class__r6   r7   rA   ^   s    z/RandSmoothFieldAdjustContrastd.set_random_state
Any | NoneNonedatar=   c                   s    t  d  | jr| j  d S r>   )r@   	randomize_do_transformr3   r5   rH   rC   r6   r7   rI   e   s    z(RandSmoothFieldAdjustContrastd.randomize"Mapping[Hashable, NdarrayOrTensor]c                 C  s   |    t|}| js@| |D ]}t|| t d||< q |S t| |D ]8\}}| j| j	|t
| j	   | || d||< qN|S N)
track_metaFrI   dictrJ   key_iteratorr   r
   	enumerater3   set_moder*   r2   r5   rH   dkeyidxr6   r6   r7   __call__k   s    z'RandSmoothFieldAdjustContrastd.__call__)NN)N)__name__
__module____qualname____doc__r   backendr   AREAr1   rA   rI   rX   __classcell__r6   r6   rC   r7   r   ,   s       c                      s   e Zd ZdZejZdejddddfddddd	d
dddd	ddZdddd d fddZ	dddd fddZ
dddddZ  ZS ) r   a  
    Dictionary version of RandSmoothFieldAdjustIntensity.

    The field is randomized once per invocation by default so the same field is applied to every selected key. The
    `mode` parameter specifying interpolation mode for the field can be a single value or a sequence of values with
    one for each key in `keys`.

    Args:
        keys: key names to apply the augment to
        spatial_size: size of input arrays, all arrays stated in `keys` must have same dimensions
        rand_size: size of the randomized field to start from
        pad: number of pixels/voxels along the edges of the field to pad with 0
        mode: interpolation mode to use when upsampling
        align_corners: if True align the corners when upsampling field
        prob: probability transform is applied
        gamma: (min, max) range of intensity multipliers
        device: Pytorch device to define field on
    r   Nr   )r   r0   r   r   r    r   r!   r"   r#   r$   r%   c
           
   
   C  sN   t | | t| | t|t| j| _t|||| jd |d||	d| _d S r/   )	r   r1   r   r   r2   r&   r*   r   r3   r4   r6   r6   r7   r1      s    z(RandSmoothFieldAdjustIntensityd.__init__r8   r9   r:   c                   s    t  || | j|| | S r>   r?   rB   rC   r6   r7   rA      s    z0RandSmoothFieldAdjustIntensityd.set_random_staterE   rF   rG   c                   s   t  d  | j  d S r>   r@   rI   r3   rK   rC   r6   r7   rI      s    z)RandSmoothFieldAdjustIntensityd.randomizerL   c                 C  s   |    t|}| js@| |D ]}t|| t d||< q |S t| |D ]8\}}| j| j	|t
| j	   | || d||< qN|S rM   rO   rT   r6   r6   r7   rX      s    z(RandSmoothFieldAdjustIntensityd.__call__)NN)N)rY   rZ   r[   r\   r   r]   r   r^   r1   rA   rI   rX   r_   r6   r6   rC   r7   r   z   s       c                      s   e Zd ZdZejZdejdddej	e
jejddf
dddd	d
dddd
ddddddZd ddd d fddZd!ddd fddZdddddZ  ZS )"r   a5  
    Dictionary version of RandSmoothDeform.

    The field is randomized once per invocation by default so the same field is applied to every selected key. The
    `field_mode` parameter specifying interpolation mode for the field can be a single value or a sequence of values
    with one for each key in `keys`. Similarly the `grid_mode` parameter can be one value or one per key.

    Args:
        keys: key names to apply the augment to
        spatial_size: input array size to which deformation grid is interpolated
        rand_size: size of the randomized field to start from
        pad: number of pixels/voxels along the edges of the field to pad with 0
        field_mode: interpolation mode to use when upsampling the deformation field
        align_corners: if True align the corners when upsampling field
        prob: probability transform is applied
        def_range: value of the deformation range in image size fractions
        grid_dtype: type for the deformation grid calculated from the field
        grid_mode: interpolation mode used for sampling input using deformation grid
        grid_padding_mode: padding mode used for sampling input using deformation grid
        grid_align_corners: if True align the corners when sampling the deformation grid
        device: Pytorch device to define field on
    r   Nr   r0   Fr   r   r    r   r!   r"   r#   strr$   )r&   r'   r(   r)   
field_moder+   r,   	def_range	grid_modegrid_padding_modegrid_align_cornersr.   c                 C  sn   t | | t| | t|t| j| _t|
t| j| _t|||| jd |d||	| jd |||d| _	d S )Nr   r0   )r(   r'   r)   rb   r+   r,   rc   
grid_dtyperd   re   rf   r.   )
r   r1   r   r   r2   r&   rb   rd   r   r3   )r5   r&   r'   r(   r)   rb   r+   r,   rc   rg   rd   re   rf   r.   r6   r6   r7   r1      s$    zRandSmoothDeformd.__init__r8   r9   r:   c                   s    t  || | j|| | S r>   r?   rB   rC   r6   r7   rA     s    z"RandSmoothDeformd.set_random_staterE   rF   rG   c                   s   t  d  | j  d S r>   r`   rK   rC   r6   r7   rI     s    zRandSmoothDeformd.randomizerL   c                 C  s   |    t|}| js@| |D ]}t|| t d||< q |S t| |D ]Z\}}| j| j	|t
| j	   | j| j|t
| j   | || d| jj||< qN|S rM   )rI   rP   rJ   rQ   r   r
   rR   r3   set_field_moderb   r2   set_grid_moderd   r.   rT   r6   r6   r7   rX     s    zRandSmoothDeformd.__call__)NN)N)rY   rZ   r[   r\   r   r]   r   r^   torchfloat32r   NEARESTr   BORDERr1   rA   rI   rX   r_   r6   r6   rC   r7   r      s$   &&   ))
__future__r   collections.abcr   r   r   typingr   numpynprj   monai.configr   r   monai.config.type_definitionsr	   monai.data.meta_objr
   Z#monai.transforms.smooth_field.arrayr   r   r   monai.transforms.transformr   r   monai.utilsr   r   r   r   r   __all__r   r   r   r   r   r   r   r   r   r6   r6   r6   r7   <module>   s4   NM\