o  iia@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.morphologyz scipy.ndimagedistance_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,_/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/apps/nuclick/transforms.pyrsrcs,eZdZdZddfd d 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. Fkeysr connectivityintallow_missing_keysboolct||||_dSN)super__init__r1)selfr0r1r3 __class__r,r-r8= zFlattenLabeld.__init__cCsXt|}|jD]"}t||tjrt||n||}tj||jd t j ||<q|S)Nr1) dictr0 isinstancetorchTensorr measurerr1astypenpuint8r9datadkeyimgr,r,r-__call__As  $zFlattenLabeld.__call__)r/F)r0rr1r2r3r4rrrrr8rK __classcell__r,r,r:r-r.3s r.cs<eZdZdZejddfdfdd ZddZddZZ 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 Fr0r centroid_keystr patch_sizetuple[int, int] | intr3r4kwargsrc s$t||||_||_||_dSr6)r7r8rPrRrT)r9r0rPrRr3rTr:r,r-r8Xs zExtractPatchd.__init__c Cst|}||j}|j|jf}|jD]2}||}||j||jdd\}}} } |dd||| | f} tdd|i|j| ||<q|S)N spatial_sizer,)r>rPrRr0bboxshaper rT) r9rGrHrroi_sizerIrJx_startx_endy_starty_endcroppedr,r,r-rKes   "zExtractPatchd.__call__c Cs|\}}|\}}tt||dd}tt||dd} ||} | |} | |kr0|} ||}| |kr:|} ||} || | | fS)Nr)r2max) r9rRrsizexymnrZr\r[r]r,r,r-rWrs zExtractPatchd.bbox) r0rrPrQrRrSr3r4rTr) rrrrrr$r8rKrWrMr,r,r:r-rNIs  rNcs:eZdZdZejejdddfdfdd ZddZZ 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 rTr0rrrQr str | Nonemin_arear2 others_valueto_binary_maskr4cs2tj|dd||_||_||_||_||_dSNF)r3)r7r8rrrirjrk)r9r0rrrirjrkr:r,r-r8s   zSplitLabeld.__init__cCsxt|}t|jdkrtddS|jD]}t||tjr"||nt||}t|}|j r=||j }d|||k<nd|||j k<t t |}|j rTd||dk<t|}d|||k<d||dk<t|r|tjt|ddd}t|d}t|tjr|tjn|}t|tjr|tjn|}t||tjr|nt|||<t||tjr|nt|||j<q|S)Nr/z8Only 'label' key is supported, more than 1 key was foundrr=)r>lenr0printr?r@rA from_numpyclonerrjr2r`rk count_nonzerorBrr typerEr)r9rGrHrIrmaskrrr,r,r-rKs2 &        $zSplitLabeld.__call__) r0rrrQrrhrir2rjr2rkr4) rrrrrr"r%r8rKrMr,r,r:r-rfsrfcs\eZdZdZddfd d Zd d ZddZ d ddZd!ddZddZ d"ddZ 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. Fr0rmin_sizer2r3r4cr5r6)r7r8rv)r9r0rvr3r:r,r-r8r<zFilterImaged.__init__cCsJt|}|jD]}t||tjrt||n||}||||<q|Sr6)r>r0r?r@rAr filterrFr,r,r-rKs  $zFilterImaged.__call__cCsJ||}||}||@}|jr|j||jdn|}|t|||gS)Nrv)filter_green_channel filter_graysrvfilter_remove_small_objectsrDdstack)r9rgbZmask_not_greenZ mask_not_grayZmask_gray_greenrsr,r,r-rws  zFilterImaged.filterTZc Cst|dddddf}||k|dk@}||}||kr8|dkr8|dur8td|d|} ||| |||}|S)Nr/rTr_) mask_percentmathceilry) r9img_npZ green_threshavoid_overmaskovermask_thresh output_typegZ gr_ch_maskmask_percentageZnew_green_threshr,r,r-rys  z!FilterImaged.filter_green_channelcCst|dddddf|dddddf|k}t|dddddf|dddddf|k}t|dddddf|dddddf|k}||@|@S)Nrr/r_)abs)r9r} toleranceZrg_diffZrb_diffZgb_diffr,r,r-rzs444zFilterImaged.filter_grayscCst|jdkr=|jddkr=|dddddf|dddddf|dddddf}dt||jd}|Sdt||jd}|S)Nr_rr/d)rmrXrDrqra)r9rZnp_sumrr,r,r-rs BzFilterImaged.mask_percent _cCsVtj|t|d}||}||kr)|dkr)|dur)t|d}|||||}|S)Nrxr/Tr_) morphologyremove_small_objectsrCr4rroundr{)r9rrvrrZrem_smrZ new_min_sizer,r,r-r{s   z(FilterImaged.filter_remove_small_objects)ruF)r0rrvr2r3r4)r~Trr4)r)rTr) rrrrr8rKrwryrzrr{rMr,r,r:r-rts  rtc @s^eZdZdZejejejdddddddf d#ddZddZ ddZ ddZ ddZ d d!Z d"S)$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?@TrrQrr drop_ratefloat jitter_ranger2gaussianr4sigma truncatedadd_exclusion_map use_distancec CsLt||||_||_||_||_||_||_||_||_ | |_ | |_ dSr6) rr8rrrrrrrrrr) r9rrrrrrrrrrr,r,r-r8s  z AddPointGuidanceSignald.__init__cCs2t|}t||jtjr||jnt||j}t||jtjr(||jnt||j}|j|d|jd}| |}|j rzt||j tjrP||j nt||j }|j |d|j|j |jd}| |}tj||d|dfdd}n tj||dfdd}t||jtjr|nt|||j<|S)Nrdtype)rrrdim)r>r?rr@rAror inclusion_mapr_apply_gaussianrr exclusion_maprrcatr )r9rGrHrrsZinc_sigrZexc_sigr,r,r-rK8s,, , $z AddPointGuidanceSignald.__call__cCJ|jr t|dkr |Std|j|jd|dd}|ddSNrr_) spatial_dimsrrrr@rqrrr unsqueezesqueezer9trbr,r,r-rM"z'AddPointGuidanceSignald._apply_gaussianc Cstdus|js:ttdrt|dk}n tt|dk}t|dkr8|j dt|}||df||dffSdSt| }t |d}t | dkd}|jj |d||t||d}tt||jd}|d|dfS)Nargwhererr/r)raprU)r rhasattrr@rrDr rmRrandintflattenexpwherechoicesumasarray unravel_indexrX transposetolist) r9rindicesindexdistance probabilityidxseedrr,r,r- _seed_pointSs   $ z#AddPointGuidanceSignald._seed_pointcCs8tj||d}||}|durd||d|df<|S)Nrr/r)r@ zeros_liker)r9rsr point_maskptr,r,r-rhs  z%AddPointGuidanceSignald.inclusion_mapc Cstj||d}tjjddg|d|gdr|S|jdd}|jdd}tt|}|D]R} tjjddg|d|gdr?q.| j \} } t t | } t t | } |rz| |j j| |d} | |j j| |d} ttd| |} ttd| |} d|| | f<q.|S)NrTFr/)rr)lowhigh)r@rrDrandomrrXrB regionpropsr rr2rfloorrrminr`) r9rrrrrmax_xmax_ystatsstatrbrcr,r,r-rps& z%AddPointGuidanceSignald.exclusion_mapN)rrQrrQrrQrrrr2rr4rrrrrr4rr4)rrrrrr r!r"r8rKrrrrr,r,r,r-r s$  rc@sPeZdZdZejejdddddfdddZddZdddZ 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 rOFrrTrrQrbb_sizer2rr4rrrrcCs.||_||_||_||_||_||_||_dSr6)rrrrrrr)r9rrrrrrrr,r,r-r8s  zAddClickSignalsd.__init__c CsBt|}t||jtjr||jnt||j}|jd}|jd}|tj j d}|d|d}}||j } | rKt | ||ftng} dd| D} dd| D} |j|| | |||jd \} } | sotd |j|| | | | ||d }| |tjj <||tjj <||tjj <t||jtjr|nt|||j<|S) NrUr)rrrr/cSg|]}|dqSrr,.0xyr,r,r- z-AddClickSignalsd.__call__..cSr)r/r,rr,r,r-rr)cxcyrbrcbbz0Failed to create patches from given click points)rJ click_maprrrrbrc)r>r?rr@rArorXgetrr&valuerrDarrayrCr2rget_clickmap_boundingboxr ValueErrorget_patches_and_signalsr(r*r)r )r9rGrHrJrbrcrtxtyposrrrrpatchesr,r,r-rKs*,   $   $zAddClickSignalsd.__call__cs`t|d}fddttD}fddttD} t|| } t| t| d|f<g} ttD]g} td| |d} td| |d}t | |}t ||}|| |krt||} |||kr~||}|| |kr|||kr| | |||gqDt d| d|d|d|d |d |d qD|| fS) Nrc(h|]}|ks|dkr|qSrr,rirrbr,r- (z.crrr,rrrcr,r-rrr/r_zIgnore smaller sized bbox (z, z ) (Min size: rb)) r@rrangermlistunionrDdeleter`rappendrn)r9rJrrrbrcrr x_del_indices y_del_indices del_indicesrrrZr\r[r]r,rrrbrcr-rs*     0z)AddClickSignalsd.get_clickmap_boundingboxc shg}fddttD} fddttD} t| | } t| t| t|D]y\} } | d}| d}| d}| d}|dd||||f}t|d}d|| | f<|||||f}| |}|j r||dk |j }|||||f}| |}| t||d|dgq5| t||dgq5t|S)Ncrrr,rrr,r-rrz;AddClickSignalsd.get_patches_and_signals..crrr,rrr,r-rrrr/r_r)rrmrrrDr enumerater@rrrrrrrrstack)r9rJrrrrrbrcrrrrr bounding_boxrZr\r[r]patchZthis_click_maprZothers_click_map other_pointsr,rr-rs.      z(AddClickSignalsd.get_patches_and_signalscCrrrrr,r,r-rrz AddClickSignalsd._apply_gaussianN)rrQrrQrr2rr4rrrrrr4)rO) rrrrrr r#r8rKrrrr,r,r,r-rs   !rc s^eZdZdZejejejejdddddej f d fdd Z ddZ d!ddZ d"ddZ 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? Fr0rrrQrrrthreshrrvr2min_holedo_reconstructionr4r3rc sHt|| ||_||_||_||_||_||_||_| |_ | |_ dSr6) r7r8rrrrrrvrrr) r9r0rrrrrrvrrr3rr:r,r-r8s zPostFilterLabeld.__init__c Cst|}||j}||j}||j}||j}|jD]%}||tj }| ||j |j |j } |j| ||||dtj ||<q|S)N)r)r>rrrrrr0rCrDrEpost_processingrrvrgen_instance_map) r9rGrHrrrbrcrIrmasksr,r,r-rK2s     "zPostFilterLabeld.__call__cCsL||k}t|jdD]}tj|||d||<tj|||d||<q |S)Nrrx)area_threshold)rrXrrremove_small_holes)r9predsrrvrr rr,r,r-r@s z PostFilterLabeld.post_processingTNc Cstj||ftjd}t|D]K\}} ||} |r"|t|kr"||nd} |r(| n|d} || d| d| d| df} t| dk| | } | || d| d| d| df<q|S)Nrr/rr_r)rDzerosuint16rrmr) r9r rrbrcrr instance_maprrsrcZthis_mapr,r,r-rGs$&z!PostFilterLabeld.gen_instance_map)r0rrrQrrQrrQrrQrrrvr2rr2rr4r3r4rrQ)rrr)TN)rrrrrr'r(r)r*r+r8rKrrrMr,r,r:r-r s  rc,eZdZdZd dfd d Zd d ZZS)AddLabelAsGuidancedz Add Label as new guidance channel Args: source: label/source key which gets added as additional guidance channel rr0rsourcerQreturnNonectj|dd||_dSrl)r7r8r)r9r0rr:r,r-r8] zAddLabelAsGuidanced.__init__cCst|}|jD]_}t||tjr||nt||}t||jtjr*||jnt||j}|dk}t|jt|jkrD|d}tj || |j gt|jdd}t||tjr`|nt |||<q|S)Nrrr) r>r0r?r@rArorrmrXrrrrr )r9rGrHrIrrr,r,r-rKas &,$"zAddLabelAsGuidanced.__call__)r)r0rrrQrrrLr,r,r:r-rUrcr)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 rr0roffsetr2rrcrrl)r7r8r)r9r0rr:r,r-r8wrzSetLabelClassd.__init__cCsZt|}|jD]#}t||tjr||nt||}tt|}||j||<q|Sr6) r>r0r?r@rAror2r`r)r9rGrHrIrrr,r,r-rK{s  &zSetLabelClassd.__call__)r)r0rrr2rrrLr,r,r:r-rorr)# __future__rrtypingrnumpyrDr@ monai.configrrmonai.networks.layersrmonai.transformsrrr monai.utilsr r r rB_rr rr.rNrfrtrrrrrr,r,r,r-s.     :CG}J