First_20200817

Author: ay

.vscode/settings.json

@@ -0,0 +1,3 @@
+{
+    "python.pythonPath": "/home/qust116-jq/anaconda3/envs/torch/bin/python"
+}

README.md

@@ -0,0 +1,19 @@
+# 语义分割学习实验-基于VOC数据集
+usage：
+1. 下载VOC数据集，将`JPEGImages` `SegmentationClass`两个文件夹放入到data文件夹下。
+2. 终端切换到目标目录，运行`python train.py -h`查看训练
+```bash
+(torch) qust116-jq@qustx-X299-WU8:~/语义分割$ python train.py -h
+usage: train.py [-h] [-m {Unet,FCN,Deeplab}] [-g GPU]
+
+choose the model
+
+optional arguments:
+  -h, --help            show this help message and exit
+  -m {Unet,FCN,Deeplab}, --model {Unet,FCN,Deeplab}
+                        输入模型名字
+  -g GPU, --gpu GPU     输入所需GPU
+```
+选择模型和GPU编号进行训练，例如运行`python train.py -m Unet -g 0`
+
+3. 预测需要手动修改`predict.py`中的模型

data/test.csv

@@ -0,0 +1,47 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import os
+
+# from model.resnet import ResNet18_OS16, ResNet34_OS16, ResNet50_OS16, ResNet101_OS16, ResNet152_OS16, ResNet18_OS8, ResNet34_OS8
+from model.resnet import ResNet18_OS8
+# from model.aspp import ASPP, ASPP_Bottleneck
+from model.aspp import ASPP
+
+class DeepLabV3(nn.Module):
+    def __init__(self,num_classes): #, model_id, project_dir
+        super(DeepLabV3, self).__init__()
+
+        self.num_classes = num_classes
+
+        # self.model_id = model_id
+        # self.project_dir = project_dir
+        # self.create_model_dirs()
+
+        self.resnet = ResNet18_OS8() # NOTE! specify the type of ResNet here
+        self.aspp = ASPP(num_classes=self.num_classes) # NOTE! if you use ResNet50-152, set self.aspp = ASPP_Bottleneck(num_classes=self.num_classes) instead
+
+    def forward(self, x):
+        # (x has shape (batch_size, 3, h, w))
+
+        h = x.size()[2]
+        w = x.size()[3]
+
+        feature_map = self.resnet(x) # (shape: (batch_size, 512, h/16, w/16)) (assuming self.resnet is ResNet18_OS16 or ResNet34_OS16. If self.resnet is ResNet18_OS8 or ResNet34_OS8, it will be (batch_size, 512, h/8, w/8). If self.resnet is ResNet50-152, it will be (batch_size, 4*512, h/16, w/16))
+
+        output = self.aspp(feature_map) # (shape: (batch_size, num_classes, h/16, w/16))
+
+        output = F.upsample(output, size=(h, w), mode="bilinear") # (shape: (batch_size, num_classes, h, w))
+
+        return output
+
+    # def create_model_dirs(self):
+    #     self.logs_dir = self.project_dir + "/training_logs"
+    #     self.model_dir = self.logs_dir + "/model_%s" % self.model_id
+    #     self.checkpoints_dir = self.model_dir + "/checkpoints"
+    #     if not os.path.exists(self.logs_dir):
+    #         os.makedirs(self.logs_dir)
+    #     if not os.path.exists(self.model_dir):
+    #         os.makedirs(self.model_dir)
+    #         os.makedirs(self.checkpoints_dir)

data/train.csv

