Method forward

axlearn/vision/detection_generator.py:297–380

Generates final detections. Args: raw_boxes: A float tensor of shape [batch, num_boxes, num_classes * 4] representing box coordinates. For boxes shared among classes `num_classes` must be 1. raw_scores: A float tensor of shape [batch, num_boxes, num_c

(
        self,
        raw_boxes: Tensor,
        raw_scores: Tensor,
        anchor_boxes: Tensor,
        image_shape: Tensor,
    )

        self.box_coder = self.config.box_coder.instantiate()

    def forward(
        self,
        raw_boxes: Tensor,
        raw_scores: Tensor,
        anchor_boxes: Tensor,
        image_shape: Tensor,
    ) -> dict[str, Tensor]:
        """Generates final detections.

        Args:
            raw_boxes: A float tensor of shape [batch, num_boxes, num_classes * 4] representing box
                coordinates. For boxes shared among classes `num_classes` must be 1.
            raw_scores: A float tensor of shape [batch, num_boxes, num_classes] representing
                logits.
            anchor_boxes: A float tensor of shape [batch, num_boxes, 4] representing anchor boxes
                corresponding to the raw boxes.
            image_shape: A tensor of shape of [batch_size, 2] with the image height and width
                that are used to clip detection boxes that exceed the image boundaries.

        Returns:
            A dictionary with the following tensors:
                detection_boxes: A float tensor of shape [batch, max_num_detections, 4]
                    representing top detected boxes in [y1, x1, y2, x2].
                detection_scores: A float tensor of shape [batch, max_num_detections] representing
                    sorted confidence scores for detected boxes. The values are between [0, 1] and
                    obtained by applying softmax on the raw scores.
                detection_classes: An int Tensor of shape [batch, max_num_detections]
                    representing classes for detected boxes.
                num_detections: An int tensor of shape [batch] only the first num_detections
                    boxes are valid proposals.

            Note that if `apply_nms` is set to False, max_num_detections is equal to
            number of input raw boxes.
        """
        normalized_scores = jax.nn.softmax(raw_scores)
        batch, num_boxes, num_coordinates = raw_boxes.shape
        raw_boxes = jnp.reshape(raw_boxes, [batch, num_boxes, num_coordinates // 4, 4])
        scores_without_bg = normalized_scores[..., 1:]
        # If raw_boxes contain per class predictions, remove boxes corresponding to background
        # class.
        if raw_boxes.shape[-2] > 1:
            raw_boxes = raw_boxes[..., 1:, :]
        # Expand anchors boxes along `num_classes` dimension to be compatible with raw_boxes.
        anchor_boxes = anchor_boxes[..., None, :]
        decoded_boxes = self.box_coder.decode(encoded_boxes=raw_boxes, anchors=anchor_boxes)
        clipped_boxes = utils_detection.clip_boxes_jax(
            decoded_boxes, image_shape[..., None, None, :]
        )
        if not self.config.apply_nms:
            num_detections = jnp.tile(decoded_boxes.shape[1], decoded_boxes.shape[0])
            # [batch, num_detections]
            detection_class_ids = jnp.argmax(scores_without_bg, axis=-1)
            num_classes_for_boxes = raw_boxes.shape[-2]
            # `num_classes_for_boxes - 1` ensures we take the only set of boxes available when
            # boxes are shared among classes.
            box_class_ids = jnp.minimum(detection_class_ids, num_classes_for_boxes - 1)
            return {
                "detection_boxes": jnp.take_along_axis(