├── .gitattributes ├── README.md ├── data_list ├── test.csv ├── train.csv └── val.csv ├── ensemble.py ├── interfere.py ├── models ├── FCN-8s_Resnet101.py ├── Unet_ResNet101.py ├── __pycache__ │ ├── Unet_ResNet101.cpython-35.pyc │ └── deeplabv3p_Xception65.cpython-35.pyc └── deeplabv3p_Xception65.py ├── train.py └── utils ├── __pycache__ ├── create_dataset.cpython-35.pyc ├── img_process.cpython-35.pyc └── lab_process.cpython-35.pyc ├── create_dataset.py ├── img_process.py ├── lab_process.py └── make_list.py /.gitattributes: -------------------------------------------------------------------------------- 1 | # Auto detect text files and perform LF normalization 2 | * text=auto 3 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Apollo--laneline-detection 2 | python pytorch opencv resnet101_unet xception_65_deeplabv3plus 3 | 4 | the project details:https://blog.csdn.net/weixin_40279184/article/details/104296607 5 | 6 | -------------------------------------------------------------------------------- /ensemble.py: -------------------------------------------------------------------------------- 1 | #coding：utf-8 2 | #@author： Jiangnan He 3 | # date： 2020.02.13 4 | #模型融合的策略是三个模型预测值取平均作为最终的输出 5 | # 预测的图片名为 test_dir + network + color_mask.jpg 6 | 7 | ''' 8 | import os 9 | import numpy as np 10 | import cv2 11 | test_dir="test_example" 12 | out1=cv2.imread(os.path.join(test_dir,"deeplabv3p","color_mask")) 13 | out2 = cv2.imread(os.path.join(test_dir, "unet", "color_mask")) 14 | img1 = out1.astype(np.float32) 15 | img2 = out2.astype(np.float32) 16 | img = (img1+img2)/2.0 17 | img = img.astype(np.uint8) 18 | cv2.imwrite('final.jpg',img) 19 | ''' 20 | 21 | 22 | #用投票的思想，少数服从多数 23 | import numpy as np 24 | import cv2 25 | import os 26 | test_dir="test_example" 27 | RESULT_PREFIXX = ['deeplabv3p', 'unet'] 28 | # each mask has 8 classes: 0~7 29 | def vote_per_image(): 30 | result_list = [] 31 | for j in range(len(RESULT_PREFIXX)): 32 | im = cv2.imread( os.path.join(test_dir,RESULT_PREFIXX[j] )+ '.jpg', 0) 33 | result_list.append(im) 34 | # each pixel 35 | height, width = result_list[0].shape 36 | vote_mask = np.zeros((height, width)) 37 | 38 | for h in range(height): 39 | for w in range(width): 40 | record = np.zeros((1, 8)) 41 | for n in range(len(result_list)): 42 | mask = result_list[n] 43 | pixel = mask[h, w] 44 | # print('pix:',pixel) 45 | record[0, pixel] += 1 46 | label = record.argmax() 47 | # print(label) 48 | vote_mask[h, w] = label 49 | cv2.imwrite('vote_mask' + '.jpg', vote_mask) 50 | vote_per_image() 51 | -------------------------------------------------------------------------------- /interfere.py: -------------------------------------------------------------------------------- 1 | #coding=utf-8 2 | #@author Jiangnan He 3 | #@date: 2020.01.11 20:10 4 | 5 | import os 6 | import cv2 7 | import torch 8 | import numpy as np 9 | from models.deeplabv3p_Xception65 import deeplabv3p 10 | from models.Unet_ResNet101 import Unet_resnet101 11 | from utils.img_process import crop_resize_data 12 | from utils.lab_process import decode_color_labels 13 | from train import network 14 | 15 | 16 | 17 | 18 | if torch.cuda.is_available(): 19 | device_id=0 20 | predict_net = 'deeplabv3p' 21 | nets = {'deeplabv3p': deeplabv3p, 'unet': Unet_resnet101} 22 | 23 | def load_model(): 24 | if torch.cuda.is_available(): 25 | 26 | net=network(predict_net).getnet().cuda(device_id) 27 | else: 28 | net=network(predict_net).getnet() 29 | net.eval() 30 | # 加载网络参数 31 | net.load_state_dict(torch.load(os.path.join(os.getcwd(),'checkpoint.pt'))) 32 | return net 33 | 34 | def img_transform(img): 35 | img = crop_resize_data(img) 36 | img = np.transpose(img, (2, 0, 1)) 37 | img = img[np.newaxis, ...].astype(np.float32) 38 | img = torch.from_numpy(img.copy()) 39 | if torch.cuda.is_available(): 40 | img = img.cuda(device=device_id) 41 | return img 42 | 43 | def get_color_mask(pred): 44 | pred = torch.softmax(pred, dim=1) 45 | pred_heatmap = torch.max(pred, dim=1) 46 | # 1,H,W,C 47 | pred = torch.argmax(pred, dim=1) 48 | pred = torch.squeeze(pred) 49 | pred = pred.detach().cpu().numpy() 50 | pred = decode_color_labels(pred) 51 | pred = np.transpose(pred, (1, 2, 0)) 52 | return pred 53 | 54 | def main(): 55 | test_dir = 'test_example' 56 | net = load_model() 57 | img_path = os.path.join(test_dir, 'test.jpg') 58 | img = cv2.imread(img_path) 59 | img = img_transform(img) 60 | pred = net(img) 61 | color_mask = get_color_mask(pred) 62 | cv2.imwrite(os.path.join(test_dir, 'color_mask.jpg'), color_mask) 63 | 64 | if __name__ == '__main__': 65 | main() 66 | -------------------------------------------------------------------------------- /models/FCN-8s_Resnet101.py: -------------------------------------------------------------------------------- 1 | # -*- coding:utf-8 -*- 2 | #@author Jiangnan He 3 | #@date 2019.12.04 4 | import torch 5 | import torch.nn as nn 6 | from torch.autograd import Variable 7 | from torchsummary import summary 8 | 9 | 10 | class Block(nn.Module): 11 | def __init__(self,in_ch,out_ch,ksize=3,stride=1,padding=1): 12 | super(Block,self).__init__() 13 | self.conv1=nn.Conv2d(in_ch,out_ch,kernel_size=ksize,stride=stride,padding=padding) 14 | self.bn=nn.BatchNorm2d(out_ch) 15 | self.relu=nn.ReLU(inplace=True) 16 | def forward(self, x): 17 | return self.relu(self.bn(self.conv1(x))) 18 | 19 | def make_layers(in_channels, layer_list): 20 | layers = [] 21 | for v in layer_list: 22 | layers += [Block(in_channels, v)] 23 | in_channels = v 24 | return nn.Sequential(*layers) 25 | 26 | #层 27 | class Layer(nn.Module): 28 | def __init__(self, in_channels, layer_list): 29 | super(Layer, self).__init__() 30 | self.layer = make_layers(in_channels, layer_list) 31 | 32 | def forward(self, x): 33 | out = self.layer(x) 34 | return out 35 | 36 | #残差块 37 | class ResBlock(nn.Module): 38 | def __init__(self,ch_list,downsample,Res):# ch_list=[in_ch,ch,out_ch] 39 | super(ResBlock,self).__init__() 40 | 41 | self.res=Res# 残差块 还是 瓶颈块的标志位 42 | self.ds = downsample # 残差块时 是否下采样 43 | 44 | #第一个1x1 卷积 45 | self.firconv1x1=Block(ch_list[0],ch_list[1],1,1,0) #第一个1x1卷积 不下采样 46 | self.firconv1x1d = Block(ch_list[0],ch_list[1], 1,2,0) #第一个1x1卷积 下采样 47 | 48 | # 3x3卷积 49 | self.conv3x3=Block(ch_list[1],ch_list[1],3,1,1) 50 | 51 | #第二个1x1卷积 52 | self.secconv1x1=Block(ch_list[1],ch_list[2],1,1,0) 53 | 54 | #skip 卷积操作 55 | self.resconv1x1d=Block(ch_list[0],ch_list[2],1,2,0) #skip下采样的1x1卷积 56 | self.resconv1x1=Block(ch_list[0],ch_list[2],1,1,0) #skip下采样的1x1卷积 57 | 58 | 59 | def forward(self,x): 60 | if self.res==True:#此时为残差块 61 | if self.ds==True:#skip有卷积操作需要下采样 62 | residual=self.resconv1x1d(x) 63 | f1=self.firconv1x1d(x) 64 | else:#skip有卷积操作不需要下采样 65 | residual = self.resconv1x1(x) 66 | f1 = self.firconv1x1(x) 67 | 68 | else:# 瓶颈块sikp无卷积操作 69 | residual=x 70 | f1 = self.firconv1x1(x) 71 | f2=self.conv3x3(f1) 72 | f3=self.secconv1x1(f2) 73 | f3+=residual 74 | return f3 75 | 76 | #Res层 77 | class ResLayer(nn.Module): 78 | def __init__(self,ch_list1,ch_list2,downsample,numBotBlock): 79 | super(ResLayer,self).__init__() 80 | self.num = numBotBlock 81 | self.resb=ResBlock(ch_list1,downsample,True) 82 | self.botb=ResBlock(ch_list2,downsample,False) 83 | self.BoB=self.make_layers() 84 | 85 | def make_layers(self): 86 | layers=[] 87 | for i in range(self.num): 88 | layers+=[self.botb] 89 | return nn.Sequential(*layers) 90 | 91 | def forward(self,x): 92 | res=self.resb(x) 93 | return self.BoB(res) 94 | 95 | 96 | #resnet101 97 | class ResNet101(nn.Module): 98 | def __init__(self): 99 | super(ResNet101,self).__init__() 100 | self.conv1=nn.Conv2d(3,64,kernel_size=3,stride=2,padding=1) #1/2 101 | self.pool=nn.MaxPool2d(kernel_size=2,stride=2) #1/4 102 | self.layer1=ResLayer([64,64,256],[256,64,256],False,4) 103 | self.layer2=ResLayer([256,128,512],[512,128,512],True,3) #1/8 ch= 512 8 104 | self.layer3=ResLayer([512,256,1024],[1024,256,1024],True,23) #1/16 ch= 1024 4 105 | self.layer4=ResLayer([1024,512,2048],[2048,512,2048],True,6) #1/32 ch=2048 106 | 107 | def forward(self, x): 108 | layer1=self.layer1(self.pool(self.conv1(x)))# 109 | layer2=self.layer2(layer1) 110 | layer3 = self.layer3(layer2) 111 | layer4 = self.layer4(layer3) 112 | return [layer2,layer3,layer4] 113 | 114 | #上采样模块 115 | class Upsample(nn.Module): 116 | def __init__(self,in_ch1,in_ch2,upsampleratio=2):#1024 ,2048 117 | super(Upsample,self).__init__() 118 | self.conv1=Block(in_ch1,in_ch2) 119 | self.conv2=Block(in_ch2,in_ch2) 120 | self.trans_conv=nn.ConvTranspose2d(in_ch2,in_ch2,upsampleratio,stride=upsampleratio) 121 | 122 | def forward(self,featrue1,featrue2): 123 | return self.conv2(self.trans_conv(self.conv2(featrue2))+self.conv1(featrue1)) 124 | 125 | class FCNDecode(nn.Module): 126 | def __init__(self,n,in_ch,out_ch,upsratio):# in_ch VGG 后通道数 out_ch 为上采样之后的通道数 127 | super(FCNDecode,self).__init__() 128 | self.conv1=Layer(in_ch,[out_ch]*n)#这里加n层Block 129 | self.trans_conv=nn.ConvTranspose2d(out_ch,out_ch,upsratio,stride=upsratio)# upsratio 上采样倍数 130 | 131 | def forward(self,x): 132 | return self.trans_conv(self.conv1(x)) 133 | 134 | class FCN(nn.Module): 135 | def __init__(self,n,in_ch,out_ch,upsratio): 136 | super(FCN,self).__init__() 137 | self.encode=ResNet101() 138 | self.ups1=Upsample(512,2048) 139 | self.ups2=Upsample(1024,2048) 140 | self.decode=FCNDecode(n,in_ch,out_ch,upsratio) 141 | self.classfier=nn.Conv2d(out_ch,10,3,padding=1) 142 | 143 | def forward(self,x): 144 | features=self.encode(x) 145 | [layer2,layer3,layer4]=features# layer2 :512*8*8 layer3 : 1024*4*4 layer4 : 2048*2*2 146 | layer3=self.ups2(layer3,layer4) 147 | out=self.ups1(layer2,layer3) 148 | return self.classfier(self.decode(out)) 149 | 150 | if __name__=="__main__": 151 | if torch.cuda.is_available(): 152 | model=FCN(4,2048,256,8).cuda() 153 | x = Variable(torch.randn(10, 3, 64, 64)).cuda() 154 | else: 155 | model=FCN(4,2048,256,8) 156 | x = Variable(torch.randn(10, 3, 128, 128)) 157 | model.eval() 158 | summary(model,(3,64,64)) 159 | 160 | -------------------------------------------------------------------------------- /models/Unet_ResNet101.py: -------------------------------------------------------------------------------- 1 | #coding:utf-8 2 | #@author: Jiangnan He 3 | #@date:2019.12.10 18:21 4 | 5 | '''''' 6 | 7 | 8 | ''' 9 | resnet101-unet 10 | 1.上采样 还是采用 转置卷积 11 | 2.进行高低层特征融合时使用了torch.cat() 替换了 FCN 中的 "+" 12 | 3.resnet101 下采样1/32 所以进行了 5次下采样 本实现进行了 3次高低层特征融合 最后得到heatmap为原图的1/4 13 | 4.与label做loss 时直接将输出 heatmap 上采样到image尺寸 14 | 15 | layer0 input 16 | ↓ 17 | layer1 conv3x3 ch=64 (1/2) 18 | ↓ 19 | layer2 reslayer(maxpool) ch=256 (1/4) --------------------------- 256+256---conv3 ch=256 20 | ↓ ↑upconv2 21 | layer3 reslayer ch=512 (1/8) ------------------------------ 512+512---conv3 ch=512 22 | ↓ ↑upconv2 23 | layer4 reslayer ch=1024 (1/16)---------------------------------- 1024+1024--conv3 ch=1024 24 | ↓ ↑upconv2 25 | layer5 reslaeyr ch=2048 (1/32)----------------------------------------- 26 | ''' 27 | import torch 28 | import torch.nn as nn 29 | from torch.autograd import Variable 30 | from torchsummary import summary 31 | from torch.nn import functional as F 32 | 33 | class Block(nn.Module): 34 | def __init__(self,in_ch,out_ch,ksize=3,stride=1,padding=1): 35 | super(Block,self).__init__() 36 | self.conv1=nn.Conv2d(in_ch,out_ch,kernel_size=ksize,stride=stride,padding=padding) 37 | self.bn=nn.BatchNorm2d(out_ch) 38 | self.relu=nn.ReLU(inplace=True) 39 | def forward(self, x): 40 | return self.relu(self.bn(self.conv1(x))) 41 | 42 | def make_layers(in_channels, layer_list): 43 | layers = [] 44 | for v in layer_list: 45 | layers += [Block(in_channels, v)] 46 | in_channels = v 47 | return nn.Sequential(*layers) 48 | 49 | #层 50 | class Layer(nn.Module): 51 | def __init__(self, in_channels, layer_list): 52 | super(Layer, self).__init__() 53 | self.layer = make_layers(in_channels, layer_list) 54 | def forward(self, x): 55 | out = self.layer(x) 56 | return out 57 | 58 | #残差块 59 | class ResBlock(nn.Module): 60 | def __init__(self,ch_list,downsample,Res):# ch_list=[in_ch,ch,out_ch] 61 | super(ResBlock,self).__init__() 62 | 63 | self.res=Res# 残差块 还是 瓶颈块的标志位 64 | self.ds = downsample # 残差块时 是否下采样 65 | #第一个1x1 卷积 66 | self.firconv1x1=Block(ch_list[0],ch_list[1],1,1,0) #第一个1x1卷积 不下采样 67 | self.firconv1x1d = Block(ch_list[0],ch_list[1], 1,2,0) #第一个1x1卷积 下采样 68 | # 3x3卷积 69 | self.conv3x3=Block(ch_list[1],ch_list[1],3,1,1) 70 | #第二个1x1卷积 71 | self.secconv1x1=Block(ch_list[1],ch_list[2],1,1,0) 72 | #skip 卷积操作 73 | self.resconv1x1d=Block(ch_list[0],ch_list[2],1,2,0) #skip下采样的1x1卷积 74 | self.resconv1x1=Block(ch_list[0],ch_list[2],1,1,0) #skip下采样的1x1卷积 75 | 76 | def forward(self,x): 77 | if self.res==True:#此时为残差块 78 | if self.ds==True:#skip有卷积操作需要下采样 79 | residual=self.resconv1x1d(x) 80 | f1=self.firconv1x1d(x) 81 | else:#skip有卷积操作不需要下采样 82 | residual = self.resconv1x1(x) 83 | f1 = self.firconv1x1(x) 84 | 85 | else:# 瓶颈块sikp无卷积操作 86 | residual=x 87 | f1 = self.firconv1x1(x) 88 | f2=self.conv3x3(f1) 89 | f3=self.secconv1x1(f2) 90 | f3+=residual 91 | return f3 92 | 93 | #Res层 94 | class ResLayer(nn.Module): 95 | def __init__(self,ch_list1,ch_list2,downsample,numBotBlock): 96 | super(ResLayer,self).__init__() 97 | self.num = numBotBlock 98 | self.resb=ResBlock(ch_list1,downsample,True) 99 | self.botb=ResBlock(ch_list2,downsample,False) 100 | self.BoB=self.make_layers() 101 | def make_layers(self): 102 | layers=[] 103 | for i in range(self.num): 104 | layers+=[self.botb] 105 | return nn.Sequential(*layers) 106 | def forward(self,x): 107 | res=self.resb(x) 108 | return self.BoB(res) 109 | 110 | #resnet101 111 | class ResNet101(nn.Module): 112 | def __init__(self): 113 | super(ResNet101,self).__init__() 114 | self.conv0=nn.Conv2d(3,64,kernel_size=3,stride=1,padding=1) 115 | self.conv1=nn.Conv2d(64,64,kernel_size=3,stride=2,padding=1) #1/2 116 | self.pool=nn.MaxPool2d(kernel_size=2,stride=2) #1/4 117 | self.layer1=ResLayer([64,64,256],[256,64,256],False,4) 118 | self.layer2=ResLayer([256,128,512],[512,128,512],True,3) #1/8 ch= 512 119 | self.layer3=ResLayer([512,256,1024],[1024,256,1024],True,23) #1/16 ch= 1024 120 | self.layer4=ResLayer([1024,512,2048],[2048,512,2048],True,6) #1/32 ch=2048 121 | #权重初始化 conv 以及BN 122 | for m in self.modules(): 123 | if isinstance(m, nn.Conv2d): 124 | nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')#mode="fan_in" weight方差在前向传播中保持不变 mode="fan_out" weight后向传播方差不变 125 | elif isinstance(m, nn.BatchNorm2d): 126 | nn.init.constant_(m.weight, 1) 127 | nn.init.constant_(m.bias, 0) 128 | 129 | def forward(self, x): 130 | layer0=self.conv0(x) #1 131 | layer1=self.conv1(layer0) #1/2 132 | layer2 =self.layer1(self.pool(layer1) ) #1/4 133 | layer3=self.layer2(layer2) #1/8 134 | layer4 = self.layer3(layer3) #1/16 135 | layer5 = self.layer4(layer4) #1/32 136 | #print(layer0.size(),layer1.size(),layer2.size(),layer3.size(),layer4.size(),layer5.size()) 137 | return [layer2,layer3,layer4,layer5] 138 | 139 | 140 | 141 | #上采样模块 142 | class Upsample(nn.Module): 143 | 144 | def __init__(self,in_ch1,in_ch2,upsampleratio=2):# layer42048 layer3 1024 145 | super(Upsample,self).__init__() 146 | self.conv1=Block(in_ch1,in_ch2)# layer4 2048-->1024 147 | self.conv2=Block(in_ch2,in_ch2)#layer3 1024-->1024 148 | self.conv3=Block(in_ch2*2,in_ch2) # cat layer3 2048--> 1024 149 | self.trans_conv=nn.ConvTranspose2d(in_ch2,in_ch2,upsampleratio,stride=upsampleratio) 150 | def forward(self,featrue1,featrue2):#layer4 ,layer3 151 | # 152 | return self.conv3( torch.cat([ self.trans_conv(self.conv1(featrue1)), self.conv2(featrue2) ],dim=1) ) 153 | 154 | class UnetDecode(nn.Module): 155 | def __init__(self,n,in_ch,out_ch,upsratio,trans_conv=False): # in_ch VGG 后通道数 out_ch 为上采样之后的通道数 156 | super(UnetDecode,self).__init__() 157 | self.conv1=Layer(in_ch,[out_ch]*n) #这里加n层Block 158 | self.trans_conv=nn.ConvTranspose2d(out_ch,out_ch,upsratio,stride=upsratio)# upsratio 上采样倍数 159 | 160 | def forward(self,x): 161 | return self.trans_conv(self.conv1(x)) 162 | 163 | class Unet_resnet101(nn.Module): 164 | def __init__(self,n,in_ch,out_ch,upsratio): 165 | super(Unet_resnet101,self).__init__() 166 | self.encode=ResNet101() 167 | self.catup1=Upsample(2048,1024) 168 | self.catup2=Upsample(1024,512) 169 | self.catup3=Upsample(512,256) 170 | self.decode=UnetDecode(n,in_ch,out_ch,upsratio) 171 | self.classfier=nn.Conv2d(256,10,3,padding=1) 172 | 173 | def forward(self,x): 174 | size=x.shape[-2:] 175 | features=self.encode(x) 176 | [layer2,layer3,layer4,layer5]=features# 177 | #layer1 256*16 layer2 :512*8*8 layer3 : 1024*4*4 layer4 : 2048*2*2 178 | layer4=self.catup1(layer5,layer4) 179 | layer3=self.catup2(layer4,layer3) 180 | layer2=self.catup3(layer3,layer2) 181 | #return self.classfier(layer2) 182 | return F.interpolate(self.classfier(layer2), size=size, mode='bilinear', align_corners=False) 183 | 184 | if __name__=="__main__": 185 | if torch.cuda.is_available(): 186 | model=Unet_resnet101(4,2048,64,4).cuda() 187 | x = Variable(torch.randn(1, 3, 64, 64)).cuda() 188 | else: 189 | model=Unet_resnet101(4,2048,64,4) 190 | x = Variable(torch.randn(1, 3, 64, 64)) 191 | model.eval() 192 | y=model(x) 193 | print("y.size",y.size()) 194 | summary(model,(3,64,64)) 195 | 196 | -------------------------------------------------------------------------------- /models/__pycache__/Unet_ResNet101.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Johnnan002/Apollo--laneline-detection/3f715540933b6ad2d71bd8cd1e08a412c53da331/models/__pycache__/Unet_ResNet101.cpython-35.pyc -------------------------------------------------------------------------------- /models/__pycache__/deeplabv3p_Xception65.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Johnnan002/Apollo--laneline-detection/3f715540933b6ad2d71bd8cd1e08a412c53da331/models/__pycache__/deeplabv3p_Xception65.cpython-35.pyc -------------------------------------------------------------------------------- /models/deeplabv3p_Xception65.py: -------------------------------------------------------------------------------- 1 | #coding:utf-8 2 | #@author: Jiangnan He 3 | #@date:2019.12.17 15:03 4 | '' 5 | ''' 6 | 本实现为deeplabv3+，backbone为Xception_65 7 | ''' 8 | 9 | import torch.nn as nn 10 | import torch 11 | from torchsummary import summary 12 | from torch.nn import functional as F 13 | #实现深度分离卷积 将卷积分程in_ch 组进行卷积 接着用1x1 卷积进行融合 14 | class depthwiseconv(nn.Module): 15 | def __init__(self, in_ch, out_ch,ksize=3,stride=1,padding=1): 16 | super(depthwiseconv, self).__init__() 17 | self.depth_conv = nn.Conv2d( 18 | in_channels=in_ch, 19 | out_channels=in_ch, 20 | kernel_size=ksize, 21 | stride=stride, 22 | padding=padding, 23 | groups=in_ch 24 | ) 25 | self.point_conv = nn.Conv2d( 26 | in_channels=in_ch, 27 | out_channels=out_ch, 28 | kernel_size=1, 29 | stride=1, 30 | padding=0, 31 | groups=1 32 | ) 33 | def forward(self, x): 34 | return self.point_conv(self.depth_conv(x)) 35 | 36 | #块 37 | class Block(nn.Module): 38 | def __init__(self,in_ch,out_ch,ksize=3,stride=1,padding=1,sep=True): 39 | super(Block,self).__init__() 40 | #是否为深度可分离卷积模块标志位 41 | self.issepconv = sep 42 | self.conv=nn.Conv2d(in_ch,out_ch,kernel_size=ksize,stride=stride,padding=padding,bias=False) 43 | self.sepconv = depthwiseconv(in_ch, out_ch, ksize, stride, padding) 44 | self.Bn=nn.BatchNorm2d(out_ch) 45 | self.relu=nn.ReLU(inplace=True) 46 | 47 | def forward(self,x): 48 | if self.issepconv==True: 49 | return self.relu(self.Bn(self.sepconv(x))) 50 | else: 51 | return self.relu(self.Bn(self.conv(x))) 52 | 53 | 54 | #层 55 | class Layer(nn.Module): 56 | def __init__(self,in_ch,out_ch,ksize,stride,padding=1,upsample=True):# 64 128 3 1 1 57 | super(Layer,self).__init__() 58 | self.upsample=upsample#是否进行下采样 59 | self.block1=Block(in_ch,out_ch,ksize,stride,padding) 60 | self.block2 = Block(out_ch, out_ch, ksize, stride, padding) 61 | self.block3 = Block(out_ch, out_ch, ksize, stride, padding)#不进行下采样时ksize=3 padding=1 same 62 | #下采样时 63 | self.block4=Block(out_ch, out_ch, ksize, stride=2, padding=1)#进行下采样时 64 | self.skipconv=nn.Conv2d(in_ch,out_ch,kernel_size=1,stride=2,padding=0)#padding设置注意 65 | self.skipconv1=nn.Conv2d(in_ch,out_ch,kernel_size=1,stride=1,padding=0)#padding设置注意 66 | if self.upsample==True:#进行下采样时 67 | layer = [self.block1, self.block2, self.block4] 68 | else:#不进行下采样时的组合 69 | layer=[self.block1,self.block2,self.block3] 70 | self.layer = self.layer=nn.Sequential(*layer) 71 | 72 | def forward(self,x): 73 | if self.upsample==True: 74 | out=self.layer(x) 75 | skipout=self.skipconv(x) 76 | return torch.cat([out,skipout],dim=1) 77 | else: 78 | out=self.layer(x) 79 | skipout=self.skipconv1(x) 80 | return torch.cat([out,skipout],dim=1) 81 | 82 | 83 | 84 | #模型 Xception 作为backbone 85 | class DCNN(nn.Module): 86 | def __init__(self): 87 | super(DCNN,self).__init__() 88 | self.conv1=Block(3,32,3,2,1,False)#conv 1/2 89 | self.conv2=Block(32,64,3,1,1,False)#conv 90 | self.layer1=Layer(64,128,3,1,1,True)#128 1/4 91 | self.layer2=Layer(256,256,3,1,1,True) #512 /8 92 | self.layer3 =Layer( 512,728, 3, 1, 1, True)# 1456 /16 93 | self.BlockLayer=Layer(1456,728,3,1,1,False)#repeat 16times 94 | self.layer4 = self.make_layer(16) 95 | self.layer5=Layer(1456,1024,3,1,1,True) #1/32 96 | self.conv3=Block(2048,1536,3,1,1,True)#sepconv 97 | self.conv4 = Block(1536, 1536, 3, 1, 1, True) # sepconv 98 | self.conv5=Block(1536, 2048, 3, 1, 1, True) # sepconv 99 | 100 | def make_layer(self,numlayer): 101 | layer=[] 102 | for i in range(numlayer): 103 | layer+=[self.BlockLayer] 104 | return nn.Sequential(*layer) 105 | 106 | def forward(self, x): 107 | #entry flow 108 | conv1=self.conv1(x) 109 | conv2=self.conv2(conv1) 110 | layer1=self.layer1(conv2)# 111 | layer2=self.layer2(layer1) 112 | layer3 = self.layer3(layer2) 113 | #middlle flow 114 | layer4 = self.layer4(layer3) 115 | #exit flow 116 | layer5 = self.layer5(layer4) 117 | conv3=self.conv3(layer5) 118 | conv4=self.conv4(conv3)# 119 | conv5=self.conv5(conv4) 120 | # print("DCNN ",conv1.size(), conv2.size(),layer1.size(),layer2.size(),layer3.size() ,layer4.size(), layer5.size(), conv3.size(),conv4.size(),conv5.size()) 121 | return layer2,conv5 122 | 123 | class ASPPConv(nn.Sequential): 124 | def __init__(self, in_channels, out_channels, dilation): 125 | modules = [ 126 | nn.Conv2d(in_channels, out_channels, 3, padding=dilation, dilation=dilation, bias=False), 127 | nn.BatchNorm2d(out_channels), 128 | nn.ReLU() 129 | ] 130 | super(ASPPConv, self).__init__(*modules) 131 | 132 | class ASPPPooling(nn.Sequential): 133 | def __init__(self, in_channels, out_channels): 134 | super(ASPPPooling, self).__init__( 135 | nn.AdaptiveAvgPool2d(1), 136 | nn.Conv2d(in_channels, out_channels, 1, bias=False), 137 | nn.BatchNorm2d(out_channels), 138 | nn.ReLU()) 139 | 140 | def forward(self, x): 141 | size = x.shape[-2:] 142 | for mod in self: 143 | x = mod(x) 144 | return F.interpolate(x, size=size, mode='bilinear', align_corners=False) 145 | 146 | 147 | class ASPP(nn.Module): 148 | def __init__(self, in_channels): 149 | super(ASPP, self).__init__() 150 | out_channels = 256 151 | modules = [] 152 | modules.append(nn.Sequential( 153 | nn.Conv2d(in_channels, out_channels, 1, bias=False), 154 | nn.BatchNorm2d(out_channels), 155 | nn.ReLU())) 156 | 157 | rate1, rate2, rate3 = (6,12,18) 158 | modules.append(ASPPConv(in_channels, out_channels, rate1)) 159 | modules.append(ASPPConv(in_channels, out_channels, rate2)) 160 | modules.append(ASPPConv(in_channels, out_channels, rate3)) 161 | modules.append(ASPPPooling(in_channels, out_channels)) 162 | self.convs = nn.ModuleList(modules) 163 | 164 | self.project = nn.Sequential( 165 | nn.Conv2d(5 * out_channels, out_channels, 1, bias=False), 166 | nn.BatchNorm2d(out_channels), 167 | nn.ReLU(), 168 | nn.Dropout(0.5)) 169 | 170 | def forward(self, x): 171 | res = [] 172 | for conv in self.convs: 173 | res.append(conv(x)) 174 | res = torch.cat(res, dim=1) 175 | return self.project(res) 176 | 177 | 178 | class deeplabv3p(nn.Module): 179 | def __init__(self): 180 | super(deeplabv3p,self).__init__() 181 | self.DCNN=DCNN() 182 | self.ASPP=ASPP(2048) 183 | self.conv1=Block(512,256,ksize=1,padding=0,sep=False) 184 | self.conv2=Block(2*256,8,ksize=3,padding=1,sep=False)# 原本为 3x3 这里测试模型 输入很小 所以改为1x1 185 | self.conv3 = Block(256, 256, ksize=1, padding=0, sep=False) 186 | 187 | 188 | #初始化参数 189 | for m in self.modules(): 190 | if isinstance(m, nn.Conv2d): 191 | nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')#mode="fan_in" weight方差在前向传播中保持不变 mode="fan_out" weight后向传播方差不变 192 | elif isinstance(m, nn.BatchNorm2d): 193 | nn.init.constant_(m.weight, 1) 194 | nn.init.constant_(m.bias, 0) 195 | 196 | def forward(self,x): 197 | lowfeature,out=self.DCNN(x) 198 | lowfeature=self.conv1(lowfeature)# 199 | #print(lowfeature.size(),'lowfeaturesize1') 200 | f1=self.conv3(self.ASPP(out)) 201 | #print(f1.size(),'f1') 202 | hightfeature=F.interpolate(f1,scale_factor=4, mode='bilinear', align_corners=False) 203 | #print(hightfeature.size(),"hightfeatsize") 204 | out1=torch.cat([lowfeature,hightfeature],dim=1) 205 | return F.interpolate(self.conv2(out1),scale_factor=8, mode='bilinear', align_corners=False) 206 | 207 | if __name__=="__main__": 208 | if torch.cuda.is_available(): 209 | x=torch.randn(1,3,128, 384).cuda() 210 | model=deeplabv3p().cuda() 211 | else: 212 | x=torch.randn(1,3,128, 384) 213 | model=deeplabv3p() 214 | model.eval() 215 | y=model(x) 216 | print(y.size(),"y.size") 217 | 218 | 219 | 220 | ''' 221 | summary(model, (3, 128, 384)) 222 | ''' 223 | 224 | 225 | 226 | 227 | 228 | 229 | 230 | -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | #coding=utf-8 2 | #@author Jiangnan He 3 | #@date: 2019.12.29 20:10 4 | 5 | ''' 6 | 实现用来训练模型 7 | 本实现有3个模型 deeplabv3+ heatmap为原图的1/2 unet-resnet101 heatmap为原图的1/4 解决方案通过双线性插值到原图大小再和label 做loss 8 | 训练模型 1. metrics, 9 | 2.loss, 10 | 3.选择选了网络, 11 | 4.lr, save model, 多卡训练 , earlystop 12 | ''' 13 | 14 | import os 15 | import torch 16 | import torch.nn as nn 17 | import numpy as np 18 | import torch.nn.functional as F 19 | from models.deeplabv3p_Xception65 import deeplabv3p 20 | from models.Unet_ResNet101 import Unet_resnet101 21 | from utils.create_dataset import DatasetFromCSV 22 | from torchvision import transforms 23 | from torch.utils.data.dataset import Dataset 24 | from time import time 25 | from utils.img_process import Process_dataset 26 | from utils.img_process import ToTensor 27 | from tqdm import tqdm 28 | 29 | 30 | #========================================== 31 | #loss: bce, dice, focal_loss , lovasz_loss, 32 | #========================================== 33 | 34 | #bce_loss 35 | class MySoftmaxCrossEntropyLoss(nn.Module): 36 | def __init__(self, nbclasses): 37 | super(MySoftmaxCrossEntropyLoss, self).__init__() 38 | self.nbclasses = nbclasses 39 | 40 | def forward(self, inputs, target): 41 | if inputs.dim() > 2: 42 | inputs = inputs.view(inputs.size(0), inputs.size(1), -1) # N,C,H,W => N,C,H*W 43 | inputs = inputs.transpose(1, 2) # N,C,H*W => N,H*W,C 44 | inputs = inputs.contiguous().view(-1, self.nbclasses) # N,H*W,C => N*H*W,C 45 | target = target.view(-1) 46 | return nn.CrossEntropyLoss(reduction="mean")(inputs, target) 47 | 48 | 49 | #dice_loss 50 | class BinaryDiceLoss(nn.Module): 51 | """Dice loss of binary class 52 | Args: 53 | smooth: A float number to smooth loss, and avoid NaN error, default: 1 54 | p: Denominator value: \sum{x^p} + \sum{y^p}, default: 2 55 | predict: A tensor of shape [N, *] 56 | target: A tensor of shape same with predict 57 | reduction: Reduction method to apply, return mean over batch if 'mean', 58 | return sum if 'sum', return a tensor of shape [N,] if 'none' 59 | Returns: 60 | Loss tensor according to arg reduction 61 | Raise: 62 | Exception if unexpected reduction 63 | """ 64 | def __init__(self, smooth=1, p=2, reduction='mean'): 65 | super(BinaryDiceLoss, self).__init__() 66 | self.smooth = smooth 67 | self.p = p 68 | self.reduction = reduction 69 | 70 | def forward(self, predict, target): 71 | assert predict.shape[0] == target.shape[0], "predict & target batch size don't match" 72 | predict = predict.contiguous().view(predict.shape[0], -1) 73 | target = target.contiguous().view(target.shape[0], -1) 74 | num = 2*torch.sum(torch.mul(predict, target), dim=1) + self.smooth 75 | den = torch.sum(predict.pow(self.p) + target.pow(self.p), dim=1) + self.smooth 76 | loss = 1 - num / den 77 | if self.reduction == 'mean': 78 | return loss.mean() 79 | elif self.reduction == 'sum': 80 | return loss.sum() 81 | elif self.reduction == 'none': 82 | return loss 83 | else: 84 | raise Exception('Unexpected reduction {}'.format(self.reduction)) 85 | 86 | 87 | class DiceLoss(nn.Module): 88 | """Dice loss, need one hot encode input 89 | Args: 90 | weight: An array of shape [num_classes,] 91 | ignore_index: class index to ignore 92 | predict: A tensor of shape [N, C, *] 93 | target: A tensor of same shape with predict 94 | other args pass to BinaryDiceLoss 95 | Return: 96 | same as BinaryDiceLoss 97 | """ 98 | def __init__(self, weight=None, ignore_index=None, **kwargs): 99 | super(DiceLoss, self).__init__() 100 | self.kwargs = kwargs 101 | self.weight = weight 102 | self.ignore_index = ignore_index 103 | 104 | def forward(self, predict, target): 105 | assert predict.shape == target.shape, 'predict & target shape do not match' 106 | dice = BinaryDiceLoss(**self.kwargs) 107 | total_loss = 0 108 | predict = F.softmax(predict, dim=1) 109 | for i in range(target.shape[1]): 110 | if i != self.ignore_index: 111 | dice_loss = dice(predict[:, i], target[:, i]) 112 | if self.weight is not None: 113 | assert self.weight.shape[0] == target.shape[1], \ 114 | 'Expect weight shape [{}], get[{}]'.format(target.shape[1], self.weight.shape[0]) 115 | dice_loss *= self.weights[i] 116 | total_loss += dice_loss 117 | return total_loss/target.shape[1] 118 | 119 | 120 | #focal_loss = -(1-pt)^gamma*log(pt) 121 | class FocalLoss2d(nn.Module): 122 | def __init__(self, gamma=2, size_average=True): 123 | super(FocalLoss2d, self).__init__() 124 | self.gamma = gamma 125 | self.size_average = size_average 126 | 127 | def forward(self, logit, target, class_weight=None, type='softmax'): 128 | target = target.view(-1, 1).long() 129 | if type=='sigmoid': 130 | if class_weight is None: 131 | class_weight = [1]*2 #[0.5, 0.5] 132 | prob = F.sigmoid(logit) 133 | prob = prob.view(-1, 1) 134 | prob = torch.cat((1-prob, prob), 1) 135 | select = torch.FloatTensor(len(prob), 2).zero_().cuda() 136 | select.scatter_(1, target, 1.) 137 | elif type=='softmax': 138 | B,C,H,W = logit.size() #one-hot编码 139 | if class_weight is None: 140 | class_weight =[1]*C #[1/C]*C 141 | logit = logit.permute(0, 2, 3, 1).contiguous().view(-1, C) 142 | prob = F.softmax(logit,1) 143 | select = torch.FloatTensor(len(prob), C).zero_().cuda() 144 | select.scatter_(1, target, 1.) 145 | class_weight = torch.FloatTensor(class_weight).cuda().view(-1,1) 146 | class_weight = torch.gather(class_weight, 0, target) 147 | prob = (prob*select).sum(1).view(-1,1) 148 | prob = torch.clamp(prob,1e-8,1-1e-8) 149 | batch_loss = - class_weight *(torch.pow((1-prob), self.gamma))*prob.log() 150 | if self.size_average: 151 | loss = batch_loss.mean() 152 | else: 153 | loss = batch_loss 154 | return loss 155 | 156 | 157 | 158 | 159 | def make_one_hot(input, num_classes): 160 | """Convert class index tensor to one hot encoding tensor. 161 | Args: 162 | input: A tensor of shape [N, 1, *] 163 | num_classes: An int of number of class 164 | Returns: 165 | A tensor of shape [N, num_classes, *] 166 | """ 167 | shape = np.array(input.shape) 168 | shape[1] = num_classes 169 | shape = tuple(shape) 170 | result = torch.zeros(shape) 171 | result = result.scatter_(1, input.cpu(), 1) 172 | return result 173 | 174 | 175 | #========================================================== 176 | #metric 177 | #========================================================== 178 | 179 | def compute_iou(pred, gt, result): 180 | """ 181 | pred : [N, H, W] 182 | gt: [N, H, W] 183 | """ 184 | pred = pred.cpu().numpy() 185 | gt = gt.cpu().numpy() 186 | for i in range(8): 187 | single_gt=gt==i 188 | single_pred=pred==i 189 | temp_tp = np.sum(single_gt * single_pred) 190 | temp_ta = np.sum(single_pred) + np.sum(single_gt) - temp_tp 191 | result["TP"][i] += temp_tp 192 | result["TA"][i] += temp_ta 193 | return result #计算iou 194 | 195 | 196 | #EarlyStopping 197 | class EarlyStopping: 198 | """Early stops the training if validation loss doesn't improve after a given patience.""" 199 | def __init__(self, patience=7, verbose=False, delta=0): 200 | """ 201 | Args: 202 | patience (int): How long to wait after last time validation loss improved. 203 | Default: 7 204 | verbose (bool): If True, prints a message for each validation loss improvement. 205 | Default: False 206 | delta (float): Minimum change in the monitored quantity to qualify as an improvement. 207 | Default: 0 208 | """ 209 | self.patience = patience 210 | self.verbose = verbose 211 | self.counter = 0 212 | self.best_score = None 213 | self.early_stop = False 214 | self.val_loss_min = np.Inf 215 | self.delta = delta 216 | def __call__(self, val_loss, model): 217 | score = -val_loss 218 | if self.best_score is None:#best_score 为空 保存一次模型 219 | self.best_score = score 220 | self.save_checkpoint(val_loss, model) 221 | elif score < self.best_score + self.delta: 222 | self.counter += 1 223 | print('EarlyStopping counter: {} out of {}'.format(self.counter,self.patience)) 224 | if self.counter >= self.patience:#计数大于忍耐次数，早停 225 | self.early_stop = True 226 | else: 227 | self.best_score = score 228 | self.save_checkpoint(val_loss, model) 229 | self.counter = 0 230 | 231 | def save_checkpoint(self, val_loss, model): 232 | ''' Saves model when validation loss decrease. ''' 233 | if self.verbose: 234 | print('Validation loss decreased ({} --> {}). Saving model ...'.format(self.val_loss_min,val_loss)) 235 | torch.save(model.state_dict(), os.path.join(os.getcwd(),'checkpoint.pt'))#保存模型到磁盘 path='checkpoint.pt' 236 | self.val_loss_min = val_loss 237 | 238 | #Network 239 | class network(): 240 | def __init__(self,net): 241 | self.net=net 242 | self.model=None 243 | def getnet(self): 244 | if self.net=="deeplabv3p": 245 | self.model= deeplabv3p() 246 | elif self.net =="Unet-ResNet101": 247 | self.model=Unet_resnet101(4,2048,64,1) 248 | return self.model 249 | 250 | #train lr save_model 251 | #cycle lr cosine lr 252 | 253 | def train(): 254 | #============================================================== 255 | # check if gpu is available 256 | #============================================================== 257 | train_on_gpu=torch.cuda.is_available() 258 | if train_on_gpu: 259 | print("cuda is available,train on gpu") 260 | device_list = [0, 6] # 261 | print(device_list) 262 | else: 263 | device_list="cpu" 264 | print("cuda is not available,train on cpu") 265 | #=============================================================== 266 | # Config 267 | #=============================================================== 268 | trained_network_name="deeplabv3p" # "deeplabv3p", "Unet-ResNet101" 269 | batch_size=2 270 | print("batch_size is:"+str(batch_size)) 271 | epoch=10 272 | print("epoch is:"+str(epoch)) 273 | #trained_network_name="deeplabv3p" 274 | lr = 3e-4# Adam 最佳初始化3e-4或者5e-4 275 | numclasses = 8 276 | #=============================================================== 277 | #data 278 | #=============================================================== 279 | 280 | train_data_list="./data_list/train.csv" 281 | val_data_list="./data_list/val.csv" 282 | transform = transforms.Compose([Process_dataset(), ToTensor()]) 283 | train_data=DatasetFromCSV(train_data_list,transform) 284 | transform = transforms.Compose([ ToTensor()]) 285 | val_data=DatasetFromCSV(val_data_list,transform) 286 | train_loader = torch.utils.data.DataLoader(train_data, batch_size=batch_size) 287 | val_loader = torch.utils.data.DataLoader(val_data, batch_size=batch_size) 288 | print() 289 | #================================================================= 290 | # Adam 291 | #================================================================= 292 | model = network(trained_network_name).getnet().cuda(device=device_list[0]) 293 | # 并行处理 294 | #model = torch.nn.DataParallel(model, device_ids=device_list) 295 | 296 | opt = torch.optim.Adam(model.parameters(), lr=lr) 297 | lam = 0.5 298 | MAX_STEP = int(1e10) 299 | #cycle lr 设置 学习率的设置 周期衰减 或者周期变化挑出局部最优cosine lr 300 | scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(opt, MAX_STEP, eta_min=1e-5)#1e-6推荐 301 | 302 | #Warming up 303 | #================================================================= 304 | #Train loop 305 | #================================================================= 306 | early_stopping = EarlyStopping(patience=10, verbose=True)#1 307 | since = time() 308 | for e in range(epoch): 309 | print(e) 310 | train_loss=0.0 311 | val_loss=0.0 312 | lr=scheduler.get_lr() 313 | print('epoch:{},lr:{}'.format(e,lr)) 314 | #=============================================================== 315 | #train 316 | #=============================================================== 317 | model.train() 318 | for x_batch,y_batch in tqdm(train_loader):#训练数据 319 | for x,y in zip(x_batch,y_batch): 320 | y = torch.unsqueeze(y, 0).long() 321 | x = torch.unsqueeze(x, 0) 322 | y = torch.unsqueeze(y, 0).long() 323 | y_=make_one_hot(y,numclasses).cuda(device=device_list[0]) 324 | x=x.cuda(device=device_list[0]) 325 | y=y.cuda(device=device_list[0]) 326 | opt.zero_grad()#梯度置为0 327 | y_pre=model(x) 328 | DLoss=DiceLoss() 329 | BLoss=MySoftmaxCrossEntropyLoss(numclasses) 330 | loss=lam*DLoss(y_pre,y_)+(1-lam)*BLoss(y_pre,y)#计算loss 331 | #loss=BLoss(y_pre,y) 332 | train_loss+=loss.item() 333 | loss.backward()#求导 334 | opt.step()#更新 335 | scheduler.step(e) 336 | #======================================================================== 337 | # validata 338 | #======================================================================== 339 | model.eval()# changes the forward() behaviour of the module it is called upon. eg, it disables dropout and has batch norm use the entire population statistics 340 | torch.no_grad() # disables tracking of gradients in autograd. 341 | result = {"TP": {i: 0 for i in range(8)}, "TA": {i: 0 for i in range(8)}} 342 | for x_batch1,y_batch1 in tqdm(val_loader): 343 | for x1,y1 in zip(x_batch1,y_batch1): 344 | y1 = torch.unsqueeze(y1, 0).long() 345 | x1 = torch.unsqueeze(x1, 0) 346 | y1 = torch.unsqueeze(y1, 0).long() 347 | y1_ = make_one_hot(y1, numclasses).cuda(device=device_list[0]) 348 | x1, y1 = x1.cuda(device=device_list[0]), y1.cuda(device=device_list[0]) 349 | y_pred1 = model(x1) 350 | #计算 dice_loss ce_loss 351 | DLoss1=DiceLoss() 352 | BLoss1 = MySoftmaxCrossEntropyLoss(numclasses) 353 | loss1=lam*DLoss1(y_pred1,y1_)+(1-lam)*BLoss1(y_pred1,y1)#计算loss 354 | val_loss+=loss1.item() 355 | #计算metric 356 | pred = torch.argmax(F.softmax(y_pred1, dim=1), dim=1) 357 | result = compute_iou(pred, y1, result) # 计算iou 358 | print("Epoch:{}, test loss is {:.4f} \n".format(epoch, val_loss / len(val_loader))) 359 | for i in range(8): 360 | result_string = "{}: {:.4f} \n".format(i, result["TP"][i] / result["TA"][i]) 361 | print(result_string) 362 | #=========================================================== 363 | # Print training information 364 | #=========================================================== 365 | 366 | print_msg = ('[{epoch:>{}}/{n_epochs:>{}}] ' .format(e,epoch)+ 367 | 'train_loss: {} '.format(train_loss) + 368 | 'valid_loss: {}'.format(val_loss)) 369 | print(print_msg) 370 | 371 | #========================================================== 372 | # Early Stopping 373 | #========================================================== 374 | 375 | early_stopping(val_loss, model) # 相当于 earlystopping.__call__(valid_loss, model) 376 | if early_stopping.early_stop: 377 | print("Early stopping") 378 | break 379 | 380 | #=============================================== 381 | # load the last checkpoint with the best model 382 | #=============================================== 383 | 384 | model.load_state_dict(torch.load('checkpoint.pt')) 385 | time_elapsed = time()- since 386 | print('training_time:{}'.format(time_elapsed/60)) 387 | 388 | if __name__=="__main__": 389 | train() 390 | 391 | 392 | 393 | 394 | -------------------------------------------------------------------------------- /utils/__pycache__/create_dataset.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Johnnan002/Apollo--laneline-detection/3f715540933b6ad2d71bd8cd1e08a412c53da331/utils/__pycache__/create_dataset.cpython-35.pyc -------------------------------------------------------------------------------- /utils/__pycache__/img_process.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Johnnan002/Apollo--laneline-detection/3f715540933b6ad2d71bd8cd1e08a412c53da331/utils/__pycache__/img_process.cpython-35.pyc -------------------------------------------------------------------------------- /utils/__pycache__/lab_process.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Johnnan002/Apollo--laneline-detection/3f715540933b6ad2d71bd8cd1e08a412c53da331/utils/__pycache__/lab_process.cpython-35.pyc -------------------------------------------------------------------------------- /utils/create_dataset.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | # @author：Jiangnan He 3 | # @date： 2019.12.25 18:54 4 | # 本实现用于生成 pytorch 的数据集 5 | 6 | import torch 7 | import pandas as pd 8 | from torchvision import transforms 9 | from torch.utils.data.dataset import Dataset 10 | from utils.lab_process import encode_labels 11 | from utils.img_process import Process_dataset 12 | from utils.img_process import ToTensor 13 | import cv2 14 | 15 | 16 | class DatasetFromCSV(Dataset): 17 | def __init__(self, csv_path, transforms=None): 18 | self.data = pd.read_csv(csv_path,header=None, names=["image","label"]) 19 | self.images = self.data["image"].values[1:] 20 | self.labels = self.data["label"].values[1:] 21 | self.transforms = transforms 22 | 23 | def __len__(self): 24 | return self.labels.shape[0] 25 | 26 | def __getitem__(self, index): 27 | image=self.images[index] 28 | label=self.labels[index] 29 | image=cv2.imread(image,cv2.IMREAD_COLOR) 30 | lab_as_np=cv2.imread(label, cv2.IMREAD_GRAYSCALE) 31 | label=encode_labels(lab_as_np)#将不同灰度值标记为不同类 # 后续需要将其转为one-hot编码 32 | sample=[image,label] 33 | img_as_tensor ,lab_as_tensor= self.transforms(sample) 34 | return img_as_tensor, lab_as_tensor 35 | 36 | #测试 37 | if __name__=="__main__": 38 | batch_size =2 39 | transform = transforms.Compose([Process_dataset(),ToTensor()] ) 40 | train_data = DatasetFromCSV('../data_list/train.csv',transform) 41 | val_data = DatasetFromCSV('../data_list/val.csv', transform) 42 | test_data = DatasetFromCSV("../data_list/test.csv",transform) 43 | train_loader = torch.utils.data.DataLoader(train_data, batch_size=batch_size) 44 | val_loader = torch.utils.data.DataLoader(val_data, batch_size=batch_size) 45 | test_loader = torch.utils.data.DataLoader(test_data, batch_size=batch_size) 46 | 47 | 48 | for batch in train_loader: 49 | img, lab = batch 50 | lab=torch.unsqueeze(lab,1) 51 | print(img.shape,lab.shape) 52 | 53 | 54 | -------------------------------------------------------------------------------- /utils/img_process.py: -------------------------------------------------------------------------------- 1 | #coding: utf-8 2 | #@author jiangnan He 3 | #@date: 15:00 2019.11.13 4 | ''' 5 | APollo laneline dataset 6 | ''' 7 | 8 | ############################### config ########################## 9 | 10 | #原图尺寸3384x1710 wh 11 | #这里采用三种尺寸进行训练 768x256,1024x384,1536x512（wh） 12 | imgH=128 13 | imgW=384 14 | offset=690 15 | ############################################################### 16 | 17 | import numpy as np 18 | import cv2 19 | import torch 20 | 21 | def gamma_transform(img, gamma): 22 | gamma_table = [np.power(x / 255.0, gamma) * 255.0 for x in range(256)] 23 | gamma_table = np.round(np.array(gamma_table)).astype(np.uint8) 24 | return cv2.LUT(img, gamma_table) 25 | 26 | def random_gamma_transform(img, gamma_vari): 27 | log_gamma_vari = np.log(gamma_vari) 28 | alpha = np.random.uniform(-log_gamma_vari, log_gamma_vari) 29 | gamma = np.exp(alpha) 30 | return gamma_transform(img, gamma) 31 | 32 | def rotate(xb,yb, angle): 33 | M_rotate = cv2.getRotationMatrix2D((imgH/2, imgW/2), angle, 1) 34 | xb = cv2.warpAffine(xb, M_rotate, (imgH , imgW )) 35 | yb = cv2.warpAffine(yb, M_rotate, (imgH , imgW )) 36 | return xb, yb 37 | 38 | def blur(img): 39 | img = cv2.blur(img, (3, 3)) 40 | return img 41 | 42 | def add_noise(img): 43 | for i in range(200): # 添加点噪声 44 | temp_x = np.random.randint(0, img.shape[0]) 45 | temp_y = np.random.randint(0, img.shape[1]) 46 | img[temp_x][temp_y] = 255 47 | return img 48 | 49 | def data_augment(xb, yb): 50 | 51 | if np.random.random() < 0.25:#只对原图操作 52 | xb = random_gamma_transform(xb, 1.0) 53 | if np.random.random() < 0.25:#只对原图操作 54 | xb = blur(xb) 55 | 56 | if np.random.random() < 0.2:#只对原图操作 57 | xb = add_noise(xb) 58 | return xb, yb 59 | 60 | class Process_dataset(object): 61 | def __call__(self,sample): 62 | image, mask=sample 63 | src_img=image[offset:,:]#hw 64 | gt_img =mask[offset:,:] #先将 原图上部分的填空部分 3384x690 剪裁掉 65 | #1 图像尺度处理 将图像缩放到 训练尺寸 这里采用三种尺寸进行训练 768x256,1024x384,1536x512（wh） 66 | src_img = cv2.resize(src_img, 67 | dsize=(imgW, imgH), 68 | interpolation=cv2.INTER_LINEAR) 69 | gt_img = cv2.resize(gt_img, 70 | dsize=(imgW, imgH), 71 | interpolation=cv2.INTER_NEAREST)# 这里用最近邻处理 72 | #样本增强 73 | src_img,gt_img=data_augment(src_img,gt_img) 74 | return src_img,gt_img 75 | 76 | class ToTensor(object): 77 | def __call__(self,sample ): 78 | image, mask=sample 79 | image = np.transpose(image,(2,0,1)) 80 | image = image.astype(np.float32) 81 | mask = mask.astype(np.uint8) 82 | return torch.from_numpy(image.copy()), torch.from_numpy(mask.copy()) 83 | 84 | 85 | 86 | # crop the image to discard useless parts 87 | def crop_resize_data(image, label=None, image_size=(1024, 384), offset=690): 88 | """ 89 | Attention: 90 | h,w, c = image.shape 91 | cv2.resize(image,(w,h)) 92 | """ 93 | roi_image = image[offset:, :] 94 | if label is not None: 95 | roi_label = label[offset:, :] 96 | train_image = cv2.resize(roi_image, image_size, interpolation=cv2.INTER_LINEAR) 97 | train_label = cv2.resize(roi_label, image_size, interpolation=cv2.INTER_NEAREST) 98 | return train_image, train_label 99 | else: 100 | train_image = cv2.resize(roi_image, image_size, interpolation=cv2.INTER_LINEAR) 101 | return train_image 102 | -------------------------------------------------------------------------------- /utils/lab_process.py: -------------------------------------------------------------------------------- 1 | #coding=utf-8 2 | #@author : Jiangnan He 3 | #@data ： 2019.12.25 20：00 4 | import numpy as np 5 | def encode_labels(color_mask):#编码 灰度值转为类别 0-7 6 | encode_mask = np.zeros((color_mask.shape[0], color_mask.shape[1])) 7 | # 0 8 | encode_mask[color_mask == 0] = 0 9 | encode_mask[color_mask == 249] = 0 10 | encode_mask[color_mask == 255] = 0 11 | # 1 12 | encode_mask[color_mask == 200] = 1 13 | encode_mask[color_mask == 204] = 1 14 | encode_mask[color_mask == 213] = 0 15 | encode_mask[color_mask == 209] = 1 16 | encode_mask[color_mask == 206] = 0 17 | encode_mask[color_mask == 207] = 0 18 | # 2 19 | encode_mask[color_mask == 201] = 2 20 | encode_mask[color_mask == 203] = 2 21 | encode_mask[color_mask == 211] = 0 22 | encode_mask[color_mask == 208] = 0 23 | # 3 24 | encode_mask[color_mask == 216] = 0 25 | encode_mask[color_mask == 217] = 3 26 | encode_mask[color_mask == 215] = 0 27 | # 4 In the test, it will be ignored 28 | encode_mask[color_mask == 218] = 0 29 | encode_mask[color_mask == 219] = 0 30 | # 4 31 | encode_mask[color_mask == 210] = 4 32 | encode_mask[color_mask == 232] = 0 33 | # 5 34 | encode_mask[color_mask == 214] = 5 35 | # 6 36 | encode_mask[color_mask == 202] = 0 37 | encode_mask[color_mask == 220] = 6 38 | encode_mask[color_mask == 221] = 6 39 | encode_mask[color_mask == 222] = 6 40 | encode_mask[color_mask == 231] = 0 41 | encode_mask[color_mask == 224] = 6 42 | encode_mask[color_mask == 225] = 6 43 | encode_mask[color_mask == 226] = 6 44 | encode_mask[color_mask == 230] = 0 45 | encode_mask[color_mask == 228] = 0 46 | encode_mask[color_mask == 229] = 0 47 | encode_mask[color_mask == 233] = 0 48 | # 7 49 | encode_mask[color_mask == 205] = 7 50 | encode_mask[color_mask == 212] = 0 51 | encode_mask[color_mask == 227] = 7 52 | encode_mask[color_mask == 223] = 0 53 | encode_mask[color_mask == 250] = 7 54 | return encode_mask 55 | 56 | 57 | #将0-7 个类别翻译成 gray灰度值 58 | def decode_labels(labels):#解码 gray 59 | deocde_mask = np.zeros((labels.shape[0], labels.shape[1]), dtype='uint8') 60 | # 0 61 | deocde_mask[labels == 0] = 0 62 | # 1 63 | deocde_mask[labels == 1] = 204 64 | # 2 65 | deocde_mask[labels == 2] = 203 66 | # 3 67 | deocde_mask[labels == 3] = 217 68 | # 4 69 | deocde_mask[labels == 4] = 210 70 | # 5 71 | deocde_mask[labels == 5] = 214 72 | # 6 73 | deocde_mask[labels == 6] = 224 74 | # 7 75 | deocde_mask[labels == 7] = 227 76 | 77 | return deocde_mask 78 | 79 | 80 | def decode_color_labels(labels):#解码 rgb 81 | 82 | decode_mask = np.zeros((3, labels.shape[0], labels.shape[1]), dtype='uint8') 83 | # 0 84 | decode_mask[0][labels == 0] = 0 85 | decode_mask[1][labels == 0] = 0 86 | decode_mask[2][labels == 0] = 0 87 | # 1 88 | decode_mask[0][labels == 1] = 70 89 | decode_mask[1][labels == 1] = 130 90 | decode_mask[2][labels == 1] = 180 91 | # 2 92 | decode_mask[0][labels == 2] = 0 93 | decode_mask[1][labels == 2] = 0 94 | decode_mask[2][labels == 2] = 142 95 | # 3 96 | decode_mask[0][labels == 3] = 153 97 | decode_mask[1][labels == 3] = 153 98 | decode_mask[2][labels == 3] = 153 99 | # 4 100 | decode_mask[0][labels == 4] = 128 101 | decode_mask[1][labels == 4] = 64 102 | decode_mask[2][labels == 4] = 128 103 | # 5 104 | decode_mask[0][labels == 5] = 190 105 | decode_mask[1][labels == 5] = 153 106 | decode_mask[2][labels == 5] = 153 107 | # 6 108 | decode_mask[0][labels == 6] = 0 109 | decode_mask[1][labels == 6] = 0 110 | decode_mask[2][labels == 6] = 230 111 | # 7 112 | decode_mask[0][labels == 7] = 255 113 | decode_mask[1][labels == 7] = 128 114 | decode_mask[2][labels == 7] = 0 115 | return decode_mask -------------------------------------------------------------------------------- /utils/make_list.py: -------------------------------------------------------------------------------- 1 | #coding=utf-8 2 | #@author : Jiangnan He 3 | #@date : 2019.12.25 15:47 4 | 5 | """ 6 | img_process 生成的文件结构如下： 7 | --train_set :D:\Dataset\Apollolaneline\train_set - gt_image 8 | - src_image 9 | 10 | -- test_set :D:\Dataset\Apollolaneline\test_set - src_image 11 | """ 12 | import pandas as pd 13 | import os 14 | from sklearn.utils import shuffle 15 | import glob 16 | # train_set 共有21914 个样本 我们按照 train：val：test=3:1：1 的比例划分 训练集和 验证集 测试集 17 | datasetpath1=datasetpath='D:/Dataset/Apollolaneline/' 18 | #================================================ 19 | # make train & validation & test lists 20 | #================================================ 21 | def get_train_val_test(): 22 | img_list=[] 23 | lab_list=[] 24 | for file in glob.glob(str(datasetpath1)+"/Road*/"): 25 | if file.find("zip")!=-1: 26 | continue 27 | for i in glob.glob(str(file)+"/ColorImage_road*"): 28 | file1=os.path.join(datasetpath1,i,"ColorImage") 29 | for j in os.listdir(file1): 30 | path=os.path.join(file1,j) 31 | for k in os.listdir(path): 32 | path1=os.path.join(path,k) 33 | print(path1,len(os.listdir(path1))) 34 | for z in os.listdir(path1): 35 | image_path=os.path.join(path1,z) 36 | # D:\Dataset\Apollolaneline\Road02\ColorImage_road02\ColorImage\Record001\Camera 5\170927_063811892_Camera_5.jpg 37 | #D:\Dataset\Apollolaneline\Gray_Label\Gray_Label\Label_road02\Label\Record001\Camera 5\170927_063811892_Camera_5_bin.png 38 | print(image_path,"image_path") 39 | image_name=str(image_path.split()[1]).split('\\')[1].split(".")[0]+'_bin.png' 40 | print("image_name",image_name) 41 | #获取src_img 文件地址 42 | src_path=os.path.join(path1,z) 43 | # 获取gt_img 文件地址 44 | gt_path=os.path.join("D:/Dataset/Apollolaneline/Gray_Label/Gray_Label/","Label_"+str(image_path.split("\\")[2]).split("_")[1],"Label","/".join(image_path.split("\\")[4:6]),image_name) 45 | img_list.append(src_path) 46 | lab_list.append(gt_path) 47 | assert len(img_list)==len(lab_list) 48 | total_len=len(img_list) 49 | six_part=int(total_len*0.6) 50 | eight_part=int(total_len*0.8) 51 | all=pd.DataFrame({'image': img_list, 'label': lab_list}) 52 | all_shuffle = shuffle(all) # 打乱存储 53 | train_list=all_shuffle[:six_part] 54 | val_list=all_shuffle[six_part:eight_part] 55 | test_list=all_shuffle[eight_part:] 56 | #生成val.csv 57 | train_list.to_csv('../data_list/train.csv', index=False)#生成val的数据集csv文件 58 | # 生成train.csv 59 | val_list.to_csv('../data_list/val.csv', index=False)#生成train的数据集csv文件 60 | #生成test.csv 61 | test_list.to_csv('../data_list/test.csv', index=False) # 生成test的数据集csv文件 62 | 63 | if __name__=="__main__": 64 | get_train_val_test() 65 | 66 | 67 | --------------------------------------------------------------------------------