| |
| 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 annotator.oneformer.detectron2.config import configurable |
| from annotator.oneformer.detectron2.layers import Conv2d, ShapeSpec, get_norm |
| from annotator.oneformer.detectron2.structures import ImageList |
| from annotator.oneformer.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 |
|
|
