def __call__()

in seamseg/algos/rpn.py [0:0]

• 27 lines of code
• 7 McCabe index (conditional complexity)

    def __call__(self, boxes, scores, training):
        """Perform NMS-based selection of proposals

        Parameters
        ----------
        boxes : torch.Tensor
            Tensor of bounding boxes with shape N x M
        scores : torch.Tensor
            Tensor of bounding box scores with shape N x M x 4
        training : bool
            Switch between training and validation modes

        Returns
        -------
        proposals : PackedSequence
            Sequence of N tensors of selected bounding boxes with shape M_i x 4, entries can be None
        """
        if training:
            num_pre_nms = self.num_pre_nms_train
            num_post_nms = self.num_post_nms_train
        else:
            num_pre_nms = self.num_pre_nms_val
            num_post_nms = self.num_post_nms_val

        proposals = []
        for bbx_i, obj_i in zip(boxes, scores):
            try:
                # Optional size pre-selection
                if self.min_size > 0:
                    bbx_size = bbx_i[:, 2:] - bbx_i[:, :2]
                    valid = (bbx_size[:, 0] >= self.min_size) & (bbx_size[:, 1] >= self.min_size)

                    if valid.any().item():
                        bbx_i, obj_i = bbx_i[valid], obj_i[valid]
                    else:
                        raise Empty

                # Score pre-selection
                obj_i, idx = obj_i.topk(min(obj_i.size(0), num_pre_nms))
                bbx_i = bbx_i[idx]

                # NMS
                idx = nms(bbx_i, obj_i, self.nms_threshold, num_post_nms)
                if idx.numel() == 0:
                    raise Empty
                bbx_i = bbx_i[idx]

                proposals.append(bbx_i)
            except Empty:
                proposals.append(None)

        return PackedSequence(proposals)