|
|
| import numpy as np
|
| from typing import List
|
| import fvcore.nn.weight_init as weight_init
|
| import torch
|
| from torch import nn
|
|
|
| from detectron2.config import configurable
|
| from detectron2.layers import Conv2d, ShapeSpec, get_norm
|
| from detectron2.utils.registry import Registry
|
|
|
| __all__ = ["FastRCNNConvFCHead", "build_box_head", "ROI_BOX_HEAD_REGISTRY"]
|
|
|
| ROI_BOX_HEAD_REGISTRY = Registry("ROI_BOX_HEAD")
|
| ROI_BOX_HEAD_REGISTRY.__doc__ = """
|
| Registry for box heads, which make box predictions from per-region features.
|
|
|
| The registered object will be called with `obj(cfg, input_shape)`.
|
| """
|
|
|
|
|
|
|
|
|
|
|
| @ROI_BOX_HEAD_REGISTRY.register()
|
| class FastRCNNConvFCHead(nn.Sequential):
|
| """
|
| A head with several 3x3 conv layers (each followed by norm & relu) and then
|
| several fc layers (each followed by relu).
|
| """
|
|
|
| @configurable
|
| def __init__(
|
| self, input_shape: ShapeSpec, *, conv_dims: List[int], fc_dims: List[int], conv_norm=""
|
| ):
|
| """
|
| NOTE: this interface is experimental.
|
|
|
| Args:
|
| input_shape (ShapeSpec): shape of the input feature.
|
| conv_dims (list[int]): the output dimensions of the conv layers
|
| fc_dims (list[int]): the output dimensions of the fc layers
|
| conv_norm (str or callable): normalization for the conv layers.
|
| See :func:`detectron2.layers.get_norm` for supported types.
|
| """
|
| super().__init__()
|
| assert len(conv_dims) + len(fc_dims) > 0
|
|
|
| self._output_size = (input_shape.channels, input_shape.height, input_shape.width)
|
|
|
| self.conv_norm_relus = []
|
| for k, conv_dim in enumerate(conv_dims):
|
| conv = Conv2d(
|
| self._output_size[0],
|
| conv_dim,
|
| kernel_size=3,
|
| padding=1,
|
| bias=not conv_norm,
|
| norm=get_norm(conv_norm, conv_dim),
|
| activation=nn.ReLU(),
|
| )
|
| self.add_module("conv{}".format(k + 1), conv)
|
| self.conv_norm_relus.append(conv)
|
| self._output_size = (conv_dim, self._output_size[1], self._output_size[2])
|
|
|
| self.fcs = []
|
| for k, fc_dim in enumerate(fc_dims):
|
| if k == 0:
|
| self.add_module("flatten", nn.Flatten())
|
| fc = nn.Linear(int(np.prod(self._output_size)), fc_dim)
|
| self.add_module("fc{}".format(k + 1), fc)
|
| self.add_module("fc_relu{}".format(k + 1), nn.ReLU())
|
| self.fcs.append(fc)
|
| self._output_size = fc_dim
|
|
|
| for layer in self.conv_norm_relus:
|
| weight_init.c2_msra_fill(layer)
|
| for layer in self.fcs:
|
| weight_init.c2_xavier_fill(layer)
|
|
|
| @classmethod
|
| def from_config(cls, cfg, input_shape):
|
| num_conv = cfg.MODEL.ROI_BOX_HEAD.NUM_CONV
|
| conv_dim = cfg.MODEL.ROI_BOX_HEAD.CONV_DIM
|
| num_fc = cfg.MODEL.ROI_BOX_HEAD.NUM_FC
|
| fc_dim = cfg.MODEL.ROI_BOX_HEAD.FC_DIM
|
| return {
|
| "input_shape": input_shape,
|
| "conv_dims": [conv_dim] * num_conv,
|
| "fc_dims": [fc_dim] * num_fc,
|
| "conv_norm": cfg.MODEL.ROI_BOX_HEAD.NORM,
|
| }
|
|
|
| def forward(self, x):
|
| for layer in self:
|
| x = layer(x)
|
| return x
|
|
|
| @property
|
| @torch.jit.unused
|
| def output_shape(self):
|
| """
|
| Returns:
|
| ShapeSpec: the output feature shape
|
| """
|
| o = self._output_size
|
| if isinstance(o, int):
|
| return ShapeSpec(channels=o)
|
| else:
|
| return ShapeSpec(channels=o[0], height=o[1], width=o[2])
|
|
|
|
|
| def build_box_head(cfg, input_shape):
|
| """
|
| Build a box head defined by `cfg.MODEL.ROI_BOX_HEAD.NAME`.
|
| """
|
| name = cfg.MODEL.ROI_BOX_HEAD.NAME
|
| return ROI_BOX_HEAD_REGISTRY.get(name)(cfg, input_shape)
|
|
|
