o  iG#@sbdZddlmZddlmZddlZddlmZddlmZm Z m Z ddl m Z Gdd d Z dS) z{ Part of this script is adapted from https://github.com/pytorch/vision/blob/main/torchvision/models/detection/retinanet.py ) annotations)CallableN)Tensor) batched_nmsbox_iouclip_boxes_to_image) floor_dividec@s<eZdZdZedddddfd!ddZd"ddZd#ddZd S)$ BoxSelectora Box selector which selects the predicted boxes. The box selection is performed with the following steps: #. For each level, discard boxes with scores less than self.score_thresh. #. For each level, keep boxes with top self.topk_candidates_per_level scores. #. For the whole image, perform non-maximum suppression (NMS) on boxes, with overlapping threshold nms_thresh. #. For the whole image, keep boxes with top self.detections_per_img scores. Args: apply_sigmoid: whether to apply sigmoid to get scores from classification logits score_thresh: no box with scores less than score_thresh will be kept topk_candidates_per_level: max number of boxes to keep for each level nms_thresh: box overlapping threshold for NMS detections_per_img: max number of boxes to keep for each image Example: .. code-block:: python input_param = { "apply_sigmoid": True, "score_thresh": 0.1, "topk_candidates_per_level": 2, "nms_thresh": 0.1, "detections_per_img": 5, } box_selector = BoxSelector(**input_param) boxes = [torch.randn([3,6]), torch.randn([7,6])] logits = [torch.randn([3,3]), torch.randn([7,3])] spatial_size = (8,8,8) selected_boxes, selected_scores, selected_labels = box_selector.select_boxes_per_image( boxes, logits, spatial_size ) Tg?ig?i,box_overlap_metricr apply_sigmoidbool score_threshfloattopk_candidates_per_levelint nms_threshdetections_per_imgcCs(||_||_||_||_||_||_dS)N)r r r rrr)selfr r r rrrri/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/apps/detection/utils/box_selector.py__init__Xs   zBoxSelector.__init__logitsrreturntuple[Tensor, Tensor, Tensor]c Cs|jd}|jrt|tj}n|}||jk}||}t|d}t |j | d}|tj |\}}||}||} t ||} | || fS)a Select indices with highest scores. The indices selection is performed with the following steps: #. If self.apply_sigmoid, get scores by applying sigmoid to logits. Otherwise, use logits as scores. #. Discard indices with scores less than self.score_thresh #. Keep indices with top self.topk_candidates_per_level scores Args: logits: predicted classification logits, Tensor sized (N, num_classes) Return: - topk_idxs: selected M indices, Tensor sized (M, ) - selected_scores: selected M scores, Tensor sized (M, ) - selected_labels: selected M labels, Tensor sized (M, ) r)shaper torchsigmoidtofloat32flattenr whereminrsizetopkr) rr num_classesscores keep_idxsZflatten_topk_idxsnum_topkselected_scoresidxsselected_labels topk_idxsrrrselect_top_score_idx_per_levelis     z*BoxSelector.select_top_score_idx_per_level boxes_list list[Tensor] logits_list spatial_sizelist[int] | tuple[int]cCs"t|t|krtdt|dt|g}g}g}|dj}|dj}t||D],\} } || \} } } | | } t| |dd\} }|| || ||| |q+tj|dd}tj|dd}tj|dd}t ||||j |j |j d}|| |}|| |}||}|||fS)a6 Postprocessing to generate detection result from classification logits and boxes. The box selection is performed with the following steps: #. For each level, discard boxes with scores less than self.score_thresh. #. For each level, keep boxes with top self.topk_candidates_per_level scores. #. For the whole image, perform non-maximum suppression (NMS) on boxes, with overlapping threshold nms_thresh. #. For the whole image, keep boxes with top self.detections_per_img scores. Args: boxes_list: list of predicted boxes from a single image, each element i is a Tensor sized (N_i, 2*spatial_dims) logits_list: list of predicted classification logits from a single image, each element i is a Tensor sized (N_i, num_classes) spatial_size: spatial size of the image Return: - selected boxes, Tensor sized (P, 2*spatial_dims) - selected_scores, Tensor sized (P, ) - selected_labels, Tensor sized (P, ) zFlen(boxes_list) should equal to len(logits_list). Got len(boxes_list)=z, len(logits_list)=rT) remove_empty)dim)r max_proposals)len ValueErrordtypezipr-rappendrcatrrr rr)rr.r0r1 image_boxes image_scores image_labelsZ boxes_dtypeZ logits_dtypeboxes_per_levellogits_per_levelr,scores_per_levellabels_per_levelkeepZ image_boxes_tZimage_scores_tZimage_labels_tkeep_tselected_boxesr)r+rrrselect_boxes_per_imagesJ     z"BoxSelector.select_boxes_per_imageN) r rr r r rrrrrrr)rrrr)r.r/r0r/r1r2rr)__name__ __module__ __qualname____doc__rrr-rFrrrrr 3s& +r )rJ __future__rcollections.abcrrrmonai.data.box_utilsrrr0monai.transforms.utils_pytorch_numpy_unificationrr rrrrs"