text: # Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #    http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  import os
  from typing import List
  
  import paddle
  import numpy as np
  import paddle.nn as nn
  import paddle.nn.functional as F
  from paddlehub.module.module import moduleinfo
  import paddlehub.vision.segmentation_transforms as T
  from paddlehub.module.cv_module import ImageSegmentationModule
  
  import ocrnet_hrnetw18_cityscapes.layers as L
  from ocrnet_hrnetw18_cityscapes.hrnet import HRNet_W18
  
  
  @moduleinfo(
      name="ocrnet_hrnetw18_cityscapes",
      type="CV/semantic_segmentation",
      author="paddlepaddle",
      author_email="",
      summary="OCRNetHRNetW18 is a segmentation model pretrained by pascal voc.",
      version="1.0.0",
      meta=ImageSegmentationModule)
  class OCRNetHRNetW18(nn.Layer):
      """
      The OCRNet implementation based on PaddlePaddle.
      The original article refers to
          Yuan, Yuhui, et al. "Object-Contextual Representations for Semantic Segmentation"
          (https://arxiv.org/pdf/1909.11065.pdf)
      Args:
          num_classes (int): The unique number of target classes.
          backbone_indices (list): A list indicates the indices of output of backbone.
              It can be either one or two values, if two values, the first index will be taken as
              a deep-supervision feature in auxiliary layer; the second one will be taken as
              input of pixel representation. If one value, it is taken by both above.
          ocr_mid_channels (int, optional): The number of middle channels in OCRHead. Default: 512.
          ocr_key_channels (int, optional): The number of key channels in ObjectAttentionBlock. Default: 256.
          align_corners (bool): An argument of F.interpolate. It should be set to False when the output size of feature
              is even, e.g. 1024x512, otherwise it is True, e.g. 769x769.  Default: False.
          pretrained (str, optional): The path or url of pretrained model. Default: None.
      """
  
      def __init__(self,
                   num_classes: int = 19,
                   backbone_indices: List[int] = [0],
                   ocr_mid_channels: int = 512,
                   ocr_key_channels: int = 256,
                   align_corners: bool = False,
                   pretrained: str = None):
          super(OCRNetHRNetW18, self).__init__()
          self.backbone = HRNet_W18()
          self.backbone_indices = backbone_indices
          in_channels = [self.backbone.feat_channels[i] for i in backbone_indices]
          self.head = OCRHead(
              num_classes=num_classes,
              in_channels=in_channels,
              ocr_mid_channels=ocr_mid_channels,
              ocr_key_channels=ocr_key_channels)
          self.align_corners = align_corners
          self.transforms = T.Compose([T.Normalize()])
  
          if pretrained is not None:
              model_dict = paddle.load(pretrained)
              self.set_dict(model_dict)
              print("load custom parameters success")
  
          else:
              checkpoint = os.path.join(self.directory, 'model.pdparams')
              model_dict = paddle.load(checkpoint)
              self.set_dict(model_dict)
              print("load pretrained parameters success")
  
      def transform(self, img: np.ndarray) -> np.ndarray:
          return self.transforms(img)
  
      def forward(self, x: paddle.Tensor) -> paddle.Tensor:
          feats = self.backbone(x)
          feats = [feats[i] for i in self.backbone_indices]
          logit_list = self.head(feats)
          logit_list = [
              F.interpolate(logit, x.shape[2:], mode='bilinear', align_corners=self.align_corners) for logit in logit_list
          ]
          return logit_list
  
  
  class OCRHead(nn.Layer):
      """
      The Object contextual representation head.
      Args:
          num_classes(int): The unique number of target classes.
          in_channels(tuple): The number of input channels.
          ocr_mid_channels(int, optional): The number of middle channels in OCRHead. Default: 512.
          ocr_key_channels(int, optional): The number of key channels in ObjectAttentionBlock. Default: 256.
      """
  
      def __init__(self, num_classes: int, in_channels: int, ocr_mid_channels: int = 512, ocr_key_channels: int = 256):
          super().__init__()
  
          self.num_classes = num_classes
          self.spatial_gather = SpatialGatherBlock()
          self.spatial_ocr = SpatialOCRModule(ocr_mid_channels, ocr_key_channels, ocr_mid_channels)
  
          self.indices = [-2, -1] if len(in_channels) > 1 else [-1, -1]
  
          self.conv3x3_ocr = L.ConvBNReLU(in_channels[self.indices[1]], ocr_mid_channels, 3, padding=1)
          self.cls_head = nn.Conv2D(ocr_mid_channels, self.num_classes, 1)
          self.aux_head = nn.Sequential(
              L.ConvBNReLU(in_channels[self.indices[0]], in_channels[self.indices[0]], 1),
              nn.Conv2D(in_channels[self.indices[0]], self.num_classes, 1))
  
      def forward(self, feat_list: List[paddle.Tensor]) -> paddle.Tensor:
          feat_shallow, feat_deep = feat_list[self.indices[0]], feat_list[self.indices[1]]
  
          soft_regions = self.aux_head(feat_shallow)
          pixels = self.conv3x3_ocr(feat_deep)
  
          object_regions = self.spatial_gather(pixels, soft_regions)
          ocr = self.spatial_ocr(pixels, object_regions)
  
          logit = self.cls_head(ocr)
          return [logit, soft_regions]
  
  
  class SpatialGatherBlock(nn.Layer):
      """Aggregation layer to compute the pixel-region representation."""
  
      def forward(self, pixels: paddle.Tensor, regions: paddle.Tensor) -> paddle.Tensor:
          n, c, h, w = pixels.shape
          _, k, _, _ = regions.shape
  
          # pixels: from (n, c, h, w) to (n, h*w, c)
          pixels = paddle.reshape(pixels, (n, c, h * w))
          pixels = paddle.transpose(pixels, [0, 2, 1])
  
          # regions: from (n, k, h, w) to (n, k, h*w)
          regions = paddle.reshape(regions, (n, k, h * w))
          regions = F.softmax(regions, axis=2)
  
          # feats: from (n, k, c) to (n, c, k, 1)
          feats = paddle.bmm(regions, pixels)
          feats = paddle.transpose(feats, [0, 2, 1])
          feats = paddle.unsqueeze(feats, axis=-1)
  
          return feats
  
  
  class SpatialOCRModule(nn.Layer):
      """Aggregate the global object representation to update the representation for each pixel."""
  
      def __init__(self, in_channels: int, key_channels: int, out_channels: int, dropout_rate: float = 0.1):
          super().__init__()
  
          self.attention_block = ObjectAttentionBlock(in_channels, key_channels)
          self.conv1x1 = nn.Sequential(L.ConvBNReLU(2 * in_channels, out_channels, 1), nn.Dropout2D(dropout_rate))
  
      def forward(self, pixels: paddle.Tensor, regions: paddle.Tensor) -> paddle.Tensor:
          context = self.attention_block(pixels, regions)
          feats = paddle.concat([context, pixels], axis=1)
          feats = self.conv1x1(feats)
  
          return feats
  
  
  class ObjectAttentionBlock(nn.Layer):
      """A self-attention module."""
  
      def __init__(self, in_channels: int, key_channels: int):
          super().__init__()
  
          self.in_channels = in_channels
          self.key_channels = key_channels
  
          self.f_pixel = nn.Sequential(
              L.ConvBNReLU(in_channels, key_channels, 1), L.ConvBNReLU(key_channels, key_channels, 1))
  
          self.f_object = nn.Sequential(
              L.ConvBNReLU(in_channels, key_channels, 1), L.ConvBNReLU(key_channels, key_channels, 1))
  
          self.f_down = L.ConvBNReLU(in_channels, key_channels, 1)
  
          self.f_up = L.ConvBNReLU(key_channels, in_channels, 1)
  
      def forward(self, x: paddle.Tensor, proxy: paddle.Tensor) -> paddle.Tensor:
          n, _, h, w = x.shape
  
          # query : from (n, c1, h1, w1) to (n, h1*w1, key_channels)
          query = self.f_pixel(x)
          query = paddle.reshape(query, (n, self.key_channels, -1))
          query = paddle.transpose(query, [0, 2, 1])
  
          # key : from (n, c2, h2, w2) to (n, key_channels, h2*w2)
          key = self.f_object(proxy)
          key = paddle.reshape(key, (n, self.key_channels, -1))
  
          # value : from (n, c2, h2, w2) to (n, h2*w2, key_channels)
          value = self.f_down(proxy)
          value = paddle.reshape(value, (n, self.key_channels, -1))
          value = paddle.transpose(value, [0, 2, 1])
  
          # sim_map (n, h1*w1, h2*w2)
          sim_map = paddle.bmm(query, key)
          sim_map = (self.key_channels**-.5) * sim_map
          sim_map = F.softmax(sim_map, axis=-1)
  
          # context from (n, h1*w1, key_channels) to (n , out_channels, h1, w1)
          context = paddle.bmm(sim_map, value)
          context = paddle.transpose(context, [0, 2, 1])
          context = paddle.reshape(context, (n, self.key_channels, h, w))
          context = self.f_up(context)
  
          return context
  