o  iܔ@sddlmZddlZddlZddlZddlZddlmZmZm Z m Z ddl Z ddl Z ddlmZddlmZddlmZddlmZmZmZddlmZmZed d e\ZZeeZed d d \Z ZGdddeZ!GdddeZ"GdddeZ#GdddeZ$GdddeZ%GdddeeZ&GdddeZ'GdddeeZ(GdddeZ)Gd d!d!eZ*Gd"d#d#eZ+Gd$d%d%eeZ,Gd&d'd'eZ-dS)() annotationsN)HashableMappingSequenceSized)KeysCollection) MetaTensor)GaussianFilter) MapTransform Randomizable Transform) min_versionoptional_importzskimage.measurez0.14.2z scipy.ndimagedistance_transform_cdt)namecs:eZdZ    ddfdd ZddZdddZZS)DiscardAddGuidanced?NFkeysrnumber_intensity_chint probabilityfloat label_names Sized | Noneallow_missing_keysboolcs(t||||_||_|pg|_dS)aE Discard positive and negative points according to discard probability Args: keys: The ``keys`` parameter will be used to get and set the actual data item to transform number_intensity_ch: number of intensity channels probability: probability of discarding clicks N)super__init__rdiscard_probabilityr)selfrrrrr __class__`/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/apps/deepedit/transforms.pyr&szDiscardAddGuidanced.__init__cCs|jdkstjjddg|jd|jgdrNtjt|j|jd|jd|jdftjd }|jd |j t|jkrE|||j dd f<|Stj ||gd d }|S) NrTFrpdtyper.axis) rnprandomchoicezeroslenrshapefloat32r concatenate)r imagesignalr#r#r$_apply<s&zDiscardAddGuidanced._applydataMapping[Hashable, np.ndarray]returndict[Hashable, np.ndarray]cCs^t|}||D]#}|dkr(|||}t||tr#|||_q |||<q tdq |S)Nr6z(This transform only applies to the image)dict key_iteratorr8 isinstancerarrayprint)r r9dkey tmp_imager#r#r$__call__Is   zDiscardAddGuidanced.__call__)rrNF) rrrrrrrrrrr9r:r;r<)__name__ __module__ __qualname__rr8rE __classcell__r#r#r!r$r$s rc,eZdZ ddfd d ZdddZZS)NormalizeLabelsInDatasetdNFrrrdict[str, int] | Nonerrcst|||p i|_dS)z Normalize label values according to label names dictionary Args: keys: The ``keys`` parameter will be used to get and set the actual data item to transform label_names: all label names N)rrrr rrrr!r#r$rYs z"NormalizeLabelsInDatasetd.__init__r9r:r;r<c Cst|}||D]I}i}t||j}t|jddD]\}\}}|dkr4|||<|||||k<|dkrr.r1r3 enumerateritemsr?rr@) r r9rBrCZnew_label_nameslabelidx key_label val_labelr#r#r$rEgs   z"NormalizeLabelsInDatasetd.__call__NF)rrrrMrrrFrGrHrIrrErJr#r#r!r$rLWsrLcrK)SingleLabelSelectiondNFrrrSequence[str] | Nonerrcs@t|||p g|_ddddddddd d d d d dd|_dS)z Selects one label at a time to train the DeepEdit Args: keys: The ``keys`` parameter will be used to get and set the actual data item to transform label_names: all label names r )Zspleenz right kidneyz left kidneyZ gallbladderZ esophagusZliverZstomachZaortazinferior vena cavaZ portal_veinZ splenic_veinZpancreaszright adrenal glandzleft adrenal glandN)rrrall_label_valuesrNr!r#r$rs"  zSingleLabelSelectiond.__init__r9r:r;r<cCst|}||D]D}|dkrHtj|j}||d<d|||||j|k<|j|d}|||||dk<td|d|| q t dq |S) NrSZ current_labelgrrz Using label z with number: z(This transform only applies to the label) r=r>r.r/r0rrhindexrAmaxwarningswarn)r r9rBrCZt_labelZ max_label_valr#r#r$rEs zSingleLabelSelectiond.__call__rW)rrrrZrrrFrXr#r#r!r$rY}srYcs>eZdZdZ    ddfdd ZddZdddZZS)AddGuidanceSignalDeepEditdaB Add Guidance signal for input image. Multilabel DeepEdit Based on the "guidance" points, apply Gaussian to them and add them as new channel for input image. Args: guidance: key to store guidance. sigma: standard deviation for Gaussian kernel. number_intensity_ch: channel index. guidancer\rFrrstrsigmarrrrcs$t||||_||_||_dSN)rrrnrpr)r rrnrprrr!r#r$rs z#AddGuidanceSignalDeepEditd.__init__c Cst|jdkr dnd}t|tjr|n|}t|tr!t|n|}t|r"|dkrBtj d|jd|jd|jdftj d}ntj d|jd|jdftj d}|j}|D]u}t t |dkreqX|dkrt dtt|d|dd}t dtt|d|dd}t dtt|d|dd} d |dd||| f<qXt dtt|d|dd}t dtt|d|dd}d |dd||f<qXt |ddkr t|d} tt| j|jd } | | dd} | dd} | |d<|dt|dt |dt|d|d<|S|dkr>tj d|jd|jd|jdftj d}|Stj d|jd|jdftj d}|S) Nr\r[rr'r(r)r*rr)rp)r2r3r?r.ndarraytolistrojsonloadsr1r4anyasarrayrjminrtorchtensorr rp unsqueezesqueezedetachcpunumpy) r r6rn dimensionsr7sshapepointp1p2p3Z signal_tensorZ pt_gaussianr#r#r$ _get_signals@ ,"     6 *"z&AddGuidanceSignalDeepEditd._get_signalr9r:r;r<c Cst|}||D]K}|dkrP||}|dd|jdf}||j}|D]$}||||}tj||gdd}t||t rG|||_ q'|||<q'|St dq |S)Nr6r.r,z(This transform only applies to image key) r=r>rrnrrr.r5r?rr@rA) r r9rBrCr6rDrnrUr7r#r#r$rEs     z#AddGuidanceSignalDeepEditd.__call__)rnr\rF) rrrnrorprrrrrrF)rGrHrI__doc__rrrErJr#r#r!r$rms -rmc6eZdZdZddfd d Zd d ZdddZZS)FindAllValidSlicesDeepEditdz Find/List all valid slices in the labels. Label is assumed to be a 4D Volume with shape CHWD, where C=1. Args: sids: key to store slices indices having valid label map. sidsFrrrrrct||||_dSrqrrrr rrrr!r#r$r z$FindAllValidSlicesDeepEditd.__init__cCs`i}|dD]%}g}t|jdD]}|d||dd|fvr(||q|||<q|S)Nrr)r.rranger3appendr rSrBrrUZl_idssidr#r#r$r8s  z"FindAllValidSlicesDeepEditd._applyr9r:r;r<cCt|}||D]=}|dkrB||}|jddkrtdt|jdkr)td|||}|dur>t|r>|||j<|Stdq |SNrSrrz$Only supports single channel labels!r]z$Only supports label with shape CHWD!(This transform only applies to label key r=r>r3 ValueErrorr2r8rrrAr r9rBrCrSrr#r#r$rE    z$FindAllValidSlicesDeepEditd.__call__rFrrrrrrrFrGrHrIrrr8rErJr#r#r!r$r s  rcHeZdZdZ     ddfdd ZddZddZdddZZS)AddInitialSeedPointDeepEditdai Add random guidance as initial seed point for a given label. Note that the label is of size (C, D, H, W) or (C, H, W) The guidance is of size (2, N, # of dims) where N is number of guidance added. # of dims = 4 when C, D, H, W; # of dims = 3 when (C, H, W) Args: guidance: key to store guidance. sids: key that represents lists of valid slice indices for the given label. sid: key that represents the slice to add initial seed point. If not present, random sid will be chosen. connected_regions: maximum connected regions to use for adding initial points. rnrrr^Frrroconnected_regionsrrrc2t||||_||_t|_||_||_dSrqrrsids_keysid_keyr=rrnrr rrnrrrrr!r#r$rD   z%AddInitialSeedPointDeepEditd.__init__c Cst|jdkr dnd}dg|d|_|}|dur*|dkr*d}|dd|ftj}|dktj}|dkr@tj|t ddn|}t |dkrPt d |g}t d|dkrZdn|j dD]}}|dkr|||ktj}t|dkr|||jq`t|} t| d } t|dkd} |jj| d| | t| | d } | | } tt| |jd}| d|d<|dks|dkr||q`||d|d |d|gq`t|gS) Nr\r[r)rr.?rPzSLICES NOT FOUND FOR LABEL: rsizer&r()r2r3default_guidancer.newaxisastyper4measurerSrrjAssertionErrorrrsumrrflattenexpwhereRr0rw unravel_index transposers)r rSrrUrdims blobs_labelsZ pos_guidanceridxdistancerrTseeddstgr#r#r$r8Ts8    $     z#AddInitialSeedPointDeepEditd._applycC||jdur||j|nd}||jdur$||j|nd}|dur=|r=|dus4||vr<|jj|dd}n td|d}||j|<dSNF)replacezNot slice IDs for label: getrrrr0loggerinforr rBrUrrr#r#r$ _randomize&& z'AddInitialSeedPointDeepEditd._randomizer9r:r;r<c Cst|}||D]a}|dkrfi}|dD]E}|||t||}|dkr7d||t|d|k<nd||t|d|k<d|}t| ||j || t ||<q|||j<|Stdq |SNrSrrPrrrrr=r>rrr.copyrrtdumpsr8rrrrrsrnrAr r9rBrCZlabel_guidancesrU tmp_labelr#r#r$rEs$   z%AddInitialSeedPointDeepEditd.__call__rnrrr^F rrrnrorrorrorrrrrF rGrHrIrrr8rrErJr#r#r!r$r4s- rcsHeZdZdZ   ddfd d Zed d ZddZdddZZ S)FindDiscrepancyRegionsDeepEditdz Find discrepancy between prediction and actual during click interactions during training. Args: pred: key to prediction source. discrepancy: key to store discrepancies found between label and prediction. pred discrepancyFrrrorrcst||||_||_dSrq)rrrr)r rrrrr!r#r$rs z(FindDiscrepancyRegionsDeepEditd.__init__cCs0||}|dktj}|dktj}||gS)Nr)rr.r4)rSr disparityZ pos_disparityZ neg_disparityr#r#r$rsz)FindDiscrepancyRegionsDeepEditd.disparitycCs |||Srq)r)r rSrr#r#r$r8s z&FindDiscrepancyRegionsDeepEditd._applyr9r:r;r<c Cs6t|}||D]}|dkri}t|dD]q\}\}}|dkrOt||}d|||k<|dktj}t||j} d| | |k<| dktj} n3t||}d|||k<d|}|dktj}t||j} d| | |k<d| } | dktj} | || ||<q|||j <|St dq |S)NrSrrPrrrz*This transform only applies to 'label' key) r=r>rQrRr.rrr4rr8rrA) r r9rBrCZall_discrepancies_rUrVrSrr#r#r$rEs2      z(FindDiscrepancyRegionsDeepEditd.__call__)rrF)rrrrorrorrrF) rGrHrIrr staticmethodrr8rErJr#r#r!r$rs  rcsPeZdZdZ    ddfd d ZdddZddZddZdddZZ S)AddRandomGuidanceDeepEditdab Add random guidance based on discrepancies that were found between label and prediction. Args: guidance: key to guidance source, shape (2, N, # of dim) discrepancy: key to discrepancy map between label and prediction shape (2, C, H, W, D) or (2, C, H, W) probability: key to click/interaction probability, shape (1) rnrrFrrrorrcsBt||||_||_||_d|_d|_d|_d|_i|_ dSrq) rr guidance_keyrr_will_interactis_posis_otherrrn)r rrnrrrr!r#r$rs z#AddRandomGuidanceDeepEditd.__init__NcCs,||j}|jjddg|d|gd|_dS)NTFrr%)rrr0r)r r9rr#r#r$ randomizes "z$AddRandomGuidanceDeepEditd.randomizecCst|}t|d}t|dkd}t|dkdkrQ|jj|d||t||d}||}tt ||j  d}|d|d<|SdS)Nrrrr) rrr.rrrrr0rwrr3rrs)r rrrrTrrrr#r#r$ find_guidance s $  z(AddRandomGuidanceDeepEditd.find_guidancec Csz|d}i}t|D]C\}\}} |dkr5t|} d| | | k<| dktj} t|d| ||<q t|} d| | | k<d| } t|d| ||<q t|dkrb|||d|_ | D]T}||dkrd|_ |dkrt|} d| | ||k<| dktj} |j |||d| qft|} d| | ||k<d| } |j |||d| qfdS)NrrPrrT2) rQrRr.rrr4rrrrrrrn) r rnrrlabelsZ pos_discrZother_discrepancy_areasrrUrVrr#r#r$ add_guidances:         z'AddRandomGuidanceDeepEditd.add_guidancer9r:r;r<c Cst|}||j}||j}|||jr|dD](}||}t|tjr,| n|}t|t r8t |n|}dd|D|j |<q|dD]}||j ||||d|dqKtdd}d}g} tt|d} | | vrzn-| | |t|j | }||kr|dD] }|| vrg|j |<qtd |nt| t|dkrtd |nqj|j ||j<|S) NrcSsg|]}d|vr|qS)r)r#).0jr#r#r$ Lsz7AddRandomGuidanceDeepEditd.__call__..rSrrcrTzNumber of simulated clicks: )r=rrrrrr?r.rrrsrortrurnrr/randintr0listrr2rr) r r9rBrnrrUZtmp_guiZ num_clickscounterZ keep_guidanceZ aux_labelr#r#r$rEAsD   $    z#AddRandomGuidanceDeepEditd.__call__)rnrrF) rrrnrorrorrorrrqrF) rGrHrIrrrrrrErJr#r#r!r$rs  'rc@s8eZdZdZ    ddd d ZeddZddZdS)AddGuidanceFromPointsDeepEditdaC Add guidance based on user clicks. ONLY WORKS FOR 3D We assume the input is loaded by LoadImaged and has the shape of (H, W, D) originally. Clicks always specify the coordinates in (H, W, D) Args: ref_image: key to reference image to fetch current and original image details. guidance: output key to store guidance. meta_keys: explicitly indicate the key of the metadata dictionary of `ref_image`. for example, for data with key `image`, the metadata by default is in `image_meta_dict`. the metadata is a dictionary object which contains: filename, original_shape, etc. if None, will try to construct meta_keys by `{ref_image}_{meta_key_postfix}`. meta_key_postfix: if meta_key is None, use `{ref_image}_{meta_key_postfix}` to fetch the metadata according to the key data, default is `meta_dict`, the metadata is a dictionary object. For example, to handle key `image`, read/write affine matrices from the metadata `image_meta_dict` dictionary's `affine` field. rnN meta_dict ref_imageror dict | None meta_keys str | Nonemeta_key_postfixcCs&||_||_|p i|_||_||_dSrq)rrnrrr)r rrnrrrr#r#r$rs   z'AddGuidanceFromPointsDeepEditd.__init__cCs&t|rt||t}|SgSrq)r2r.multiplyrrrs)clicksfactorrnr#r#r$r8sz%AddGuidanceFromPointsDeepEditd._applyc Cst|}|jp|jd|j}t||jtr||jj}n||vr(||}nt|dd|vr7td|d}t ||jj dd}t ||}i}|j D]} || g} t t | t} || ||| <qU|||j<|S)NrI is not found. Please check whether it is the correct the image meta key. spatial_shapez%Missing "spatial_shape" in meta_dict!r)r=rrrr?rmetar RuntimeErrorrr3r.r@rrrrrr8rn) r r9rB meta_dict_keyroriginal_shape current_shaper all_guidancesrUrr#r#r$rEs*   z'AddGuidanceFromPointsDeepEditd.__call__)rnNNr) rrornrorrrrrro)rGrHrIrrrr8rEr#r#r#r$rns   rc@s"eZdZdZd ddZd d Zd S) $ResizeGuidanceMultipleLabelDeepEditdzA Resize the guidance based on cropped vs resized image. rnrorr;NonecCs||_||_dSrq)rnr)r rnrr#r#r$rs z-ResizeGuidanceMultipleLabelDeepEditd.__init__c Cst|}||jjdd}d}t||jtr||jj}n||vr(||}nt|d|d}t||}i}||j  D]!} t ||j | r]t ||j | | tng} | || <qB|||j <|S)Nrimage_meta_dictrr)r=rr3r?rrrr.dividernrr2rrrrs) r r9rBrrrrrrrUrnr#r#r$rEs*  "  z-ResizeGuidanceMultipleLabelDeepEditd.__call__N)rnrorror;r)rGrHrIrrrEr#r#r#r$rs  rc@seZdZdZd ddZdS) SplitPredsLabeldz; Split preds and labels for individual evaluation r9r:r;r<cCst|}||D]E}|dkrEt|dD],\}\}}|dkrC|||ddfd|d|<|d|ddfd|d|<qq |dkrNtd q |S) NrrrPr.pred_rSlabel_zThis is only for pred key)r=r>rQrRrr)r r9rBrCrTrUrr#r#r$rEs"" zSplitPredsLabeld.__call__NrF)rGrHrIrrEr#r#r#r$rsrcr)!AddInitialSeedPointMissingLabelsdaf Add random guidance as initial seed point for a given label. Note that the label is of size (C, D, H, W) or (C, H, W) The guidance is of size (2, N, # of dims) where N is number of guidance added. # of dims = 4 when C, D, H, W; # of dims = 3 when (C, H, W) Args: guidance: key to store guidance. sids: key that represents lists of valid slice indices for the given label. sid: key that represents the slice to add initial seed point. If not present, random sid will be chosen. connected_regions: maximum connected regions to use for adding initial points. rnrrr^Frrrorrrrcrrqrrr!r#r$rrz*AddInitialSeedPointMissingLabelsd.__init__c Cst|jdkr dnd}dg|d|_|}|dur*|dkr*d}|dd|ftj}|dktj}|dkr@tj|t ddn|}g}t |dkrR| |jnt d|dkrZdn|j dD]}}|dkr|||ktj}t|dkr|| |jq`t|}t|d } t|dkd} |jj| d| | t| | d } || } tt| |jd} | d| d<|dks|dkr| | q`| | d| d | d|gq`t|S) Nr\r[r)rr.rrrrr()r2r3rr.rrr4rrSrrjrrrrrrrrrr0rwrrrs)r rSrrrrZlabel_guidancerrrrTrrrr#r#r$r8s8    $     z(AddInitialSeedPointMissingLabelsd._applycCrrrrr#r#r$rDrz,AddInitialSeedPointMissingLabelsd._randomizer9r:r;r<c Cst|}||D]`}|dkrei}|dD]D}|||t||}|dkr7d||t|d|k<nd||t|d|k<d|}t| ||j | t ||<q|||j<|Stdq |Srrrr#r#r$rEOs$   z*AddInitialSeedPointMissingLabelsd.__call__rrrFrr#r#r!r$rs0 rcr) FindAllValidSlicesMissingLabelsdz Find/List all valid slices in the labels. Label is assumed to be a 4D Volume with shape CHWD, where C=1. Args: sids: key to store slices indices having valid label map. rFrrrrrcrrqrrr!r#r$rqrz)FindAllValidSlicesMissingLabelsd.__init__cCsri}|dD].}g}t|jdD]}|d||dd|fvr(||q|gkr2dgd}|||<q|S)Nrr)r.rcrrr#r#r$r8us   z'FindAllValidSlicesMissingLabelsd._applyr9r:r;r<cCrrrrr#r#r$rErz)FindAllValidSlicesMissingLabelsd.__call__rrrFrr#r#r!r$ris  r). __future__rrtloggingr/rkcollections.abcrrrrrr.ry monai.configr monai.datarmonai.networks.layersr monai.transforms.transformr r r monai.utilsr rrr getLoggerrGrrrrLrYrmrrrrrrrrrr#r#r#r$s<     3&1[+rFL*r