U
    Php&                     @  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
 d dlmZ d dlmZmZmZmZ ddgZG d	d
 d
ZG dd dZddddddZdd Zdd Zdd Ze
fddZeedhZdddddddZdS )    )annotationsN)NdarrayOrTensor)
AFFINE_TOL)allclose)LazyAttrconvert_to_numpyconvert_to_tensorlook_up_optionresamplecombine_transformsc                   @  s(   e Zd ZdZdZdd Zedd ZdS )Affinez0A class to represent an affine transform matrix.datac                 C  s
   || _ d S Nr   selfr    r   P/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/transforms/lazy/utils.py__init__!   s    zAffine.__init__c                 C  sZ   t | trdS t | trdS t| dr4t| jdk r8dS | jd dkoX| jd | jd kS )z&Check if the data is an affine matrix.TFshape   )      )
isinstancer   DisplacementFieldhasattrlenr   r   r   r   r   is_affine_shaped$   s    

zAffine.is_affine_shapedN)__name__
__module____qualname____doc__	__slots__r   staticmethodr   r   r   r   r   r      s
   r   c                   @  s(   e Zd ZdZdZdd Zedd ZdS )r   z0A class to represent a dense displacement field.r   c                 C  s
   || _ d S r   r   r   r   r   r   r   5   s    zDisplacementField.__init__c                 C  sD   t | trdS t | trdS t| dr4t| jdk r8dS t|  S )zCheck if the data is a DDF.TFr   r   )r   r   r   r   r   r   r   r   r   r   r   is_ddf_shaped8   s    

zDisplacementField.is_ddf_shapedN)r    r!   r"   r#   r$   r   r%   r&   r   r   r   r   r   0   s
   r   ztorch.Tensor)leftrightreturnc                 C  s   t | rXt |rXtt| t r&| jn| dd} tt|t rB|jn|dd}t| |S t| rt|rtt| tr~| jn| dd} tt|tr|jn|dd}| | S t	dS )znGiven transforms A and B to be applied to x, return the combined transform (AB), so that A(B(x)) becomes AB(x)Twrap_sequenceN)
r   r   r   r   r   torchmatmulr   r&   NotImplementedError)r'   r(   r   r   r   r   D   s    c                 C  s.   t | tjtjfr| S t | tr*| tj S | S )z8Extract the affine matrix from a pending transform item.)r   r,   Tensornpndarraydictr   AFFINE)pending_itemr   r   r   affine_from_pendingS   s
    

r5   c                 C  sj   t | tsi S tj| tjdtj| tjdi}tj| krL| tj |tj< tj| krf| tj |tj< |S )z-Extract kwargs from a pending transform item.N)r   r2   r   INTERP_MODEgetPADDING_MODESHAPEDTYPE)r4   retr   r   r   kwargs_from_pending\   s    
  

r<   c                 C  s   dS )zGCheck if two sets of kwargs are compatible (to be combined in `apply`).Tr   )Zkwargs_1Zkwargs_2r   r   r   is_compatible_apply_kwargsk   s    r=   c                 C  s   t | dd} | dddf }tj|t||ds6dS t| d }g dg }}t| d|d|f D ]\}}t|D ]v\}}	tj|	d|dstj|	d|dr|d }
||ks|
|kr  dS || ||
 qvtj|	d|dsv  dS qvqf|S )	a  
    Check whether the transformation matrix suggests voxel-wise interpolation.

    Returns None if the affine matrix suggests interpolation.
    Otherwise, the matrix suggests that the resampling could be achieved by simple array operations
    such as flip/permute/pad_nd/slice; in this case this function returns axes information about simple axes
    operations.

    Args:
        matrix: the affine matrix to check.
        atol: absolute tolerance for checking if the matrix is close to an integer.
    Tr*   Nr   atol   r   g        )r   r0   r   roundr   	enumerateiscloseappend)matrixr?   sndimoxoyxrycZ	y_channelr   r   r   requires_interpp   s"      
rN   r?   r   zdict | None)r   rE   kwargsc              
   C  s,  t |std|j dt| tjjr:| jr:t	
d |p@i }|D ]}t|t qF|dt}|tjd}|tj| j|tjdd}t|d }t| tj d	}tj||j}	|tjd
}
|
d
kr| n|
}t|dd}||	tj||	d  |tj|tjd}t ||d}|d
k	r|dkr|d st!"t|dd}t!#|$ }t!%|t!&t|s|'|d
t|j }|jd|d  }|dd |dd
 D  |d
|< dd t(|d
| D }|r8t)j*||d}|D ]:}|d }d|||f< || d ||df  ||df< qt!+|dksXt,d| dt%|t!-t||drt|t|krt%t|dd|r|d |_|.t)j/}|S tj0j1||||d d}|.t)j/}|d |_|S tj0j2f |}d|_3|4d  |f d|i|W  5 Q R  S Q R X d
S )a  
    Resample `data` using the affine transformation defined by ``matrix``.

    Args:
        data: input data to be resampled.
        matrix: affine transformation matrix.
        kwargs: currently supports (see also: ``monai.utils.enums.LazyAttr``)

            - "lazy_shape" for output spatial shape
            - "lazy_padding_mode"
            - "lazy_interpolation_mode" (this option might be ignored when ``mode="auto"``.)
            - "lazy_align_corners"
            - "lazy_dtype" (dtype for resampling computation; this might be ignored when ``mode="auto"``.)
            - "atol" for tolerance for matrix floating point comparison.
            - "lazy_resample_mode" for resampling backend, default to `"auto"`. Setting to other values will use the
              `monai.transforms.SpatialResample` for resampling.

    See Also:
        :py:class:`monai.transforms.SpatialResample`
    z>Calling the dense grid resample API directly not implemented, .zMdata.pending_operations is not empty, the resampling output may be incorrect.r?   autoF)dtypealign_cornersr@   )r   
track_metaNTr*   r   )spatial_size
dst_affinemodepadding_moder>   rS   c                 S  s   g | ]}|d  qS )r@   r   ).0rJ   r   r   r   
<listcomp>   s     zresample.<locals>.<listcomp>c                 S  s$   g | ]\}}|| d kr|d qS )r   r@   r   )rY   idxvalr   r   r   rZ      s      )dimsr   z4Resampling out_spatial_size should be positive, got rV   rX   )rW   img)5r   r   r.   r   r   monair   
MetaTensorpending_operationswarningswarnr	   __override_lazy_keywordsr7   r   r   RESAMPLE_MODEr:   rR   ALIGN_CORNERSr   r   get_track_metato_affine_ndaffiner9   peek_pending_shaper   utilsconvert_to_dst_typer6   r8   rN   r0   rA   argsorttolistr   arangepermuterB   r,   flipall
ValueErroreyetofloat32
transformscrop_or_pad_ndSpatialResamplelazytrace_transform)r   rE   rO   kr?   rW   Zinit_kwargsrG   r^   init_affinerU   Zout_spatial_sizeZcall_kwargsaxesZ	matrix_npfull_transposein_shaperq   fZind_fZ	resamplerr   r   r   r
      sr    



"$

)N)
__future__r   rb   numpyr0   r,   r_   monai.configr   monai.data.utilsr   Z0monai.transforms.utils_pytorch_numpy_unificationr   monai.utilsr   r   r   r	   __all__r   r   r   r5   r<   r=   rN   listrd   r
   r   r   r   r   <module>   s$   	!