@@ -0,0 +1,193 @@
+# -*- encoding: utf-8 -*-
+'''
+@File    :   unet.py
+@Time    :   2020/08/02 10:19:44
+@Author  :   AngYi
+@Contact :   [email protected]
+@Department   :  QDKD shuli
+@description : 
+'''
+
+# here put the import lib
+import torch
+from torch import nn
+from torchvision import models
+import numpy as np
+
+pretrained_model=models.vgg16(pretrained=True) #用于 FCN32x FCN16x FCN8x
+pretrained_net = models.resnet34(pretrained=True) #用于 FCN8s
+
+def bilinear_kernel(in_channels, out_channels, kernel_size):
+    '''
+    return a bilinear filter tensor
+    双线性卷积核，用于反卷积
+    '''
+    factor = (kernel_size + 1) // 2
+    if kernel_size % 2 == 1:
+        center = factor - 1
+    else:
+        center = factor - 0.5
+    og = np.ogrid[:kernel_size, :kernel_size]
+    filt = (1 - abs(og[0] - center) / factor) * (1 - abs(og[1] - center) / factor)
+    weight = np.zeros((in_channels, out_channels, kernel_size, kernel_size), dtype='float32')
+    weight[range(in_channels), range(out_channels), :, :] = filt
+    return torch.from_numpy(weight)
+
+class FCN32s(nn.Module):
+    def __init__(self,num_classes):
+        super(FCN32s, self).__init__()
+
+        self.feature=pretrained_model.features
+
+        self.conv=nn.Conv2d(512,num_classes, kernel_size=1, stride=1, padding=0)
+        self.upsample32x=nn.Sequential(
+            nn.ConvTranspose2d(num_classes, num_classes, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1),
+            nn.ConvTranspose2d(num_classes, num_classes, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1),
+            nn.ConvTranspose2d(num_classes, num_classes, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1),
+            nn.ConvTranspose2d(num_classes, num_classes, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1),
+            nn.ConvTranspose2d(num_classes, num_classes, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1),
+        )
+        
+        for m in self.modules():
+            if isinstance(m,nn.ConvTranspose2d):
+                m.weight.data.copy_(bilinear_kernel(int(m.in_channels),int(m.out_channels),m.kernel_size[0]))
+
+    def forward(self,x):
+        x=self.feature(x) # 1/32
+        x=self.conv(x)
+        x=self.upsample32x(x)
+        return x
+
+class FCN16s(nn.Module):
+    def __init__(self,num_classes):
+        super(FCN16s, self).__init__()
+
+        self.feature_1=nn.Sequential(*list(pretrained_model.features.children())[:24])
+        self.feature_2=nn.Sequential(*list(pretrained_model.features.children())[24:])
+
+        self.conv_1=nn.Conv2d(512,num_classes,kernel_size=1,stride=1,padding=0)
+        self.conv_2=nn.Conv2d(512, num_classes, kernel_size=1, stride=1, padding=0)
+
+        self.upsample2x=nn.ConvTranspose2d(num_classes,num_classes,kernel_size=3,stride=2,padding=1,output_padding=1,dilation=1)
+        self.upsample16x=nn.Sequential(
+            nn.ConvTranspose2d(num_classes, num_classes, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1),
+            nn.ConvTranspose2d(num_classes, num_classes, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1),
+            nn.ConvTranspose2d(num_classes, num_classes, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1),
+            nn.ConvTranspose2d(num_classes, num_classes, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1),
+        )
+
+        for m in self.modules():
+            if isinstance(m,nn.ConvTranspose2d):
+                m.weight.data.copy_(bilinear_kernel(m.in_channels,m.out_channels,m.kernel_size[0]))
+
+    def forward(self, x):
+        x1=self.feature_1(x)
+        x2=self.feature_2(x1)
+
+        x1=self.conv_1(x1)
+        x2=self.conv_2(x2)
+        x2=self.upsample2x(x2)
+        x2+=x1
+
+        x2=self.upsample16x(x2)
+        return x2
+
+
+class FCN8s(nn.Module):
+    def __init__(self, num_classes):
+        super(FCN8s, self).__init__()
+
+        self.feature_1 = nn.Sequential(*list(pretrained_model.features.children())[:17])
+        self.feature_2 = nn.Sequential(*list(pretrained_model.features.children())[17:24])
+        self.feature_3 = nn.Sequential(*list(pretrained_model.features.children())[24:])
+
+        self.conv_1 = nn.Conv2d(512,num_classes, kernel_size=1, stride=1, padding=0)
+        self.conv_2=nn.Conv2d(256,num_classes,kernel_size=1, stride=1, padding=0)
+        self.conv_3=nn.Conv2d(512,num_classes, kernel_size=1, stride=1, padding=0)
+
+        self.upsample2x_1 = nn.ConvTranspose2d(num_classes, num_classes, kernel_size=3, stride=2, padding=1, output_padding=1, dilation=1)
+        self.upsample2x_2 = nn.ConvTranspose2d(num_classes, num_classes, kernel_size=3, stride=2, padding=1, output_padding=1,dilation=1)
+        self.upsample8x = nn.Sequential(
+            nn.ConvTranspose2d(num_classes, num_classes, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1),
+            nn.ConvTranspose2d(num_classes, num_classes, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1),
+            nn.ConvTranspose2d(num_classes, num_classes, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1),
+        )
+
+        for m in self.modules():
+            if isinstance(m, nn.ConvTranspose2d):
+                m.weight.data=bilinear_kernel(m.in_channels, m.out_channels, m.kernel_size[0])
+
+    def forward(self, x):
+        x1 = self.feature_1(x)
+        x2 = self.feature_2(x1)
+        x3=self.feature_3(x2)
+
+        x2=self.conv_1(x2)
+        x3=self.conv_3(x3)
+        x3=self.upsample2x_1(x3)
+        x3+=x2
+
+        x1=self.conv_2(x1)
+        x3=self.upsample2x_2(x3)
+        x3+=x1
+
+        x3=self.upsample8x(x3)
+        return x3
+
+class FCN8x(nn.Module):
+    def __init__(self, num_classes):
+        super(FCN8x, self).__init__()
+
+        self.stage1 = nn.Sequential(*list(pretrained_net.children())[:-4])  # 第一段
+        self.stage2 = list(pretrained_net.children())[-4]  # 第二段
+        self.stage3 = list(pretrained_net.children())[-3]  # 第三段
+
+        self.scores1 = nn.Conv2d(512, num_classes, 1)
+        self.scores2 = nn.Conv2d(256, num_classes, 1)
+        self.scores3 = nn.Conv2d(128, num_classes, 1)
+
+        self.upsample_8x = nn.ConvTranspose2d(num_classes, num_classes, 16, 8, 4, bias=False)
+        self.upsample_8x.weight.data = bilinear_kernel(num_classes, num_classes, 16)  # 使用双线性 kernel
+
+        self.upsample_4x = nn.ConvTranspose2d(num_classes, num_classes, 4, 2, 1, bias=False)
+        self.upsample_4x.weight.data = bilinear_kernel(num_classes, num_classes, 4)  # 使用双线性 kernel
+
+        self.upsample_2x = nn.ConvTranspose2d(num_classes, num_classes, 4, 2, 1, bias=False)
+        self.upsample_2x.weight.data = bilinear_kernel(num_classes, num_classes, 4)  # 使用双线性 kernel
+
+    def forward(self, x):
+        x = self.stage1(x)
+        s1 = x  # 1/8
+
+        x = self.stage2(x)
+        s2 = x  # 1/16
+
+        x = self.stage3(x)
+        s3 = x  # 1/32
+
+        s3 = self.scores1(s3)
+        s3 = self.upsample_2x(s3)
+        s2 = self.scores2(s2)
+        s2 = s2 + s3
+
+        s1 = self.scores3(s1)
+        s2 = self.upsample_4x(s2)
+        s = s1 + s2
+
+        s = self.upsample_8x(s2)
+        return s
+
+
+
+
+
+
+
+if __name__ == "__main__":
+    pass
+    from torchsummary import summary
+    fcn = FCN32s(3)
+    # fcn.cuda(1)
+    # summary(fcn,(3,128,128))
+    # pretrained_model.cuda(1)
+    # summary(pretrained_model,(3,128,128))

