Repository: Project-MONAI/MONAI
File: monai/apps/detection/networks/retinanet_network.py (Python)

Code before:
        """
        # compute features maps list from the input images.
        features = self.feature_extractor(images)
        if isinstance(features, Tensor):
            feature_maps = [features]
        elif torch.jit.isinstance(features, Dict[str, Tensor]):
            feature_maps = list(features.values())
        else:
            feature_maps = list(features)

        if not isinstance(feature_maps[0], Tensor):
            raise ValueError("feature_extractor output format must be Tensor, Dict[str, Tensor], or Sequence[Tensor].")

        # compute classification and box regression maps from the feature maps
        # expandable for mask prediction in the future

        if not output_namedtuple:
            # output dict
            head_outputs: dict[str, list[Tensor]] = {self.cls_key: self.classification_head(feature_maps)}
            head_outputs[self.box_reg_key] = self.regression_head(feature_maps)
        else:
            # output namedtuple
            Model_output = namedtuple('Model_output', ["classification", "box_regression"])
            head_outputs = Model_output(self.classification_head(feature_maps), self.regression_head(feature_maps))

        return head_outputs


def resnet_fpn_feature_extractor(
    backbone: resnet.ResNet,
    spatial_dims: int,
    pretrained_backbone: bool = False,
    returned_layers: Sequence[int] = (1, 2, 3),
    trainable_backbone_layers: int | None = None,
) -> BackboneWithFPN:
    """
    Constructs a feature extractor network with a ResNet-FPN backbone, used as feature_extractor in RetinaNet.

    Reference: `"Focal Loss for Dense Object Detection" <https://arxiv.org/abs/1708.02002>`_.

    The returned feature_extractor network takes an image tensor as inputs,
    and outputs a dictionary that maps string to the extracted feature maps (Tensor).

    The input to the returned feature_extractor is expected to be a list of tensors,
    each of shape ``[C, H, W]`` or ``[C, H, W, D]``,
    one for each image. Different images can have different sizes.


    Args:
        backbone: a ResNet model, used as backbone.
        spatial_dims: number of spatial dimensions of the images. We support both 2D and 3D images.

Review comment (suggestion):
Hi @Can-Zhao ,
      Since the TensorRT model doesn't have any of [these attributes](https://github.com/Project-MONAI/MONAI/blob/01b2162e1bfce5d68531ba773174bea0a3bc3641/monai/apps/detection/networks/retinanet_detector.py#L193), is it possible to add these parameters to the `__init__` function and use them if not provided by the `network`?

Thanks,
Bin

Code after:
            the predicted classification maps and second N tensors being the predicted box regression maps.

        """
        # compute features maps list from the input images.
        features = self.feature_extractor(images)
        if isinstance(features, Tensor):
            feature_maps = [features]
        elif torch.jit.isinstance(features, Dict[str, Tensor]):
            feature_maps = list(features.values())
        else:
            feature_maps = list(features)

        if not isinstance(feature_maps[0], Tensor):
            raise ValueError("feature_extractor output format must be Tensor, Dict[str, Tensor], or Sequence[Tensor].")

        # compute classification and box regression maps from the feature maps
        # expandable for mask prediction in the future

        if not self.use_list_output:
            # output dict
            head_outputs = {self.cls_key: self.classification_head(feature_maps)}
            head_outputs[self.box_reg_key] = self.regression_head(feature_maps)
            return head_outputs
        else:
            # output list of tensor, first half is classification, second half is box regression
            head_outputs_sequence = self.classification_head(feature_maps) + self.regression_head(feature_maps)
            return head_outputs_sequence


def resnet_fpn_feature_extractor(
    backbone: resnet.ResNet,
    spatial_dims: int,
    pretrained_backbone: bool = False,
    returned_layers: Sequence[int] = (1, 2, 3),
    trainable_backbone_layers: int | None = None,
) -> BackboneWithFPN:
    """
    Constructs a feature extractor network with a ResNet-FPN backbone, used as feature_extractor in RetinaNet.

    Reference: `"Focal Loss for Dense Object Detection" <https://arxiv.org/abs/1708.02002>`_.

    The returned feature_extractor network takes an image tensor as inputs,
    and outputs a dictionary that maps string to the extracted feature maps (Tensor).

    The input to the returned feature_extractor is expected to be a list of tensors,
    each of shape ``[C, H, W]`` or ``[C, H, W, D]``,
    one for each image. Different images can have different sizes.


    Args:
        backbone: a ResNet model, used as backbone.