o  i"@sddlmZddlZddlmZddlmZddlZddlm Z ddl m Z e ddd\Z Z d gZ    d0d1d"d Zd2d'd(Zd3d4d*d+Z d5d6d.d/ZdS)7) annotationsN)Sequence)Any) MetaTensor)optional_importtqdm)namepoint_based_window_infererTinputstorch.Tensor | MetaTensorroi_size Sequence[int] predictortorch.nn.Module point_coords torch.Tensor point_labels class_vectortorch.Tensor | None prompt_class prev_mask torch.Tensor | MetaTensor | None point_startint center_onlyboolmarginkwargsrreturnc # Ks|jddks tdt|jdkstdtt||\} } |t| d| d| dg|j }|d urBtt||dnd }d }|d|d D]}t |d|d| jd | | d \}}t |d|d| jd| | d \}}t |d |d | jd | | d \}}t t|D]}t t|D]}t t|D]}||||||||||||f\}}}}}}t d t d t t |t |t t |t |t t |t |g}| |} || f|||||g|d| }!|d ur!tjd|!jd| jd | jd| jd gdd}tjd|!jd| jd | jd| jd gdd}"|||!d7<d|"|<qqqqN||"}|d d d d | d| jd | d| d | jd| d| d| jd | df}|"d d d d | d| jd | d| d | jd| d| d| jd | df}"|d ur|d d d d | d| jd | d| d | jd| d| d| jd | df}|d}||"dk||"dk<t|tjrt|}t|dst||jd|jd}|S)a Point-based window inferer that takes an input image, a set of points, and a model, and returns a segmented image. The inferer algorithm crops the input image into patches that centered at the point sets, which is followed by patch inference and average output stitching, and finally returns the segmented mask. Args: inputs: [1CHWD], input image to be processed. 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: the model. For vista3D, the output is [B, 1, H, W, D] which needs to be transposed to [1, B, H, W, D]. Add transpose=True in kwargs for vista3d. point_coords: [B, N, 3]. Point coordinates for B foreground objects, each has N points. point_labels: [B, N]. Point labels. 0/1 means negative/positive points for regular supported or zero-shot classes. 2/3 means negative/positive points for special supported classes (e.g. tumor, vessel). class_vector: [B]. Used for class-head automatic segmentation. Can be None value. prompt_class: [B]. The same as class_vector representing the point class and inform point head about supported class or zeroshot, not used for automatic segmentation. If None, point head is default to supported class segmentation. prev_mask: [1, B, H, W, D]. The value is before sigmoid. An optional tensor of previously segmented masks. point_start: only use points starting from this number. All points before this number is used to generate prev_mask. This is used to avoid re-calculating the points in previous iterations if given prev_mask. center_only: for each point, only crop the patch centered at this point. If false, crop 3 patches for each point. margin: if center_only is false, this value is the distance between point to the patch boundary. Returns: stitched_output: [1, B, H, W, D]. The value is before sigmoid. Notice: The function only supports SINGLE OBJECT INFERENCE with B=1. rz(Only supports single object point click.r zInput image should be 5D.iN)rr)rrrr patch_coordsrcpu)devicemetaaffine)r,r+)shape ValueErrorlen_pad_previous_maskcopydeepcopytorchtensortor(_get_window_idxrangeslicerzeros isinstanceTensorrhasattrr+)#r r rrrrrrrrrrimagepadZstitched_outputpZlx_Zrx_Zly_Zry_Zlz_Zrz_ijklxrxlyrylzrz unravel_sliceZ batch_imageoutputZ stitched_maskrK\/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/apps/vista3d/inferer.pyr s-( $$$4 &&  VV   r?roistuple[int, int]cCsr||ddkrd|}}||fS||d|kr$|||}}||fSt||dt||d}}||fS)z(Helper function to get the window index.r$r)r)r?rMrNleftrightrKrKrL_get_window_idx_cs "rRtuple[list[int], list[int]]c Csht|||\}}|r|g|gfStd|||}t||||}||||g} ||||g} | | fS)zGet the window index.r)rRmaxmin) r?rMrNrrrPrQZ left_mostZ right_most left_list right_listrKrKrLr6s r6padvalue+tuple[torch.Tensor | MetaTensor, list[int]]cCs|g}tt|jdddD]}t||d|j|d}|d}||||gq t|r:tjjj ||d|d}||fS)zHelper function to pad inputs.r r%r$rconstant)r>modevalue) r7r/r-rTextendanyr3nn functionalr>)r r rXpad_sizerBdiffhalfrKrKrLr0sr0)NNNrTr )r r r rrrrrrrrrrrrrrrrrrrrrrr)r?rrMrrNrrrO)Tr ) r?rrMrrNrrrrrrrS)r)r r r rrXrrrY) __future__rr1collections.abcrtypingrr3monai.data.meta_tensorr monai.utilsrr___all__r rRr6r0rKrKrKrLs(       s