Create AugmentCE2P.py

3rdparty/SCHP/networks/AugmentCE2P.py

@@ -0,0 +1,481 @@
+#!/usr/bin/env python
+# -*- encoding: utf-8 -*-
+
+"""
+@Author  :   Peike Li
+@Contact :   peike.li@yahoo.com
+@File    :   AugmentCE2P.py
+@Time    :   8/4/19 3:35 PM
+@Desc    :
+@License :   This source code is licensed under the license found in the
+             LICENSE file in the root directory of this source tree.
+"""
+
+import torch
+import torch.nn as nn
+
+from torch.nn import BatchNorm2d, functional as F, LeakyReLU
+
+affine_par = True
+pretrained_settings = {
+    "resnet101": {
+        "imagenet": {
+            "input_space": "BGR",
+            "input_size": [3, 224, 224],
+            "input_range": [0, 1],
+            "mean": [0.406, 0.456, 0.485],
+            "std": [0.225, 0.224, 0.229],
+            "num_classes": 1000,
+        }
+    },
+}
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    "3x3 convolution with padding"
+    return nn.Conv2d(
+        in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False
+    )
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(
+        self,
+        inplanes,
+        planes,
+        stride=1,
+        dilation=1,
+        downsample=None,
+        fist_dilation=1,
+        multi_grid=1,
+    ):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+        self.bn1 = BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(
+            planes,
+            planes,
+            kernel_size=3,
+            stride=stride,
+            padding=dilation * multi_grid,
+            dilation=dilation * multi_grid,
+            bias=False,
+        )
+        self.bn2 = BatchNorm2d(planes)
+        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
+        self.bn3 = BatchNorm2d(planes * 4)
+        self.relu = nn.ReLU(inplace=False)
+        self.relu_inplace = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.dilation = dilation
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out = out + residual
+        out = self.relu_inplace(out)
+
+        return out
+
+
+class PSPModule(nn.Module):
+    """
+    Reference:
+        Zhao, Hengshuang, et al. *"Pyramid scene parsing network."*
+    """
+
+    def __init__(self, features, out_features=512, sizes=(1, 2, 3, 6)):
+        super(PSPModule, self).__init__()
+
+        self.stages = []
+        self.stages = nn.ModuleList(
+            [self._make_stage(features, out_features, size) for size in sizes]
+        )
+        self.bottleneck = nn.Sequential(
+            nn.Conv2d(
+                features + len(sizes) * out_features,
+                out_features,
+                kernel_size=3,
+                padding=1,
+                dilation=1,
+                bias=False,
+            ),
+            BatchNorm2d(out_features),
+            LeakyReLU(),
+        )
+
+    def _make_stage(self, features, out_features, size):
+        prior = nn.AdaptiveAvgPool2d(output_size=(size, size))
+        conv = nn.Conv2d(features, out_features, kernel_size=1, bias=False)
+        return nn.Sequential(
+            prior,
+            conv,
+            # bn
+            BatchNorm2d(out_features),
+            LeakyReLU(),
+        )
+
+    def forward(self, feats):
+        h, w = feats.size(2), feats.size(3)
+        priors = [
+            F.interpolate(
+                input=stage(feats), size=(h, w), mode="bilinear", align_corners=True
+            )
+            for stage in self.stages
+        ] + [feats]
+        bottle = self.bottleneck(torch.cat(priors, 1))
+        return bottle
+
+
+class ASPPModule(nn.Module):
+    """
+    Reference:
+        Chen, Liang-Chieh, et al. *"Rethinking Atrous Convolution for Semantic Image Segmentation."*
+    """
+
+    def __init__(
+        self, features, inner_features=256, out_features=512, dilations=(12, 24, 36)
+    ):
+        super(ASPPModule, self).__init__()
+
+        self.conv1 = nn.Sequential(
+            nn.AdaptiveAvgPool2d((1, 1)),
+            nn.Conv2d(
+                features,
+                inner_features,
+                kernel_size=1,
+                padding=0,
+                dilation=1,
+                bias=False,
+            ),
+            #    InPlaceABNSync(inner_features)
+            BatchNorm2d(inner_features),
+            LeakyReLU(),
+        )
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(
+                features,
+                inner_features,
+                kernel_size=1,
+                padding=0,
+                dilation=1,
+                bias=False,
+            ),
+            BatchNorm2d(inner_features),
+            LeakyReLU(),
+        )
+        self.conv3 = nn.Sequential(
+            nn.Conv2d(
+                features,
+                inner_features,
+                kernel_size=3,
+                padding=dilations[0],
+                dilation=dilations[0],
+                bias=False,
+            ),
+            BatchNorm2d(inner_features),
+            LeakyReLU(),
+        )
+        self.conv4 = nn.Sequential(
+            nn.Conv2d(
+                features,
+                inner_features,
+                kernel_size=3,
+                padding=dilations[1],
+                dilation=dilations[1],
+                bias=False,
+            ),
+            BatchNorm2d(inner_features),
+            LeakyReLU(),
+        )
+        self.conv5 = nn.Sequential(
+            nn.Conv2d(
+                features,
+                inner_features,
+                kernel_size=3,
+                padding=dilations[2],
+                dilation=dilations[2],
+                bias=False,
+            ),
+            BatchNorm2d(inner_features),
+            LeakyReLU(),
+        )
+
+        self.bottleneck = nn.Sequential(
+            nn.Conv2d(
+                inner_features * 5,
+                out_features,
+                kernel_size=1,
+                padding=0,
+                dilation=1,
+                bias=False,
+            ),
+            BatchNorm2d(inner_features),
+            LeakyReLU(),
+            nn.Dropout2d(0.1),
+        )
+
+    def forward(self, x):
+        _, _, h, w = x.size()
+
+        feat1 = F.interpolate(
+            self.conv1(x), size=(h, w), mode="bilinear", align_corners=True
+        )
+
+        feat2 = self.conv2(x)
+        feat3 = self.conv3(x)
+        feat4 = self.conv4(x)
+        feat5 = self.conv5(x)
+        out = torch.cat((feat1, feat2, feat3, feat4, feat5), 1)
+
+        bottle = self.bottleneck(out)
+        return bottle
+
+
+class Edge_Module(nn.Module):
+    """
+    Edge Learning Branch
+    """
+
+    def __init__(self, in_fea=[256, 512, 1024], mid_fea=256, out_fea=2):
+        super(Edge_Module, self).__init__()
+
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(
+                in_fea[0], mid_fea, kernel_size=1, padding=0, dilation=1, bias=False
+            ),
+            BatchNorm2d(mid_fea),
+            LeakyReLU(),
+        )
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(
+                in_fea[1], mid_fea, kernel_size=1, padding=0, dilation=1, bias=False
+            ),
+            BatchNorm2d(mid_fea),
+            LeakyReLU(),
+        )
+        self.conv3 = nn.Sequential(
+            nn.Conv2d(
+                in_fea[2], mid_fea, kernel_size=1, padding=0, dilation=1, bias=False
+            ),
+            BatchNorm2d(mid_fea),
+            LeakyReLU(),
+        )
+        self.conv4 = nn.Conv2d(
+            mid_fea, out_fea, kernel_size=3, padding=1, dilation=1, bias=True
+        )
+        # self.conv5 = nn.Conv2d(out_fea * 3, out_fea, kernel_size=1, padding=0, dilation=1, bias=True)
+
+    def forward(self, x1, x2, x3):
+        _, _, h, w = x1.size()
+
+        edge1_fea = self.conv1(x1)
+        # edge1 = self.conv4(edge1_fea)
+        edge2_fea = self.conv2(x2)
+        edge2 = self.conv4(edge2_fea)
+        edge3_fea = self.conv3(x3)
+        edge3 = self.conv4(edge3_fea)
+
+        edge2_fea = F.interpolate(
+            edge2_fea, size=(h, w), mode="bilinear", align_corners=True
+        )
+        edge3_fea = F.interpolate(
+            edge3_fea, size=(h, w), mode="bilinear", align_corners=True
+        )
+        edge2 = F.interpolate(edge2, size=(h, w), mode="bilinear", align_corners=True)
+        edge3 = F.interpolate(edge3, size=(h, w), mode="bilinear", align_corners=True)
+
+        # edge = torch.cat([edge1, edge2, edge3], dim=1)
+        edge_fea = torch.cat([edge1_fea, edge2_fea, edge3_fea], dim=1)
+        # edge = self.conv5(edge)
+
+        # return edge, edge_fea
+        return edge_fea
+
+
+class Decoder_Module(nn.Module):
+    """
+    Parsing Branch Decoder Module.
+    """
+
+    def __init__(self, num_classes):
+        super(Decoder_Module, self).__init__()
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(512, 256, kernel_size=1, padding=0, dilation=1, bias=False),
+            BatchNorm2d(256),
+            LeakyReLU(),
+        )
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(
+                256, 48, kernel_size=1, stride=1, padding=0, dilation=1, bias=False
+            ),
+            BatchNorm2d(48),
+            LeakyReLU(),
+        )
+        self.conv3 = nn.Sequential(
+            nn.Conv2d(304, 256, kernel_size=1, padding=0, dilation=1, bias=False),
+            BatchNorm2d(256),
+            LeakyReLU(),
+            nn.Conv2d(256, 256, kernel_size=1, padding=0, dilation=1, bias=False),
+            BatchNorm2d(256),
+            LeakyReLU(),
+        )
+
+        # self.conv4 = nn.Conv2d(256, num_classes, kernel_size=1, padding=0, dilation=1, bias=True)
+
+    def forward(self, xt, xl):
+        _, _, h, w = xl.size()
+        xt = F.interpolate(
+            self.conv1(xt), size=(h, w), mode="bilinear", align_corners=True
+        )
+        xl = self.conv2(xl)
+        x = torch.cat([xt, xl], dim=1)
+        x = self.conv3(x)
+        # seg = self.conv4(x)
+        # return seg, x
+        return x
+
+
+class ResNet(nn.Module):
+    def __init__(self, block, layers, num_classes):
+        self.inplanes = 128
+        super(ResNet, self).__init__()
+        self.conv1 = conv3x3(3, 64, stride=2)
+        self.bn1 = BatchNorm2d(64)
+        self.relu1 = nn.ReLU(inplace=False)
+        self.conv2 = conv3x3(64, 64)
+        self.bn2 = BatchNorm2d(64)
+        self.relu2 = nn.ReLU(inplace=False)
+        self.conv3 = conv3x3(64, 128)
+        self.bn3 = BatchNorm2d(128)
+        self.relu3 = nn.ReLU(inplace=False)
+
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+        self.layer4 = self._make_layer(
+            block, 512, layers[3], stride=1, dilation=2, multi_grid=(1, 1, 1)
+        )
+
+        self.context_encoding = PSPModule(2048, 512)
+
+        self.edge = Edge_Module()
+        self.decoder = Decoder_Module(num_classes)
+
+        self.fushion = nn.Sequential(
+            nn.Conv2d(1024, 256, kernel_size=1, padding=0, dilation=1, bias=False),
+            BatchNorm2d(256),
+            LeakyReLU(),
+            nn.Dropout2d(0.1),
+            nn.Conv2d(
+                256, num_classes, kernel_size=1, padding=0, dilation=1, bias=True
+            ),
+        )
+
+    def _make_layer(self, block, planes, blocks, stride=1, dilation=1, multi_grid=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(
+                    self.inplanes,
+                    planes * block.expansion,
+                    kernel_size=1,
+                    stride=stride,
+                    bias=False,
+                ),
+                BatchNorm2d(planes * block.expansion, affine=affine_par),
+            )
+
+        layers = []
+        generate_multi_grid = lambda index, grids: (
+            grids[index % len(grids)] if isinstance(grids, tuple) else 1
+        )
+        layers.append(
+            block(
+                self.inplanes,
+                planes,
+                stride,
+                dilation=dilation,
+                downsample=downsample,
+                multi_grid=generate_multi_grid(0, multi_grid),
+            )
+        )
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(
+                block(
+                    self.inplanes,
+                    planes,
+                    dilation=dilation,
+                    multi_grid=generate_multi_grid(i, multi_grid),
+                )
+            )
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.relu1(self.bn1(self.conv1(x)))
+        x = self.relu2(self.bn2(self.conv2(x)))
+        x = self.relu3(self.bn3(self.conv3(x)))
+        x = self.maxpool(x)
+        x2 = self.layer1(x)
+        x3 = self.layer2(x2)
+        x4 = self.layer3(x3)
+        x5 = self.layer4(x4)
+        x = self.context_encoding(x5)
+        # parsing_result, parsing_fea = self.decoder(x, x2)
+        parsing_fea = self.decoder(x, x2)
+        # Edge Branch
+        # edge_result, edge_fea = self.edge(x2, x3, x4)
+        edge_fea = self.edge(x2, x3, x4)
+        # Fusion Branch
+        x = torch.cat([parsing_fea, edge_fea], dim=1)
+        fusion_result = self.fushion(x)
+        # return [[parsing_result, fusion_result], [edge_result]]
+        return fusion_result
+
+
+def initialize_pretrained_model(
+    model, settings, pretrained="./models/resnet101-imagenet.pth"
+):
+    model.input_space = settings["input_space"]
+    model.input_size = settings["input_size"]
+    model.input_range = settings["input_range"]
+    model.mean = settings["mean"]
+    model.std = settings["std"]
+
+    if pretrained is not None:
+        saved_state_dict = torch.load(pretrained)
+        new_params = model.state_dict().copy()
+        for i in saved_state_dict:
+            i_parts = i.split(".")
+            if not i_parts[0] == "fc":
+                new_params[".".join(i_parts[0:])] = saved_state_dict[i]
+        model.load_state_dict(new_params)
+
+
+def resnet101(num_classes=20, pretrained="./models/resnet101-imagenet.pth"):
+    model = ResNet(Bottleneck, [3, 4, 23, 3], num_classes)
+    settings = pretrained_settings["resnet101"]["imagenet"]
+    initialize_pretrained_model(model, settings, pretrained)
+    return model
+