o  iR ã@s ddlmZddlZddlZddlmZmZddlmZddl Z ddl Z ddl m Z dZ dZdd gd gd gd œZd gZd,dd„Z        d-d.d+d „ZdS)/é)Ú annotationsN)ÚCallableÚSequence)ÚAny)ÚTensorT)éééééé9é€r r rr )ééé„Úsample_prompt_pairsÚidÚintÚreturnútuple[int, int]cCs|tvrtrdSdS)N)éé)rr)Ú SPECIAL_INDEXÚENABLE_SPECIAL)r©rú\/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/apps/vista3d/sampler.pyÚ_get_point_label"s rréFçš™™™™™É?ÚlabelsrÚ label_setú Sequence[int]Ú max_promptú int | NoneÚmax_forepromptÚmax_backpromptÚ max_pointÚinclude_backgroundÚboolÚdrop_label_probÚfloatÚdrop_point_probÚ point_samplerúCallable | NoneÚpoint_sampler_kwargsrúAtuple[Tensor | None, Tensor | None, Tensor | None, Tensor | None]c ! Ks |jddks tdƒ‚|d}|j} | ¡ ¡ ¡ ¡} |r-tt| ƒt| ƒt|ƒƒ} ntt| ƒt| ƒt|ƒdhƒ} tt|ƒt| ƒƒ} |dur]t | ƒ|kr]t   | ¡| d|…} |durrt | ƒ|krrt   | ¡| d|…} |dur™t | ƒt | ƒ|kr™t | ƒ|krt   | |¡} g} n t   | |t | ƒ¡} g}g}| durZt |tt t jd|dd¡ƒdƒ}t |tt t jd|dd¡ƒƒ}| D]m}t|ƒ\}}|t|ƒk}|}t |¡}t |¡}t t |ƒ|ƒ}t t |ƒ|ƒ}| t t j||dt j||dtjgd ¢| d g||||¡¡| t |g||g|d g||||¡ | ¡¡qÈ| D] }| t ||d ¡ | ¡¡| t ||¡ | ¡d¡q8n1| | fi| ¤Ž\}}| D]$}| t t |dƒd ¡ | ¡¡| t t |dƒ¡ | ¡d¡qft | ƒdkr t | ƒdkr d \}}}}n`t | | ¡ d ¡ | ¡ ¡}t |¡}t |¡}t |¡}t  dd¡|krót | ƒdkród}t | ƒ} |dt |ƒ| …}|dt |ƒ| …}|dt |ƒ| …}n t  dd¡|krd}d}||||fS)aÈ Sample training pairs for VISTA3D training. Args: labels: [1, 1, H, W, D], ground truth labels. label_set: the label list for the specific dataset. Note if 0 is included in label_set, it will be added into automatic branch training. Recommend removing 0 from label_set for multi-partially-labeled-dataset training, and adding 0 for finetuning specific dataset. The reason is region with 0 in one partially labeled dataset may contain foregrounds in another dataset. max_prompt: int, max number of total prompt, including foreground and background. max_foreprompt: int, max number of prompt from foreground. max_backprompt: int, max number of prompt from background. max_point: maximum number of points for each object. include_background: if include 0 into training prompt. If included, background 0 is treated the same as foreground and points will be sampled. Can be true only if user want to segment background 0 with point clicks, otherwise always be false. drop_label_prob: probability to drop label prompt. drop_point_prob: probability to drop point prompt. point_sampler: sampler to augment masks with supervoxel. point_sampler_kwargs: arguments for point_sampler. Returns: tuple: - label_prompt (Tensor | None): Tensor of shape [B, 1] containing the classes used for training automatic segmentation. - point (Tensor | None): Tensor of shape [B, N, 3] representing the corresponding points for each class. Note that background label prompts require matching points as well (e.g., [0, 0, 0] is used). - point_label (Tensor | None): Tensor of shape [B, N] representing the corresponding point labels for each point (negative or positive). -1 is used for padding the background label prompt and will be ignored. - prompt_class (Tensor | None): Tensor of shape [B, 1], exactly the same as label_prompt for label indexing during training. If label_prompt is None, prompt_class is used to identify point classes. rrzonly support batch size 1)rrNr)ÚmuÚsigma)Úk)rrr)Údeviceéÿÿÿÿr)NNNN) ÚshapeÚ ValueErrorr3ÚuniqueÚcpuÚnumpyÚtolistÚlistÚsetÚlenÚrandomÚshuffleÚsampleÚminrÚnpÚabsÚgaussrÚtorchÚnonzeroÚappendÚstackÚchoicesÚtensorÚtoÚzerosÚ unsqueezeÚlongÚcopyÚdeepcopyÚuniform)!rr r"r$r%r&r'r)r+r,r.r3Ú unique_labelsZbackground_labelsÚ_pointÚ _point_labelÚnum_pÚnum_nrZneg_idZpos_idÚplabelsÚnlabelsÚ plabelpointsÚ nlabelpointsZnum_paZnum_naÚ_Z label_promptÚpointÚ point_labelÚ prompt_classÚpadrrrr)s”4"         &"      ÿ"þÿÿ0ÿÿ"ý &     )rrrr)NNrrFrrN)rrr r!r"r#r$r#r%rr&rr'r(r)r*r+r*r,r-r.rrr/)Ú __future__rrOr>Úcollections.abcrrÚtypingrr9rBrErrrZ MERGE_LISTÚ__all__rrrrrrÚs2   ý  ö