o  iB6@sNdZddlmZddlZddlZddlmZGdddZGdddeZdS) z The functions in this script are adapted from nnDetection, https://github.com/MIC-DKFZ/nnDetection/blob/main/nndet/core/boxes/sampler.py ) annotationsN)Tensorc@s&eZdZdZddddZdddZdS)HardNegativeSamplerBasea Base class of hard negative sampler. Hard negative sampler is used to suppress false positive rate in classification tasks. During training, it select negative samples with high prediction scores. The training workflow is described as the follows: 1) forward network and get prediction scores (classification prob/logits) for all the samples; 2) use hard negative sampler to choose negative samples with high prediction scores and some positive samples; 3) compute classification loss for the selected samples; 4) do back propagation. Args: pool_size: when we need ``num_neg`` hard negative samples, they will be randomly selected from ``num_neg * pool_size`` negative samples with the highest prediction scores. Larger ``pool_size`` gives more randomness, yet selects negative samples that are less 'hard', i.e., negative samples with lower prediction scores. pool_sizefloatreturnNonecCs ||_dS)Nr)selfrr r/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/apps/detection/utils/hard_negative_sampler.py__init__<s z HardNegativeSamplerBase.__init__negativernum_negintfg_probsc Cs||kr tdt||j}t||}||tjj|ddd\}}||}tj ||j dd|}||} tj |tj d} d| | <| S) a; Select hard negative samples. Args: negative: indices of all the negative samples, sized (P,), where P is the number of negative samples num_neg: number of negative samples to sample fg_probs: maximum foreground prediction scores (probability) across all the classes for each sample, sized (A,), where A is the number of samples. Returns: binary mask of negative samples to choose, sized (A,), where A is the number of samples in one image zSThe number of negative samples should not be larger than the number of all samples.rT)dimsorteddeviceNdtype) numel ValueErrorrrmintotorchfloat32topkrandpermr zeros_likeuint8) r rrrpool_Znegative_idx_poolZ hard_negativeperm2Zselected_neg_idxneg_maskr r r select_negatives?s z(HardNegativeSamplerBase.select_negativesN)r)rrrr )rrrrrrrr)__name__ __module__ __qualname____doc__rr(r r r r r(s rcsbeZdZdZ d(d)fd d Zd*ddZd+ddZd,ddZd-dd Zd.d#d$Z d/d&d'Z Z S)0HardNegativeSamplera HardNegativeSampler is used to suppress false positive rate in classification tasks. During training, it selects negative samples with high prediction scores. The training workflow is described as the follows: 1) forward network and get prediction scores (classification prob/logits) for all the samples; 2) use hard negative sampler to choose negative samples with high prediction scores and some positive samples; 3) compute classification loss for the selected samples; 4) do back propagation. Args: batch_size_per_image: number of training samples to be randomly selected per image positive_fraction: percentage of positive elements in the selected samples min_neg: minimum number of negative samples to select if possible. pool_size: when we need ``num_neg`` hard negative samples, they will be randomly selected from ``num_neg * pool_size`` negative samples with the highest prediction scores. Larger ``pool_size`` gives more randomness, yet selects negative samples that are less 'hard', i.e., negative samples with lower prediction scores. rrbatch_size_per_imagerpositive_fractionrmin_negrrr cs.tj|d||_||_||_tddS)Nr z,Sampling hard negatives on a per batch basis)superrr0r.r/logginginfo)r r.r/r0r __class__r r rxs zHardNegativeSampler.__init__ target_labels list[Tensor]concat_fg_probsr!tuple[list[Tensor], list[Tensor]]cCs&dd|D}||d}|||S)a Select positives and hard negatives from list samples per image. Hard negative sampler will be applied to each image independently. Args: target_labels: list of labels per image. For image i in the batch, target_labels[i] is a Tensor sized (A_i,), where A_i is the number of samples in image i. Positive samples have positive labels, negative samples have label 0. concat_fg_probs: concatenated maximum foreground probability for all the images, sized (R,), where R is the sum of all samples inside one batch, i.e., R = A_0 + A_1 + ... Returns: - list of binary mask for positive samples - list of binary mask for negative samples Example: .. code-block:: python sampler = HardNegativeSampler( batch_size_per_image=6, positive_fraction=0.5, min_neg=1, pool_size=2 ) # two images with different number of samples target_labels = [ torch.tensor([0,1]), torch.tensor([1,0,2,1])] concat_fg_probs = torch.rand(6) pos_idx_list, neg_idx_list = sampler(target_labels, concat_fg_probs) cSsg|]}|jdqS)r)shape).0Zsamples_in_imager r r sz0HardNegativeSampler.__call__..r)splitselect_samples_img_list)r r6r8samples_per_imagerr r r __call__s  zHardNegativeSampler.__call__rc Csvg}g}t|t|krtdt|dt|dt||D]\}}|||\}}||||q ||fS)a% Select positives and hard negatives from list samples per image. Hard negative sampler will be applied to each image independently. Args: target_labels: list of labels per image. For image i in the batch, target_labels[i] is a Tensor sized (A_i,), where A_i is the number of samples in image i. Positive samples have positive labels, negative samples have label 0. fg_probs: list of maximum foreground probability per images, For image i in the batch, target_labels[i] is a Tensor sized (A_i,), where A_i is the number of samples in image i. Returns: - list of binary mask for positive samples - list binary mask for negative samples Example: .. code-block:: python sampler = HardNegativeSampler( batch_size_per_image=6, positive_fraction=0.5, min_neg=1, pool_size=2 ) # two images with different number of samples target_labels = [ torch.tensor([0,1]), torch.tensor([1,0,2,1])] fg_probs = [ torch.rand(2), torch.rand(4)] pos_idx_list, neg_idx_list = sampler.select_samples_img_list(target_labels, fg_probs) zDRequire len(target_labels) == len(fg_probs). Got len(target_labels)=z, len(fg_probs)=.)lenrzipselect_samples_per_imgappend) r r6rpos_idxneg_idxlabels_per_imgfg_probs_per_imgpos_idx_per_image_maskneg_idx_per_image_maskr r r r>s"  z+HardNegativeSampler.select_samples_img_listrHrItuple[Tensor, Tensor]c Csv||kr tdt|dkd}t|dkd}||}||||}|||}||||}||fS)at Select positives and hard negatives from samples. Args: labels_per_img: labels, sized (A,). Positive samples have positive labels, negative samples have label 0. fg_probs_per_img: maximum foreground probability, sized (A,) Returns: - binary mask for positive samples, sized (A,) - binary mask for negative samples, sized (A,) Example: .. code-block:: python sampler = HardNegativeSampler( batch_size_per_image=6, positive_fraction=0.5, min_neg=1, pool_size=2 ) # two images with different number of samples target_labels = torch.tensor([1,0,2,1]) fg_probs = torch.rand(4) pos_idx, neg_idx = sampler.select_samples_per_img(target_labels, fg_probs) zHlabels_per_img and fg_probs_per_img should have same number of elements.rr)rrrwhere get_num_posselect_positives get_num_negr() r rHrIpositivernum_posrJrrKr r r rDs  z*HardNegativeSampler.select_samples_per_imgrQ torch.TensorcCs"t|j|j}t||}|S)z Number of positive samples to draw Args: positive: indices of positive samples Returns: number of positive sample )rr.r/rr)r rQrRr r r rNs zHardNegativeSampler.get_num_posrrRcCs>ttd|tddt|j}t|t||j}|S)a Sample enough negatives to fill up ``self.batch_size_per_image`` Args: negative: indices of positive samples num_pos: number of positive samples to draw Returns: number of negative samples rg?)rmaxabsrr/rrr0)r rrRrr r r rPs$ zHardNegativeSampler.get_num_neglabelscCsX||kr tdtj||jdd|}||}tj|tjd}d||<|S)a Select positive samples Args: positive: indices of positive samples, sized (P,), where P is the number of positive samples num_pos: number of positive samples to sample labels: labels for all samples, sized (A,), where A is the number of samples. Returns: binary mask of positive samples to choose, sized (A,), where A is the number of samples in one image zSThe number of positive samples should not be larger than the number of all samples.rNrr)rrrr!rr"r#)r rQrRrVperm1pos_idx_per_imagerJr r r rOsz$HardNegativeSampler.select_positives)rr) r.rr/rr0rrrrr )r6r7r8rrr9)r6r7rr7rr9)rHrrIrrrL)rQrSrr)rrSrRrrr)rQrrRrrVrrr) r)r*r+r,rr@r>rDrNrPrO __classcell__r r r4r r-cs 0 ' r-)r, __future__rr2rrrr-r r r r s  ;