o  i`1@sdZddlmZddlZddlZddlZddlmZ ddl m Z ddl m Z ddlmZddlmZdd lmZdd lmZmZdd lmZmZmZmZmZgd Zd dZd.ddZd.ddZ ej!fd/ddZ"d0d d!Z#ej!d"fd1d(d)Z$d2d,d-Z%dS)3zA A collection of "functional" transforms for spatial operations. ) annotationsN)pad) NdarrayTensor)get_track_meta) MetaTensor) to_affine_nd)TraceableTransform)convert_pad_modecreate_translate)PytorchPadModeconvert_to_dst_typeconvert_to_numpyconvert_to_tensor ensure_tuple)pad_ndpad_func crop_funccrop_or_pad_ndcCs8|dus|dvr tjS|dvrtjS|dvrtjStjS)zSget the most similar mode of `pad` from ``padding_mode`` of the spatial resampling.N)zerosconstantz grid-constant) reflectionreflectmirrorz grid-mirror)wrapz grid-wrap)r CONSTANTREFLECTCIRCULAR REPLICATE) padding_modere/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/transforms/croppad/functional.py_convert_pt_pad_mode"sr!imgr pad_widthlist[tuple[int, int]]modestrreturnc Kst|tjr#|jrtd|jd|jd|d| }nt |}t ||dj}|dkr>d|vr>|d|d<t j||fd |i|}t||d d S) NzPadding: moving img z from cuda to cpu for dtype=z mode=.dstr%rvalueconstant_valuesr%r*r) isinstancetorchTensoris_cudawarningswarnshapedtypedetachcpunumpynpasarrayr r+poprr )r"r#r%kwargsZimg_nprrr _np_pad-s " r=cKst|}t||dj}|dkr d|vr |}|d|d<n|}dd|ddDddd}t|d |fd |i|d }t ||d d S) Nr)rr,r+cSs$g|]}|dddD]}|q qS)Nr).0sublistvalrrr C$z_pt_pad..r>rr%r-) r/ as_tensorr r+copyr;pad_pt unsqueezesqueezer )r"r#r%r<Zimg_pt_kwargsZ pt_pad_widthrrr _pt_pad;s  "rKto_padc s|dvrt|f||d|Sz!t}|dvr%|jtjtjtjtjhvr%t}||f||d|WStt t fyzBt t sKt fdddDr\t|f||d|WYdSt|jd|d|d|d|jdt |tjrw|jnd dww) a Pad `img` for a given an amount of padding in each dimension. `torch.nn.functional.pad` is used unless the mode or kwargs are not available in torch, in which case `np.pad` will be used. Args: img: data to be transformed, assuming `img` is channel-first and padding doesn't apply to the channel dim. to_pad: the amount to be padded in each dimension [(low_H, high_H), (low_W, high_W), ...]. default to `self.to_pad`. mode: available modes: (Numpy) {``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``, ``"mean"``, ``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``} (PyTorch) {``"constant"``, ``"reflect"``, ``"replicate"``, ``"circular"``}. One of the listed string values or a user supplied function. Defaults to ``"constant"``. See also: https://numpy.org/doc/stable/reference/generated/numpy.pad.html https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html kwargs: other arguments for the `np.pad` or `torch.pad` function. note that `np.pad` treats channel dimension as the first dimension. >emptymaximummean linear_rampminimum symmetricmedian)r#r%>rcircularrr replicateedgec3s|] }|tvVqdS)N)r&)r?kerrrr ls  zpad_nd..) supportedzunexpected keywordZ implementedr+N )r=r5r/int16int64booluint8rK ValueError TypeError RuntimeErrorr.NotImplementedErroranyr4r0device)r"rLr%r<_padrrXr rIs0 ":r torch.Tensor spatial_sizetuple[int, ...]c Kst|jd}tt|t|dd}tt||}ttjdd|Dddi}| t|df}t |t |t |d df}|j dd |jdd } } d gtd gd d f\} } } }t| | |jdd D]S\}}}| p~|d kp~||dk|p|d kp||dk} }| |d krd nt| ||dkrd nt||dfg7} | ttt|d t|d| dd g7} qn| rt|}t|| fd|i|}|r|| }|S)a Crop or pad using the translation matrix and spatial size. The translation coefficients are rounded to the nearest integers. For a more generic implementation, please see :py:class:`monai.transforms.SpatialResample`. Args: img: data to be transformed, assuming `img` is channel-first and padding doesn't apply to the channel dim. translation_mat: the translation matrix to be applied to the image. A translation matrix generated by, for example, :py:func:`monai.transforms.utils.create_translate`. The translation coefficients are rounded to the nearest integers. spatial_size: the spatial size of the output image. mode: the padding mode. kwargs: other arguments for the `np.pad` or `torch.pad` function. rDT) wrap_sequencecSsg|]}d|dgqS)g?r)r?xrrr rBsz"crop_or_pad_nd..indexingijr>N)axisrrFrr%)lenr4r9roundrr rFr:meshgridreshapefloor concatenate ones_likeminmaxslicezipintr!r)r"Ztranslation_matrir%r<ndim matrix_npccZsrc_ccZ src_startZsrc_endrLZto_cropdo_padZdo_cropsesp_moderrr rus$ & 2<0rFtuple[tuple[int, int]]transform_infodictlazyr_c Ks||d}t|tr|n|jdd}t|tr|nd}t|} | rgdd|D} t| t|jkrH| dgt|jt| 7} dd| ddD} t || } d dt || ddD} n|} t j t |dt d t jd } tj|| | ||||d }tt|tr|n|td }|rt|tr||S|S| rt|| |fi|n|}t|td }t|tr||S|S)a7 Functional implementation of padding a MetaTensor. This function operates eagerly or lazily according to ``lazy`` (default ``False``). `torch.nn.functional.pad` is used unless the mode or kwargs are not available in torch, in which case `np.pad` will be used. Args: img: data to be transformed, assuming `img` is channel-first and padding doesn't apply to the channel dim. to_pad: the amount to be padded in each dimension [(low_H, high_H), (low_W, high_W), ...]. note that it including channel dimension. transform_info: a dictionary with the relevant information pertaining to an applied transform. mode: available modes: (Numpy) {``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``, ``"mean"``, ``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``} (PyTorch) {``"constant"``, ``"reflect"``, ``"replicate"``, ``"circular"``}. One of the listed string values or a user supplied function. Defaults to ``"constant"``. See also: https://numpy.org/doc/stable/reference/generated/numpy.pad.html https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html lazy: a flag indicating whether the operation should be performed in a lazy fashion or not. transform_info: a dictionary with the relevant information pertaining to an applied transform. kwargs: other arguments for the `np.pad` or `torch.pad` function. note that `np.pad` treats channel dimension as the first dimension. )paddedr%rDNcSs$g|]}t|dt|dfqSrrD)r|)r?prrr rBrCzpad_func..rpcSsg|]}|d qS)rr)r?rrrr rBscSs g|] \}\}}|||qSrr)r?drrrrr rBs r7)rfr5sp_sizeaffine extra_info orig_sizerr track_meta)r.rpeek_pending_shaper4peek_pending_rankr9r:rerqr r{r/eyer|rffloat64rtrack_transform_metarrErcopy_meta_fromr)r"rLrr%rr<rimg_size spatial_rankrZ to_pad_listto_shiftxformr4 meta_infooutrrr rs8     rslicestuple[slice, ...]c CsZt|tr |n|jdd}t|tr|nd}tddt|dd|D}d| i}g}t t |ddD]!\} } | j dur\| | j dkrW|| | j n| j q@| dq@ddt|dd|D} tj|| t||||||d } tt|tr|n|td } |rt| tr| | S| S| |} t| tr| | S| S) aI Functional implementation of cropping a MetaTensor. This function operates eagerly or lazily according to ``lazy`` (default ``False``). Args: img: data to be transformed, assuming `img` is channel-first and cropping doesn't apply to the channel dim. slices: the crop slices computed based on specified `center & size` or `start & end` or `slices`. lazy: a flag indicating whether the operation should be performed in a lazy fashion or not. transform_info: a dictionary with the relevant information pertaining to an applied transform. rDNrcSs0g|]\}}||d|||dgqSrindicesr?rorrr rBs0zcrop_func..croppedrcSs,g|]\}}||d||dqS)rDrrrrrr rBs,rr)r.rrr4rr9r:r{flattentolist enumeraterstartappendrrr rrErr)r"rrrrrrrrirr4rrrrr rs0 " &  r)r"rr#r$r%r&r'r)r"rrLr$r%r&r'r)r"rhrirjr%r&) r"rhrLrrrr%r&rr_r'rh) r"rhrrrr_rrr'rh)&__doc__ __future__rr2r8r9r/torch.nn.functionalrrGmonai.config.type_definitionsrmonai.data.meta_objrmonai.data.meta_tensorrmonai.data.utilsrmonai.transforms.inversermonai.transforms.utilsr r monai.utilsr r r rr__all__r!r=rKrrrrrrrrr s0         ,& >