o  i)@sddlmZddlZddlmZddlZddlZddlm Z m Z ddl m Z ddl mZddlmZddlmZgd ZdddZdddZGdddeZGdddeZGdddeZdS)) annotationsN)Sequence) DtypeLikeKeysCollection) MapLabelValue) MapTransform)$keep_components_with_positive_points)look_up_option)VistaPreTransformdVistaPostTransformdRelabeld labels_dict dict | NonereturndictcCs"i}|durdd|D}|S)z#get the label name to index mappingNcSsi|] \}}|t|qS)lowerint).0kvrr_/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/apps/vista3d/transforms.py !sz._get_name_to_index_mapping..)items)r name_to_index_mappingrrr_get_name_to_index_mappingsrr label_prompt list | NonecCs|durWt|trWg}|D]}t|tr&|s&||vr&t|||<q |D]+}t|ttfrO|t|trI|||rFt|ndnt|q)||q)|S|S)zconvert the label name to indexNr) isinstanceliststrisdigitrlenrappendget)rrZconverted_label_promptlrrr_convert_name_to_index%s 0 r&cs0eZdZ    ddfdd ZddZZS)r FuNkeysrallow_missing_keysbool special_index Sequence[int]r rsubclassrNonecs(t||||_||_t||_dS)a Pre-transform for Vista3d. It performs two functionalities: 1. If label prompt shows the points belong to special class (defined by special index, e.g. tumors, vessels), convert point labels from 0 (negative), 1 (positive) to special 2 (negative), 3 (positive). 2. If label prompt is within the keys in subclass, convert the label prompt to its subclasses defined by subclass[key]. e.g. "lung" label is converted to ["left lung", "right lung"]. The `label_prompt` is a list of int values of length [B] and `point_labels` is a list of length B, where each element is an int value of length [B, N]. Args: keys: keys of the corresponding items to be transformed. special_index: the index that defines the special class. subclass: a dictionary that maps a label prompt to its subclasses. allow_missing_keys: don't raise exception if key is missing. N)super__init__r1r3rr)selfr.r/r1r r3 __class__rrr6:szVistaPreTransformd.__init__cCs|dd}|dd}t|j|}zm|jdurO|durOg}ttt|j}tt |D]}|||vrC| |jt ||q.| ||q.||d<|durz|dur}|d|j vrst|}d||dk<d||dk<|}||d<W|SW|SW|StytdY|Sw)Nr point_labelsrzDVistaPreTransformd failed to transform label prompt or point labels.)r$r&rr3rmaprr.ranger"extendr r#r1nparraytolist Exceptionwarningswarn)r7datarr:Z _label_promptZ subclass_keysirrr__call__[s8          zVistaPreTransformd.__call__)Fr'NN) r.rr/r0r1r2r rr3rrr4__name__ __module__ __qualname__r6rI __classcell__rrr8rr 9s!r cs(eZdZd d fdd Zd d ZZS)r Fr.rr/r0rr4cst||dS)as Post-transform for Vista3d. It converts the model output logits into final segmentation masks. If `label_prompt` is None, the output will be thresholded to be sequential indexes [0,1,2,...], else the indexes will be [0, label_prompt[0], label_prompt[1], ...]. If `label_prompt` is None while `points` are provided, the model will perform postprocess to remove regions that does not contain positive points. Args: keys: keys of the corresponding items to be transformed. dataset_transforms: a dictionary specifies the transform for corresponding dataset: key: dataset name, value: list of data transforms. dataset_key: key to get the dataset name from the data dictionary, default to "dataset_name". allow_missing_keys: don't raise exception if key is missing. N)r5r6)r7r.r/r8rrr6|szVistaPostTransformd.__init__c CsL|jD]}||vr||}|jd}|j}|dddur=|dddur=t|d|d||d|dd}d||dk<|dkrbtj|dkdd d }| dd d }d||<nd ||dk<d|vr|ddur|d 7}|d|}t d|dD]}|d } ||d|j ||| k<qd||d k<|||<q|S) z)data["label_prompt"] should not contain 0rrNpointsr:) point_coordsr:gr=T)dimkeepdimg?g?) r.shapedevicer$r unsqueezetotorchallargmaxfloatr?dtype) r7rGr.predZ object_numrTZis_bkrrHfracrrrrIs:       zVistaPostTransformd.__call__)F)r.rr/r0rr4rJrrr8rr {sr cs0eZdZejddfdfdd ZddZZS)r dataset_nameFr.rlabel_mappings dict[str, list[tuple[int, int]]]r[r dataset_keyr r/r0rr4csXt||i|_||_|D]\}}tdd|Ddd|D|d|j|<qdS)aS Remap the voxel labels in the input data dictionary based on the specified mapping. This list of local -> global label mappings will be applied to each input `data[keys]`. if `data[dataset_key]` is not in `label_mappings`, label_mappings['default']` will be used. if `label_mappings[data[dataset_key]]` is None, no relabeling will be performed. Args: keys: keys of the corresponding items to be transformed. label_mappings: a dictionary specifies how local dataset class indices are mapped to the global class indices. The dictionary keys are dataset names and the values are lists of list of (local label, global label) pairs. This list of local -> global label mappings will be applied to each input `data[keys]`. If `data[dataset_key]` is not in `label_mappings`, label_mappings['default']` will be used. if `label_mappings[data[dataset_key]]` is None, no relabeling will be performed. Please set `label_mappings={}` to completely skip this transform. dtype: convert the output data to dtype, default to float32. dataset_key: key to get the dataset name from the data dictionary, default to "dataset_name". allow_missing_keys: don't raise exception if key is missing. cSg|]}t|dqS)rrrpairrrr z%Relabeld.__init__..cSrb)r=rcrdrrrrfrg) orig_labels target_labelsr[N)r5r6mappersrarr)r7r.r_r[rar/namemappingr8rrr6s  zRelabeld.__init__cCsVt|}||jd}t||jdd}|dur|S||D] }|||||<q|S)Ndefault)rm)rr$rar rj key_iterator)r7rGdr^_mkeyrrrrIszRelabeld.__call__) r.rr_r`r[rrar r/r0rr4)rKrLrMrAint16r6rIrNrrr8rr s &r )r rrr)rrrrrr) __future__rrEcollections.abcrnumpyrArW monai.configrrmonai.transformsrmonai.transforms.transformrmonai.transforms.utilsr monai.utilsr __all__rr&r r r rrrrs        B6