o
    ip&                     @  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Zdd Zdd Zdd Ze
fddZh eedZ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                   @  (   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  
   || _ d S Nr   selfr    r   ]/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/transforms/lazy/utils.py__init__!      
zAffine.__init__c                 C  sZ   t | trdS t | trdS t| drt| jdk rdS | jd dv o,| 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          r   c                   @  r   )r   z0A class to represent a dense displacement field.r   c                 C  r   r   r   r   r   r   r   r   5   r   zDisplacementField.__init__c                 C  sD   t | trdS t | trdS t| drt| jdk rd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   r)   r   lefttorch.Tensorrightreturnc                 C  s   t | r,t |r,tt| t r| jn| dd} tt|t r!|jn|dd}t| |S t| rVt|rVtt| tr?| jn| dd} tt|trM|jn|dd}| | S t	)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_sequence)
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   r1   Tensornpndarraydictr   AFFINE)pending_itemr   r   r   affine_from_pendingS   s
   

r:   c                 C  sj   t | tsi S tj| tjdtj| tjdi}tj| v r&| tj |tj< tj| v r3| tj |tj< |S )z-Extract kwargs from a pending transform item.N)r   r7   r   INTERP_MODEgetPADDING_MODESHAPEDTYPE)r9   retr   r   r   kwargs_from_pending\   s   


rA   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   rB   c                 C  s   t | dd} | dddf }tj|t||dsdS t| d }g dg}}t| d|d|f D ]D\}}t|D ];\}}	tj|	d|dsOtj|	d|drj|d }
||v s[|
|v r_  dS || ||
 q;tj|	d|dsv  dS q;q3|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   r5   r   roundr!   	enumerateiscloseappend)matrixrD   sndimoxoyxrycZ	y_channelr   r   r   requires_interpp   s&     
	rS   rD   r   rJ   r   kwargsdict | Nonec                 C  s:  t |std|j dt| tjjr| jrt	
d |p i }|D ]}t|t q#|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
u rn| n|
}t|dd}||	tj||	d  |tj|tjd}t ||d}|d
urp|dkrp|d spt!"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 }|rt)j*||d}|D ]}|d }d|||f< || d ||df  ||df< qt!+|dks*t,d| dt%|t!-t||drWt|t|krWt%t|dd|rW|d |_|.t)j/}|S tj0j1||||d d}|.t)j/}|d |_|S tj0j2di |}d|_3|4d |dd|i|W  d
   S 1 sw   Y  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.rD   autoF)dtypealign_cornersrE   )r   
track_metaNTr/   r   )spatial_size
dst_affinemodepadding_moderC   rY   c                 S  s   g | ]}|d  qS )rE   r   ).0rO   r   r   r   
<listcomp>   s    zresample.<locals>.<listcomp>c                 S  s$   g | ]\}}|| d kr|d qS )r   rE   r   )r_   idxvalr   r   r   r`      s   $ )dimsr   z4Resampling out_spatial_size should be positive, got r\   r^   )r]   imgr   )5r   r"   r3   r   r   monair   
MetaTensorpending_operationswarningswarnr	   __override_lazy_keywordsr<   r   r   RESAMPLE_MODEr?   rX   ALIGN_CORNERSr!   r   get_track_metato_affine_ndaffiner>   peek_pending_shaper   utilsconvert_to_dst_typer;   r=   rS   r5   rF   argsorttolistr   arangepermuterG   r1   flipall
ValueErroreyetofloat32
transformscrop_or_pad_ndSpatialResamplelazytrace_transform)r   rJ   rT   krD   r]   Zinit_kwargsrL   rd   init_affiner[   Zout_spatial_sizeZcall_kwargsaxesZ	matrix_npZfull_transposein_shaperw   fZind_f	resamplerr   r   r   r
      sr   



""

&)r+   r,   r-   r,   r.   r,   r   )r   r,   rJ   r   rT   rU   )
__future__r   rh   numpyr5   r1   re   monai.configr   monai.data.utilsr   Z0monai.transforms.utils_pytorch_numpy_unificationr   monai.utilsr   r   r   r	   __all__r   r   r   r:   rA   rB   rS   listrj   r
   r   r   r   r   <module>   s&   
	!