configurable for rcnn
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
import logging
import numpy as np
+from typing import Optional, Tuple
import torch
from torch import nn
+from detectron2.config import configurable
from detectron2.data.detection_utils import convert_image_to_rgb
from detectron2.structures import ImageList
from detectron2.utils.events import get_event_storage
from detectron2.utils.logger import log_first_n
-from ..backbone import build_backbone
+from ..backbone import Backbone, build_backbone
from ..postprocessing import detector_postprocess
from ..proposal_generator import build_proposal_generator
from ..roi_heads import build_roi_heads
@@ -27,18 +29,59 @@ class GeneralizedRCNN(nn.Module):
3. Per-region feature extraction and prediction
"""
- def __init__(self, cfg):
- super().__init__()
-
- self.backbone = build_backbone(cfg)
- self.proposal_generator = build_proposal_generator(cfg, self.backbone.output_shape())
- self.roi_heads = build_roi_heads(cfg, self.backbone.output_shape())
- self.vis_period = cfg.VIS_PERIOD
- self.input_format = cfg.INPUT.FORMAT
+ @configurable
+ def __init__(
+ self,
+ *,
+ backbone: Backbone,
+ proposal_generator: nn.Module,
+ roi_heads: nn.Module,
+ pixel_mean: Tuple[float],
+ pixel_std: Tuple[float],
+ input_format: Optional[str] = None,
+ vis_period: int = 0,
+ ):
+ """
+ NOTE: this interface is experimental.
- assert len(cfg.MODEL.PIXEL_MEAN) == len(cfg.MODEL.PIXEL_STD)
- self.register_buffer("pixel_mean", torch.Tensor(cfg.MODEL.PIXEL_MEAN).view(-1, 1, 1))
- self.register_buffer("pixel_std", torch.Tensor(cfg.MODEL.PIXEL_STD).view(-1, 1, 1))
+ Args:
+ backbone: a backbone module, must follow detectron2's backbone interface
+ proposal_generator: a module that generates proposals using backbone features
+ roi_heads: a ROI head that performs per-region computation
+ pixel_mean, pixel_std: list or tuple with #channels element,
+ representing the per-channel mean and std to be used to normalize
+ the input image
+ input_format: describe the meaning of channels of input. Needed by visualization
+ vis_period: the period to run visualization. Set to 0 to disable.
+ """
+ super().__init__()
+ self.backbone = backbone
+ self.proposal_generator = proposal_generator
+ self.roi_heads = roi_heads
+
+ self.input_format = input_format
+ self.vis_period = vis_period
+ if vis_period > 0:
+ assert input_format is not None, "input_format is required for visualization!"
+
+ self.register_buffer("pixel_mean", torch.Tensor(pixel_mean).view(-1, 1, 1))
+ self.register_buffer("pixel_std", torch.Tensor(pixel_std).view(-1, 1, 1))
+ assert (
+ self.pixel_mean.shape == self.pixel_std.shape
+ ), f"{self.pixel_mean} and {self.pixel_std} have different shapes!"
+
+ @classmethod
+ def from_config(cls, cfg):
+ backbone = build_backbone(cfg)
+ return {
+ "backbone": backbone,
+ "proposal_generator": build_proposal_generator(cfg, backbone.output_shape()),
+ "roi_heads": build_roi_heads(cfg, backbone.output_shape()),
+ "input_format": cfg.INPUT.FORMAT,
+ "vis_period": cfg.VIS_PERIOD,
+ "pixel_mean": cfg.MODEL.PIXEL_MEAN,
+ "pixel_std": cfg.MODEL.PIXEL_STD,
+ }
@property
def device(self):
