U Phta@sBddlmZddlZddlmZddlZddlZddlm Z m Z ddl m Z ddl mZmZmZddlmZmZmZed\ZZed \ZZed d d \ZZGd ddeZGdddeZGdddeZGdddeZGdddeZGdddeeZGdddeZGdddeZ GdddeZ!Gdd d eZ"dS)!) annotationsN)Any)KeysCollectionNdarrayOrTensor)GaussianFilter) MapTransform Randomizable SpatialPad)StrEnumconvert_to_numpyoptional_importzskimage.measurezskimage.morphologyzscipy.ndimage.morphologydistance_transform_cdt)namec@s@eZdZdZdZdZdZdZdZdZ dZ d Z d Z d Z d Zd ZdS) NuclickKeysz& Keys for nuclick transforms. imagelabelothers foregroundcentroid mask_valuelocation nuc_pointsbounding_boxes img_height img_width pred_classesN)__name__ __module__ __qualname____doc__IMAGELABELOTHERS FOREGROUNDCENTROID MASK_VALUELOCATION NUC_POINTSBOUNDING_BOXES IMG_HEIGHT IMG_WIDTH PRED_CLASSESr,r,R/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/apps/nuclick/transforms.pyrsrcs4eZdZdZd ddddfdd Zd d ZZS) FlattenLabelda FlattenLabeld creates labels per closed object contour (defined by a connectivity). For e.g if there are 12 small regions of 1's it will delineate them into 12 different label classes Args: connectivity: Max no. of orthogonal hops to consider a pixel/voxel as a neighbor. Refer skimage.measure.label allow_missing_keys: don't raise exception if key is missing. Frintbool)keys connectivityallow_missing_keyscst||||_dSN)super__init__r3)selfr2r3r4 __class__r,r-r7=szFlattenLabeld.__init__cCsXt|}|jD]D}t||tjr.t||n||}tj||jd t j ||<q|S)Nr3) dictr2 isinstancetorchTensorr measurerr3astypenpuint8r8datadkeyimgr,r,r-__call__As  $zFlattenLabeld.__call__)r/Frrrrr7rI __classcell__r,r,r9r-r.3s r.csHeZdZdZejddfdddddd fd d Zd d ZddZZ S) ExtractPatchda Extracts a patch from the given image and label, however it is based on the centroid location. The centroid location is a 2D coordinate (H, W). The extracted patch is extracted around the centroid, if the centroid is towards the edge, the centroid will not be the center of the image as the patch will be extracted from the edges onwards Args: keys: image, label centroid_key: key where the centroid values are stored, defaults to ``"centroid"`` patch_size: size of the extracted patch allow_missing_keys: don't raise exception if key is missing. pad_kwargs: other arguments for the SpatialPad transform Frstrztuple[int, int] | intr1r)r2 centroid_key patch_sizer4kwargsc s$t||||_||_||_dSr5)r6r7rOrPrQ)r8r2rOrPr4rQr9r,r-r7XszExtractPatchd.__init__c Cst|}||j}|j|jf}|jD]d}||}||j||jdd\}}} } |dd||| | f} tfd|i|j| ||<q$|S)N spatial_size)r<rOrPr2bboxshaper rQ) r8rErFrroi_sizerGrHx_startx_endy_starty_endcroppedr,r,r-rIes   "zExtractPatchd.__call__c Cs|\}}|\}}tt||dd}tt||dd} ||} | |} | |kr`|} ||}| |krt|} ||} || | | fS)Nr)r0max) r8rPrsizexymnrWrYrXrZr,r,r-rTrszExtractPatchd.bbox) rrrrrr$r7rIrTrKr,r,r9r-rLIs  rLcsHeZdZdZejejdddfdddddd d fd d Zd dZZ S) SplitLabeldap Extracts a single label from all the given classes, the single label is defined by mask_value, the remaining labels are kept in others Args: label: key of the label source others: other labels storage key, defaults to ``"others"`` mask_value: the mask_value that will be kept for binarization of the label, defaults to ``"mask_value"`` min_area: The smallest allowable object size. others_value: Value/class for other nuclei; Use this to separate core nuclei vs others. to_binary_mask: Convert mask to binary; Set it false to restore original class values rTrrNz str | Noner0r1)r2rrmin_area others_valueto_binary_maskcs2tj|dd||_||_||_||_||_dSNF)r4)r6r7rrrerfrg)r8r2rrrerfrgr9r,r-r7s zSplitLabeld.__init__cCst|}t|jdkr"tddS|jD]R}t||tjrF||n t||}t|}|j r|||j }d|||k<nd|||j k<t t |}|j rd||dk<t|}d|||k<d||dk<t|rtjt|ddd}t|d}t|tjr|tjn|}t|tjr2|tjn|}t||tjrL|nt|||<t||tjrn|nt|||j<q(|S)Nr/z8Only 'label' key is supported, more than 1 key was foundrr;)r<lenr2printr=r>r? from_numpyclonerrfr0r]rg count_nonzeror@rr typerCr)r8rErFrGrmaskrrr,r,r-rIs2 &       "&zSplitLabeld.__call__) rrrrrr"r%r7rIrKr,r,r9r-rcsrccsbeZdZdZdddddfdd Zd d Zd d ZdddZdddZddZ dddZ Z S) FilterImageda Filters Green and Gray channel of the image using an allowable object size, this pre-processing transform is specific towards NuClick training process. More details can be referred in this paper Koohbanani, Navid Alemi, et al. "NuClick: a deep learning framework for interactive segmentation of microscopic images." Medical Image Analysis 65 (2020): 101771. Args: min_size: The smallest allowable object size allow_missing_keys: don't raise exception if key is missing. Frr0r1)r2min_sizer4cst||||_dSr5)r6r7rr)r8r2rrr4r9r,r-r7szFilterImaged.__init__cCsJt|}|jD]6}t||tjr.t||n||}||||<q|Sr5)r<r2r=r>r?r filterrDr,r,r-rIs  $zFilterImaged.__call__cCsJ||}||}||@}|jr2|j||jdn|}|t|||gS)Nrr)filter_green_channel filter_graysrrfilter_remove_small_objectsrBdstack)r8rgbZmask_not_greenZ mask_not_grayZmask_gray_greenror,r,r-rss  zFilterImaged.filterTZc Cst|dddddf}||k|dk@}||}||krp|dkrp|dkrptd|d|} ||| |||}|S)Nr/rTr\) mask_percentmathceilru) r8img_npZ green_threshavoid_overmaskovermask_thresh output_typegZ gr_ch_maskmask_percentageZnew_green_threshr,r,r-rus z!FilterImaged.filter_green_channelcCst|dddddf|dddddf|k}t|dddddf|dddddf|k}t|dddddf|dddddf|k}||@|@S)Nrr/r\)abs)r8ry toleranceZrg_diffZrb_diffZgb_diffr,r,r-rvs444zFilterImaged.filter_grayscCst|jdkrx|jddkrx|dddddf|dddddf|dddddf}dt||jd}ndt||jd}|S)Nr\rr/d)rirUrBrmr^)r8rZnp_sumrr,r,r-r}s BzFilterImaged.mask_percent _cCsVtj|t|d}||}||krR|dkrR|dkrRt|d}|||||}|S)Nrtr/Tr\) morphologyremove_small_objectsrAr1r}roundrw)r8rrrrrZrem_smrZ new_min_sizer,r,r-rws   z(FilterImaged.filter_remove_small_objects)rqF)rzTr{r1)r)rTr) rrrrr7rIrsrurvr}rwrKr,r,r9r-rps  rpc @steZdZdZejejejdddddddf dddd d d d d d d d d dZddZ ddZ ddZ ddZ ddZ dS)AddPointGuidanceSignalda Adds Guidance Signal to the input image Args: image: key of source image, defaults to ``"image"`` label: key of source label, defaults to ``"label"`` others: source others (other labels from the binary mask which are not being used for training) defaults to ``"others"`` drop_rate: probability of dropping the signal, defaults to ``0.5`` jitter_range: noise added to the points in the point mask for exclusion mask, defaults to ``3`` gaussian: add gaussian sigma: sigma value for gaussian truncated: spreads how many stds for gaussian add_exclusion_map: add exclusion map/signal g?rF?@TrNfloatr0r1) rrr drop_rate jitter_rangegaussiansigma truncatedadd_exclusion_map use_distancec CsLt||||_||_||_||_||_||_||_||_ | |_ | |_ dSr5) rr7rrrrrrrrrr) r8rrrrrrrrrrr,r,r-r7s z AddPointGuidanceSignald.__init__cCs4t|}t||jtjr$||jnt||j}t||jtjrP||jnt||j}|j|d|jd}| |}|j rt||j tjr||j nt||j }|j |d|j|j |jd}| |}tj||d|dfdd}ntj||dfdd}t||jtjr"|nt|||j<|S)Nrdtype)rrrdim)r<r=rr>r?rkr inclusion_mapr_apply_gaussianrr exclusion_maprrcatr )r8rErFrroZinc_sigrZexc_sigr,r,r-rI8s$,, , &z AddPointGuidanceSignald.__call__cCsJ|jrt|dkr|Std|j|jd|dd}|ddSNrr\) spatial_dimsrrrr>rmrrr unsqueezesqueezer8tr_r,r,r-rMs"z'AddPointGuidanceSignald._apply_gaussianc Cstdks|jstttdr(t|dk}ntt|dk}t|dkrp|j dt|}||df||dffSdSt| }t |d}t | dkd}|jj |d||t||d}tt||jd}|d|dfS)Nargwhererr/r)r^prR)r rhasattrr>rrBr riRrandintflattenexpwherechoicesumasarray unravel_indexrU transposetolist) r8rindicesindexdistance probabilityidxseedrr,r,r- _seed_pointSs   $ z#AddPointGuidanceSignald._seed_pointcCs8tj||d}||}|dk r4d||d|df<|S)Nrr/r)r> zeros_liker)r8ror point_maskptr,r,r-rhs  z%AddPointGuidanceSignald.inclusion_mapc Cstj||d}tjjddg|d|gdr.|S|jdd}|jdd}tt|}|D]} tjjddg|d|gdr~q\| j \} } t t | } t t | } |r| |j j| |d} | |j j| |d} ttd| |} ttd| |} d|| | f<q\|S)NrTFr/)rr)lowhigh)r>rrBrandomrrUr@ regionpropsr rr0r~floorrrminr]) r8rrrrrmax_xmax_ystatsstatr_r`r,r,r-rps& z%AddPointGuidanceSignald.exclusion_mapN)rrrrrr r!r"r7rIrrrrr,r,r,r-r s""rc @s`eZdZdZejejdddddfdddd d d d d d d ZddZdddZ ddZ ddZ dS)AddClickSignalsda Adds Click Signal to the input image Args: image: source image, defaults to ``"image"`` foreground: 2D click indices as list, defaults to ``"foreground"`` bb_size: single integer size, defines a bounding box like (bb_size, bb_size) gaussian: add gaussian sigma: sigma value for gaussian truncated: spreads how many stds for gaussian add_exclusion_map: add exclusion map/signal rMFrrTrNr0r1rrrbb_sizerrrrcCs.||_||_||_||_||_||_||_dSr5r)r8rrrrrrrr,r,r-r7s zAddClickSignalsd.__init__c CsDt|}t||jtjr$||jnt||j}|jd}|jd}|tj j d}|d|d}}||j } | rt | ||ftng} dd| D} dd| D} |j|| | |||jd \} } | std |j|| | | | ||d }| |tjj <||tjj <||tjj <t||jtjr2|nt|||j<|S) NrRr)rrrr/cSsg|] }|dqSrr,.0xyr,r,r- sz-AddClickSignalsd.__call__..cSsg|] }|dqS)r/r,rr,r,r-rs)cxcyr_r`bbz0Failed to create patches from given click points)rH click_maprrrr_r`)r<r=rr>r?rkrUgetrr&valuerrBarrayrAr0rget_clickmap_boundingboxr ValueErrorget_patches_and_signalsr(r*r)r )r8rErFrHr_r`rtxtyposrrrrpatchesr,r,r-rIs6,   $   &zAddClickSignalsd.__call__csdt|d}fddttD}fddttD} t|| } t| t| d|f<g} ttD]} td| |d} td| |d}t | |}t ||}|| |kr||} |||kr||}|| |kr,|||kr,| | |||gqt d| d|d|d|d |d |d q|| fS) Nrcs(h|] }|ks |dkr|qSrr,rirr_r,r- s z.cs(h|] }|ks |dkr|qSrr,rrr`r,r-rs r/r\zIgnore smaller sized bbox (z, z ) (Min size: r_)) r>rrangerilistunionrBdeleter]rappendrj)r8rHrrr_r`rr x_del_indices y_del_indices del_indicesrrrWrYrXrZr,rrr_r`r-rs*     0z)AddClickSignalsd.get_clickmap_boundingboxc sjg}fddttD} fddttD} t| | } t| t| t|D]\} } | d}| d}| d}| d}|dd||||f}t|d}d|| | f<|||||f}| |}|j rF||dk |j }|||||f}| |}| t||d|dgqj| t||dgqjt|S)Ncs(h|] }|ks |dkr|qSrr,rrr,r-rs z;AddClickSignalsd.get_patches_and_signals..cs(h|] }|ks |dkr|qSrr,rrr,r-rs rr/r\r)rrirrrBr enumerater>rrrrnrrrstack)r8rHrrrrr_r`rrrrr bounding_boxrWrYrXrZpatchZthis_click_maprZothers_click_map other_pointsr,rr-rs.     z(AddClickSignalsd.get_patches_and_signalscCsJ|jrt|dkr|Std|j|jd|dd}|ddSrrrr,r,r-rs"z AddClickSignalsd._apply_gaussianN)rM) rrrrrr r#r7rIrrrr,r,r,r-rs !rc sveZdZdZejejejejdddddej f ddddddd d d d dd fd d Z ddZ dddZ dddZ ZS)PostFilterLabeldaV Performs Filtering of Labels on the predicted probability map Args: thresh: probability threshold for classifying a pixel as a mask min_size: min_size objects that will be removed from the image, refer skimage remove_small_objects min_hole: min_hole that will be removed from the image, refer skimage remove_small_holes do_reconstruction: Boolean Flag, Perform a morphological reconstruction of an image, refer skimage allow_missing_keys: don't raise exception if key is missing. pred_classes: List of Predicted class for each instance Q? FrrNrr0r1) r2rrrrthreshrrmin_holedo_reconstructionr4rc sHt|| ||_||_||_||_||_||_||_| |_ | |_ dSr5) r6r7rrrrrrrrrr) r8r2rrrrrrrrrr4rr9r,r-r7szPostFilterLabeld.__init__c Cst|}||j}||j}||j}||j}|jD]J}||tj }| ||j |j |j } |j| ||||dtj ||<q8|S)N)r)r<rrrrrr2rArBrCpost_processingrrrrgen_instance_map) r8rErFrrr_r`rGrmasksr,r,r-rI2s     "zPostFilterLabeld.__call__cCsL||k}t|jdD]0}tj|||d||<tj|||d||<q|S)Nrrt)area_threshold)rrUrrremove_small_holes)r8predsrrrrrrr,r,r-r@s z PostFilterLabeld.post_processingTNc Cstj||ftjd}t|D]\}} ||} |rD|t|krD||nd} |rP| n|d} || d| d| d| df} t| dk| | } | || d| d| d| df<q|S)Nrr/rr\r)rBzerosuint16rrir) r8rrr_r`rr instance_maprrorcZthis_mapr,r,r-rGs$&z!PostFilterLabeld.gen_instance_map)rrr)TN)rrrrrr'r(r)r*r+r7rIrrrKr,r,r9r-r s( rcs4eZdZdZd ddddfdd Zd d ZZS) AddLabelAsGuidancedz Add Label as new guidance channel Args: source: label/source key which gets added as additional guidance channel rrrNNone)r2sourcereturncstj|dd||_dSrh)r6r7r )r8r2r r9r,r-r7]szAddLabelAsGuidanced.__init__cCst|}|jD]}t||tjr*||n t||}t||jtjrT||jnt||j}|dk}t|jt|jkr|d}tj || |j gt|jdd}t||tjr|nt |||<q|S)Nrrr) r<r2r=r>r?rkr rirUrrnrr )r8rErFrGrrr,r,r-rIas &,$"zAddLabelAsGuidanced.__call__)rrJr,r,r9r-rUsrcs4eZdZdZd ddddfdd Zd d ZZS) SetLabelClassdz Assign class value from the labelmap. This converts multi-dimension tensor to single scalar tensor. Args: offset: offset value to be added to the mask value to determine the final class rrr0r )r2offsetr cstj|dd||_dSrh)r6r7r )r8r2r r9r,r-r7wszSetLabelClassd.__init__cCsZt|}|jD]F}t||tjr*||n t||}tt|}||j||<q|Sr5) r<r2r=r>r?rkr0r]r )r8rErFrGrrr,r,r-rI{s  &zSetLabelClassd.__call__)rrJr,r,r9r-r osr )# __future__rr~typingrnumpyrBr> monai.configrrmonai.networks.layersrmonai.transformsrrr monai.utilsr r r r@_rr rr.rLrcrprrrrr r,r,r,r- s,     :CG}J