交通事故检测代码
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.

323 lines
13KB

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