o / iQ @sddlmZddlZddlmZmZmZmZddlm Z ddl Z ddl Z ddl mmZddlmZddlmZmZmZddlmZmZmZmZmZmZmZmZm Z e ddd \Z!Z"d Z#d gZ$d ej%d ej%dddddddddf dDd4d Z&d5d6Z'd7d8Z(dEd>d?Z)d@dAZ*dFdBdCZ+dS)G) annotationsN)CallableIterableMappingSequence)Any) MetaTensor)compute_importance_mapdense_patch_slicesget_valid_patch_size) BlendModePytorchPadModeconvert_data_typeconvert_to_dst_type ensure_tupleensure_tuple_repfall_back_tuplelook_up_optionoptional_importtqdm)namez nearest-exactsliding_window_inferenceg?g?gFinputstorch.Tensor | MetaTensorroi_sizeSequence[int] | int sw_batch_sizeint predictorNCallable[..., torch.Tensor | Sequence[torch.Tensor] | dict[Any, torch.Tensor]]overlapSequence[float] | floatmodeBlendMode | str sigma_scale padding_modePytorchPadMode | strcvalfloat sw_devicetorch.device | str | Nonedeviceprogressboolroi_weight_maptorch.Tensor | None process_fnCallable | None buffer_steps int | None buffer_dim with_coordargsrkwargsreturnAtorch.Tensor | tuple[torch.Tensor, ...] | dict[Any, torch.Tensor]cQ s |duo|dk}tjd}|r1|| ks||kr)td| d|d|d|dkr1||7}t||}|D]}|dksB|dkrJtd |dq8j}j^}}| pXj} | p]j} |d dd}ttrttgj d d }t t j d ddt tfddt|D}g}ttjdddD]}t|dj|d}|d}||||gqt|rtj|t|t|ddurtj|t|t|dt|||}t||| dt|} |sd }!td| |}"n-t||||\}#}$}"t jd}!}%|$d|#D]}&|&d|%kr&d }!n|&d}%qt|}'|'kr>| dur>| }(n>zt|'})t|)||| |d}(t|(j|kr[| s[|(d}(Wnty{}*zt d|)d|d|d| d |*d}*~*wwt |(t j | |dd}(gggddf\}+},}-}.}/| rt!|"n|"D]:}0t|0t"|0||r|$|.dn| }1fdd|1D}2|dkrt #fdd|2D$| }3durt #fdd|2D$| }4|4|d <n|2d$| }3dur|2d$| }4|4|d <|r||3|2g|Ri|}5n ||3g|Ri|}5t%|5\}6}7| r1| |7|3|(\}7}8n|(}8t|8j|kr?|8d}8|8j$|| d }8|r|$|.dd\}9}:|-st|7djd}t&|};|:|9|;|<t j'd|g|;|| d!g}-t(|7d|2D].\}<}=|=|dj)|9}>t*|>|>||=|d<t*dd|=d<|-d|=|<|87<q{|/t|27}/|/|$|.dkrqnt&|7}-tt|-D]}?|-|?j}@|@d|@dd}A}Bd}C|s|Bkrd"dt(|BD}Ctj+|8|Bt,d#}8t|+|?krh||Ag}D|D|Crd$dt(||CDnt&|7}D|!rt j-nt j'}E|+.|E|D|| d |,.t j'ddg|Ddd|| d |8$| }FD]'}G|CdurStd%dt(|G|CD}G|,dt*dt*dg|GR|F7<q@|rt*dgtj}Ht*|9|:|H|d<|.|#}It*|I|Id|Hd<|!r|+d|Hj/|-d|!d&q|+d|H|-dj$| d'7<q|-|?|89<|-|?$| |-|?<t0|2|C|+|?|-|?qg}-|r|.d7}.q|!rt j12tt|+D]}?|+|?|,d<qt|rb|3d(|it4|+D]Z\}?}Jd)dt(|JjddD}Kg}Lt|D]0}M||Md}Nt*t5t6||Md|K|Nt5t6||Md|N|K|N}O|L7d|Oq |Jt*dt*dg|LR|+|?<qt8|+|6}P|durwt9|P|| d'd}P|PSt9|P| d'd}P|PS)*a Sliding window inference on `inputs` with `predictor`. The outputs of `predictor` could be a tensor, a tuple, or a dictionary of tensors. Each output in the tuple or dict value is allowed to have different resolutions with respect to the input. e.g., the input patch spatial size is [128,128,128], the output (a tuple of two patches) patch sizes could be ([128,64,256], [64,32,128]). In this case, the parameter `overlap` and `roi_size` need to be carefully chosen to ensure the output ROI is still an integer. If the predictor's input and output spatial sizes are not equal, we recommend choosing the parameters so that `overlap*roi_size*output_size/input_size` is an integer (for each spatial dimension). When roi_size is larger than the inputs' spatial size, the input image are padded during inference. To maintain the same spatial sizes, the output image will be cropped to the original input size. Args: inputs: input image to be processed (assuming NCHW[D]) roi_size: the spatial window size for inferences. When its components have None or non-positives, the corresponding inputs dimension will be used. if the components of the `roi_size` are non-positive values, the transform will use the corresponding components of img size. For example, `roi_size=(32, -1)` will be adapted to `(32, 64)` if the second spatial dimension size of img is `64`. sw_batch_size: the batch size to run window slices. predictor: given input tensor ``patch_data`` in shape NCHW[D], The outputs of the function call ``predictor(patch_data)`` should be a tensor, a tuple, or a dictionary with Tensor values. Each output in the tuple or dict value should have the same batch_size, i.e. NM'H'W'[D']; where H'W'[D'] represents the output patch's spatial size, M is the number of output channels, N is `sw_batch_size`, e.g., the input shape is (7, 1, 128,128,128), the output could be a tuple of two tensors, with shapes: ((7, 5, 128, 64, 256), (7, 4, 64, 32, 128)). In this case, the parameter `overlap` and `roi_size` need to be carefully chosen to ensure the scaled output ROI sizes are still integers. If the `predictor`'s input and output spatial sizes are different, we recommend choosing the parameters so that ``overlap*roi_size*zoom_scale`` is an integer for each dimension. overlap: Amount of overlap between scans along each spatial dimension, defaults to ``0.25``. mode: {``"constant"``, ``"gaussian"``} How to blend output of overlapping windows. Defaults to ``"constant"``. - ``"constant``": gives equal weight to all predictions. - ``"gaussian``": gives less weight to predictions on edges of windows. sigma_scale: the standard deviation coefficient of the Gaussian window when `mode` is ``"gaussian"``. Default: 0.125. Actual window sigma is ``sigma_scale`` * ``dim_size``. When sigma_scale is a sequence of floats, the values denote sigma_scale at the corresponding spatial dimensions. padding_mode: {``"constant"``, ``"reflect"``, ``"replicate"``, ``"circular"``} Padding mode for ``inputs``, when ``roi_size`` is larger than inputs. Defaults to ``"constant"`` See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html cval: fill value for 'constant' padding mode. Default: 0 sw_device: device for the window data. By default the device (and accordingly the memory) of the `inputs` is used. Normally `sw_device` should be consistent with the device where `predictor` is defined. device: device for the stitched output prediction. By default the device (and accordingly the memory) of the `inputs` is used. If for example set to device=torch.device('cpu') the gpu memory consumption is less and independent of the `inputs` and `roi_size`. Output is on the `device`. progress: whether to print a `tqdm` progress bar. roi_weight_map: pre-computed (non-negative) weight map for each ROI. If not given, and ``mode`` is not `constant`, this map will be computed on the fly. process_fn: process inference output and adjust the importance map per window buffer_steps: the number of sliding window iterations along the ``buffer_dim`` to be buffered on ``sw_device`` before writing to ``device``. (Typically, ``sw_device`` is ``cuda`` and ``device`` is ``cpu``.) default is None, no buffering. For the buffer dim, when spatial size is divisible by buffer_steps*roi_size, (i.e. no overlapping among the buffers) non_blocking copy may be automatically enabled for efficiency. buffer_dim: the spatial dimension along which the buffers are created. 0 indicates the first spatial dimension. Default is -1, the last spatial dimension. with_coord: whether to pass the window coordinates to ``predictor``. Default is False. If True, the signature of ``predictor`` should be ``predictor(patch_data, patch_coord, ...)``. args: optional args to be passed to ``predictor``. kwargs: optional keyword args to be passed to ``predictor``. Note: - input must be channel-first and have a batch dim, supports N-D sliding window. Nrzbuffer_dim must be in [z, z], got .z"overlap must be >= 0 and < 1, got conditionF) copy_attrT) wrap_sequencec3s"|] }t||VqdSN)max.0i) image_size_rV/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/inferers/utils.py  z+sliding_window_inference..r)padr#value) return_slice)r#r%r,dtype)NNz patch size z, mode=z, sigma_scale=z , device=z^ Seems to be OOM. Please try smaller patch size or mode='constant' instead of mode='gaussian'.)r,rNcs:g|]}t||dtdgt|qS)r=N)slicelist)rDidx)num_winslicesrGrH s,z,sliding_window_inference..cg|]}|qSrGrGrDZ win_slice)rrGrHrTcrUrGrGrV)r>rGrHrTrW)rNr,)sizerNr,cSsg|] \}}|t|qSrG)r))rDZout_w_iZin_w_irGrGrHrT)r#cSsg|] \}}t||qSrG)r)rD_i_zrGrGrHrTrYcss2|]\}}tt|j|t|j|VqdSrA)rOrstartstop)rD_siZz_srGrGrHrIs0) non_blocking)r,pad_sizecSsg|]\}}||qSrGrG)rDZ_shape_dZ _roi_size_drGrGrHrT:s):lenshape ValueErrorrrNr,pop isinstancercopy_meta_fromrtorchTensorrtuplerangerBextendanyFrKrr _get_scan_intervalr _create_buffered_slicescuda is_availabler rr Exception RuntimeErrorrmincatto_flatten_structrPzeroszipr\rO interpolate _nearest_modeemptyappendcopy__compute_coordscurrent_stream synchronizeupdate enumeraterroundinsert _pack_structr)Qrrrrr!r#r%r&r(r*r,r-r/r1r3r5r6r7r8bufferednum_spatial_dimso compute_dtype batch_size_Z temp_meta image_sizer`kdiffhalf scan_intervalZ total_slicesr_ windows_range n_per_batchb_slices_ssxZvalid_patch_sizeZimportance_map_Z valid_p_sizeeZoutput_image_listZcount_map_listZsw_device_bufferb_sb_iZslice_gZ slice_rangeZ unravel_sliceZwin_dataZ win_conditionZ seg_prob_out dict_keysZ seg_tupleZw_tc_startc_endsp_sizepsoffsetssb_shapeZseg_chnsZ seg_shapez_scale output_shape new_tensorZw_t_Z__sZo_sliceZimg_boutput_iZ zoom_scaleZ final_slicingspsi slice_dimZ final_outputrG)r>rFrrRrrSrHr*sN_            &        &&  *$   ""  cs`t|}|tj|dd|dfdd}dd|D}|dd|f}tj|dddfddd\}}}t|}dg|ddtt|t|d |d kr\ |d td fd dt |D} g} t | D]/\} } || dkr| | d j t|nddf} || j d t|d f}| | j | |fqut j| } || | fS) zrearrange slices for bufferingNr mergesort)kindcSsg|] }tdd|DqS)css"|] }t|d|dVqdS)rr=N)rO)rDcrGrGrHrIRrJz5_create_buffered_slices...)rirCrGrGrHrTRsz+_create_buffered_slices..T) return_counts return_indexrr=c sJg|]!}tD]}t|d||d|dqqS)rr=)rj)rDbrErrrrGrHrT[s.)npasarrayargsortuniquecumsumtolistrtrarr}rjrr] itertoolschain)rSrrr5r3Z slices_nprZ_b_lensZb_endsrr_s_rZs_sZs_erGrrHroNs(  "" *  roc Cst||D]9\}}t|}|r6tdt|D]}tt||j||dt||j||d||<q|||7<qdS)zKsliding window batch spatial scaling indexing for multi-resolution outputs.r;N)ryrPrjrarOrr\r])coordsroutpatch original_idxrZidx_zmaxisrGrGrHris0 rr Sequence[int]rSequence[float]tuple[int, ...]cCst||krtdt|d|dt||kr&tdt|d|dg}tt||D]+\}}||||krE|t||q/t||d|}||dkrW|ndq/t|S)z Compute scan interval according to the image size, roi size and overlap. Scan interval will be `int((1 - overlap) * roi_size)`, if interval is 0, use 1 instead to make sure sliding window works. zlen(image_size) z different from spatial dims r<zlen(roi_size) r=r)rarcryrjr}rri)rrrr!rrErintervalrGrGrHrnus  rncsbd}ttjrf}||fSttr)t}tfdd|D}||fSt}||fS)Nc3s|]}|VqdSrArG)rDrseg_outrGrHrIsz"_flatten_struct..)rergrhrsortedkeysrir)rrZ seg_probsrGrrHrws  rwcCs@|dur tt||St|ttfrt|dkr|dSt|S)Nr=r)dictryrerPrirar)rrrGrGrHrs r)(rrrrrrrr r!r"r#r$r%r"r&r'r(r)r*r+r,r+r-r.r/r0r1r2r3r4r5rr6r.r7rr8rr9r:) rrrrrrr!rr9rrA), __future__rrcollections.abcrrrrtypingrnumpyrrgtorch.nn.functionalnn functionalrmmonai.data.meta_tensorrmonai.data.utilsr r r monai.utilsr r rrrrrrrrrr{__all__CONSTANTrrorrnrwrrGrGrGrHsD   ,  &