o # i}8@s&ddlmZddlmZmZddlZddlZddlm Z ddl m Z m Z ddl mZddlmZddlmZdd lmZmZmZmZerKddlZd Zned \ZZgd Zdefd.ddZdejfd/ddZd0d1ddZd2ddZ dd Z!d!d"Z"d3d$d%Z#d3d&d'Z$d3d(d)Z%d1d*d+Z&d,d-Z'dS)4) annotations) TYPE_CHECKINGcastN) DtypeLike) ITKReader ITKWriter) MetaTensor)orientation_ras_lps)EnsureChannelFirst)MetaKeys SpaceKeysconvert_to_dst_typeoptional_importTitk)itk_image_to_metatensormetatensor_to_itk_imageitk_to_monai_affinemonai_to_itk_affineget_itk_image_centermonai_to_itk_ddf channel_dimstr | int | NonedtypeDtypeLike | torch.dtypereturnrcCsPtd|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"]/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/data/itk_torch_bridge.pyr+s   r meta_tensor int | NonercKs||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  r torch.TensorcCst||j}|rt||}n tj|dtjd}tj|dtjd}tj|tjd|d|d|f<tj|tjd|d||f<|rPt||\}} | ||}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_imager:Zreference_affine_matrix affine_matrix offset_matrixinverse_offset_matrixdirection_matrixinverse_direction_matrixspacing_matrixinverse_spacing_matrixr"r"r#ris        rc Cst|t|\}}|||}t|\}}|||}|r+t||\}}|||}|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)r9rBrAr@r:numpytolist) rrFrErKrLrIrJrGrHr:rCrDr"r"r#rs     rcCsVt|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 GetDirectionrN)r image_sizespacingorigincenterr"r"r#rs  rc Cs|}|}|}t|}t|}t|}t|jddd}t|}t|} t|} t|dkoUt| dkoUt| dk} t ||ogt ||ogt ||} | snt d| stt ddS)Nrz9ITK-MONAI bridge: non-zero ITK region indices encounteredz9ITK-MONAI bridge: ITK regions should be of the same shape) rRZGetBufferedRegionZGetRequestedRegionrParrayrSshapeZGetIndexall 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#r9s,(   r9!tuple[torch.Tensor, torch.Tensor]cCs|j}tt|t|}tj|dtjd}tj|tjd|d||f<tj|dtjd}tj|tjd |d||f<||fSNr7r8)r:rPrQrr<r=r>r?)rrEr:offsetrGrHr"r"r#r@sr@cCsx|j}tj|tjd}tj|dtjd}tj|dtjd}t|D]\}}||||f<d||||f<q%||fS)Nr8r7)r:rPrQrUr>r<r= enumerate)rr:rYrKrLier"r"r#rBs rBcCs|j}t|}tj|dtjd}tj|tjd|d|d|f<t|}tj|dtjd}tj|tjd|d|d|f<||fSrd) r:rarray_from_matrixrWr<r=r>r?ZGetInverseDirection)rr: directionrIZinverse_directionrJr"r"r#rA s  rAcs|jfddt|D\}}fddt|D\}}fddt|D\}}fddt|D\}} ||||} d} t|tj} ||||| | d} | t|t|7} t || | d}|S) Nc3(|]}|ddfVqdSNrM.0mr:r"r# &z9_compute_reference_space_affine_matrix..c3rkrlrmrnrqr"r#rrrsc3rkrlrmrnrqr"r#rrrsc3rkrlrmrnrqr"r#rrs r\rO)rCrD) r:rBrArPrQrRrSrTrVr)rZ ref_imagerKZinv_spacing_matrixZref_spacing_matrixZref_inv_spacing_matrixrIZinv_direction_matrixZref_direction_matrixZref_inv_direction_matrixrCZ pixel_offsetrXrDZref_affine_matrixr"rqr#r;s&   r;c 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 r7r.Nr\r8)ndminz ij,...j->...i)ttype)r: transposetuplelistrangerPrQrUr>r]rWeinsumastyperTrFVectorImageGetImageFromArray SetSpacing SetOriginrV SetDirection) rddfr:rYrjZvector_component_typeZvector_pixel_typeZdisplacement_field_typeZdisplacement_fieldr"r"r#r6s "r)rrrrrr)r$rrr%rr)NN)rr6rl)rrc)( __future__rtypingrrrMrPr<monai.config.type_definitionsr monai.datarrmonai.data.meta_tensorrmonai.data.utilsr monai.transformsr monai.utilsr r r rrhas_itk__all__floatrrTrrrrr9r@rBrAr;rr"r"r"r#s8        $ 4%