| | |
| | import numpy as np |
| | from typing import Callable, Dict, Optional, Tuple, Union |
| | import fvcore.nn.weight_init as weight_init |
| | import torch |
| | from torch import nn |
| | from torch.nn import functional as F |
| |
|
| | from detectron2.config import configurable |
| | from detectron2.layers import Conv2d, ShapeSpec, get_norm |
| | from detectron2.structures import ImageList |
| | from detectron2.utils.registry import Registry |
| |
|
| | from ..backbone import Backbone, build_backbone |
| | from ..postprocessing import sem_seg_postprocess |
| | from .build import META_ARCH_REGISTRY |
| |
|
| | __all__ = [ |
| | "SemanticSegmentor", |
| | "SEM_SEG_HEADS_REGISTRY", |
| | "SemSegFPNHead", |
| | "build_sem_seg_head", |
| | ] |
| |
|
| |
|
| | SEM_SEG_HEADS_REGISTRY = Registry("SEM_SEG_HEADS") |
| | SEM_SEG_HEADS_REGISTRY.__doc__ = """ |
| | Registry for semantic segmentation heads, which make semantic segmentation predictions |
| | from feature maps. |
| | """ |
| |
|
| |
|
| | @META_ARCH_REGISTRY.register() |
| | class SemanticSegmentor(nn.Module): |
| | """ |
| | Main class for semantic segmentation architectures. |
| | """ |
| |
|
| | @configurable |
| | def __init__( |
| | self, |
| | *, |
| | backbone: Backbone, |
| | sem_seg_head: nn.Module, |
| | pixel_mean: Tuple[float], |
| | pixel_std: Tuple[float], |
| | ): |
| | """ |
| | Args: |
| | backbone: a backbone module, must follow detectron2's backbone interface |
| | sem_seg_head: a module that predicts semantic segmentation from backbone features |
| | 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 |
| | """ |
| | super().__init__() |
| | self.backbone = backbone |
| | self.sem_seg_head = sem_seg_head |
| | self.register_buffer("pixel_mean", torch.tensor(pixel_mean).view(-1, 1, 1), False) |
| | self.register_buffer("pixel_std", torch.tensor(pixel_std).view(-1, 1, 1), False) |
| |
|
| | @classmethod |
| | def from_config(cls, cfg): |
| | backbone = build_backbone(cfg) |
| | sem_seg_head = build_sem_seg_head(cfg, backbone.output_shape()) |
| | return { |
| | "backbone": backbone, |
| | "sem_seg_head": sem_seg_head, |
| | "pixel_mean": cfg.MODEL.PIXEL_MEAN, |
| | "pixel_std": cfg.MODEL.PIXEL_STD, |
| | } |
| |
|
| | @property |
| | def device(self): |
| | return self.pixel_mean.device |
| |
|
| | def forward(self, batched_inputs): |
| | """ |
| | Args: |
| | batched_inputs: a list, batched outputs of :class:`DatasetMapper`. |
| | Each item in the list contains the inputs for one image. |
| | |
| | For now, each item in the list is a dict that contains: |
| | |
| | * "image": Tensor, image in (C, H, W) format. |
| | * "sem_seg": semantic segmentation ground truth |
| | * Other information that's included in the original dicts, such as: |
| | "height", "width" (int): the output resolution of the model (may be different |
| | from input resolution), used in inference. |
| | |
| | |
| | Returns: |
| | list[dict]: |
| | Each dict is the output for one input image. |
| | The dict contains one key "sem_seg" whose value is a |
| | Tensor that represents the |
| | per-pixel segmentation prediced by the head. |
| | The prediction has shape KxHxW that represents the logits of |
| | each class for each pixel. |
| | """ |
| | images = [x["image"].to(self.device) for x in batched_inputs] |
| | images = [(x - self.pixel_mean) / self.pixel_std for x in images] |
| | images = ImageList.from_tensors( |
| | images, |
| | self.backbone.size_divisibility, |
| | padding_constraints=self.backbone.padding_constraints, |
| | ) |
| |
|
| | features = self.backbone(images.tensor) |
| |
|
| | if "sem_seg" in batched_inputs[0]: |
| | targets = [x["sem_seg"].to(self.device) for x in batched_inputs] |
| | targets = ImageList.from_tensors( |
| | targets, |
| | self.backbone.size_divisibility, |
| | self.sem_seg_head.ignore_value, |
| | self.backbone.padding_constraints, |
| | ).tensor |
| | else: |
| | targets = None |
| | results, losses = self.sem_seg_head(features, targets) |
| |
|
| | if self.training: |
| | return losses |
| |
|
| | processed_results = [] |
| | for result, input_per_image, image_size in zip(results, batched_inputs, images.image_sizes): |
| | height = input_per_image.get("height", image_size[0]) |
| | width = input_per_image.get("width", image_size[1]) |
| | r = sem_seg_postprocess(result, image_size, height, width) |
| | processed_results.append({"sem_seg": r}) |
| | return processed_results |
| |
|
| |
|
| | def build_sem_seg_head(cfg, input_shape): |
| | """ |
| | Build a semantic segmentation head from `cfg.MODEL.SEM_SEG_HEAD.NAME`. |
| | """ |
| | name = cfg.MODEL.SEM_SEG_HEAD.NAME |
| | return SEM_SEG_HEADS_REGISTRY.get(name)(cfg, input_shape) |
| |
|
| |
|
| | @SEM_SEG_HEADS_REGISTRY.register() |
| | class SemSegFPNHead(nn.Module): |
| | """ |
| | A semantic segmentation head described in :paper:`PanopticFPN`. |
| | It takes a list of FPN features as input, and applies a sequence of |
| | 3x3 convs and upsampling to scale all of them to the stride defined by |
| | ``common_stride``. Then these features are added and used to make final |
| | predictions by another 1x1 conv layer. |
| | """ |
| |
|
| | @configurable |
| | def __init__( |
| | self, |
| | input_shape: Dict[str, ShapeSpec], |
| | *, |
| | num_classes: int, |
| | conv_dims: int, |
| | common_stride: int, |
| | loss_weight: float = 1.0, |
| | norm: Optional[Union[str, Callable]] = None, |
| | ignore_value: int = -1, |
| | ): |
| | """ |
| | NOTE: this interface is experimental. |
| | |
| | Args: |
| | input_shape: shapes (channels and stride) of the input features |
| | num_classes: number of classes to predict |
| | conv_dims: number of output channels for the intermediate conv layers. |
| | common_stride: the common stride that all features will be upscaled to |
| | loss_weight: loss weight |
| | norm (str or callable): normalization for all conv layers |
| | ignore_value: category id to be ignored during training. |
| | """ |
| | super().__init__() |
| | input_shape = sorted(input_shape.items(), key=lambda x: x[1].stride) |
| | if not len(input_shape): |
| | raise ValueError("SemSegFPNHead(input_shape=) cannot be empty!") |
| | self.in_features = [k for k, v in input_shape] |
| | feature_strides = [v.stride for k, v in input_shape] |
| | feature_channels = [v.channels for k, v in input_shape] |
| |
|
| | self.ignore_value = ignore_value |
| | self.common_stride = common_stride |
| | self.loss_weight = loss_weight |
| |
|
| | self.scale_heads = [] |
| | for in_feature, stride, channels in zip( |
| | self.in_features, feature_strides, feature_channels |
| | ): |
| | head_ops = [] |
| | head_length = max(1, int(np.log2(stride) - np.log2(self.common_stride))) |
| | for k in range(head_length): |
| | norm_module = get_norm(norm, conv_dims) |
| | conv = Conv2d( |
| | channels if k == 0 else conv_dims, |
| | conv_dims, |
| | kernel_size=3, |
| | stride=1, |
| | padding=1, |
| | bias=not norm, |
| | norm=norm_module, |
| | activation=F.relu, |
| | ) |
| | weight_init.c2_msra_fill(conv) |
| | head_ops.append(conv) |
| | if stride != self.common_stride: |
| | head_ops.append( |
| | nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False) |
| | ) |
| | self.scale_heads.append(nn.Sequential(*head_ops)) |
| | self.add_module(in_feature, self.scale_heads[-1]) |
| | self.predictor = Conv2d(conv_dims, num_classes, kernel_size=1, stride=1, padding=0) |
| | weight_init.c2_msra_fill(self.predictor) |
| |
|
| | @classmethod |
| | def from_config(cls, cfg, input_shape: Dict[str, ShapeSpec]): |
| | return { |
| | "input_shape": { |
| | k: v for k, v in input_shape.items() if k in cfg.MODEL.SEM_SEG_HEAD.IN_FEATURES |
| | }, |
| | "ignore_value": cfg.MODEL.SEM_SEG_HEAD.IGNORE_VALUE, |
| | "num_classes": cfg.MODEL.SEM_SEG_HEAD.NUM_CLASSES, |
| | "conv_dims": cfg.MODEL.SEM_SEG_HEAD.CONVS_DIM, |
| | "common_stride": cfg.MODEL.SEM_SEG_HEAD.COMMON_STRIDE, |
| | "norm": cfg.MODEL.SEM_SEG_HEAD.NORM, |
| | "loss_weight": cfg.MODEL.SEM_SEG_HEAD.LOSS_WEIGHT, |
| | } |
| |
|
| | def forward(self, features, targets=None): |
| | """ |
| | Returns: |
| | In training, returns (None, dict of losses) |
| | In inference, returns (CxHxW logits, {}) |
| | """ |
| | x = self.layers(features) |
| | if self.training: |
| | return None, self.losses(x, targets) |
| | else: |
| | x = F.interpolate( |
| | x, scale_factor=self.common_stride, mode="bilinear", align_corners=False |
| | ) |
| | return x, {} |
| |
|
| | def layers(self, features): |
| | for i, f in enumerate(self.in_features): |
| | if i == 0: |
| | x = self.scale_heads[i](features[f]) |
| | else: |
| | x = x + self.scale_heads[i](features[f]) |
| | x = self.predictor(x) |
| | return x |
| |
|
| | def losses(self, predictions, targets): |
| | predictions = predictions.float() |
| | predictions = F.interpolate( |
| | predictions, |
| | scale_factor=self.common_stride, |
| | mode="bilinear", |
| | align_corners=False, |
| | ) |
| | loss = F.cross_entropy( |
| | predictions, targets, reduction="mean", ignore_index=self.ignore_value |
| | ) |
| | losses = {"loss_sem_seg": loss * self.loss_weight} |
| | return losses |
| |
|
