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