diff --git a/README.md b/README.md index f63d402..833bc3d 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,20 @@ # pytorch-lanenet 效果很好的lanenet网络，主干网络基于bisenetv2并对主干网络做了修改，效果远好于bisnetv2 + 可直接训练自己的数据应用于生产 + +inspired by https://github.com/MaybeShewill-CV/lanenet-lane-detection and https://github.com/leonfrank/lanenet-danet-pytorch + +Using Bisenetv2 as Encoder. + +install(安装) + + +python setup.py install + +Usage(使用) + +Train(训练) + + + diff --git a/lanenet/config.py b/lanenet/config.py index 7807a23..4483c8f 100644 --- a/lanenet/config.py +++ b/lanenet/config.py @@ -1,8 +1,30 @@ # coding: utf-8 +# number of unique pixel values n_labels=6 +#embeding特征数 no_of_instances=6 +#分割目标的种类数 只分割车道线值为2 num_classes=2 +#device_ids 中第一个gpu号 gpu_no='cuda:0' +# 使用的gpu device_ids = [0, 1,2,3,4,5,6] +lr=0.001 +epochs=2000 +bs=16 +show_interval=30 +save_interval=3 # 1-train，2-val -is_training=1 \ No newline at end of file +is_training=1 +# 模型保存位置 +save_path='./checkpoints' +#loss权重设置详见compute_loss +w1=0.25 +#loss权重设置详见compute_loss +w2=0.25 +#loss权重设置详见compute_loss +w3=0.25 +#loss权重设置详见compute_loss +w4=0.25 +train_dataset_file = '/workspace/all/index/train1' +val_dataset_file = '/workspace/all/index/val1' \ No newline at end of file diff --git a/lanenet/model/BiseNet.py b/lanenet/model/BiseNet.py deleted file mode 100644 index 322c8cc..0000000 --- a/lanenet/model/BiseNet.py +++ /dev/null @@ -1,248 +0,0 @@ -import torch -import torch.nn as nn -import torch.nn.init as init -import torch.nn.functional as F -from torch.utils import model_zoo -from torchvision import models - -class conv2d(nn.Module): - def __init__(self,in_dim,out_dim,k,pad,stride,groups = 1,bias=False,use_bn = True,use_rl = True): - super(conv2d,self).__init__() - self.use_bn = use_bn - self.use_rl = use_rl - self.conv = nn.Conv2d(in_dim,out_dim,k,padding=pad,stride=stride, groups=groups,bias=bias) - self.bn = nn.BatchNorm2d(out_dim) - self.relu = nn.ReLU(inplace=True) - def forward(self,bottom): - if self.use_bn and self.use_rl: - return self.relu(self.bn(self.conv(bottom))) - elif self.use_bn: - return self.bn(self.conv(bottom)) - else: - return self.conv(bottom) - -class StemBlock(nn.Module): - def __init__(self): - super(StemBlock,self).__init__() - self.conv1 = conv2d(3,16,3,1,2) - self.conv_1x1 = conv2d(16,8,1,0,1) - self.conv_3x3 = conv2d(8,16,3,1,2) - self.mpooling = nn.MaxPool2d(3,2,1) - self.conv2 = conv2d(32,16,3,1,1) - def forward(self,bottom): - base = self.conv1(bottom) - conv_1 = self.conv_1x1(base) - conv_3 = self.conv_3x3(conv_1) - pool = self.mpooling(base) - cat = torch.cat([conv_3,pool],1) - res = self.conv2(cat) - return res - -class ContextEmbeddingBlock(nn.Module): - def __init__(self,in_dim): - super(ContextEmbeddingBlock,self).__init__() - self.gap = nn.AdaptiveAvgPool2d(1)#1 - self.bn1 = nn.BatchNorm2d(in_dim) - self.conv1 = conv2d(in_dim,in_dim,1,0,1) - self.conv2 = conv2d(in_dim,in_dim,3,1,1,use_bn = False,use_rl = False) - def forward(self,bottom): - gap = self.gap(bottom) - bn = self.bn1(gap) - conv1 = self.conv1(bn) - feat = bottom+conv1 - res = self.conv2(feat) - return res - -class GatherExpansion(nn.Module): - def __init__(self,in_dim,out_dim,stride = 1,exp = 6): - super(GatherExpansion,self).__init__() - exp_dim = in_dim*exp - self.stride = stride - self.conv1 = conv2d(in_dim,exp_dim,3,1,1) - self.dwconv2 = conv2d(exp_dim,exp_dim,3,1,1,exp_dim,use_rl = False) - self.conv_11 = conv2d(exp_dim,out_dim,1,0,1,use_rl = False) - - self.dwconv1 = conv2d(exp_dim,exp_dim,3,1,2,exp_dim,use_rl = False) - self.dwconv3 = conv2d(in_dim,in_dim,3,1,2,in_dim,use_rl = False) - self.conv_12 = conv2d(in_dim,out_dim,1,0,1,use_rl = False) - self.relu = nn.ReLU(inplace=True) - def forward(self,bottom): - base = self.conv1(bottom) - if self.stride == 2: - base = self.dwconv1(base) - bottom = self.dwconv3(bottom) - bottom = self.conv_12(bottom) - x = self.dwconv2(base) - x = self.conv_11(x) - res = self.relu(x+bottom) - return res - -class BGA(nn.Module): - def __init__(self,in_dim): - super(BGA,self).__init__() - self.in_dim = in_dim - self.db_dwconv = conv2d(in_dim,in_dim,3,1,1,in_dim,use_rl=False) - self.db_conv1x1 = conv2d(in_dim,in_dim,1,0,1,use_rl=False,use_bn=False) - self.db_conv = conv2d(in_dim,in_dim,3,1,2,use_rl=False) - self.db_apooling = nn.AvgPool2d(3,2,1) - - self.sb_dwconv = conv2d(in_dim,in_dim,3,1,1,in_dim,use_rl=False) - self.sb_conv1x1 = conv2d(in_dim,in_dim,1,0,1,use_rl=False,use_bn=False) - self.sb_conv = conv2d(in_dim,in_dim,3,1,1,use_rl=False) - self.sb_sigmoid = nn.Sigmoid() - - self.conv = conv2d(in_dim,in_dim,3,1,1,use_rl=False) - def forward(self,db,sb): - db_dwc = self.db_dwconv(db) - db_out = self.db_conv1x1(db_dwc)# - db_conv = self.db_conv(db) - db_pool = self.db_apooling(db_conv) - - sb_dwc = self.sb_dwconv(sb) - sb_out = self.sb_sigmoid(self.sb_conv1x1(sb_dwc))# - sb_conv = self.sb_conv(sb) - sb_up = self.sb_sigmoid(F.interpolate(sb_conv, size=db_out.size()[2:], mode="bilinear",align_corners=True)) - db_l = db_out*sb_up - sb_r = F.interpolate(sb_out*db_pool, size=db_out.size()[2:], mode="bilinear",align_corners=True) - res = self.conv(db_l+sb_r) - return res - -class SegHead(nn.Module): - def __init__(self,in_dim,out_dim,cls,size=[512,1024]): - super(SegHead,self).__init__() - self.size = size - self.conv = conv2d(in_dim,out_dim,3,1,1) - self.cls = conv2d(out_dim,cls,1,0,1,use_bn=False,use_rl=False) - def forward(self,feat): - x = self.conv(feat) - x = self.cls(x) - pred = F.interpolate(x, size=self.size, mode="bilinear",align_corners=True) - return pred - - -class DetailedBranch(nn.Module): - def __init__(self): - super(DetailedBranch,self).__init__() - self.s1_conv1 = conv2d(3,64,3,1,2) - self.s1_conv2 = conv2d(64,64,3,1,1) - - self.s2_conv1 = conv2d(64,64,3,1,2) - self.s2_conv2 = conv2d(64,64,3,1,1) - self.s2_conv3 = conv2d(64,64,3,1,1) - - self.s3_conv1 = conv2d(64,128,3,1,2) - self.s3_conv2 = conv2d(128,128,3,1,1) - self.s3_conv3 = conv2d(128,128,3,1,1) - def forward(self,bottom): - s1_1 = self.s1_conv1(bottom) - s1_2 = self.s1_conv2(s1_1) - - s2_1 = self.s2_conv1(s1_2) - s2_2 = self.s2_conv2(s2_1) - s2_3 = self.s2_conv3(s2_2) - - s3_1 = self.s3_conv1(s2_3) - s3_2 = self.s3_conv2(s3_1) - s3_3 = self.s3_conv3(s3_2) - return s3_3 - -class SemanticBranch(nn.Module): - def __init__(self, cls): - super(SemanticBranch,self).__init__() - self.stem = StemBlock() - self.s3_ge1 = GatherExpansion(16,32,2) - self.s3_ge2 = GatherExpansion(32,32) - - self.s4_ge1 = GatherExpansion(32,64,2) - self.s4_ge2 = GatherExpansion(64,64) - - self.s5_ge1 = GatherExpansion(64,128,2) - self.s5_ge2 = GatherExpansion(128,128) - self.s5_ge3 = GatherExpansion(128,128) - self.s5_ge4 = GatherExpansion(128,128) - self.s5_ge5 = GatherExpansion(128,128,exp=1) - if self.training: - self.seghead1 = SegHead(16,16,cls) - self.seghead2 = SegHead(32,32,cls) - self.seghead3 = SegHead(64,64,cls) - self.seghead4 = SegHead(128,128,cls) - - self.ceb = ContextEmbeddingBlock(128) - - def forward(self,bottom): - stg12 = self.stem(bottom) - #print(stg12.size()) - stg3 = self.s3_ge1(stg12) - stg3 = self.s3_ge2(stg3) - #print(stg3.size()) - stg4 = self.s4_ge1(stg3) - stg4 = self.s4_ge2(stg4) - #print(stg4.size()) - stg5 = self.s5_ge1(stg4) - stg5 = self.s5_ge2(stg5) - stg5 = self.s5_ge3(stg5) - stg5 = self.s5_ge4(stg5) - stg5 = self.s5_ge5(stg5) - #print(stg5.size()) - out = self.ceb(stg5) - if self.training: - seghead1 = self.seghead1(stg12) - seghead2 = self.seghead2(stg3) - seghead3 = self.seghead3(stg4) - seghead4 = self.seghead4(stg5) - return out,seghead1,seghead2,seghead3,seghead4 - else: - return out - - -class BiSeNet(nn.Module): - def __init__(self,cls): - super(BiSeNet, self).__init__() - self.db = DetailedBranch() - self.sb = SemanticBranch(cls) - self.bga = BGA(128) - self.seghead = SegHead(128,128,cls) - self._init_params() - self.criterion = nn.CrossEntropyLoss(ignore_index=255) - def _init_params(self): - for m in self.modules(): - if isinstance(m, nn.Conv2d): - nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') - if m.bias is not None: - nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.BatchNorm2d): - nn.init.constant_(m.weight, 1) - nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.BatchNorm1d): - nn.init.constant_(m.weight, 1) - nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.Linear): - nn.init.normal_(m.weight, 0, 0.01) - if m.bias is not None: - nn.init.constant_(m.bias, 0) - def forward(self,data,y=None): - db = self.db(data) - if self.training: - sb,head1,head2,head3,head4 = self.sb(data) - else: - sb = self.sb(data) - bga = self.bga(db,sb) - pred = self.seghead(bga) - if self.training: - main_loss = self.criterion(pred, y) - aux1_loss = self.criterion(head1, y) - aux2_loss = self.criterion(head2, y) - aux3_loss = self.criterion(head3, y) - aux4_loss = self.criterion(head4, y) - return pred.max(1)[1],main_loss,(aux1_loss,aux2_loss,aux3_loss,aux4_loss) - return pred - -if __name__ == '__main__': - import os - os.environ["CUDA_VISIBLE_DEVICES"] = '0' - input = torch.rand(4, 3, 720, 960).cuda() - model = BiSeNet(11,False).cuda() - model.eval() - print(model) - output = model(input) - print('BiSeNet', output.size()) \ No newline at end of file diff --git a/lanenet/model/BiseNet_v2-2021-04-21.py b/lanenet/model/BiseNet_v2-2021-04-21.py deleted file mode 100644 index a7e7520..0000000 --- a/lanenet/model/BiseNet_v2-2021-04-21.py +++ /dev/null @@ -1,334 +0,0 @@ -import torch -import torch.nn as nn -import torch.nn.init as init -import torch.nn.functional as F -from torch.utils import model_zoo -from torchvision import models -from lanenet import config -import collections - -class conv2d(nn.Module): - def __init__(self,in_dim,out_dim,k,pad,stride,groups = 1,bias=False,use_bn = True,use_rl = True): - super(conv2d,self).__init__() - self.use_bn = use_bn - self.use_rl = use_rl - self.conv = nn.Conv2d(in_dim,out_dim,k,padding=pad,stride=stride, groups=groups,bias=bias) - self.bn = nn.BatchNorm2d(out_dim) - self.relu = nn.ReLU(inplace=True) - def forward(self,bottom): - if self.use_bn and self.use_rl: - return self.relu(self.bn(self.conv(bottom))) - elif self.use_bn: - return self.bn(self.conv(bottom)) - else: - return self.conv(bottom) - -class SegHead(nn.Module): - def __init__(self,in_dim,out_dim,cls,size=[720,1280]): - super(SegHead,self).__init__() - self.size = size - self.conv = conv2d(in_dim,out_dim,3,1,1) - self.cls = conv2d(out_dim,cls,1,0,1,use_bn=False,use_rl=False) - def forward(self,feat): - x = self.conv(feat) - x = self.cls(x) - pred = F.interpolate(x, size=self.size, mode="bilinear",align_corners=True) - return pred - -class StemBlock(nn.Module): - def __init__(self): - super(StemBlock,self).__init__() - self.conv1 = conv2d(3,16,3,1,2) - self.conv_1x1 = conv2d(16,8,1,0,1) - self.conv_3x3 = conv2d(8,16,3,1,2) - self.mpooling = nn.MaxPool2d(3,2,1) - self.conv2 = conv2d(32,16,3,1,1) - def forward(self,bottom): - base = self.conv1(bottom) - conv_1 = self.conv_1x1(base) - conv_3 = self.conv_3x3(conv_1) - pool = self.mpooling(base) - cat = torch.cat([conv_3,pool],1) - res = self.conv2(cat) - return res - -class ContextEmbeddingBlock(nn.Module): - def __init__(self,in_dim): - super(ContextEmbeddingBlock,self).__init__() - self.gap = nn.AdaptiveAvgPool2d(1)#1 - self.bn1 = nn.BatchNorm2d(in_dim) - self.conv1 = conv2d(in_dim,in_dim,1,0,1) - self.conv2 = conv2d(in_dim,in_dim,3,1,1,use_bn = False,use_rl = False) - def forward(self,bottom): - gap = self.gap(bottom) - bn = self.bn1(gap) - conv1 = self.conv1(bn) - feat = bottom+conv1 - res = self.conv2(feat) - return res - -class GatherExpansion(nn.Module): - def __init__(self,in_dim,out_dim,stride = 1,exp = 6): - super(GatherExpansion,self).__init__() - exp_dim = in_dim*exp - self.stride = stride - self.conv1 = conv2d(in_dim,exp_dim,3,1,1) - self.dwconv2 = conv2d(exp_dim,exp_dim,3,1,1,exp_dim,use_rl = False) - self.conv_11 = conv2d(exp_dim,out_dim,1,0,1,use_rl = False) - - self.dwconv1 = conv2d(exp_dim,exp_dim,3,1,2,exp_dim,use_rl = False) - self.dwconv3 = conv2d(in_dim,in_dim,3,1,2,in_dim,use_rl = False) - self.conv_12 = conv2d(in_dim,out_dim,1,0,1,use_rl = False) - self.relu = nn.ReLU(inplace=True) - def forward(self,bottom): - base = self.conv1(bottom) - if self.stride == 2: - base = self.dwconv1(base) - bottom = self.dwconv3(bottom) - bottom = self.conv_12(bottom) - x = self.dwconv2(base) - x = self.conv_11(x) - res = self.relu(x+bottom) - return res - -class BGA(nn.Module): - def __init__(self,in_dim): - super(BGA,self).__init__() - self.in_dim = in_dim - self.db_dwconv = conv2d(in_dim,in_dim,3,1,1,in_dim,use_rl=False) - self.db_conv1x1 = conv2d(in_dim,in_dim,1,0,1,use_rl=False,use_bn=False) - self.db_conv = conv2d(in_dim,in_dim,3,1,2,use_rl=False) - self.db_apooling = nn.AvgPool2d(3,2,1) - - self.sb_dwconv = conv2d(in_dim,in_dim,3,1,1,in_dim,use_rl=False) - self.sb_conv1x1 = conv2d(in_dim,in_dim,1,0,1,use_rl=False,use_bn=False) - self.sb_conv = conv2d(in_dim,in_dim,3,1,1,use_rl=False) - self.sb_sigmoid = nn.Sigmoid() - - self.conv = conv2d(in_dim,in_dim,3,1,1,use_rl=False) - def forward(self,db,sb): - db_dwc = self.db_dwconv(db) - db_out = self.db_conv1x1(db_dwc)# - db_conv = self.db_conv(db) - db_pool = self.db_apooling(db_conv) - - sb_dwc = self.sb_dwconv(sb) - sb_out = self.sb_sigmoid(self.sb_conv1x1(sb_dwc))# - sb_conv = self.sb_conv(sb) - sb_up = self.sb_sigmoid(F.interpolate(sb_conv, size=db_out.size()[2:], mode="bilinear",align_corners=True)) - db_l = db_out*sb_up - sb_r = F.interpolate(sb_out*db_pool, size=db_out.size()[2:], mode="bilinear",align_corners=True) - res = self.conv(db_l+sb_r) - return res - -class DetailedBranch(nn.Module): - def __init__(self): - super(DetailedBranch,self).__init__() - self.s1_conv1 = conv2d(3,64,3,1,2) - self.s1_conv2 = conv2d(64,64,3,1,1) - - self.s2_conv1 = conv2d(64,64,3,1,2) - self.s2_conv2 = conv2d(64,64,3,1,1) - self.s2_conv3 = conv2d(64,64,3,1,1) - - self.s3_conv1 = conv2d(64,128,3,1,2) - self.s3_conv2 = conv2d(128,128,3,1,1) - self.s3_conv3 = conv2d(128,128,3,1,1) - def forward(self,bottom): - - detail_stage_outputs = collections.OrderedDict() - - s1_1 = self.s1_conv1(bottom) - s1_2 = self.s1_conv2(s1_1) - - detail_stage_outputs["stg1"] = s1_2 - - s2_1 = self.s2_conv1(s1_2) - s2_2 = self.s2_conv2(s2_1) - s2_3 = self.s2_conv3(s2_2) - - detail_stage_outputs["stg2"] = s2_3 - - s3_1 = self.s3_conv1(s2_3) - s3_2 = self.s3_conv2(s3_1) - s3_3 = self.s3_conv3(s3_2) - - detail_stage_outputs["stg3"] = s3_3 - - return { - 'out': s3_3, - 'detail_stage_outputs': detail_stage_outputs - } - -class SemanticBranch(nn.Module): - def __init__(self): - super(SemanticBranch,self).__init__() - self.stem = StemBlock() - self.s3_ge1 = GatherExpansion(16,32,2) - self.s3_ge2 = GatherExpansion(32,32) - - self.s4_ge1 = GatherExpansion(32,64,2) - self.s4_ge2 = GatherExpansion(64,64) - - self.s5_ge1 = GatherExpansion(64,128,2) - self.s5_ge2 = GatherExpansion(128,128) - self.s5_ge3 = GatherExpansion(128,128) - self.s5_ge4 = GatherExpansion(128,128) - self.s5_ge5 = GatherExpansion(128,128,exp=1) - - self.ceb = ContextEmbeddingBlock(128) - - def forward(self,bottom): - seg_stage_outputs = collections.OrderedDict() - - stg1 = self.stem(bottom) - #print(stg12.size()) - seg_stage_outputs["stg1"] = stg1 - - stg3 = self.s3_ge1(stg1) - stg3 = self.s3_ge2(stg3) - #print(stg3.size()) - seg_stage_outputs["stg3"] = stg3 - - stg4 = self.s4_ge1(stg3) - stg4 = self.s4_ge2(stg4) - - seg_stage_outputs["stg4"] = stg4 - #print(stg4.size()) - stg5 = self.s5_ge1(stg4) - stg5 = self.s5_ge2(stg5) - stg5 = self.s5_ge3(stg5) - stg5 = self.s5_ge4(stg5) - stg5 = self.s5_ge5(stg5) - - seg_stage_outputs["stg5"] = stg5 - #print(stg5.size()) - out = self.ceb(stg5) - - return { - 'out': out, - 'seg_stage_outputs': seg_stage_outputs - } - -class InstanceSegmentationBranch(nn.Module): - def __init__(self): - super(InstanceSegmentationBranch,self).__init__() - self.bsconv1 = conv2d(128,256,3,1,1,use_rl=True) - self.bsconv2 = conv2d(256,128,1,0,1,use_rl=True) - def forward(self,data): - input_tensor_size=list(data.size()) - tmp_size=input_tensor_size[2:] - out_put_tensor_size=tuple([int(tmp * 8) for tmp in tmp_size]) - conv1_out=self.bsconv1(data) - conv2_out=self.bsconv2(conv1_out) - isb_out = F.interpolate(conv2_out, size=out_put_tensor_size, mode="bilinear",align_corners=True) - return isb_out - -class BinarySegmentationBranch(nn.Module): - - def __init__(self): - super(BinarySegmentationBranch,self).__init__() - - self.bsconv1_pre = conv2d(128,32,3,1,1,use_rl=True) - self.bsconv1_pre1 = conv2d(32,64,3,1,1,use_rl=True) - self.bsconv1_pre2 = conv2d(144,64,3,1,1,use_rl=True) - self.bsconv1_pre3 = conv2d(64,128,3,1,1,use_rl=True) - self.bsconv1_pre4 = conv2d(192,64,3,1,1,use_rl=True) - self.bsconv1_pre5 = conv2d(64,128,3,1,1,use_rl=True) - self.bsconv3 = conv2d(128,config.num_classes,1,0,1,use_rl=False,use_bn=True) - - - # self.bsconv1 = conv2d(128,256,3,1,1,use_rl=True) - # self.bsconv2 = conv2d(256,128,1,0,1,use_rl=True) - # self.bsconv3 = conv2d(128,config.num_classes,1,0,1,use_rl=False,use_bn=True) - - def forward(self,data,seg_stage_outputs,detail_stage_outputs): - - input_tensor_size=list(data.size()) - tmp_size=input_tensor_size[2:] - output_stage2_size=[int(tmp * 2) for tmp in tmp_size] - output_stage1_size=[int(tmp * 4) for tmp in tmp_size] - out_put_tensor_size=[int(tmp * 8) for tmp in tmp_size] - - out1=self.bsconv1_pre(data) - out2=self.bsconv1_pre1(out1) - output_stage2_tensor = F.interpolate(out2, size=tuple(output_stage2_size), mode="bilinear",align_corners=True) - # output_stage2_tensor = tf.concat([output_stage2_tensor, detail_stage_outputs['stage_2'], semantic_stage_outputs['stage_1']], axis=-1, name='stage_2_concate_features') - output_stage2_tensor=torch.cat([output_stage2_tensor,detail_stage_outputs['stg2'], seg_stage_outputs['stg1']],1) - output_stage2_tensor=self.bsconv1_pre2(output_stage2_tensor) - output_stage2_tensor=self.bsconv1_pre3(output_stage2_tensor) - output_stage1_tensor = F.interpolate(output_stage2_tensor, size=tuple(output_stage1_size), mode="bilinear",align_corners=True) - #output_stage1_tensor = tf.concat([output_stage1_tensor, detail_stage_outputs['stage_1']], axis=-1, name='stage_1_concate_features') - output_stage1_tensor=torch.cat([output_stage1_tensor,detail_stage_outputs['stg1']],1) - output_stage1_tensor=self.bsconv1_pre4(output_stage1_tensor) - output_stage1_tensor=self.bsconv1_pre5(output_stage1_tensor) - output_stage1_tensor=self.bsconv3(output_stage1_tensor) - - # conv_out1=self.bsconv1(data) - # conv_out2=self.bsconv2(conv_out1) - # conv_out3=self.bsconv3(conv_out2) - - # print("--------------------conv_out3 size--------------------") - # print(list(conv_out3.size())) - # print("--------------------bga size--------------------") - - # print("--------------------out_put_tensor_size--------------------") - # print(out_put_tensor_size) - # print(tuple(out_put_tensor_size)) - # print("--------------------out_put_tensor_size--------------------") - bsb_out = F.interpolate(output_stage1_tensor, size=tuple(out_put_tensor_size), mode="bilinear",align_corners=True) - return bsb_out - -class BiSeNet(nn.Module): - def __init__(self): - super(BiSeNet, self).__init__() - self.db = DetailedBranch() - self.sb = SemanticBranch() - self.bga = BGA(128) - self._init_params() - self.criterion = nn.CrossEntropyLoss(ignore_index=255) - self.binarySegmentationBranch=BinarySegmentationBranch() - self.instanceSegmentationBranch=InstanceSegmentationBranch() - def _init_params(self): - for m in self.modules(): - if isinstance(m, nn.Conv2d): - nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') - if m.bias is not None: - nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.BatchNorm2d): - nn.init.constant_(m.weight, 1) - nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.BatchNorm1d): - nn.init.constant_(m.weight, 1) - nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.Linear): - nn.init.normal_(m.weight, 0, 0.01) - if m.bias is not None: - nn.init.constant_(m.bias, 0) - def forward(self,data,y=None): - db = self.db(data) - sb = self.sb(data) - bga = self.bga(db["out"],sb["out"]) - # print("--------------------bga size--------------------") - # print(bga.size()) - # print("--------------------bga size--------------------") - bsb_res=self.binarySegmentationBranch(bga,sb["seg_stage_outputs"],db["detail_stage_outputs"]) - isb_res=self.instanceSegmentationBranch(bga) - return { - 'instance_seg_logits': isb_res, - # 'binary_seg_pred': binary_seg_ret, - 'binary_seg_logits': bsb_res - } - -if __name__ == '__main__': - import os - os.environ["CUDA_VISIBLE_DEVICES"] = '5' - input = torch.rand(1, 3, 256, 512).cuda() - model = BiSeNet().cuda() - model.eval() - print(model) - output = model(input) - print('+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++') - print('BiSeNet_v2', output["instance_seg_logits"].size()) - print('BiSeNet_v2', output["binary_seg_logits"].size()) - print('+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++') diff --git a/lanenet/model/BiseNet_v2-bak.py b/lanenet/model/BiseNet_v2-bak.py deleted file mode 100644 index a7e7520..0000000 --- a/lanenet/model/BiseNet_v2-bak.py +++ /dev/null @@ -1,334 +0,0 @@ -import torch -import torch.nn as nn -import torch.nn.init as init -import torch.nn.functional as F -from torch.utils import model_zoo -from torchvision import models -from lanenet import config -import collections - -class conv2d(nn.Module): - def __init__(self,in_dim,out_dim,k,pad,stride,groups = 1,bias=False,use_bn = True,use_rl = True): - super(conv2d,self).__init__() - self.use_bn = use_bn - self.use_rl = use_rl - self.conv = nn.Conv2d(in_dim,out_dim,k,padding=pad,stride=stride, groups=groups,bias=bias) - self.bn = nn.BatchNorm2d(out_dim) - self.relu = nn.ReLU(inplace=True) - def forward(self,bottom): - if self.use_bn and self.use_rl: - return self.relu(self.bn(self.conv(bottom))) - elif self.use_bn: - return self.bn(self.conv(bottom)) - else: - return self.conv(bottom) - -class SegHead(nn.Module): - def __init__(self,in_dim,out_dim,cls,size=[720,1280]): - super(SegHead,self).__init__() - self.size = size - self.conv = conv2d(in_dim,out_dim,3,1,1) - self.cls = conv2d(out_dim,cls,1,0,1,use_bn=False,use_rl=False) - def forward(self,feat): - x = self.conv(feat) - x = self.cls(x) - pred = F.interpolate(x, size=self.size, mode="bilinear",align_corners=True) - return pred - -class StemBlock(nn.Module): - def __init__(self): - super(StemBlock,self).__init__() - self.conv1 = conv2d(3,16,3,1,2) - self.conv_1x1 = conv2d(16,8,1,0,1) - self.conv_3x3 = conv2d(8,16,3,1,2) - self.mpooling = nn.MaxPool2d(3,2,1) - self.conv2 = conv2d(32,16,3,1,1) - def forward(self,bottom): - base = self.conv1(bottom) - conv_1 = self.conv_1x1(base) - conv_3 = self.conv_3x3(conv_1) - pool = self.mpooling(base) - cat = torch.cat([conv_3,pool],1) - res = self.conv2(cat) - return res - -class ContextEmbeddingBlock(nn.Module): - def __init__(self,in_dim): - super(ContextEmbeddingBlock,self).__init__() - self.gap = nn.AdaptiveAvgPool2d(1)#1 - self.bn1 = nn.BatchNorm2d(in_dim) - self.conv1 = conv2d(in_dim,in_dim,1,0,1) - self.conv2 = conv2d(in_dim,in_dim,3,1,1,use_bn = False,use_rl = False) - def forward(self,bottom): - gap = self.gap(bottom) - bn = self.bn1(gap) - conv1 = self.conv1(bn) - feat = bottom+conv1 - res = self.conv2(feat) - return res - -class GatherExpansion(nn.Module): - def __init__(self,in_dim,out_dim,stride = 1,exp = 6): - super(GatherExpansion,self).__init__() - exp_dim = in_dim*exp - self.stride = stride - self.conv1 = conv2d(in_dim,exp_dim,3,1,1) - self.dwconv2 = conv2d(exp_dim,exp_dim,3,1,1,exp_dim,use_rl = False) - self.conv_11 = conv2d(exp_dim,out_dim,1,0,1,use_rl = False) - - self.dwconv1 = conv2d(exp_dim,exp_dim,3,1,2,exp_dim,use_rl = False) - self.dwconv3 = conv2d(in_dim,in_dim,3,1,2,in_dim,use_rl = False) - self.conv_12 = conv2d(in_dim,out_dim,1,0,1,use_rl = False) - self.relu = nn.ReLU(inplace=True) - def forward(self,bottom): - base = self.conv1(bottom) - if self.stride == 2: - base = self.dwconv1(base) - bottom = self.dwconv3(bottom) - bottom = self.conv_12(bottom) - x = self.dwconv2(base) - x = self.conv_11(x) - res = self.relu(x+bottom) - return res - -class BGA(nn.Module): - def __init__(self,in_dim): - super(BGA,self).__init__() - self.in_dim = in_dim - self.db_dwconv = conv2d(in_dim,in_dim,3,1,1,in_dim,use_rl=False) - self.db_conv1x1 = conv2d(in_dim,in_dim,1,0,1,use_rl=False,use_bn=False) - self.db_conv = conv2d(in_dim,in_dim,3,1,2,use_rl=False) - self.db_apooling = nn.AvgPool2d(3,2,1) - - self.sb_dwconv = conv2d(in_dim,in_dim,3,1,1,in_dim,use_rl=False) - self.sb_conv1x1 = conv2d(in_dim,in_dim,1,0,1,use_rl=False,use_bn=False) - self.sb_conv = conv2d(in_dim,in_dim,3,1,1,use_rl=False) - self.sb_sigmoid = nn.Sigmoid() - - self.conv = conv2d(in_dim,in_dim,3,1,1,use_rl=False) - def forward(self,db,sb): - db_dwc = self.db_dwconv(db) - db_out = self.db_conv1x1(db_dwc)# - db_conv = self.db_conv(db) - db_pool = self.db_apooling(db_conv) - - sb_dwc = self.sb_dwconv(sb) - sb_out = self.sb_sigmoid(self.sb_conv1x1(sb_dwc))# - sb_conv = self.sb_conv(sb) - sb_up = self.sb_sigmoid(F.interpolate(sb_conv, size=db_out.size()[2:], mode="bilinear",align_corners=True)) - db_l = db_out*sb_up - sb_r = F.interpolate(sb_out*db_pool, size=db_out.size()[2:], mode="bilinear",align_corners=True) - res = self.conv(db_l+sb_r) - return res - -class DetailedBranch(nn.Module): - def __init__(self): - super(DetailedBranch,self).__init__() - self.s1_conv1 = conv2d(3,64,3,1,2) - self.s1_conv2 = conv2d(64,64,3,1,1) - - self.s2_conv1 = conv2d(64,64,3,1,2) - self.s2_conv2 = conv2d(64,64,3,1,1) - self.s2_conv3 = conv2d(64,64,3,1,1) - - self.s3_conv1 = conv2d(64,128,3,1,2) - self.s3_conv2 = conv2d(128,128,3,1,1) - self.s3_conv3 = conv2d(128,128,3,1,1) - def forward(self,bottom): - - detail_stage_outputs = collections.OrderedDict() - - s1_1 = self.s1_conv1(bottom) - s1_2 = self.s1_conv2(s1_1) - - detail_stage_outputs["stg1"] = s1_2 - - s2_1 = self.s2_conv1(s1_2) - s2_2 = self.s2_conv2(s2_1) - s2_3 = self.s2_conv3(s2_2) - - detail_stage_outputs["stg2"] = s2_3 - - s3_1 = self.s3_conv1(s2_3) - s3_2 = self.s3_conv2(s3_1) - s3_3 = self.s3_conv3(s3_2) - - detail_stage_outputs["stg3"] = s3_3 - - return { - 'out': s3_3, - 'detail_stage_outputs': detail_stage_outputs - } - -class SemanticBranch(nn.Module): - def __init__(self): - super(SemanticBranch,self).__init__() - self.stem = StemBlock() - self.s3_ge1 = GatherExpansion(16,32,2) - self.s3_ge2 = GatherExpansion(32,32) - - self.s4_ge1 = GatherExpansion(32,64,2) - self.s4_ge2 = GatherExpansion(64,64) - - self.s5_ge1 = GatherExpansion(64,128,2) - self.s5_ge2 = GatherExpansion(128,128) - self.s5_ge3 = GatherExpansion(128,128) - self.s5_ge4 = GatherExpansion(128,128) - self.s5_ge5 = GatherExpansion(128,128,exp=1) - - self.ceb = ContextEmbeddingBlock(128) - - def forward(self,bottom): - seg_stage_outputs = collections.OrderedDict() - - stg1 = self.stem(bottom) - #print(stg12.size()) - seg_stage_outputs["stg1"] = stg1 - - stg3 = self.s3_ge1(stg1) - stg3 = self.s3_ge2(stg3) - #print(stg3.size()) - seg_stage_outputs["stg3"] = stg3 - - stg4 = self.s4_ge1(stg3) - stg4 = self.s4_ge2(stg4) - - seg_stage_outputs["stg4"] = stg4 - #print(stg4.size()) - stg5 = self.s5_ge1(stg4) - stg5 = self.s5_ge2(stg5) - stg5 = self.s5_ge3(stg5) - stg5 = self.s5_ge4(stg5) - stg5 = self.s5_ge5(stg5) - - seg_stage_outputs["stg5"] = stg5 - #print(stg5.size()) - out = self.ceb(stg5) - - return { - 'out': out, - 'seg_stage_outputs': seg_stage_outputs - } - -class InstanceSegmentationBranch(nn.Module): - def __init__(self): - super(InstanceSegmentationBranch,self).__init__() - self.bsconv1 = conv2d(128,256,3,1,1,use_rl=True) - self.bsconv2 = conv2d(256,128,1,0,1,use_rl=True) - def forward(self,data): - input_tensor_size=list(data.size()) - tmp_size=input_tensor_size[2:] - out_put_tensor_size=tuple([int(tmp * 8) for tmp in tmp_size]) - conv1_out=self.bsconv1(data) - conv2_out=self.bsconv2(conv1_out) - isb_out = F.interpolate(conv2_out, size=out_put_tensor_size, mode="bilinear",align_corners=True) - return isb_out - -class BinarySegmentationBranch(nn.Module): - - def __init__(self): - super(BinarySegmentationBranch,self).__init__() - - self.bsconv1_pre = conv2d(128,32,3,1,1,use_rl=True) - self.bsconv1_pre1 = conv2d(32,64,3,1,1,use_rl=True) - self.bsconv1_pre2 = conv2d(144,64,3,1,1,use_rl=True) - self.bsconv1_pre3 = conv2d(64,128,3,1,1,use_rl=True) - self.bsconv1_pre4 = conv2d(192,64,3,1,1,use_rl=True) - self.bsconv1_pre5 = conv2d(64,128,3,1,1,use_rl=True) - self.bsconv3 = conv2d(128,config.num_classes,1,0,1,use_rl=False,use_bn=True) - - - # self.bsconv1 = conv2d(128,256,3,1,1,use_rl=True) - # self.bsconv2 = conv2d(256,128,1,0,1,use_rl=True) - # self.bsconv3 = conv2d(128,config.num_classes,1,0,1,use_rl=False,use_bn=True) - - def forward(self,data,seg_stage_outputs,detail_stage_outputs): - - input_tensor_size=list(data.size()) - tmp_size=input_tensor_size[2:] - output_stage2_size=[int(tmp * 2) for tmp in tmp_size] - output_stage1_size=[int(tmp * 4) for tmp in tmp_size] - out_put_tensor_size=[int(tmp * 8) for tmp in tmp_size] - - out1=self.bsconv1_pre(data) - out2=self.bsconv1_pre1(out1) - output_stage2_tensor = F.interpolate(out2, size=tuple(output_stage2_size), mode="bilinear",align_corners=True) - # output_stage2_tensor = tf.concat([output_stage2_tensor, detail_stage_outputs['stage_2'], semantic_stage_outputs['stage_1']], axis=-1, name='stage_2_concate_features') - output_stage2_tensor=torch.cat([output_stage2_tensor,detail_stage_outputs['stg2'], seg_stage_outputs['stg1']],1) - output_stage2_tensor=self.bsconv1_pre2(output_stage2_tensor) - output_stage2_tensor=self.bsconv1_pre3(output_stage2_tensor) - output_stage1_tensor = F.interpolate(output_stage2_tensor, size=tuple(output_stage1_size), mode="bilinear",align_corners=True) - #output_stage1_tensor = tf.concat([output_stage1_tensor, detail_stage_outputs['stage_1']], axis=-1, name='stage_1_concate_features') - output_stage1_tensor=torch.cat([output_stage1_tensor,detail_stage_outputs['stg1']],1) - output_stage1_tensor=self.bsconv1_pre4(output_stage1_tensor) - output_stage1_tensor=self.bsconv1_pre5(output_stage1_tensor) - output_stage1_tensor=self.bsconv3(output_stage1_tensor) - - # conv_out1=self.bsconv1(data) - # conv_out2=self.bsconv2(conv_out1) - # conv_out3=self.bsconv3(conv_out2) - - # print("--------------------conv_out3 size--------------------") - # print(list(conv_out3.size())) - # print("--------------------bga size--------------------") - - # print("--------------------out_put_tensor_size--------------------") - # print(out_put_tensor_size) - # print(tuple(out_put_tensor_size)) - # print("--------------------out_put_tensor_size--------------------") - bsb_out = F.interpolate(output_stage1_tensor, size=tuple(out_put_tensor_size), mode="bilinear",align_corners=True) - return bsb_out - -class BiSeNet(nn.Module): - def __init__(self): - super(BiSeNet, self).__init__() - self.db = DetailedBranch() - self.sb = SemanticBranch() - self.bga = BGA(128) - self._init_params() - self.criterion = nn.CrossEntropyLoss(ignore_index=255) - self.binarySegmentationBranch=BinarySegmentationBranch() - self.instanceSegmentationBranch=InstanceSegmentationBranch() - def _init_params(self): - for m in self.modules(): - if isinstance(m, nn.Conv2d): - nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') - if m.bias is not None: - nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.BatchNorm2d): - nn.init.constant_(m.weight, 1) - nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.BatchNorm1d): - nn.init.constant_(m.weight, 1) - nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.Linear): - nn.init.normal_(m.weight, 0, 0.01) - if m.bias is not None: - nn.init.constant_(m.bias, 0) - def forward(self,data,y=None): - db = self.db(data) - sb = self.sb(data) - bga = self.bga(db["out"],sb["out"]) - # print("--------------------bga size--------------------") - # print(bga.size()) - # print("--------------------bga size--------------------") - bsb_res=self.binarySegmentationBranch(bga,sb["seg_stage_outputs"],db["detail_stage_outputs"]) - isb_res=self.instanceSegmentationBranch(bga) - return { - 'instance_seg_logits': isb_res, - # 'binary_seg_pred': binary_seg_ret, - 'binary_seg_logits': bsb_res - } - -if __name__ == '__main__': - import os - os.environ["CUDA_VISIBLE_DEVICES"] = '5' - input = torch.rand(1, 3, 256, 512).cuda() - model = BiSeNet().cuda() - model.eval() - print(model) - output = model(input) - print('+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++') - print('BiSeNet_v2', output["instance_seg_logits"].size()) - print('BiSeNet_v2', output["binary_seg_logits"].size()) - print('+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++') diff --git a/lanenet/model/BiseNet_v2.py b/lanenet/model/BiseNet_v2.py deleted file mode 100644 index 3c73cff..0000000 --- a/lanenet/model/BiseNet_v2.py +++ /dev/null @@ -1,336 +0,0 @@ -import torch -import torch.nn as nn -import torch.nn.init as init -import torch.nn.functional as F -from torch.utils import model_zoo -from torchvision import models -from lanenet import config -import collections - -class conv2d(nn.Module): - def __init__(self,in_dim,out_dim,k,pad,stride,groups = 1,bias=False,use_bn = True,use_rl = True): - super(conv2d,self).__init__() - self.use_bn = use_bn - self.use_rl = use_rl - self.conv = nn.Conv2d(in_dim,out_dim,k,padding=pad,stride=stride, groups=groups,bias=bias) - self.bn = nn.BatchNorm2d(out_dim) - self.relu = nn.ReLU(inplace=True) - def forward(self,bottom): - if self.use_bn and self.use_rl: - return self.relu(self.bn(self.conv(bottom))) - elif self.use_bn: - return self.bn(self.conv(bottom)) - else: - return self.conv(bottom) - -class SegHead(nn.Module): - def __init__(self,in_dim,out_dim,cls,size=[720,1280]): - super(SegHead,self).__init__() - self.size = size - self.conv = conv2d(in_dim,out_dim,3,1,1) - self.cls = conv2d(out_dim,cls,1,0,1,use_bn=False,use_rl=False) - def forward(self,feat): - x = self.conv(feat) - x = self.cls(x) - pred = F.interpolate(x, size=self.size, mode="bilinear",align_corners=True) - return pred - -class StemBlock(nn.Module): - def __init__(self): - super(StemBlock,self).__init__() - self.conv1 = conv2d(3,16,3,1,2) - self.conv_1x1 = conv2d(16,32,1,0,1) - self.conv_3x3 = conv2d(32,32,3,1,2) - self.mpooling = nn.MaxPool2d(3,2,1) - self.conv2 = conv2d(48,32,3,1,1) - def forward(self,bottom): - base = self.conv1(bottom) - conv_1 = self.conv_1x1(base) - conv_3 = self.conv_3x3(conv_1) - pool = self.mpooling(base) - cat = torch.cat([conv_3,pool],1) - res = self.conv2(cat) - return res - -class ContextEmbeddingBlock(nn.Module): - def __init__(self,in_dim): - super(ContextEmbeddingBlock,self).__init__() - self.gap = nn.AdaptiveAvgPool2d(1)#1 - # self.gap = nn.AvgPool2d(3,1,1) - self.bn1 = nn.BatchNorm2d(in_dim) - self.conv1 = conv2d(in_dim,in_dim,1,0,1) - self.conv2 = conv2d(in_dim,in_dim,3,1,1,use_bn = False,use_rl = False) - def forward(self,bottom): - gap = self.gap(bottom) - # print(gap) - bn = self.bn1(gap) - conv1 = self.conv1(bn) - feat = bottom+conv1 - res = self.conv2(feat) - return res - -class GatherExpansion(nn.Module): - def __init__(self,in_dim,out_dim,stride = 1,exp = 6): - super(GatherExpansion,self).__init__() - exp_dim = in_dim*exp - self.stride = stride - self.conv1 = conv2d(in_dim,exp_dim,3,1,1) - self.dwconv2 = conv2d(exp_dim,exp_dim,3,1,1,exp_dim,use_rl = False) - self.conv_11 = conv2d(exp_dim,out_dim,1,0,1,use_rl = False) - - self.dwconv1 = conv2d(exp_dim,exp_dim,3,1,2,exp_dim,use_rl = False) - self.dwconv3 = conv2d(in_dim,in_dim,3,1,2,in_dim,use_rl = False) - self.conv_12 = conv2d(in_dim,out_dim,1,0,1,use_rl = False) - self.relu = nn.ReLU(inplace=True) - def forward(self,bottom): - base = self.conv1(bottom) - if self.stride == 2: - base = self.dwconv1(base) - bottom = self.dwconv3(bottom) - bottom = self.conv_12(bottom) - x = self.dwconv2(base) - x = self.conv_11(x) - res = self.relu(x+bottom) - return res - -class BGA(nn.Module): - def __init__(self,in_dim): - super(BGA,self).__init__() - self.in_dim = in_dim - self.db_dwconv = conv2d(in_dim,in_dim,3,1,1,in_dim,use_rl=False) - self.db_conv1x1 = conv2d(in_dim,in_dim,1,0,1,use_rl=False,use_bn=False) - self.db_conv = conv2d(in_dim,in_dim,3,1,2,use_rl=False) - self.db_apooling = nn.AvgPool2d(3,2,1) - - self.sb_dwconv = conv2d(in_dim,in_dim,3,1,1,in_dim,use_rl=False) - self.sb_conv1x1 = conv2d(in_dim,in_dim,1,0,1,use_rl=False,use_bn=False) - self.sb_conv = conv2d(in_dim,in_dim,3,1,1,use_rl=False) - self.sb_sigmoid = nn.Sigmoid() - - self.conv = conv2d(in_dim,in_dim,3,1,1,use_rl=False) - def forward(self,db,sb): - db_dwc = self.db_dwconv(db) - db_out = self.db_conv1x1(db_dwc)# - db_conv = self.db_conv(db) - db_pool = self.db_apooling(db_conv) - - sb_dwc = self.sb_dwconv(sb) - sb_out = self.sb_sigmoid(self.sb_conv1x1(sb_dwc))# - sb_conv = self.sb_conv(sb) - sb_up = self.sb_sigmoid(F.interpolate(sb_conv, size=db_out.size()[2:], mode="bilinear",align_corners=True)) - db_l = db_out*sb_up - sb_r = F.interpolate(sb_out*db_pool, size=db_out.size()[2:], mode="bilinear",align_corners=True) - res = self.conv(db_l+sb_r) - return res - -class DetailedBranch(nn.Module): - def __init__(self): - super(DetailedBranch,self).__init__() - self.s1_conv1 = conv2d(3,64,3,1,2) - self.s1_conv2 = conv2d(64,64,3,1,1) - - self.s2_conv1 = conv2d(64,64,3,1,2) - self.s2_conv2 = conv2d(64,64,3,1,1) - self.s2_conv3 = conv2d(64,64,3,1,1) - - self.s3_conv1 = conv2d(64,128,3,1,2) - self.s3_conv2 = conv2d(128,128,3,1,1) - self.s3_conv3 = conv2d(128,128,3,1,1) - def forward(self,bottom): - - detail_stage_outputs = collections.OrderedDict() - - s1_1 = self.s1_conv1(bottom) - s1_2 = self.s1_conv2(s1_1) - - detail_stage_outputs["stg1"] = s1_2 - - s2_1 = self.s2_conv1(s1_2) - s2_2 = self.s2_conv2(s2_1) - s2_3 = self.s2_conv3(s2_2) - - detail_stage_outputs["stg2"] = s2_3 - - s3_1 = self.s3_conv1(s2_3) - s3_2 = self.s3_conv2(s3_1) - s3_3 = self.s3_conv3(s3_2) - - detail_stage_outputs["stg3"] = s3_3 - - return { - 'out': s3_3, - 'detail_stage_outputs': detail_stage_outputs - } - -class SemanticBranch(nn.Module): - def __init__(self): - super(SemanticBranch,self).__init__() - self.stem = StemBlock() - self.s3_ge1 = GatherExpansion(32,32,2) - self.s3_ge2 = GatherExpansion(32,32) - - self.s4_ge1 = GatherExpansion(32,64,2) - self.s4_ge2 = GatherExpansion(64,64) - - self.s5_ge1 = GatherExpansion(64,128,2) - self.s5_ge2 = GatherExpansion(128,128) - self.s5_ge3 = GatherExpansion(128,128) - self.s5_ge4 = GatherExpansion(128,128) - self.s5_ge5 = GatherExpansion(128,128,exp=1) - - self.ceb = ContextEmbeddingBlock(128) - - def forward(self,bottom): - seg_stage_outputs = collections.OrderedDict() - - stg1 = self.stem(bottom) - #print(stg12.size()) - seg_stage_outputs["stg1"] = stg1 - - stg3 = self.s3_ge1(stg1) - stg3 = self.s3_ge2(stg3) - #print(stg3.size()) - seg_stage_outputs["stg3"] = stg3 - - stg4 = self.s4_ge1(stg3) - stg4 = self.s4_ge2(stg4) - - seg_stage_outputs["stg4"] = stg4 - #print(stg4.size()) - stg5 = self.s5_ge1(stg4) - stg5 = self.s5_ge2(stg5) - stg5 = self.s5_ge3(stg5) - stg5 = self.s5_ge4(stg5) - stg5 = self.s5_ge5(stg5) - - seg_stage_outputs["stg5"] = stg5 - #print(stg5.size()) - out = self.ceb(stg5) - - return { - 'out': out, - 'seg_stage_outputs': seg_stage_outputs - } - -class InstanceSegmentationBranch(nn.Module): - def __init__(self): - super(InstanceSegmentationBranch,self).__init__() - self.bsconv1 = conv2d(128,256,3,1,1,use_rl=True) - self.bsconv2 = conv2d(256,128,1,0,1,use_rl=True) - def forward(self,data): - input_tensor_size=list(data.size()) - tmp_size=input_tensor_size[2:] - out_put_tensor_size=tuple([int(tmp * 8) for tmp in tmp_size]) - conv1_out=self.bsconv1(data) - conv2_out=self.bsconv2(conv1_out) - isb_out = F.interpolate(conv2_out, size=out_put_tensor_size, mode="bilinear",align_corners=True) - return isb_out - -class BinarySegmentationBranch(nn.Module): - - def __init__(self): - super(BinarySegmentationBranch,self).__init__() - - self.bsconv1_pre = conv2d(128,256,3,1,1,use_rl=True) - self.bsconv1_pre1 = conv2d(256,256,3,1,1,use_rl=True) - self.bsconv1_pre2 = conv2d(352,256,3,1,1,use_rl=True) - self.bsconv1_pre3 = conv2d(256,256,3,1,1,use_rl=True) - self.bsconv1_pre4 = conv2d(320,256,3,1,1,use_rl=True) - self.bsconv1_pre5 = conv2d(256,128,3,1,1,use_rl=True) - self.bsconv3 = conv2d(128,config.num_classes,1,0,1,use_rl=False,use_bn=True) - - - # self.bsconv1 = conv2d(128,256,3,1,1,use_rl=True) - # self.bsconv2 = conv2d(256,128,1,0,1,use_rl=True) - # self.bsconv3 = conv2d(128,config.num_classes,1,0,1,use_rl=False,use_bn=True) - - def forward(self,data,seg_stage_outputs,detail_stage_outputs): - - input_tensor_size=list(data.size()) - tmp_size=input_tensor_size[2:] - output_stage2_size=[int(tmp * 2) for tmp in tmp_size] - output_stage1_size=[int(tmp * 4) for tmp in tmp_size] - out_put_tensor_size=[int(tmp * 8) for tmp in tmp_size] - - out1=self.bsconv1_pre(data) - # channel=256 - out2=self.bsconv1_pre1(out1) - - output_stage2_tensor = F.interpolate(out2, size=tuple(output_stage2_size), mode="bilinear",align_corners=True) - # output_stage2_tensor = tf.concat([output_stage2_tensor, detail_stage_outputs['stage_2'], semantic_stage_outputs['stage_1']], axis=-1, name='stage_2_concate_features') - output_stage2_tensor=torch.cat([output_stage2_tensor,detail_stage_outputs['stg2'], seg_stage_outputs['stg1']],1) - output_stage2_tensor=self.bsconv1_pre2(output_stage2_tensor) - output_stage2_tensor=self.bsconv1_pre3(output_stage2_tensor) - # channel=256 - output_stage1_tensor = F.interpolate(output_stage2_tensor, size=tuple(output_stage1_size), mode="bilinear",align_corners=True) - #output_stage1_tensor = tf.concat([output_stage1_tensor, detail_stage_outputs['stage_1']], axis=-1, name='stage_1_concate_features') - output_stage1_tensor=torch.cat([output_stage1_tensor,detail_stage_outputs['stg1']],1) - output_stage1_tensor=self.bsconv1_pre4(output_stage1_tensor) - output_stage1_tensor=self.bsconv1_pre5(output_stage1_tensor) - output_stage1_tensor=self.bsconv3(output_stage1_tensor) - - # conv_out1=self.bsconv1(data) - # conv_out2=self.bsconv2(conv_out1) - # conv_out3=self.bsconv3(conv_out2) - - # print("--------------------conv_out3 size--------------------") - # print(list(conv_out3.size())) - # print("--------------------bga size--------------------") - - # print("--------------------out_put_tensor_size--------------------") - # print(out_put_tensor_size) - # print(tuple(out_put_tensor_size)) - # print("--------------------out_put_tensor_size--------------------") - bsb_out = F.interpolate(output_stage1_tensor, size=tuple(out_put_tensor_size), mode="bilinear",align_corners=True) - return bsb_out - -class BiSeNet(nn.Module): - def __init__(self): - super(BiSeNet, self).__init__() - self.db = DetailedBranch() - self.sb = SemanticBranch() - self.bga = BGA(128) - self._init_params() - self.criterion = nn.CrossEntropyLoss(ignore_index=255) - self.binarySegmentationBranch=BinarySegmentationBranch() - self.instanceSegmentationBranch=InstanceSegmentationBranch() - def _init_params(self): - for m in self.modules(): - if isinstance(m, nn.Conv2d): - nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') - if m.bias is not None: - nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.BatchNorm2d): - nn.init.constant_(m.weight, 1) - nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.BatchNorm1d): - nn.init.constant_(m.weight, 1) - nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.Linear): - nn.init.normal_(m.weight, 0, 0.01) - if m.bias is not None: - nn.init.constant_(m.bias, 0) - def forward(self,data,y=None): - db = self.db(data) - sb = self.sb(data) - bga = self.bga(db["out"],sb["out"]) - bsb_res=self.binarySegmentationBranch(bga,sb["seg_stage_outputs"],db["detail_stage_outputs"]) - isb_res=self.instanceSegmentationBranch(bga) - return { - 'instance_seg_logits': isb_res, - # 'binary_seg_pred': binary_seg_ret, - 'binary_seg_logits': bsb_res - } - -if __name__ == '__main__': - import os - os.environ["CUDA_VISIBLE_DEVICES"] = '5' - input = torch.rand(1, 3, 256, 512) - model = BiSeNet() - model.eval() - print(model) - output = model(input) - print('+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++') - print('BiSeNet_v2', output["instance_seg_logits"].size()) - print('BiSeNet_v2', output["binary_seg_logits"].size()) - print('+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++') diff --git a/lanenet/model/BiseNet_v2_1.py b/lanenet/model/BiseNet_v2_1.py deleted file mode 100644 index 888b5e9..0000000 --- a/lanenet/model/BiseNet_v2_1.py +++ /dev/null @@ -1,377 +0,0 @@ -import torch -import torch.nn as nn -import torch.nn.init as init -import torch.nn.functional as F -from torch.utils import model_zoo -from torchvision import models -from lanenet import config -import collections - -class conv2d(nn.Module): - def __init__(self,in_dim,out_dim,k,pad,stride,groups = 1,bias=False,use_bn = True,use_rl = True): - super(conv2d,self).__init__() - self.use_bn = use_bn - self.use_rl = use_rl - self.conv = nn.Conv2d(in_dim,out_dim,k,padding=pad,stride=stride, groups=groups,bias=bias) - self.bn = nn.BatchNorm2d(out_dim) - self.relu = nn.ReLU(inplace=True) - def forward(self,bottom): - if self.use_bn and self.use_rl: - return self.relu(self.bn(self.conv(bottom))) - elif self.use_bn: - return self.bn(self.conv(bottom)) - else: - return self.conv(bottom) - -class SegHead(nn.Module): - def __init__(self,in_dim,out_dim,cls,size=[720,1280]): - super(SegHead,self).__init__() - self.size = size - self.conv = conv2d(in_dim,out_dim,3,1,1) - self.cls = conv2d(out_dim,cls,1,0,1,use_bn=False,use_rl=False) - def forward(self,feat): - x = self.conv(feat) - x = self.cls(x) - pred = F.interpolate(x, size=self.size, mode="bilinear",align_corners=True) - return pred - -class StemBlock(nn.Module): - def __init__(self): - super(StemBlock,self).__init__() - self.conv1 = conv2d(3,16,3,1,2) - self.conv_1x1 = conv2d(16,32,1,0,1) - self.conv_3x3 = conv2d(32,32,3,1,2) - self.mpooling = nn.MaxPool2d(3,2,1) - self.conv2 = conv2d(48,32,3,1,1) - def forward(self,bottom): - base = self.conv1(bottom) - conv_1 = self.conv_1x1(base) - conv_3 = self.conv_3x3(conv_1) - pool = self.mpooling(base) - cat = torch.cat([conv_3,pool],1) - res = self.conv2(cat) - return res - -class ContextEmbeddingBlock(nn.Module): - def __init__(self,in_dim): - super(ContextEmbeddingBlock,self).__init__() - self.gap = nn.AdaptiveAvgPool2d(1)#1 - # self.gap = nn.AvgPool2d(3,1,1) - self.bn1 = nn.BatchNorm2d(in_dim) - self.conv1 = conv2d(in_dim,in_dim,1,0,1) - self.conv2 = conv2d(in_dim,in_dim,3,1,1,use_bn = False,use_rl = False) - def forward(self,bottom): - gap = self.gap(bottom) - # print(gap) - bn = self.bn1(gap) - conv1 = self.conv1(bn) - feat = bottom+conv1 - res = self.conv2(feat) - return res - -class GatherExpansion(nn.Module): - def __init__(self,in_dim,out_dim,stride = 1,exp = 6): - super(GatherExpansion,self).__init__() - exp_dim = in_dim*exp - self.stride = stride - self.conv1 = conv2d(in_dim,exp_dim,3,1,1) - self.dwconv2 = conv2d(exp_dim,exp_dim,3,1,1,exp_dim,use_rl = False) - self.conv_11 = conv2d(exp_dim,out_dim,1,0,1,use_rl = False) - - self.dwconv1 = conv2d(exp_dim,exp_dim,3,1,2,exp_dim,use_rl = False) - self.dwconv3 = conv2d(in_dim,in_dim,3,1,2,in_dim,use_rl = False) - self.conv_12 = conv2d(in_dim,out_dim,1,0,1,use_rl = False) - self.relu = nn.ReLU(inplace=True) - def forward(self,bottom): - base = self.conv1(bottom) - if self.stride == 2: - base = self.dwconv1(base) - bottom = self.dwconv3(bottom) - bottom = self.conv_12(bottom) - x = self.dwconv2(base) - x = self.conv_11(x) - res = self.relu(x+bottom) - return res - -class BGA(nn.Module): - def __init__(self,in_dim): - super(BGA,self).__init__() - self.in_dim = in_dim - self.db_dwconv = conv2d(in_dim,in_dim,3,1,1,in_dim,use_rl=False) - self.db_conv1x1 = conv2d(in_dim,in_dim,1,0,1,use_rl=False,use_bn=False) - self.db_conv = conv2d(in_dim,in_dim,3,1,2,use_rl=False) - self.db_apooling = nn.AvgPool2d(3,2,1) - - self.sb_dwconv = conv2d(in_dim,in_dim,3,1,1,in_dim,use_rl=False) - self.sb_conv1x1 = conv2d(in_dim,in_dim,1,0,1,use_rl=False,use_bn=False) - self.sb_conv = conv2d(in_dim,in_dim,3,1,1,use_rl=False) - self.sb_sigmoid = nn.Sigmoid() - - self.conv = conv2d(in_dim,in_dim,3,1,1,use_rl=False) - def forward(self,db,sb): - db_dwc = self.db_dwconv(db) - db_out = self.db_conv1x1(db_dwc)# - db_conv = self.db_conv(db) - db_pool = self.db_apooling(db_conv) - - sb_dwc = self.sb_dwconv(sb) - sb_out = self.sb_sigmoid(self.sb_conv1x1(sb_dwc))# - sb_conv = self.sb_conv(sb) - sb_up = self.sb_sigmoid(F.interpolate(sb_conv, size=db_out.size()[2:], mode="bilinear",align_corners=True)) - db_l = db_out*sb_up - sb_r = F.interpolate(sb_out*db_pool, size=db_out.size()[2:], mode="bilinear",align_corners=True) - res = self.conv(db_l+sb_r) - return res - -class DetailedBranch(nn.Module): - def __init__(self): - super(DetailedBranch,self).__init__() - self.s1_conv1 = conv2d(3,64,3,1,2) - self.s1_conv2 = conv2d(64,64,3,1,1) - - self.s2_conv1 = conv2d(64,64,3,1,2) - self.s2_conv2 = conv2d(64,64,3,1,1) - self.s2_conv3 = conv2d(64,64,3,1,1) - - self.s3_conv1 = conv2d(64,128,3,1,2) - self.s3_conv2 = conv2d(128,128,3,1,1) - self.s3_conv3 = conv2d(128,128,3,1,1) - def forward(self,bottom): - - detail_stage_outputs = collections.OrderedDict() - - s1_1 = self.s1_conv1(bottom) - s1_2 = self.s1_conv2(s1_1) - - detail_stage_outputs["stg1"] = s1_2 - - s2_1 = self.s2_conv1(s1_2) - s2_2 = self.s2_conv2(s2_1) - s2_3 = self.s2_conv3(s2_2) - - detail_stage_outputs["stg2"] = s2_3 - - s3_1 = self.s3_conv1(s2_3) - s3_2 = self.s3_conv2(s3_1) - s3_3 = self.s3_conv3(s3_2) - - detail_stage_outputs["stg3"] = s3_3 - - return { - 'out': s3_3, - 'detail_stage_outputs': detail_stage_outputs - } - -class SemanticBranch(nn.Module): - def __init__(self): - super(SemanticBranch,self).__init__() - self.stem = StemBlock() - self.s3_ge1 = GatherExpansion(32,32,2) - self.s3_ge2 = GatherExpansion(32,32) - - self.s4_ge1 = GatherExpansion(32,64,2) - self.s4_ge2 = GatherExpansion(64,64) - - self.s5_ge1 = GatherExpansion(64,128,2) - self.s5_ge2 = GatherExpansion(128,128) - self.s5_ge3 = GatherExpansion(128,128) - self.s5_ge4 = GatherExpansion(128,128) - self.s5_ge5 = GatherExpansion(128,128,exp=1) - - self.ceb = ContextEmbeddingBlock(128) - - if config.is_training==1: - self.seghead1 = SegHead(32,32,config.num_classes) - self.seghead3 = SegHead(32,32,config.num_classes) - self.seghead4 = SegHead(64,64,config.num_classes) - self.seghead5 = SegHead(128,128,config.num_classes) - - def forward(self,bottom): - seg_stage_outputs = collections.OrderedDict() - - stg1 = self.stem(bottom) - #print(stg12.size()) - seg_stage_outputs["stg1"] = stg1 - - stg3 = self.s3_ge1(stg1) - stg3 = self.s3_ge2(stg3) - #print(stg3.size()) - seg_stage_outputs["stg3"] = stg3 - - stg4 = self.s4_ge1(stg3) - stg4 = self.s4_ge2(stg4) - - seg_stage_outputs["stg4"] = stg4 - #print(stg4.size()) - stg5 = self.s5_ge1(stg4) - stg5 = self.s5_ge2(stg5) - stg5 = self.s5_ge3(stg5) - stg5 = self.s5_ge4(stg5) - stg5 = self.s5_ge5(stg5) - - seg_stage_outputs["stg5"] = stg5 - #print(stg5.size()) - out = self.ceb(stg5) - - if self.training: - seghead1 = self.seghead1(stg1) - seghead2 = self.seghead3(stg3) - seghead3 = self.seghead4(stg4) - seghead4 = self.seghead5(stg5) - - return { - 'out': out, - 'seg_stage_outputs': seg_stage_outputs, - 'seghead1':seghead1, - 'seghead2':seghead2, - 'seghead3':seghead3, - 'seghead4':seghead4 - } - else: - return { - 'out': out, - 'seg_stage_outputs': seg_stage_outputs - } - -class InstanceSegmentationBranch(nn.Module): - def __init__(self): - super(InstanceSegmentationBranch,self).__init__() - self.bsconv1 = conv2d(128,256,3,1,1,use_rl=True) - self.bsconv2 = conv2d(256,128,1,0,1,use_rl=True) - def forward(self,data): - input_tensor_size=list(data.size()) - tmp_size=input_tensor_size[2:] - out_put_tensor_size=tuple([int(tmp * 8) for tmp in tmp_size]) - conv1_out=self.bsconv1(data) - conv2_out=self.bsconv2(conv1_out) - isb_out = F.interpolate(conv2_out, size=out_put_tensor_size, mode="bilinear",align_corners=True) - return isb_out - -class BinarySegmentationBranch(nn.Module): - - def __init__(self): - super(BinarySegmentationBranch,self).__init__() - - self.bsconv1_pre = conv2d(128,128,3,1,1,use_rl=True) # - # 融合1/2 特征 - self.bsconv1_pre2 = conv2d((128+(128-32)),32,3,1,1,use_rl=True) # - # 融合1/4 特征 - self.bsconv1_pre4 = conv2d(128,16,3,1,1,use_rl=True)# - # 融合1/8 特征 - self.bsconv3 = conv2d(16,config.num_classes,1,0,1,use_rl=False,use_bn=True) - - - # self.bsconv1 = conv2d(128,256,3,1,1,use_rl=True) - # self.bsconv2 = conv2d(256,128,1,0,1,use_rl=True) - # self.bsconv3 = conv2d(128,config.num_classes,1,0,1,use_rl=False,use_bn=True) - - def forward(self,data,seg_stage_outputs,detail_stage_outputs): - - input_tensor_size=list(data.size()) - tmp_size=input_tensor_size[2:] - output_stage2_size=[int(tmp * 2) for tmp in tmp_size] - output_stage1_size=[int(tmp * 4) for tmp in tmp_size] - out_put_tensor_size=[int(tmp * 8) for tmp in tmp_size] - - out1=self.bsconv1_pre(data) - # channel=256 - # out2=self.bsconv1_pre1(out1) - - output_stage2_tensor = F.interpolate(out1, size=tuple(output_stage2_size), mode="bilinear",align_corners=True) - # output_stage2_tensor = tf.concat([output_stage2_tensor, detail_stage_outputs['stage_2'], semantic_stage_outputs['stage_1']], axis=-1, name='stage_2_concate_features') - output_stage2_tensor=torch.cat([output_stage2_tensor,detail_stage_outputs['stg2'], seg_stage_outputs['stg1']],1) - - output_stage2_tensor=self.bsconv1_pre2(output_stage2_tensor) - # output_stage2_tensor=self.bsconv1_pre3(output_stage2_tensor) - # channel=256 - output_stage1_tensor = F.interpolate(output_stage2_tensor, size=tuple(output_stage1_size), mode="bilinear",align_corners=True) - #output_stage1_tensor = tf.concat([output_stage1_tensor, detail_stage_outputs['stage_1']], axis=-1, name='stage_1_concate_features') - output_stage1_tensor=torch.cat([output_stage1_tensor,detail_stage_outputs['stg1']],1) - - print(output_stage1_tensor.size()) - - output_stage1_tensor=self.bsconv1_pre4(output_stage1_tensor) - # output_stage1_tensor=self.bsconv1_pre5(output_stage1_tensor) - output_stage1_tensor=self.bsconv3(output_stage1_tensor) - - # conv_out1=self.bsconv1(data) - # conv_out2=self.bsconv2(conv_out1) - # conv_out3=self.bsconv3(conv_out2) - - # print("--------------------conv_out3 size--------------------") - # print(list(conv_out3.size())) - # print("--------------------bga size--------------------") - - # print("--------------------out_put_tensor_size--------------------") - # print(out_put_tensor_size) - # print(tuple(out_put_tensor_size)) - # print("--------------------out_put_tensor_size--------------------") - bsb_out = F.interpolate(output_stage1_tensor, size=tuple(out_put_tensor_size), mode="bilinear",align_corners=True) - return bsb_out - -class BiSeNet(nn.Module): - def __init__(self): - super(BiSeNet, self).__init__() - self.db = DetailedBranch() - self.sb = SemanticBranch() - self.bga = BGA(128) - if config.is_training==1: - self.seghead = SegHead(128,128,config.num_classes) - self.criterion = nn.CrossEntropyLoss(ignore_index=255) - self._init_params() - self.binarySegmentationBranch=BinarySegmentationBranch() - self.instanceSegmentationBranch=InstanceSegmentationBranch() - def _init_params(self): - for m in self.modules(): - if isinstance(m, nn.Conv2d): - nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') - if m.bias is not None: - nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.BatchNorm2d): - nn.init.constant_(m.weight, 1) - nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.BatchNorm1d): - nn.init.constant_(m.weight, 1) - nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.Linear): - nn.init.normal_(m.weight, 0, 0.01) - if m.bias is not None: - nn.init.constant_(m.bias, 0) - def forward(self,data,y=None): - db = self.db(data) - sb = self.sb(data) - bga = self.bga(db["out"],sb["out"]) - bsb_res=self.binarySegmentationBranch(bga,sb["seg_stage_outputs"],db["detail_stage_outputs"]) - isb_res=self.instanceSegmentationBranch(bga) - if config.is_training==1: - aux1_loss = self.criterion(sb["seghead1"], y) - aux2_loss = self.criterion(sb["seghead2"], y) - aux3_loss = self.criterion(sb["seghead3"], y) - aux4_loss = self.criterion(sb["seghead4"], y) - return { - 'instance_seg_logits': isb_res, - # 'binary_seg_pred': binary_seg_ret, - 'binary_seg_logits': bsb_res, - 'aux1_loss':aux1_loss, - 'aux2_loss':aux2_loss, - 'aux3_loss':aux3_loss, - 'aux4_loss':aux4_loss - } - return { - 'instance_seg_logits': isb_res, - # 'binary_seg_pred': binary_seg_ret, - 'binary_seg_logits': bsb_res - } - -if __name__ == '__main__': - import os - os.environ["CUDA_VISIBLE_DEVICES"] = '5' - input = torch.rand(1, 3, 256, 512) - model = BiSeNet() - model.eval() - print(model) - output = model(input) - print('+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++') - print('BiSeNet_v2', output["instance_seg_logits"].size()) - print('BiSeNet_v2', output["binary_seg_logits"].size()) - print('+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++') diff --git a/lanenet/model/bisenetv2.py b/lanenet/model/bisenetv2.py deleted file mode 100644 index b4cbbcd..0000000 --- a/lanenet/model/bisenetv2.py +++ /dev/null @@ -1,408 +0,0 @@ - -import torch -import torch.nn as nn -import torch.nn.functional as F - - -class ConvBNReLU(nn.Module): - - def __init__(self, in_chan, out_chan, ks=3, stride=1, padding=1, - dilation=1, groups=1, bias=False): - super(ConvBNReLU, self).__init__() - self.conv = nn.Conv2d( - in_chan, out_chan, kernel_size=ks, stride=stride, - padding=padding, dilation=dilation, - groups=groups, bias=bias) - self.bn = nn.BatchNorm2d(out_chan) - self.relu = nn.ReLU(inplace=True) - - def forward(self, x): - feat = self.conv(x) - feat = self.bn(feat) - feat = self.relu(feat) - return feat - - -class UpSample(nn.Module): - - def __init__(self, n_chan, factor=2): - super(UpSample, self).__init__() - out_chan = n_chan * factor * factor - self.proj = nn.Conv2d(n_chan, out_chan, 1, 1, 0) - self.up = nn.PixelShuffle(factor) - self.init_weight() - - def forward(self, x): - feat = self.proj(x) - feat = self.up(feat) - return feat - - def init_weight(self): - nn.init.xavier_normal_(self.proj.weight, gain=1.) - - -class DetailBranch(nn.Module): - - def __init__(self): - super(DetailBranch, self).__init__() - self.S1 = nn.Sequential( - ConvBNReLU(3, 64, 3, stride=2), - ConvBNReLU(64, 64, 3, stride=1), - ) - self.S2 = nn.Sequential( - ConvBNReLU(64, 64, 3, stride=2), - ConvBNReLU(64, 64, 3, stride=1), - ConvBNReLU(64, 64, 3, stride=1), - ) - self.S3 = nn.Sequential( - ConvBNReLU(64, 128, 3, stride=2), - ConvBNReLU(128, 128, 3, stride=1), - ConvBNReLU(128, 128, 3, stride=1), - ) - - def forward(self, x): - feat = self.S1(x) - feat = self.S2(feat) - feat = self.S3(feat) - return feat - - -class StemBlock(nn.Module): - - def __init__(self): - super(StemBlock, self).__init__() - self.conv = ConvBNReLU(3, 16, 3, stride=2) - self.left = nn.Sequential( - ConvBNReLU(16, 8, 1, stride=1, padding=0), - ConvBNReLU(8, 16, 3, stride=2), - ) - self.right = nn.MaxPool2d( - kernel_size=3, stride=2, padding=1, ceil_mode=False) - self.fuse = ConvBNReLU(32, 16, 3, stride=1) - - def forward(self, x): - feat = self.conv(x) - feat_left = self.left(feat) - feat_right = self.right(feat) - feat = torch.cat([feat_left, feat_right], dim=1) - feat = self.fuse(feat) - return feat - - -class CEBlock(nn.Module): - - def __init__(self): - super(CEBlock, self).__init__() - self.bn = nn.BatchNorm2d(128) - self.conv_gap = ConvBNReLU(128, 128, 1, stride=1, padding=0) - #TODO: in paper here is naive conv2d, no bn-relu - self.conv_last = ConvBNReLU(128, 128, 3, stride=1) - - def forward(self, x): - feat = torch.mean(x, dim=(2, 3), keepdim=True) - feat = self.bn(feat) - feat = self.conv_gap(feat) - feat = feat + x - feat = self.conv_last(feat) - return feat - - -class GELayerS1(nn.Module): - - def __init__(self, in_chan, out_chan, exp_ratio=6): - super(GELayerS1, self).__init__() - mid_chan = in_chan * exp_ratio - self.conv1 = ConvBNReLU(in_chan, in_chan, 3, stride=1) - self.dwconv = nn.Sequential( - nn.Conv2d( - in_chan, mid_chan, kernel_size=3, stride=1, - padding=1, groups=in_chan, bias=False), - nn.BatchNorm2d(mid_chan), - nn.ReLU(inplace=True), # not shown in paper - ) - self.conv2 = nn.Sequential( - nn.Conv2d( - mid_chan, out_chan, kernel_size=1, stride=1, - padding=0, bias=False), - nn.BatchNorm2d(out_chan), - ) - self.conv2[1].last_bn = True - self.relu = nn.ReLU(inplace=True) - - def forward(self, x): - feat = self.conv1(x) - feat = self.dwconv(feat) - feat = self.conv2(feat) - feat = feat + x - feat = self.relu(feat) - return feat - - -class GELayerS2(nn.Module): - - def __init__(self, in_chan, out_chan, exp_ratio=6): - super(GELayerS2, self).__init__() - mid_chan = in_chan * exp_ratio - self.conv1 = ConvBNReLU(in_chan, in_chan, 3, stride=1) - self.dwconv1 = nn.Sequential( - nn.Conv2d( - in_chan, mid_chan, kernel_size=3, stride=2, - padding=1, groups=in_chan, bias=False), - nn.BatchNorm2d(mid_chan), - ) - self.dwconv2 = nn.Sequential( - nn.Conv2d( - mid_chan, mid_chan, kernel_size=3, stride=1, - padding=1, groups=mid_chan, bias=False), - nn.BatchNorm2d(mid_chan), - nn.ReLU(inplace=True), # not shown in paper - ) - self.conv2 = nn.Sequential( - nn.Conv2d( - mid_chan, out_chan, kernel_size=1, stride=1, - padding=0, bias=False), - nn.BatchNorm2d(out_chan), - ) - self.conv2[1].last_bn = True - self.shortcut = nn.Sequential( - nn.Conv2d( - in_chan, in_chan, kernel_size=3, stride=2, - padding=1, groups=in_chan, bias=False), - nn.BatchNorm2d(in_chan), - nn.Conv2d( - in_chan, out_chan, kernel_size=1, stride=1, - padding=0, bias=False), - nn.BatchNorm2d(out_chan), - ) - self.relu = nn.ReLU(inplace=True) - - def forward(self, x): - feat = self.conv1(x) - feat = self.dwconv1(feat) - feat = self.dwconv2(feat) - feat = self.conv2(feat) - shortcut = self.shortcut(x) - feat = feat + shortcut - feat = self.relu(feat) - return feat - - -class SegmentBranch(nn.Module): - - def __init__(self): - super(SegmentBranch, self).__init__() - self.S1S2 = StemBlock() - self.S3 = nn.Sequential( - GELayerS2(16, 32), - GELayerS1(32, 32), - ) - self.S4 = nn.Sequential( - GELayerS2(32, 64), - GELayerS1(64, 64), - ) - self.S5_4 = nn.Sequential( - GELayerS2(64, 128), - GELayerS1(128, 128), - GELayerS1(128, 128), - GELayerS1(128, 128), - ) - self.S5_5 = CEBlock() - - def forward(self, x): - feat2 = self.S1S2(x) - feat3 = self.S3(feat2) - feat4 = self.S4(feat3) - feat5_4 = self.S5_4(feat4) - feat5_5 = self.S5_5(feat5_4) - return feat2, feat3, feat4, feat5_4, feat5_5 - - -class BGALayer(nn.Module): - - def __init__(self): - super(BGALayer, self).__init__() - self.left1 = nn.Sequential( - nn.Conv2d( - 128, 128, kernel_size=3, stride=1, - padding=1, groups=128, bias=False), - nn.BatchNorm2d(128), - nn.Conv2d( - 128, 128, kernel_size=1, stride=1, - padding=0, bias=False), - ) - self.left2 = nn.Sequential( - nn.Conv2d( - 128, 128, kernel_size=3, stride=2, - padding=1, bias=False), - nn.BatchNorm2d(128), - nn.AvgPool2d(kernel_size=3, stride=2, padding=1, ceil_mode=False) - ) - self.right1 = nn.Sequential( - nn.Conv2d( - 128, 128, kernel_size=3, stride=1, - padding=1, bias=False), - nn.BatchNorm2d(128), - ) - self.right2 = nn.Sequential( - nn.Conv2d( - 128, 128, kernel_size=3, stride=1, - padding=1, groups=128, bias=False), - nn.BatchNorm2d(128), - nn.Conv2d( - 128, 128, kernel_size=1, stride=1, - padding=0, bias=False), - ) - self.up1 = nn.Upsample(scale_factor=4) - self.up2 = nn.Upsample(scale_factor=4) - ##TODO: does this really has no relu? - self.conv = nn.Sequential( - nn.Conv2d( - 128, 128, kernel_size=3, stride=1, - padding=1, bias=False), - nn.BatchNorm2d(128), - nn.ReLU(inplace=True), # not shown in paper - ) - - def forward(self, x_d, x_s): - dsize = x_d.size()[2:] - left1 = self.left1(x_d) - left2 = self.left2(x_d) - right1 = self.right1(x_s) - right2 = self.right2(x_s) - right1 = self.up1(right1) - left = left1 * torch.sigmoid(right1) - right = left2 * torch.sigmoid(right2) - right = self.up2(right) - out = self.conv(left + right) - return out - - - -class SegmentHead(nn.Module): - - def __init__(self, in_chan, mid_chan, n_classes, up_factor=8, aux=True): - super(SegmentHead, self).__init__() - self.conv = ConvBNReLU(in_chan, mid_chan, 3, stride=1) - self.drop = nn.Dropout(0.1) - self.up_factor = up_factor - - out_chan = n_classes * up_factor * up_factor - if aux: - self.conv_out = nn.Sequential( - ConvBNReLU(mid_chan, up_factor * up_factor, 3, stride=1), - nn.Conv2d(up_factor * up_factor, out_chan, 1, 1, 0), - nn.PixelShuffle(up_factor) - ) - else: - self.conv_out = nn.Sequential( - nn.Conv2d(mid_chan, out_chan, 1, 1, 0), - nn.PixelShuffle(up_factor) - ) - - def forward(self, x): - feat = self.conv(x) - feat = self.drop(feat) - feat = self.conv_out(feat) - return feat - - -class BiSeNetV2(nn.Module): - - def __init__(self, n_classes, output_aux=True): - super(BiSeNetV2, self).__init__() - self.output_aux = output_aux - self.detail = DetailBranch() - self.segment = SegmentBranch() - self.bga = BGALayer() - - ## TODO: what is the number of mid chan ? - self.head = SegmentHead(128, 1024, n_classes, up_factor=8, aux=False) - if self.output_aux: - self.aux2 = SegmentHead(16, 128, n_classes, up_factor=4) - self.aux3 = SegmentHead(32, 128, n_classes, up_factor=8) - self.aux4 = SegmentHead(64, 128, n_classes, up_factor=16) - self.aux5_4 = SegmentHead(128, 128, n_classes, up_factor=32) - - self.init_weights() - - def forward(self, x): - size = x.size()[2:] - feat_d = self.detail(x) - feat2, feat3, feat4, feat5_4, feat_s = self.segment(x) - feat_head = self.bga(feat_d, feat_s) - - logits = self.head(feat_head) - if self.output_aux: - logits_aux2 = self.aux2(feat2) - logits_aux3 = self.aux3(feat3) - logits_aux4 = self.aux4(feat4) - logits_aux5_4 = self.aux5_4(feat5_4) - return logits, logits_aux2, logits_aux3, logits_aux4, logits_aux5_4 - pred = logits.argmax(dim=1) - return pred - - def init_weights(self): - for name, module in self.named_modules(): - if isinstance(module, (nn.Conv2d, nn.Linear)): - nn.init.kaiming_normal_(module.weight, mode='fan_out') - if not module.bias is None: nn.init.constant_(module.bias, 0) - elif isinstance(module, nn.modules.batchnorm._BatchNorm): - if hasattr(module, 'last_bn') and module.last_bn: - nn.init.zeros_(module.weight) - else: - nn.init.ones_(module.weight) - nn.init.zeros_(module.bias) - - -if __name__ == "__main__": - # x = torch.randn(16, 3, 1024, 2048) - # detail = DetailBranch() - # feat = detail(x) - # print('detail', feat.size()) - # - # x = torch.randn(16, 3, 1024, 2048) - # stem = StemBlock() - # feat = stem(x) - # print('stem', feat.size()) - # - # x = torch.randn(16, 128, 16, 32) - # ceb = CEBlock() - # feat = ceb(x) - # print(feat.size()) - # - # x = torch.randn(16, 32, 16, 32) - # ge1 = GELayerS1(32, 32) - # feat = ge1(x) - # print(feat.size()) - # - # x = torch.randn(16, 16, 16, 32) - # ge2 = GELayerS2(16, 32) - # feat = ge2(x) - # print(feat.size()) - # - # left = torch.randn(16, 128, 64, 128) - # right = torch.randn(16, 128, 16, 32) - # bga = BGALayer() - # feat = bga(left, right) - # print(feat.size()) - # - # x = torch.randn(16, 128, 64, 128) - # head = SegmentHead(128, 128, 19) - # logits = head(x) - # print(logits.size()) - # - # x = torch.randn(16, 3, 1024, 2048) - # segment = SegmentBranch() - # feat = segment(x)[0] - # print(feat.size()) - # - x = torch.randn(16, 3, 1024, 2048) - model = BiSeNetV2(n_classes=19) - outs = model(x) - for out in outs: - print(out.size()) - # print(logits.size()) - - # for name, param in model.named_parameters(): - # if len(param.size()) == 1: - # print(name) \ No newline at end of file diff --git a/lanenet/model/decoders.py b/lanenet/model/decoders.py deleted file mode 100644 index 5418bfd..0000000 --- a/lanenet/model/decoders.py +++ /dev/null @@ -1,175 +0,0 @@ -import torch -import torch.nn as nn -from lanenet import config - -DEVICE = torch.device(config.gpu_no if torch.cuda.is_available() else 'cpu') -from .blocks import * - -class ESPNetDecoder(): - def __init__(self): - - # light-weight decoder - self.level3_C = C(128 + 3, classes, 1, 1) - self.br = nn.BatchNorm2d(classes, eps=1e-03) - self.conv = CBR(19 + classes, classes, 3, 1) - - self.up_l3 = nn.Sequential( - nn.ConvTranspose2d(classes, classes, 2, stride=2, padding=0, output_padding=0, bias=False)) - self.combine_l2_l3 = nn.Sequential(BR(2 * classes), - DilatedParllelResidualBlockB(2 * classes, classes, add=False)) - - self.up_l2 = nn.Sequential( - nn.ConvTranspose2d(classes, classes, 2, stride=2, padding=0, output_padding=0, bias=False), BR(classes)) - - self.classifier = nn.ConvTranspose2d(classes, classes, 2, stride=2, padding=0, output_padding=0, bias=False) - - def forward(self, input): - ''' - :param input: RGB image - :return: transformed feature map - ''' - output0 = self.modules[0](input) - inp1 = self.modules[1](input) - inp2 = self.modules[2](input) - - output0_cat = self.modules[3](torch.cat([output0, inp1], 1)) - output1_0 = self.modules[4](output0_cat) # down-sampled - - for i, layer in enumerate(self.modules[5]): - if i == 0: - output1 = layer(output1_0) - else: - output1 = layer(output1) - - output1_cat = self.modules[6](torch.cat([output1, output1_0, inp2], 1)) - - output2_0 = self.modules[7](output1_cat) # down-sampled - for i, layer in enumerate(self.modules[8]): - if i == 0: - output2 = layer(output2_0) - else: - output2 = layer(output2) - - output2_cat = self.modules[9](torch.cat([output2_0, output2], 1)) # concatenate for feature map width expansion - - output2_c = self.up_l3(self.br(self.modules[10](output2_cat))) # RUM - - output1_C = self.level3_C(output1_cat) # project to C-dimensional space - comb_l2_l3 = self.up_l2(self.combine_l2_l3(torch.cat([output1_C, output2_c], 1))) # RUM - - concat_features = self.conv(torch.cat([comb_l2_l3, output0_cat], 1)) - - classifier = self.classifier(concat_features) - return classifier - - -class ENetDecoder(): - def __init__(self): - # Stage 4 - Decoder - self.upsample4_0 = UpsamplingBottleneck( - 128, 64, padding=1, dropout_prob=0.1, relu=decoder_relu) - self.regular4_1 = RegularBottleneck( - 64, padding=1, dropout_prob=0.1, relu=decoder_relu) - self.regular4_2 = RegularBottleneck( - 64, padding=1, dropout_prob=0.1, relu=decoder_relu) - - # Stage 5 - Decoder - self.upsample5_0 = UpsamplingBottleneck( - 64, 16, padding=1, dropout_prob=0.1, relu=decoder_relu) - self.regular5_1 = RegularBottleneck( - 16, padding=1, dropout_prob=0.1, relu=decoder_relu) - self.transposed_conv = nn.ConvTranspose2d( - 16, - num_classes, - kernel_size=3, - stride=2, - padding=1, - output_padding=1, - bias=False) - - def forward(self, x): - # Initial block - x = self.initial_block(x) - - # Stage 1 - Encoder - x, max_indices1_0 = self.downsample1_0(x) - x = self.regular1_1(x) - x = self.regular1_2(x) - x = self.regular1_3(x) - x = self.regular1_4(x) - - # Stage 2 - Encoder - x, max_indices2_0 = self.downsample2_0(x) - x = self.regular2_1(x) - x = self.dilated2_2(x) - x = self.asymmetric2_3(x) - x = self.dilated2_4(x) - x = self.regular2_5(x) - x = self.dilated2_6(x) - x = self.asymmetric2_7(x) - x = self.dilated2_8(x) - - # Stage 3 - Encoder - x = self.regular3_0(x) - x = self.dilated3_1(x) - x = self.asymmetric3_2(x) - x = self.dilated3_3(x) - x = self.regular3_4(x) - x = self.dilated3_5(x) - x = self.asymmetric3_6(x) - x = self.dilated3_7(x) - - # Stage 4 - Decoder - x = self.upsample4_0(x, max_indices2_0) - x = self.regular4_1(x) - x = self.regular4_2(x) - - # Stage 5 - Decoder - x = self.upsample5_0(x, max_indices1_0) - x = self.regular5_1(x) - x = self.transposed_conv(x) - - return x - - -class FCNDecoder(nn.Module): - def __init__(self, decode_layers, decode_channels=[], decode_last_stride=8): - super(FCNDecoder, self).__init__() - - self._decode_channels = [512, 256] - self._out_channel = 64 - self._decode_layers = decode_layers - - self._conv_layers = [] - for _ch in self._decode_channels: - self._conv_layers.append(nn.Conv2d(_ch, self._out_channel, kernel_size=1, bias=False).to(DEVICE)) - - self._conv_final = nn.Conv2d(self._out_channel, 2, kernel_size=1, bias=False) - self._deconv = nn.ConvTranspose2d(self._out_channel, self._out_channel, kernel_size=4, stride=2, padding=1, - bias=False) - - self._deconv_final = nn.ConvTranspose2d(self._out_channel, self._out_channel, kernel_size=16, - stride=decode_last_stride, - padding=4, bias=False) - - def forward(self, encode_data): - ret = {} - input_tensor = encode_data[self._decode_layers[0]] - input_tensor.to(DEVICE) - score = self._conv_layers[0](input_tensor) - for i, layer in enumerate(self._decode_layers[1:]): - deconv = self._deconv(score) - - input_tensor = encode_data[layer] - score = self._conv_layers[i](input_tensor) - - fused = torch.add(deconv, score) - score = fused - - deconv_final = self._deconv_final(score) - score_final = self._conv_final(deconv_final) - - - ret['logits'] = score_final - ret['deconv'] = deconv_final - return ret diff --git a/lanenet/model/encoders.py b/lanenet/model/encoders.py deleted file mode 100644 index f4e99fd..0000000 --- a/lanenet/model/encoders.py +++ /dev/null @@ -1,184 +0,0 @@ -# coding: utf-8 -""" -Shared encoders (U-net). -""" -import torch -import torch.nn as nn -from collections import OrderedDict - -import torchvision.models as models - -from .blocks import RegularBottleneck, DownsamplingBottleneck, InitialBlock, InputProjectionA, \ - DilatedParallelResidualBlockB, DownSamplerB, C, CBR, BR - - -class VGGEncoder(nn.Module): - """ - Simple VGG Encoder - """ - - def __init__(self, num_blocks, in_channels, out_channels): - super(VGGEncoder, self).__init__() - - self.pretrained_modules = models.vgg16(pretrained=True).features - - self.num_blocks = num_blocks - self._in_channels = in_channels - self._out_channels = out_channels - self._conv_reps = [2, 2, 3, 3, 3] - self.net = nn.Sequential() - self.pretrained_net = nn.Sequential() - - for i in range(num_blocks): - self.net.add_module("block" + str(i + 1), self._encode_block(i + 1)) - self.pretrained_net.add_module("block" + str(i + 1), self._encode_pretrained_block(i + 1)) - - def _encode_block(self, block_id, kernel_size=3, stride=1): - out_channels = self._out_channels[block_id - 1] - padding = (kernel_size - 1) // 2 - seq = nn.Sequential() - - for i in range(self._conv_reps[block_id - 1]): - if i == 0: - in_channels = self._in_channels[block_id - 1] - else: - in_channels = out_channels - seq.add_module("conv_{}_{}".format(block_id, i + 1), - nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding)) - seq.add_module("bn_{}_{}".format(block_id, i + 1), nn.BatchNorm2d(out_channels)) - seq.add_module("relu_{}_{}".format(block_id, i + 1), nn.ReLU()) - seq.add_module("maxpool" + str(block_id), nn.MaxPool2d(kernel_size=2, stride=2)) - return seq - - def _encode_pretrained_block(self, block_id): - seq = nn.Sequential() - for i in range(0, self._conv_reps[block_id - 1], 4): - seq.add_module("conv_{}_{}".format(block_id, i + 1), self.pretrained_modules[i]) - seq.add_module("relu_{}_{}".format(block_id, i + 2), self.pretrained_modules[i + 1]) - seq.add_module("conv_{}_{}".format(block_id, i + 3), self.pretrained_modules[i + 2]) - seq.add_module("relu_{}_{}".format(block_id, i + 4), self.pretrained_modules[i + 3]) - seq.add_module("maxpool" + str(block_id), self.pretrained_modules[i + 4]) - return seq - - def forward(self, input_tensor): - ret = OrderedDict() - # 5 stage of encoding - X = input_tensor - for i, block in enumerate(self.net): - pool = block(X) - ret["pool" + str(i + 1)] = pool - - X = pool - return ret - - -class ESPNetEncoder(nn.Module): - """ - ESPNet-C encoder - """ - - def __init__(self, classes=20, p=5, q=3): - ''' - :param classes: number of classes in the dataset. Default is 20 for the cityscapes - :param p: depth multiplier - :param q: depth multiplier - ''' - super().__init__() - self.level1 = CBR(3, 16, 3, 2) - self.sample1 = InputProjectionA(1) - self.sample2 = InputProjectionA(2) - - self.b1 = BR(16 + 3) - self.level2_0 = DownSamplerB(16 + 3, 64) - - self.level2 = nn.ModuleList() - for i in range(0, p): - self.level2.append(DilatedParallelResidualBlockB(64, 64)) - self.b2 = BR(128 + 3) - - self.level3_0 = DownSamplerB(128 + 3, 128) - self.level3 = nn.ModuleList() - for i in range(0, q): - self.level3.append(DilatedParallelResidualBlockB(128, 128)) - self.b3 = BR(256) - - self.classifier = C(256, classes, 1, 1) - - def forward(self, input): - ''' - :param input: Receives the input RGB image - :return: the transformed feature map with spatial dimensions 1/8th of the input image - ''' - output0 = self.level1(input) - inp1 = self.sample1(input) - inp2 = self.sample2(input) - - output0_cat = self.b1(torch.cat([output0, inp1], 1)) - output1_0 = self.level2_0(output0_cat) # down-sampled - - for i, layer in enumerate(self.level2): - if i == 0: - output1 = layer(output1_0) - else: - output1 = layer(output1) - - output1_cat = self.b2(torch.cat([output1, output1_0, inp2], 1)) - - output2_0 = self.level3_0(output1_cat) # down-sampled - for i, layer in enumerate(self.level3): - if i == 0: - output2 = layer(output2_0) - else: - output2 = layer(output2) - - output2_cat = self.b3(torch.cat([output2_0, output2], 1)) - - classifier = self.classifier(output2_cat) - - return classifier - - -class ENetEncoder(nn.Module): - """ - ENET Encoder - """ - - def __init__(self, num_classes, encoder_relu=False, decoder_relu=True): - super().__init__() - - def forward(self, input): - self.initial_block = InitialBlock(3, 16, padding=1, relu=encoder_relu) - - # Stage 1 - Encoder - self.downsample1_0 = DownsamplingBottleneck(16, 64, padding=1, return_indices=True, dropout_prob=0.01, - relu=encoder_relu) - self.regular1_1 = RegularBottleneck(64, padding=1, dropout_prob=0.01, relu=encoder_relu) - self.regular1_2 = RegularBottleneck(64, padding=1, dropout_prob=0.01, relu=encoder_relu) - self.regular1_3 = RegularBottleneck(64, padding=1, dropout_prob=0.01, relu=encoder_relu) - self.regular1_4 = RegularBottleneck(64, padding=1, dropout_prob=0.01, relu=encoder_relu) - - # Stage 2 - Encoder - self.downsample2_0 = DownsamplingBottleneck(64, 128, padding=1, return_indices=True, dropout_prob=0.1, - relu=encoder_relu) - self.regular2_1 = RegularBottleneck(128, padding=1, dropout_prob=0.1, relu=encoder_relu) - self.dilated2_2 = RegularBottleneck(128, dilation=2, padding=2, dropout_prob=0.1, relu=encoder_relu) - self.asymmetric2_3 = RegularBottleneck(128, kernel_size=5, padding=2, asymmetric=True, dropout_prob=0.1, - relu=encoder_relu) - self.dilated2_4 = RegularBottleneck(128, dilation=4, padding=4, dropout_prob=0.1, relu=encoder_relu) - self.regular2_5 = RegularBottleneck(128, padding=1, dropout_prob=0.1, relu=encoder_relu) - self.dilated2_6 = RegularBottleneck(128, dilation=8, padding=8, dropout_prob=0.1, relu=encoder_relu) - self.asymmetric2_7 = RegularBottleneck(128, kernel_size=5, asymmetric=True, padding=2, dropout_prob=0.1, - relu=encoder_relu) - self.dilated2_8 = RegularBottleneck(128, dilation=16, padding=16, dropout_prob=0.1, relu=encoder_relu) - - # Stage 3 - Encoder - self.regular3_0 = RegularBottleneck(128, padding=1, dropout_prob=0.1, relu=encoder_relu) - self.dilated3_1 = RegularBottleneck(128, dilation=2, padding=2, dropout_prob=0.1, relu=encoder_relu) - self.asymmetric3_2 = RegularBottleneck(128, kernel_size=5, padding=2, asymmetric=True, dropout_prob=0.1, - relu=encoder_relu) - self.dilated3_3 = RegularBottleneck(128, dilation=4, padding=4, dropout_prob=0.1, relu=encoder_relu) - self.regular3_4 = RegularBottleneck(128, padding=1, dropout_prob=0.1, relu=encoder_relu) - self.dilated3_5 = RegularBottleneck(128, dilation=8, padding=8, dropout_prob=0.1, relu=encoder_relu) - self.asymmetric3_6 = RegularBottleneck(128, kernel_size=5, asymmetric=True, padding=2, dropout_prob=0.1, - relu=encoder_relu) - self.dilated3_7 = RegularBottleneck(128, dilation=16, padding=16, dropout_prob=0.1, relu=encoder_relu) diff --git a/lanenet/model/lanenet_postprocess.py b/lanenet/model/lanenet_postprocess.py deleted file mode 100644 index 41690f1..0000000 --- a/lanenet/model/lanenet_postprocess.py +++ /dev/null @@ -1,431 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding: utf-8 -*- -# @Time : 18-5-30 上午10:04 -# @Author : MaybeShewill-CV -# @Site : https://github.com/MaybeShewill-CV/lanenet-lane-detection -# @File : lanenet_postprocess.py -# @IDE: PyCharm Community Edition -""" -LaneNet model post process -""" -import os.path as ops -import math - -import cv2 -import glog as log -import numpy as np -from sklearn.cluster import DBSCAN -from sklearn.preprocessing import StandardScaler - - -def _morphological_process(image, kernel_size=5): - """ - morphological process to fill the hole in the binary segmentation result - :param image: - :param kernel_size: - :return: - """ - if len(image.shape) == 3: - raise ValueError('Binary segmentation result image should be a single channel image') - - if image.dtype is not np.uint8: - image = np.array(image, np.uint8) - - kernel = cv2.getStructuringElement(shape=cv2.MORPH_ELLIPSE, ksize=(kernel_size, kernel_size)) - - # close operation fille hole - closing = cv2.morphologyEx(image, cv2.MORPH_CLOSE, kernel, iterations=1) - - return closing - - -def _connect_components_analysis(image): - """ - connect components analysis to remove the small components - :param image: - :return: - """ - if len(image.shape) == 3: - gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) - else: - gray_image = image - - return cv2.connectedComponentsWithStats(gray_image, connectivity=8, ltype=cv2.CV_32S) - - -class _LaneFeat(object): - """ - - """ - def __init__(self, feat, coord, class_id=-1): - """ - lane feat object - :param feat: lane embeddng feats [feature_1, feature_2, ...] - :param coord: lane coordinates [x, y] - :param class_id: lane class id - """ - self._feat = feat - self._coord = coord - self._class_id = class_id - - @property - def feat(self): - """ - - :return: - """ - return self._feat - - @feat.setter - def feat(self, value): - """ - - :param value: - :return: - """ - if not isinstance(value, np.ndarray): - value = np.array(value, dtype=np.float64) - - if value.dtype != np.float32: - value = np.array(value, dtype=np.float64) - - self._feat = value - - @property - def coord(self): - """ - - :return: - """ - return self._coord - - @coord.setter - def coord(self, value): - """ - - :param value: - :return: - """ - if not isinstance(value, np.ndarray): - value = np.array(value) - - if value.dtype != np.int32: - value = np.array(value, dtype=np.int32) - - self._coord = value - - @property - def class_id(self): - """ - - :return: - """ - return self._class_id - - @class_id.setter - def class_id(self, value): - """ - - :param value: - :return: - """ - if not isinstance(value, np.int64): - raise ValueError('Class id must be integer') - - self._class_id = value - - -class _LaneNetCluster(object): - """ - Instance segmentation result cluster - """ - - def __init__(self, cfg): - """ - - """ - self._color_map = [np.array([255, 0, 0]), - np.array([0, 255, 0]), - np.array([0, 0, 255]), - np.array([125, 125, 0]), - np.array([0, 125, 125]), - np.array([125, 0, 125]), - np.array([50, 100, 50]), - np.array([100, 50, 100])] - self._cfg = cfg - - def _embedding_feats_dbscan_cluster(self, embedding_image_feats): - """ - dbscan cluster - :param embedding_image_feats: - :return: - """ - db = DBSCAN(eps=self._cfg.POSTPROCESS.DBSCAN_EPS, min_samples=self._cfg.POSTPROCESS.DBSCAN_MIN_SAMPLES) - try: - features = StandardScaler().fit_transform(embedding_image_feats) - db.fit(features) - except Exception as err: - log.error(err) - ret = { - 'origin_features': None, - 'cluster_nums': 0, - 'db_labels': None, - 'unique_labels': None, - 'cluster_center': None - } - return ret - db_labels = db.labels_ - unique_labels = np.unique(db_labels) - - num_clusters = len(unique_labels) - cluster_centers = db.components_ - - ret = { - 'origin_features': features, - 'cluster_nums': num_clusters, - 'db_labels': db_labels, - 'unique_labels': unique_labels, - 'cluster_center': cluster_centers - } - - return ret - - @staticmethod - def _get_lane_embedding_feats(binary_seg_ret, instance_seg_ret): - """ - get lane embedding features according the binary seg result - :param binary_seg_ret: - :param instance_seg_ret: - :return: - """ - idx = np.where(binary_seg_ret == 255) - lane_embedding_feats = instance_seg_ret[idx] - # idx_scale = np.vstack((idx[0] / 256.0, idx[1] / 512.0)).transpose() - # lane_embedding_feats = np.hstack((lane_embedding_feats, idx_scale)) - lane_coordinate = np.vstack((idx[1], idx[0])).transpose() - - assert lane_embedding_feats.shape[0] == lane_coordinate.shape[0] - - ret = { - 'lane_embedding_feats': lane_embedding_feats, - 'lane_coordinates': lane_coordinate - } - - return ret - - def apply_lane_feats_cluster(self, binary_seg_result, instance_seg_result): - """ - - :param binary_seg_result: - :param instance_seg_result: - :return: - """ - # get embedding feats and coords - get_lane_embedding_feats_result = self._get_lane_embedding_feats( - binary_seg_ret=binary_seg_result, - instance_seg_ret=instance_seg_result - ) - - # dbscan cluster - dbscan_cluster_result = self._embedding_feats_dbscan_cluster( - embedding_image_feats=get_lane_embedding_feats_result['lane_embedding_feats'] - ) - - mask = np.zeros(shape=[binary_seg_result.shape[0], binary_seg_result.shape[1], 3], dtype=np.uint8) - db_labels = dbscan_cluster_result['db_labels'] - unique_labels = dbscan_cluster_result['unique_labels'] - coord = get_lane_embedding_feats_result['lane_coordinates'] - - if db_labels is None: - return None, None - - lane_coords = [] - - for index, label in enumerate(unique_labels.tolist()): - if label == -1: - continue - idx = np.where(db_labels == label) - pix_coord_idx = tuple((coord[idx][:, 1], coord[idx][:, 0])) - mask[pix_coord_idx] = self._color_map[index] - lane_coords.append(coord[idx]) - - return mask, lane_coords - - -class LaneNetPostProcessor(object): - """ - lanenet post process for lane generation - """ - def __init__(self, cfg, ipm_remap_file_path='./data/tusimple_ipm_remap.yml'): - """ - - :param ipm_remap_file_path: ipm generate file path - """ - assert ops.exists(ipm_remap_file_path), '{:s} not exist'.format(ipm_remap_file_path) - - self._cfg = cfg - self._cluster = _LaneNetCluster(cfg=cfg) - self._ipm_remap_file_path = ipm_remap_file_path - - remap_file_load_ret = self._load_remap_matrix() - self._remap_to_ipm_x = remap_file_load_ret['remap_to_ipm_x'] - self._remap_to_ipm_y = remap_file_load_ret['remap_to_ipm_y'] - - self._color_map = [np.array([255, 0, 0]), - np.array([0, 255, 0]), - np.array([0, 0, 255]), - np.array([125, 125, 0]), - np.array([0, 125, 125]), - np.array([125, 0, 125]), - np.array([50, 100, 50]), - np.array([100, 50, 100])] - - def _load_remap_matrix(self): - """ - - :return: - """ - fs = cv2.FileStorage(self._ipm_remap_file_path, cv2.FILE_STORAGE_READ) - - remap_to_ipm_x = fs.getNode('remap_ipm_x').mat() - remap_to_ipm_y = fs.getNode('remap_ipm_y').mat() - - ret = { - 'remap_to_ipm_x': remap_to_ipm_x, - 'remap_to_ipm_y': remap_to_ipm_y, - } - - fs.release() - - return ret - - def postprocess(self, binary_seg_result, instance_seg_result=None, - min_area_threshold=100, source_image=None, - data_source='tusimple'): - """ - - :param binary_seg_result: - :param instance_seg_result: - :param min_area_threshold: - :param source_image: - :param data_source: - :return: - """ - # convert binary_seg_result - binary_seg_result = np.array(binary_seg_result * 255, dtype=np.uint8) - - # apply image morphology operation to fill in the hold and reduce the small area - morphological_ret = _morphological_process(binary_seg_result, kernel_size=5) - - connect_components_analysis_ret = _connect_components_analysis(image=morphological_ret) - - labels = connect_components_analysis_ret[1] - stats = connect_components_analysis_ret[2] - for index, stat in enumerate(stats): - if stat[4] <= min_area_threshold: - idx = np.where(labels == index) - morphological_ret[idx] = 0 - - # apply embedding features cluster - mask_image, lane_coords = self._cluster.apply_lane_feats_cluster( - binary_seg_result=morphological_ret, - instance_seg_result=instance_seg_result - ) - - if mask_image is None: - return { - 'mask_image': None, - 'fit_params': None, - 'source_image': None, - } - - # lane line fit - fit_params = [] - src_lane_pts = [] # lane pts every single lane - for lane_index, coords in enumerate(lane_coords): - if data_source == 'tusimple': - tmp_mask = np.zeros(shape=(720, 1280), dtype=np.uint8) - tmp_mask[tuple((np.int_(coords[:, 1] * 720 / 256), np.int_(coords[:, 0] * 1280 / 512)))] = 255 - else: - raise ValueError('Wrong data source now only support tusimple') - tmp_ipm_mask = cv2.remap( - tmp_mask, - self._remap_to_ipm_x, - self._remap_to_ipm_y, - interpolation=cv2.INTER_NEAREST - ) - nonzero_y = np.array(tmp_ipm_mask.nonzero()[0]) - nonzero_x = np.array(tmp_ipm_mask.nonzero()[1]) - - fit_param = np.polyfit(nonzero_y, nonzero_x, 2) - fit_params.append(fit_param) - - [ipm_image_height, ipm_image_width] = tmp_ipm_mask.shape - plot_y = np.linspace(10, ipm_image_height, ipm_image_height - 10) - fit_x = fit_param[0] * plot_y ** 2 + fit_param[1] * plot_y + fit_param[2] - # fit_x = fit_param[0] * plot_y ** 3 + fit_param[1] * plot_y ** 2 + fit_param[2] * plot_y + fit_param[3] - - lane_pts = [] - for index in range(0, plot_y.shape[0], 5): - src_x = self._remap_to_ipm_x[ - int(plot_y[index]), int(np.clip(fit_x[index], 0, ipm_image_width - 1))] - if src_x <= 0: - continue - src_y = self._remap_to_ipm_y[ - int(plot_y[index]), int(np.clip(fit_x[index], 0, ipm_image_width - 1))] - src_y = src_y if src_y > 0 else 0 - - lane_pts.append([src_x, src_y]) - - src_lane_pts.append(lane_pts) - - # tusimple test data sample point along y axis every 10 pixels - source_image_width = source_image.shape[1] - for index, single_lane_pts in enumerate(src_lane_pts): - single_lane_pt_x = np.array(single_lane_pts, dtype=np.float32)[:, 0] - single_lane_pt_y = np.array(single_lane_pts, dtype=np.float32)[:, 1] - if data_source == 'tusimple': - start_plot_y = 240 - end_plot_y = 720 - else: - raise ValueError('Wrong data source now only support tusimple') - step = int(math.floor((end_plot_y - start_plot_y) / 10)) - for plot_y in np.linspace(start_plot_y, end_plot_y, step): - diff = single_lane_pt_y - plot_y - fake_diff_bigger_than_zero = diff.copy() - fake_diff_smaller_than_zero = diff.copy() - fake_diff_bigger_than_zero[np.where(diff <= 0)] = float('inf') - fake_diff_smaller_than_zero[np.where(diff > 0)] = float('-inf') - idx_low = np.argmax(fake_diff_smaller_than_zero) - idx_high = np.argmin(fake_diff_bigger_than_zero) - - previous_src_pt_x = single_lane_pt_x[idx_low] - previous_src_pt_y = single_lane_pt_y[idx_low] - last_src_pt_x = single_lane_pt_x[idx_high] - last_src_pt_y = single_lane_pt_y[idx_high] - - if previous_src_pt_y < start_plot_y or last_src_pt_y < start_plot_y or \ - fake_diff_smaller_than_zero[idx_low] == float('-inf') or \ - fake_diff_bigger_than_zero[idx_high] == float('inf'): - continue - - interpolation_src_pt_x = (abs(previous_src_pt_y - plot_y) * previous_src_pt_x + - abs(last_src_pt_y - plot_y) * last_src_pt_x) / \ - (abs(previous_src_pt_y - plot_y) + abs(last_src_pt_y - plot_y)) - interpolation_src_pt_y = (abs(previous_src_pt_y - plot_y) * previous_src_pt_y + - abs(last_src_pt_y - plot_y) * last_src_pt_y) / \ - (abs(previous_src_pt_y - plot_y) + abs(last_src_pt_y - plot_y)) - - if interpolation_src_pt_x > source_image_width or interpolation_src_pt_x < 10: - continue - - lane_color = self._color_map[index].tolist() - cv2.circle(source_image, (int(interpolation_src_pt_x), - int(interpolation_src_pt_y)), 5, lane_color, -1) - ret = { - 'mask_image': mask_image, - 'fit_params': fit_params, - 'source_image': source_image, - } - - return ret diff --git a/lanenet/model/model-old.py b/lanenet/model/model-old.py deleted file mode 100644 index e3f5f26..0000000 --- a/lanenet/model/model-old.py +++ /dev/null @@ -1,109 +0,0 @@ -# coding: utf-8 -""" -LaneNet model -https://arxiv.org/pdf/1807.01726.pdf -""" -import torch -import torch.nn as nn -import torch.nn.functional as F - -from lanenet.model.loss import DiscriminativeLoss -from lanenet.model.encoders import VGGEncoder -from lanenet.model.decoders import ESPNetDecoder, FCNDecoder -from lanenet import config - -DEVICE = torch.device(config.gpu_no if torch.cuda.is_available() else 'cpu') - - -class LaneNet1(nn.Module): - def __init__(self, arch="VGG"): - super(LaneNet1, self).__init__() - # no of instances for segmentation - self.no_of_instances = 5 - encode_num_blocks = 5 - in_channels = [3, 64, 128, 256, 512] - out_channels = in_channels[1:] + [512] - self._arch = arch - if self._arch == 'VGG': - self._encoder = VGGEncoder(encode_num_blocks, in_channels, out_channels) - self._encoder.to(DEVICE) - - decode_layers = ["pool5", "pool4", "pool3"] - decode_channels = out_channels[:-len(decode_layers) - 1:-1] - decode_last_stride = 8 - # self._decoder = ESPNetDecoder() - self._decoder = FCNDecoder(decode_layers, decode_channels, decode_last_stride) - self._decoder.to(DEVICE) - elif self._arch == 'ESPNet': - raise NotImplementedError - elif self._arch == 'ENNet': - raise NotImplementedError - - self._pix_layer = nn.Conv2d(in_channels=64, out_channels=self.no_of_instances, kernel_size=1, bias=False).to( - DEVICE) - self.relu = nn.ReLU().to(DEVICE) - - def forward(self, input_tensor): - encode_ret = self._encoder(input_tensor) - decode_ret = self._decoder(encode_ret) - - decode_logits = decode_ret['logits'] - - decode_logits = decode_logits.to(DEVICE) - - binary_seg_ret = torch.argmax(F.softmax(decode_logits, dim=1), dim=1, keepdim=True) - - decode_deconv = decode_ret['deconv'] - pix_embedding = self.relu(self._pix_layer(decode_deconv)) - - return { - 'instance_seg_logits': pix_embedding, - 'binary_seg_pred': binary_seg_ret, - 'binary_seg_logits': decode_logits - } -if __name__ == '__main__': - import os - os.environ["CUDA_VISIBLE_DEVICES"] = '5' - input = torch.rand(1, 3, 256, 512).cuda() - model = LaneNet1().cuda() - model.eval() - print(model) - output = model(input) - binary_seg_pred=output["binary_seg_pred"].squeeze(0) - binary_seg_pred=binary_seg_pred.squeeze(0) - - instance_seg_logits=output["instance_seg_logits"].squeeze(0) - instance_seg_logits=instance_seg_logits.permute(1, 2, 0) - - print('+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++') - print('BiSeNet_v2', binary_seg_pred.size()) - print('BiSeNet_v2', instance_seg_logits.size()) - print('+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++') - -def compute_loss(net_output, binary_label, instance_label): - k_binary = 0.7 - k_instance = 0.3 - k_dist = 1.0 - - ce_loss_fn = nn.CrossEntropyLoss() - binary_seg_logits = net_output["binary_seg_logits"] - binary_loss = ce_loss_fn(binary_seg_logits, binary_label) - - pix_embedding = net_output["instance_seg_logits"] - ds_loss_fn = DiscriminativeLoss(0.5, 1.5, 1.0, 1.0, 0.001) - var_loss, dist_loss, reg_loss = ds_loss_fn(pix_embedding, instance_label) - binary_loss = binary_loss * k_binary - instance_loss = var_loss * k_instance - dist_loss = dist_loss * k_dist - total_loss = binary_loss + instance_loss + dist_loss - out = net_output["binary_seg_pred"] - iou = 0 - batch_size = out.size()[0] - for i in range(batch_size): - PR = out[i].squeeze(0).nonzero().size()[0] - GT = binary_label[i].nonzero().size()[0] - TP = (out[i].squeeze(0) * binary_label[i]).nonzero().size()[0] - union = PR + GT - TP - iou += TP / union - iou = iou / batch_size - return total_loss, binary_loss, instance_loss, out, iou diff --git a/lanenet/online_test_video.py b/lanenet/online_test_video.py index 599f9a6..dccfa75 100644 --- a/lanenet/online_test_video.py +++ b/lanenet/online_test_video.py @@ -8,9 +8,10 @@ from lanenet.dataloader.transformers import Rescale from lanenet.model.model import LaneNet import torch.nn as nn +import argparse import os -DEVICE = torch.device('cuda:3' if torch.cuda.is_available() else 'cpu') -os.environ["CUDA_VISIBLE_DEVICES"] = '3' +DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') +os.environ["CUDA_VISIBLE_DEVICES"] = '0' def compose_img(image_data, out, i=0): oridata=image_data[i].cpu().numpy() @@ -38,7 +39,16 @@ def compose_img(image_data, out, i=0): return val_gt if __name__ == '__main__': - model_path = './checkpoints-combine-new1/83_checkpoint.pth' + parser = argparse.ArgumentParser() + parser.add_argument("--weight", help="./checkpoints/83_checkpoint.pth") + parser.add_argument("--video", help="./checkpoints/83_checkpoint.pth") + args=parser.parse_args() + model_path = args.weight + + # sourceFileName='ch0_20200318140335_20200318140435' + # video_path = os.path.join("/workspace/lanenet-lane-detection-11000", sourceFileName+'.mp4') + video_path = args.video + gpu = True if not torch.cuda.is_available(): gpu = False @@ -61,8 +71,6 @@ def compose_img(image_data, out, i=0): model.eval() - sourceFileName='ch0_20200318140335_20200318140435' - video_path = os.path.join("/workspace/lanenet-lane-detection-11000", sourceFileName+'.mp4') times=0 frameFrequency=1 camera = cv2.VideoCapture(video_path) diff --git a/lanenet/train.py b/lanenet/train.py index 37b64d6..557d6c5 100644 --- a/lanenet/train.py +++ b/lanenet/train.py @@ -97,7 +97,7 @@ def train(train_loader, model, optimizer, epoch,w1,w2,w3,w4): batch_time.update(time.time() - end) end = time.time() - if step % 30 == 0: + if step % config.show_interval == 0: print( "Epoch {ep} Step {st} |({batch}/{size})| ETA: {et:.2f}|Total loss:{tot:.5f}|Binary loss:{bin:.5f}|Instance loss:{ins:.5f}|IoU:{iou:.5f}".format( ep=epoch + 1, @@ -131,35 +131,35 @@ def save_model(save_path, epoch, model): def main(): - args = parse_args() + # args = parse_args() - save_path = args.save - w1 = args.w1 - w2 = args.w2 - w3 = args.w3 - w4 = args.w4 + save_path = config.save_path + w1 = config.w1 + w2 = config.w2 + w3 = config.w3 + w4 = config.w4 if not os.path.isdir(save_path): os.makedirs(save_path) - train_dataset_file = '/workspace/all/index/train1' - val_dataset_file = '/workspace/all/index/val1' + train_dataset_file = config.train_dataset_file + val_dataset_file = config.val_dataset_file train_dataset = LaneDataSet(train_dataset_file, transform=None) # train_dataset = LaneDataSet(train_dataset_file, transform=transforms.Compose([Rescale((1280, 720))])) - train_loader = DataLoader(train_dataset, batch_size=args.bs, shuffle=True,num_workers=24,pin_memory=True,drop_last=True) + train_loader = DataLoader(train_dataset, batch_size=config.bs, shuffle=True,num_workers=4,pin_memory=True,drop_last=True) - if args.val: - val_dataset = LaneDataSet(val_dataset_file, transform=None) - # val_dataset = LaneDataSet(val_dataset_file, transform=transforms.Compose([Rescale((1280, 720))])) - val_loader = DataLoader(val_dataset, batch_size=args.bs, shuffle=True,num_workers=24,pin_memory=True,drop_last=True) + # if args.val: + val_dataset = LaneDataSet(val_dataset_file, transform=None) + # val_dataset = LaneDataSet(val_dataset_file, transform=transforms.Compose([Rescale((1280, 720))])) + val_loader = DataLoader(val_dataset, batch_size=config.bs, shuffle=True,num_workers=4,pin_memory=True,drop_last=True) model = LaneNet() model = nn.DataParallel(model, device_ids=config.device_ids) model.to(DEVICE) - optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) - print(f"{args.epochs} epochs {len(train_dataset)} training samples\n") + optimizer = torch.optim.Adam(model.parameters(), lr=config.lr) + print(f"{config.epochs} epochs {len(train_dataset)} training samples\n") log_model="/workspace/pytorch-lanenet-master/checkpoints-combine-new1/83_checkpoint_state.pth" # 如果有保存的模型，则加载模型，并在其基础上继续训练 if os.path.exists(log_model): @@ -177,12 +177,12 @@ def main(): start_epoch = 0 print('no model,will start train from 0 epoche') - for epoch in range(start_epoch, args.epochs): + for epoch in range(start_epoch, config.epochs): print(f"Epoch {epoch}") train_iou = train(train_loader, model, optimizer, epoch,w1,w2,w3,w4) - if args.val: - val_iou = test(val_loader, model, epoch) - if (epoch+1) % 3 == 0: + # if args.val: + val_iou = test(val_loader, model, epoch) + if (epoch+1) % config.save_interval == 0: save_model(save_path, epoch, model) save_state_name = os.path.join(save_path, f'{epoch}_checkpoint_state.pth') checkpoint = { @@ -192,8 +192,8 @@ def main(): } torch.save(checkpoint, save_state_name) print(f"Train IoU : {train_iou}") - if args.val: - print(f"Val IoU : {val_iou}") + # if args.val: + print(f"Val IoU : {val_iou}") if __name__ == '__main__': diff --git a/lanenet/utils/postprocess.py b/lanenet/utils/postprocess.py index 39f23f4..01487ae 100644 --- a/lanenet/utils/postprocess.py +++ b/lanenet/utils/postprocess.py @@ -2,7 +2,7 @@ from sklearn.cluster import MeanShift, estimate_bandwidth -def embedding_post_process(embedding, bin_seg, band_width=1.5, max_num_lane=4): +def embedding_post_process(embedding, bin_seg, band_width=1.5, max_num_lane=6): """ First use mean shift to find dense cluster center.