U Ph}8@sRddlmZddlmZmZddlZddlZddlm Z ddl m Z m Z ddl mZddlmZddlmZdd lmZmZmZmZerddlZd Zn ed \ZZd d ddddgZdefdddddd Zdejfdddddd Zd,ddddZd-ddZ ddZ!d d!Z"d"dd#d$Z#d"dd%d&Z$d"dd'd(Z%ddd)d*Z&d+dZ'dS).) annotations) TYPE_CHECKINGcastN) DtypeLike) ITKReader ITKWriter) MetaTensor)orientation_ras_lps)EnsureChannelFirst)MetaKeys SpaceKeysconvert_to_dst_typeoptional_importTitkitk_image_to_metatensormetatensor_to_itk_imageitk_to_monai_affinemonai_to_itk_affineget_itk_image_centermonai_to_itk_ddfzstr | int | NonezDtypeLike | torch.dtyper) channel_dimdtypereturncCsPtd|d}||\}}t|||dd}t||}t|d|}tt|S)a| Converts an ITK image to a MetaTensor object. Args: image: The ITK image to be converted. channel_dim: the channel dimension of the input image, default is None. This is used to set original_channel_dim in the metadata, EnsureChannelFirst reads this field. If None, the channel_dim is inferred automatically. If the input array doesn't have a channel dim, this value should be ``'no_channel'``. dtype: output dtype, defaults to the Python built-in `float`. Returns: A MetaTensor object containing the array data and metadata in ChannelFirst format. F)affine_lps_to_rasr)dstrr)r)rget_datar rensure_torch_and_prune_metar r)imagerrreaderZ image_array meta_dataZ metatensorr P/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/data/itk_torch_bridge.pyr+s   z int | Noner) meta_tensorrrcKs||jtjtjtjkr<|}t|j |_ tj|jtj<n|}t |dd}|j |j |dd|j |j|j|j d|j|dS)a Converts a MetaTensor object to an ITK image. Expects the MetaTensor to be in ChannelFirst format. Args: meta_tensor: The MetaTensor to be converted. channel_dim: channel dimension of the data array, defaults to ``0`` (Channel-first). ``None`` indicates no channel dimension. This is used to create a Vector Image if it is not ``None``. dtype: output data type, defaults to `np.float32`. kwargs: additional keyword arguments. Currently `itk.GetImageFromArray` will get ``ttype`` from this dictionary. Returns: The ITK image. See also: :py:func:`ITKWriter.create_backend_obj` F) output_dtyperT) data_arrayrsqueeze_end_dims)raffinerrkwargs)metagetr SPACEr LPSRAScloner r&rset_data_arraydatacreate_backend_objdata_objrr#)r"rrr'Z _meta_tensorwriterr r r!rEs  z torch.Tensor)rcCst||j}|rt||}ntj|dtjd}tj|dtjd}tj|tjd|d|d|f<tj|tjd|d||f<|rt||\}} | ||}t|\} } | || }t |\} } | || }||S)aa Converts an ITK affine matrix (2x2 for 2D or 3x3 for 3D matrix and translation vector) to a MONAI affine matrix. Args: image: The ITK image object. This is used to extract the spacing and direction information. matrix: The 2x2 or 3x3 ITK affine matrix. translation: The 2-element or 3-element ITK affine translation vector. center_of_rotation: The center of rotation. If provided, the affine matrix will be adjusted to account for the difference between the center of the image and the center of rotation. reference_image: The coordinate space that matrix and translation were defined in respect to. If not supplied, the coordinate space of image is used. Returns: A 4x4 MONAI affine matrix. rN) _assert_itk_regions_match_arrayndim&_compute_reference_space_affine_matrixtorcheyefloat64tensor_compute_offset_matrix_compute_direction_matrix_compute_spacing_matrix)rmatrix translationcenter_of_rotationZreference_imager6Zreference_affine_matrix affine_matrix offset_matrixinverse_offset_matrixdirection_matrixinverse_direction_matrixspacing_matrixinverse_spacing_matrixr r r!ris        c Cst|t|\}}|||}t|\}}|||}|rVt||\}}|||}|j} |d| d| f} |d| | f} | | fS)a Converts a MONAI affine matrix to an ITK affine matrix (2x2 for 2D or 3x3 for 3D matrix and translation vector). See also 'itk_to_monai_affine'. Args: image: The ITK image object. This is used to extract the spacing and direction information. affine_matrix: The 3x3 for 2D or 4x4 for 3D MONAI affine matrix. center_of_rotation: The center of rotation. If provided, the affine matrix will be adjusted to account for the difference between the center of the image and the center of rotation. Returns: The ITK matrix and the translation vector. N)r5r>r=r<r6numpytolist) rrBrArGrHrErFrCrDr6r?r@r r r!rs     cCsVt|tj}t|}t|}||dd||}|S)a Calculates the center of the ITK image based on its origin, size, and spacing. This center is equivalent to the implicit image center that MONAI uses. Args: image: The ITK image. Returns: The center of the image as a list of coordinates. g?) npasarrayGetLargestPossibleRegionGetSizefloat32 GetSpacing GetOrigin GetDirectionrJ)r image_sizespacingorigincenterr r r!rs c Cs|}|}|}t|}t|}t|}t|jddd}t|}t|} t|} t|dkot| dkot| dk} t ||ot ||ot ||} | st d| st ddS)Nrz9ITK-MONAI bridge: non-zero ITK region indices encounteredz9ITK-MONAI bridge: ITK regions should be of the same shape) rNZGetBufferedRegionZGetRequestedRegionrLarrayrOshapeZGetIndexall array_equalAssertionError) rZlargest_regionZbuffered_regionZrequested_regionZlargest_region_sizeZbuffered_region_sizeZrequested_region_size array_sizeZlargest_region_indexZbuffered_region_indexZrequested_region_indexZindices_are_zerosZ sizes_matchr r r!r5s*(   r5z!tuple[torch.Tensor, torch.Tensor]cCs|j}tt|t|}tj|dtjd}tj|tjd|d||f<tj|dtjd}tj|tjd |d||f<||fSNr3r4)r6rLrMrr8r9r:r;)rrAr6offsetrCrDr r r!r<sr<cCsx|j}tj|tjd}tj|dtjd}tj|dtjd}t|D]$\}}||||f<d||||f<qJ||fS)Nr4r3)r6rLrMrQr:r8r9 enumerate)rr6rUrGrHier r r!r>s r>cCs|j}t|}tj|dtjd}tj|tjd|d|d|f<t|}tj|dtjd}tj|tjd|d|d|f<||fSr_) r6rarray_from_matrixrSr8r9r:r;ZGetInverseDirection)rr6 directionrEZinverse_directionrFr r r!r= s  r=cs|jfddt|D\}}fddt|D\}}fddt|D\}}fddt|D\}} ||||} d} t|tj} ||||| | d} | t|t|7} t || | d}|S) Nc3s&|]}|ddfVqdSNrI.0mr6r r! sz9_compute_reference_space_affine_matrix..c3s&|]}|ddfVqdSrfrgrhrkr r!rlsc3s&|]}|ddfVqdSrfrgrhrkr r!rlsc3s&|]}|ddfVqdSrfrgrhrkr r!rlsrXrK)r?r@) r6r>r=rLrMrNrOrPrRr)rZ ref_imagerGZinv_spacing_matrixZref_spacing_matrixZref_inv_spacing_matrixrEZinv_direction_matrixZref_direction_matrixZref_inv_direction_matrixr?Z pixel_offsetrTr@Zref_affine_matrixr rkr!r7s&   r7c Cs|j}|tttd|ddg}|ddddf}tj|tjd}|tj ||dd9}tj| tjd}tj d||tjd tj }tj}tj||f}tj||f}tj||d }|||||| |S) a? converting the dense displacement field from the MONAI space to the ITK Args: image: itk image of array shape 2D: (H, W) or 3D: (D, H, W) ddf: numpy array of shape 2D: (2, H, W) or 3D: (3, D, H, W) Returns: displacement_field: itk image of the corresponding displacement field r3r.NrXr4)ndminz ij,...j->...i)ttype)r6 transposetuplelistrangerLrMrQr:rYrSeinsumastyperPrFVectorImageGetImageFromArray SetSpacing SetOriginrR SetDirection) rZddfr6rUreZvector_component_typeZvector_pixel_typeZdisplacement_field_typeZdisplacement_fieldr r r!r6s ")NN)N)( __future__rtypingrrrIrLr8monai.config.type_definitionsr monai.datarrmonai.data.meta_tensorrmonai.data.utilsr monai.transformsr monai.utilsr r r rrhas_itk__all__floatrrPrrrrr5r<r>r=r7rr r r r! sF       $4 %