add DUC, GCN

Author: ycszen

Diff for: README.md

@@ -0,0 +1,8 @@
+# Pytorch for Semantic Segmentation
+A repository contains some exiting networks and some experimental networks for semantic segmentation.
++ [x] ResNet(FCN)
+  - [x] ResNet-50
+  - [x] ResNet-101
+  - [ ] Wide-ResNet
++ [x] DUC(*Understanding Convolution for Semantic Segmentation*)[pdf](https://arxiv.org/abs/1702.08502)
++ [x] GCN(*Large Kernel Matters -- Improve Semantic Segmentation by Global Convolutional Network*)[pdf](https://arxiv.org/abs/1703.02719)

Diff for: duc.py

@@ -0,0 +1,104 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.nn.init as init
+import torch.utils.model_zoo as model_zoo
+from torchvision import models
+
+import math
+
+
+class DUC(nn.Module):
+    def __init__(self, inplanes, planes, upscale_factor=2):
+        super(DUC, self).__init__()
+        self.relu = nn.ReLU()
+        self.conv = nn.Conv2d(inplanes, planes, kernel_size=3,
+                              padding=1)
+        self.bn = nn.BatchNorm2d(planes)
+        self.pixel_shuffle = nn.PixelShuffle(upscale_factor)
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        x = self.relu(x)
+        x = self.pixel_shuffle(x)
+        return x
+
+class FCN(nn.Module):
+    def __init__(self, num_classes):
+        super(FCN, self).__init__()
+
+        self.num_classes = num_classes
+
+        resnet = models.resnet50(pretrained=True)
+
+        self.conv1 = resnet.conv1
+        self.bn0 = resnet.bn1
+        self.relu = resnet.relu
+        self.maxpool = resnet.maxpool
+
+        self.layer1 = resnet.layer1
+        self.layer2 = resnet.layer2
+        self.layer3 = resnet.layer3
+        self.layer4 = resnet.layer4
+
+        self.duc1 = DUC(2048, 2048*2)
+        self.duc2 = DUC(1024, 1024*2)
+        self.duc3 = DUC(512, 512*2)
+        self.duc4 = DUC(128, 128*2)
+        self.duc5 = DUC(64, 64*2)
+
+        self.out1 = self._classifier(1024)
+        self.out2 = self._classifier(512)
+        self.out3 = self._classifier(128)
+        self.out4 = self._classifier(64)
+        self.out5 = self._classifier(32)
+
+        self.transformer = nn.Conv2d(320, 128, kernel_size=1)
+
+    def _classifier(self, inplanes):
+        if inplanes == 32:
+            return nn.Sequential(
+                nn.Conv2d(inplanes, self.num_classes, 1),
+                nn.Conv2d(self.num_classes, self.num_classes,
+                          kernel_size=3, padding=1)
+            )
+        return nn.Sequential(
+            nn.Conv2d(inplanes, inplanes/2, 3, padding=1, bias=False),
+            nn.BatchNorm2d(inplanes/2, momentum=.95),
+            nn.ReLU(inplace=True),
+            nn.Dropout(.1),
+            nn.Conv2d(inplanes/2, self.num_classes, 1),
+        )
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn0(x)
+        x = self.relu(x)
+        conv_x = x
+        x = self.maxpool(x)
+        pool_x = x
+
+        fm1 = self.layer1(x)
+        fm2 = self.layer2(fm1)
+        fm3 = self.layer3(fm2)
+        fm4 = self.layer4(fm3)
+
+        dfm1 = fm3 + self.duc1(fm4)
+        out16 = self.out1(dfm1)
+
+        dfm2 = fm2 + self.duc2(dfm1)
+        out8 = self.out2(dfm2)
+
+        dfm3 = fm1 + self.duc3(dfm2)
+
+        dfm3_t = self.transformer(torch.cat((dfm3, pool_x), 1))
+        out4 = self.out3(dfm3_t)
+
+        dfm4 = conv_x + self.duc4(dfm3_t)
+        out2 = self.out4(dfm4)
+
+        dfm5 = self.duc5(dfm4)
+        out = self.out5(dfm5)
+
+        return out, out2, out4, out8, out16

Diff for: gcn.py

@@ -0,0 +1,137 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.nn.init as init
+import torch.utils.model_zoo as model_zoo
+from torchvision import models
+
+import math
+
+
+class GCN(nn.Module):
+    def __init__(self, inplanes, planes, ks=7):
+        super(GCN, self).__init__()
+        self.conv_l1 = nn.Conv2d(inplanes, planes, kernel_size=(ks, 1),
+                                 padding=(ks/2, 0))
+
+        self.conv_l2 = nn.Conv2d(planes, planes, kernel_size=(1, ks),
+                                 padding=(0, ks/2))
+        self.conv_r1 = nn.Conv2d(inplanes, planes, kernel_size=(1, ks),
+                                 padding=(0, ks/2))
+        self.conv_r2 = nn.Conv2d(planes, planes, kernel_size=(ks, 1),
+                                 padding=(ks/2, 0))
+
+    def forward(self, x):
+        x_l = self.conv_l1(x)
+        x_l = self.conv_l2(x_l)
+
+        x_r = self.conv_r1(x)
+        x_r = self.conv_r2(x_r)
+
+        x = x_l + x_r
+
+        return x
+
+
+class Refine(nn.Module):
+    def __init__(self, planes):
+        super(Refine, self).__init__()
+        self.bn = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv1 = nn.Conv2d(planes, planes, kernel_size=3, padding=1)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, padding=1)
+
+    def forward(self, x):
+        residual = x
+        x = self.bn(x)
+        x = self.relu(x)
+        x = self.conv1(x)
+        x = self.bn(x)
+        x = self.relu(x)
+        x = self.conv2(x)
+
+        out = residual + x
+        return out
+
+
+class FCN(nn.Module):
+    def __init__(self, num_classes):
+        super(FCN, self).__init__()
+
+        self.num_classes = num_classes
+
+        resnet = models.resnet50(pretrained=True)
+
+        self.conv1 = resnet.conv1
+        self.bn0 = resnet.bn1
+        self.relu = resnet.relu
+        self.maxpool = resnet.maxpool
+
+        self.layer1 = resnet.layer1
+        self.layer2 = resnet.layer2
+        self.layer3 = resnet.layer3
+        self.layer4 = resnet.layer4
+
+        self.gcn1 = GCN(2048, self.num_classes)
+        self.gcn2 = GCN(1024, self.num_classes)
+        self.gcn3 = GCN(512, self.num_classes)
+        self.gcn4 = GCN(64, self.num_classes)
+        self.gcn5 = GCN(64, self.num_classes)
+
+        self.refine1 = Refine(self.num_classes)
+        self.refine2 = Refine(self.num_classes)
+        self.refine3 = Refine(self.num_classes)
+        self.refine4 = Refine(self.num_classes)
+        self.refine5 = Refine(self.num_classes)
+        self.refine6 = Refine(self.num_classes)
+        self.refine7 = Refine(self.num_classes)
+        self.refine8 = Refine(self.num_classes)
+        self.refine9 = Refine(self.num_classes)
+        self.refine10 = Refine(self.num_classes)
+
+        self.out0 = self._classifier(2048)
+        self.out1 = self._classifier(1024)
+        self.out2 = self._classifier(512)
+        self.out_e = self._classifier(256)
+        self.out3 = self._classifier(64)
+        self.out4 = self._classifier(64)
+        self.out5 = self._classifier(32)
+
+        self.transformer = nn.Conv2d(256, 64, kernel_size=1)
+
+    def _classifier(self, inplanes):
+        return nn.Sequential(
+            nn.Conv2d(inplanes, inplanes, 3, padding=1, bias=False),
+            nn.BatchNorm2d(inplanes/2),
+            nn.ReLU(inplace=True),
+            nn.Dropout(.1),
+            nn.Conv2d(inplanes/2, self.num_classes, 1),
+        )
+
+    def forward(self, x):
+        input = x
+        x = self.conv1(x)
+        x = self.bn0(x)
+        x = self.relu(x)
+        conv_x = x
+        x = self.maxpool(x)
+        pool_x = x
+
+        fm1 = self.layer1(x)
+        fm2 = self.layer2(fm1)
+        fm3 = self.layer3(fm2)
+        fm4 = self.layer4(fm3)
+
+        gcfm1 = self.refine1(self.gcn1(fm4))
+        gcfm2 = self.refine2(self.gcn2(fm3))
+        gcfm3 = self.refine3(self.gcn3(fm2))
+        gcfm4 = self.refine4(self.gcn4(pool_x))
+        gcfm5 = self.refine5(self.gcn5(conv_x))
+
+        fs1 = self.refine6(F.upsample_bilinear(gcfm1, fm3.size()[2:]) + gcfm2)
+        fs2 = self.refine7(F.upsample_bilinear(fs1, fm2.size()[2:]) + gcfm3)
+        fs3 = self.refine8(F.upsample_bilinear(fs2, pool_x.size()[2:]) + gcfm4)
+        fs4 = self.refine9(F.upsample_bilinear(fs3, conv_x.size()[2:]) + gcfm5)
+        out = self.refine10(F.upsample_bilinear(fs4, input.size()[2:]))
+
+        return out, fs4, fs3, fs2, fs1, gcfm1

Diff for: tester.py

@@ -0,0 +1,64 @@
+import torch
+from torch.utils import data
+import torch.optim as optim
+from torch.autograd import Variable
+from transform import Colorize
+from torchvision.transforms import ToPILImage, Compose, ToTensor, CenterCrop
+from transform import Scale
+# from resnet import FCN
+from upsample import FCN
+# from gcn import FCN
+from datasets import VOCTestSet
+from PIL import Image
+import numpy as np
+from tqdm import tqdm
+
+
+label_transform = Compose([Scale((256, 256), Image.BILINEAR), ToTensor()])
+batch_size = 1
+dst = VOCTestSet("./data", transform=label_transform)
+
+testloader = data.DataLoader(dst, batch_size=batch_size,
+                             num_workers=8)
+
+
+model = torch.nn.DataParallel(FCN(22), device_ids=[0, 1, 2, 3])
+# model = FCN(22)
+model.cuda()
+model.load_state_dict(torch.load("./pth/fcn-deconv-40.pth"))
+model.eval()
+
+
+# 10 13 48 86 101
+img = Image.open("./data/VOC2012test/JPEGImages/2008_000101.jpg").convert("RGB")
+original_size = img.size
+img.save("original.png")
+img = img.resize((256, 256), Image.BILINEAR)
+img = ToTensor()(img)
+img = Variable(img).unsqueeze(0)
+outputs = model(img)
+# 22 256 256
+for i, output in enumerate(outputs):
+    output = output[0].data.max(0)[1]
+    output = Colorize()(output)
+    output = np.transpose(output.numpy(), (1, 2, 0))
+    img = Image.fromarray(output, "RGB")
+    if i == 0:
+        img = img.resize(original_size, Image.NEAREST)
+    img.save("test-%d.png" % i)
+
+'''
+
+for index, (imgs, name, size) in tqdm(enumerate(testloader)):
+    imgs = Variable(imgs.cuda())
+    outputs = model(imgs)
+
+    output = outputs[0][0].data.max(0)[1]
+    output = Colorize()(output)
+    print(output)
+    output = np.transpose(output.numpy(), (1, 2, 0))
+    img = Image.fromarray(output, "RGB")
+    # img = Image.fromarray(output[0].cpu().numpy(), "P")
+    img = img.resize((size[0].numpy(), size[1].numpy()), Image.NEAREST)
+    img.save("./results/VOC2012/Segmentation/comp5_test_cls/%s.png" % name)
+'''