| | |
| | 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) |
| |
|