data/val.csv

@@ -0,0 +1,118 @@
+""" Parts of the U-Net model """
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class DoubleConv(nn.Module):
+    """(convolution => [BN] => ReLU) * 2"""
+
+    def __init__(self, in_channels, out_channels, mid_channels=None):
+        super().__init__()
+        if not mid_channels:
+            mid_channels = out_channels
+        self.double_conv = nn.Sequential(
+            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1),
+            nn.BatchNorm2d(mid_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True)
+        )
+
+    def forward(self, x):
+        return self.double_conv(x)
+
+
+class Down(nn.Module):
+    """Downscaling with maxpool then double conv"""
+
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.maxpool_conv = nn.Sequential(
+            nn.MaxPool2d(2),
+            DoubleConv(in_channels, out_channels)
+        )
+
+    def forward(self, x):
+        return self.maxpool_conv(x)
+
+
+class Up(nn.Module):
+    """Upscaling then double conv"""
+
+    def __init__(self, in_channels, out_channels, bilinear=True):
+        super().__init__()
+
+        # if bilinear, use the normal convolutions to reduce the number of channels
+        if bilinear:
+            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+            self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
+        else:
+            self.up = nn.ConvTranspose2d(in_channels , in_channels // 2, kernel_size=2, stride=2)
+            self.conv = DoubleConv(in_channels, out_channels)
+
+
+    def forward(self, x1, x2):
+        x1 = self.up(x1)
+        # input is CHW
+        diffY = x2.size()[2] - x1.size()[2]
+        diffX = x2.size()[3] - x1.size()[3]
+
+        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
+                        diffY // 2, diffY - diffY // 2])
+        # if you have padding issues, see
+        # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
+        # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd
+        x = torch.cat([x2, x1], dim=1)
+        return self.conv(x)
+
+
+class OutConv(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(OutConv, self).__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+""" Full assembly of the parts to form the complete network """
+
+import torch.nn.functional as F
+
+
+
+class UNet(nn.Module):
+    def __init__(self, n_channels, n_classes, bilinear=True):
+        super(UNet, self).__init__()
+        self.n_channels = n_channels
+        self.n_classes = n_classes
+        self.bilinear = bilinear
+
+        self.inc = DoubleConv(n_channels, 64)
+        self.down1 = Down(64, 128)
+        self.down2 = Down(128, 256)
+        self.down3 = Down(256, 512)
+        factor = 2 if bilinear else 1
+        self.down4 = Down(512, 1024 // factor)
+        self.up1 = Up(1024, 512 // factor, bilinear)
+        self.up2 = Up(512, 256 // factor, bilinear)
+        self.up3 = Up(256, 128 // factor, bilinear)
+        self.up4 = Up(128, 64, bilinear)
+        self.outc = OutConv(64, n_classes)
+
+    def forward(self, x):
+        x1 = self.inc(x)
+        x2 = self.down1(x1)
+        x3 = self.down2(x2)
+        x4 = self.down3(x3)
+        x5 = self.down4(x4)
+        x = self.up1(x5, x4)
+        x = self.up2(x, x3)
+        x = self.up3(x, x2)
+        x = self.up4(x, x1)
+        logits = self.outc(x)
+        return logits
+