高速公路违停检测
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

299 line
12KB

  1. """Bilateral Segmentation Network"""
  2. import torch
  3. import torch.nn as nn
  4. import torch.nn.functional as F
  5. import numpy as np
  6. from core.models.base_models.resnet import resnet18,resnet50
  7. from core.nn import _ConvBNReLU
  8. __all__ = ['BiSeNet', 'get_bisenet', 'get_bisenet_resnet18_citys']
  9. class BiSeNet(nn.Module):
  10. def __init__(self, nclass, backbone='resnet18', aux=False, jpu=False, pretrained_base=True, **kwargs):
  11. super(BiSeNet, self).__init__()
  12. self.aux = aux
  13. self.spatial_path = SpatialPath(3, 128, **kwargs)
  14. self.context_path = ContextPath(backbone, pretrained_base, **kwargs)
  15. self.ffm = FeatureFusion(256, 256, 4, **kwargs)
  16. self.head = _BiSeHead(256, 64, nclass, **kwargs)
  17. if aux:
  18. self.auxlayer1 = _BiSeHead(128, 256, nclass, **kwargs)
  19. self.auxlayer2 = _BiSeHead(128, 256, nclass, **kwargs)
  20. self.__setattr__('exclusive',
  21. ['spatial_path', 'context_path', 'ffm', 'head', 'auxlayer1', 'auxlayer2'] if aux else [
  22. 'spatial_path', 'context_path', 'ffm', 'head'])
  23. def forward(self, x,outsize=None,test_flag=False):
  24. size = x.size()[2:]
  25. spatial_out = self.spatial_path(x)
  26. context_out = self.context_path(x)
  27. fusion_out = self.ffm(spatial_out, context_out[-1])
  28. outputs = []
  29. x = self.head(fusion_out)
  30. x = F.interpolate(x, size, mode='bilinear', align_corners=True)
  31. if outsize:
  32. print('######using torch resize#######',outsize)
  33. x = F.interpolate(x, outsize, mode='bilinear', align_corners=True)
  34. outputs.append(x)
  35. if self.aux:
  36. auxout1 = self.auxlayer1(context_out[0])
  37. auxout1 = F.interpolate(auxout1, size, mode='bilinear', align_corners=True)
  38. outputs.append(auxout1)
  39. auxout2 = self.auxlayer2(context_out[1])
  40. auxout2 = F.interpolate(auxout2, size, mode='bilinear', align_corners=True)
  41. outputs.append(auxout2)
  42. if test_flag:
  43. outputs = [torch.argmax(outputx ,axis=1) for outputx in outputs]
  44. #return tuple(outputs)
  45. return outputs[0]
  46. class BiSeNet_MultiOutput(nn.Module):
  47. def __init__(self, nclass, backbone='resnet18', aux=False, jpu=False, pretrained_base=True, **kwargs):
  48. super(BiSeNet_MultiOutput, self).__init__()
  49. self.aux = aux
  50. self.spatial_path = SpatialPath(3, 128, **kwargs)
  51. self.context_path = ContextPath(backbone, pretrained_base, **kwargs)
  52. self.ffm = FeatureFusion(256, 256, 4, **kwargs)
  53. assert isinstance(nclass,list)
  54. self.outCnt = len(nclass)
  55. for ii,nclassii in enumerate(nclass):
  56. setattr(self,'head%d'%(ii) , _BiSeHead(256, 64, nclassii, **kwargs))
  57. if aux:
  58. self.auxlayer1 = _BiSeHead(128, 256, nclass, **kwargs)
  59. self.auxlayer2 = _BiSeHead(128, 256, nclass, **kwargs)
  60. self.__setattr__('exclusive',
  61. ['spatial_path', 'context_path', 'ffm', 'head', 'auxlayer1', 'auxlayer2'] if aux else [
  62. 'spatial_path', 'context_path', 'ffm', 'head'])
  63. def forward(self, x,outsize=None,test_flag=False,smooth_kernel=0):
  64. size = x.size()[2:]
  65. spatial_out = self.spatial_path(x)
  66. context_out = self.context_path(x)
  67. fusion_out = self.ffm(spatial_out, context_out[-1])
  68. outputs = []
  69. for ii in range(self.outCnt):
  70. x = getattr(self,'head%d'%(ii))(fusion_out)
  71. x = F.interpolate(x, size, mode='bilinear', align_corners=True)
  72. outputs.append(x)
  73. if self.aux:
  74. auxout1 = self.auxlayer1(context_out[0])
  75. auxout1 = F.interpolate(auxout1, size, mode='bilinear', align_corners=True)
  76. outputs.append(auxout1)
  77. auxout2 = self.auxlayer2(context_out[1])
  78. auxout2 = F.interpolate(auxout2, size, mode='bilinear', align_corners=True)
  79. outputs.append(auxout2)
  80. if test_flag:
  81. outputs = [torch.argmax(outputx ,axis=1) for outputx in outputs]
  82. if smooth_kernel>0:
  83. gaussian_kernel = torch.from_numpy(np.ones((1,1,smooth_kernel,smooth_kernel)) )
  84. pad = int((smooth_kernel - 1)/2)
  85. if not gaussian_kernel.is_cuda:
  86. gaussian_kernel = gaussian_kernel.to(x.device)
  87. #print(gaussian_kernel.dtype,gaussian_kernel,outputs[0].dtype)
  88. outputs = [ x.unsqueeze(1).double() for x in outputs]
  89. outputs = [torch.conv2d(x, gaussian_kernel, padding=pad) for x in outputs ]
  90. outputs = [ x.squeeze(1).long() for x in outputs]
  91. #return tuple(outputs)
  92. return outputs
  93. class _BiSeHead(nn.Module):
  94. def __init__(self, in_channels, inter_channels, nclass, norm_layer=nn.BatchNorm2d, **kwargs):
  95. super(_BiSeHead, self).__init__()
  96. self.block = nn.Sequential(
  97. _ConvBNReLU(in_channels, inter_channels, 3, 1, 1, norm_layer=norm_layer),
  98. nn.Dropout(0.1),
  99. nn.Conv2d(inter_channels, nclass, 1)
  100. )
  101. def forward(self, x):
  102. x = self.block(x)
  103. return x
  104. class SpatialPath(nn.Module):
  105. """Spatial path"""
  106. def __init__(self, in_channels, out_channels, norm_layer=nn.BatchNorm2d, **kwargs):
  107. super(SpatialPath, self).__init__()
  108. inter_channels = 64
  109. self.conv7x7 = _ConvBNReLU(in_channels, inter_channels, 7, 2, 3, norm_layer=norm_layer)
  110. self.conv3x3_1 = _ConvBNReLU(inter_channels, inter_channels, 3, 2, 1, norm_layer=norm_layer)
  111. self.conv3x3_2 = _ConvBNReLU(inter_channels, inter_channels, 3, 2, 1, norm_layer=norm_layer)
  112. self.conv1x1 = _ConvBNReLU(inter_channels, out_channels, 1, 1, 0, norm_layer=norm_layer)
  113. def forward(self, x):
  114. x = self.conv7x7(x)
  115. x = self.conv3x3_1(x)
  116. x = self.conv3x3_2(x)
  117. x = self.conv1x1(x)
  118. return x
  119. class _GlobalAvgPooling(nn.Module):
  120. def __init__(self, in_channels, out_channels, norm_layer, **kwargs):
  121. super(_GlobalAvgPooling, self).__init__()
  122. self.gap = nn.Sequential(
  123. nn.AdaptiveAvgPool2d(1),
  124. nn.Conv2d(in_channels, out_channels, 1, bias=False),
  125. norm_layer(out_channels),
  126. nn.ReLU(True)
  127. )
  128. def forward(self, x):
  129. size = x.size()[2:]
  130. pool = self.gap(x)
  131. out = F.interpolate(pool, size, mode='bilinear', align_corners=True)
  132. return out
  133. class AttentionRefinmentModule(nn.Module):
  134. def __init__(self, in_channels, out_channels, norm_layer=nn.BatchNorm2d, **kwargs):
  135. super(AttentionRefinmentModule, self).__init__()
  136. self.conv3x3 = _ConvBNReLU(in_channels, out_channels, 3, 1, 1, norm_layer=norm_layer)
  137. self.channel_attention = nn.Sequential(
  138. nn.AdaptiveAvgPool2d(1),
  139. _ConvBNReLU(out_channels, out_channels, 1, 1, 0, norm_layer=norm_layer),
  140. nn.Sigmoid()
  141. )
  142. def forward(self, x):
  143. x = self.conv3x3(x)
  144. attention = self.channel_attention(x)
  145. x = x * attention
  146. return x
  147. class ContextPath(nn.Module):
  148. def __init__(self, backbone='resnet18', pretrained_base=True, norm_layer=nn.BatchNorm2d, **kwargs):
  149. super(ContextPath, self).__init__()
  150. if backbone == 'resnet18':
  151. pretrained = resnet18(pretrained=pretrained_base, **kwargs)
  152. elif backbone=='resnet50':
  153. pretrained = resnet50(pretrained=pretrained_base, **kwargs)
  154. else:
  155. raise RuntimeError('unknown backbone: {}'.format(backbone))
  156. self.conv1 = pretrained.conv1
  157. self.bn1 = pretrained.bn1
  158. self.relu = pretrained.relu
  159. self.maxpool = pretrained.maxpool
  160. self.layer1 = pretrained.layer1
  161. self.layer2 = pretrained.layer2
  162. self.layer3 = pretrained.layer3
  163. self.layer4 = pretrained.layer4
  164. inter_channels = 128
  165. self.global_context = _GlobalAvgPooling(512, inter_channels, norm_layer)
  166. self.arms = nn.ModuleList(
  167. [AttentionRefinmentModule(512, inter_channels, norm_layer, **kwargs),
  168. AttentionRefinmentModule(256, inter_channels, norm_layer, **kwargs)]
  169. )
  170. self.refines = nn.ModuleList(
  171. [_ConvBNReLU(inter_channels, inter_channels, 3, 1, 1, norm_layer=norm_layer),
  172. _ConvBNReLU(inter_channels, inter_channels, 3, 1, 1, norm_layer=norm_layer)]
  173. )
  174. def forward(self, x):
  175. x = self.conv1(x)
  176. x = self.bn1(x)
  177. x = self.relu(x)
  178. x = self.maxpool(x)
  179. x = self.layer1(x)
  180. context_blocks = []
  181. context_blocks.append(x)
  182. x = self.layer2(x)
  183. context_blocks.append(x)
  184. c3 = self.layer3(x)
  185. context_blocks.append(c3)
  186. c4 = self.layer4(c3)
  187. context_blocks.append(c4)
  188. context_blocks.reverse()
  189. global_context = self.global_context(c4)
  190. last_feature = global_context
  191. context_outputs = []
  192. for i, (feature, arm, refine) in enumerate(zip(context_blocks[:2], self.arms, self.refines)):
  193. feature = arm(feature)
  194. feature += last_feature
  195. last_feature = F.interpolate(feature, size=context_blocks[i + 1].size()[2:],
  196. mode='bilinear', align_corners=True)
  197. last_feature = refine(last_feature)
  198. context_outputs.append(last_feature)
  199. return context_outputs
  200. class FeatureFusion(nn.Module):
  201. def __init__(self, in_channels, out_channels, reduction=1, norm_layer=nn.BatchNorm2d, **kwargs):
  202. super(FeatureFusion, self).__init__()
  203. self.conv1x1 = _ConvBNReLU(in_channels, out_channels, 1, 1, 0, norm_layer=norm_layer, **kwargs)
  204. self.channel_attention = nn.Sequential(
  205. nn.AdaptiveAvgPool2d(1),
  206. _ConvBNReLU(out_channels, out_channels // reduction, 1, 1, 0, norm_layer=norm_layer),
  207. _ConvBNReLU(out_channels // reduction, out_channels, 1, 1, 0, norm_layer=norm_layer),
  208. nn.Sigmoid()
  209. )
  210. def forward(self, x1, x2):
  211. fusion = torch.cat([x1, x2], dim=1)
  212. out = self.conv1x1(fusion)
  213. attention = self.channel_attention(out)
  214. out = out + out * attention
  215. return out
  216. def get_bisenet(dataset='citys', backbone='resnet18', pretrained=False, root='~/.torch/models',
  217. pretrained_base=True, **kwargs):
  218. acronyms = {
  219. 'pascal_voc': 'pascal_voc',
  220. 'pascal_aug': 'pascal_aug',
  221. 'ade20k': 'ade',
  222. 'coco': 'coco',
  223. 'citys': 'citys',
  224. }
  225. from ..data.dataloader import datasets
  226. model = BiSeNet(datasets[dataset].NUM_CLASS, backbone=backbone, pretrained_base=pretrained_base, **kwargs)
  227. if pretrained:
  228. from .model_store import get_model_file
  229. device = torch.device(kwargs['local_rank'])
  230. model.load_state_dict(torch.load(get_model_file('bisenet_%s_%s' % (backbone, acronyms[dataset]), root=root),
  231. map_location=device))
  232. return model
  233. def get_bisenet_resnet18_citys(**kwargs):
  234. return get_bisenet('citys', 'resnet18', **kwargs)
  235. if __name__ == '__main__':
  236. # img = torch.randn(2, 3, 224, 224)
  237. # model = BiSeNet(19, backbone='resnet18')
  238. # print(model.exclusive)
  239. input = torch.rand(2, 3, 224, 224)
  240. model = BiSeNet(4, pretrained_base=True)
  241. # target = torch.zeros(4, 512, 512).cuda()
  242. # model.eval()
  243. # print(model)
  244. loss = model(input)
  245. print(loss, loss.shape)
  246. # from torchsummary import summary
  247. #
  248. # summary(model, (3, 224, 224)) # 打印表格,按顺序输出每层的输出形状和参数
  249. import torch
  250. from thop import profile
  251. from torchsummary import summary
  252. flop, params = profile(model, input_size=(1, 3, 512, 512))
  253. print('flops:{:.3f}G\nparams:{:.3f}M'.format(flop / 1e9, params / 1e6))