Spaces:
Running
Running
File size: 3,605 Bytes
4a3ab35 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 |
import torch
import torch.nn as nn
from nets.CSPdarknet53_tiny import darknet53_tiny
from nets.attention import cbam_block, eca_block, se_block, CA_Block
attention_block = [se_block, cbam_block, eca_block, CA_Block]
#-------------------------------------------------#
# 卷积块 -> 卷积 + 标准化 + 激活函数
# Conv2d + BatchNormalization + LeakyReLU
#-------------------------------------------------#
class BasicConv(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size, stride=1):
super(BasicConv, self).__init__()
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, kernel_size//2, bias=False)
self.bn = nn.BatchNorm2d(out_channels)
self.activation = nn.LeakyReLU(0.1)
def forward(self, x):
x = self.conv(x)
x = self.bn(x)
x = self.activation(x)
return x
#---------------------------------------------------#
# 卷积 + 上采样
#---------------------------------------------------#
class Upsample(nn.Module):
def __init__(self, in_channels, out_channels):
super(Upsample, self).__init__()
self.upsample = nn.Sequential(
BasicConv(in_channels, out_channels, 1),
nn.Upsample(scale_factor=2, mode='nearest')
)
def forward(self, x,):
x = self.upsample(x)
return x
#---------------------------------------------------#
# 最后获得yolov4的输出
#---------------------------------------------------#
def yolo_head(filters_list, in_filters):
m = nn.Sequential(
BasicConv(in_filters, filters_list[0], 3),
nn.Conv2d(filters_list[0], filters_list[1], 1),
)
return m
#---------------------------------------------------#
# yolo_body
#---------------------------------------------------#
class YoloBodytiny(nn.Module):
def __init__(self, anchors_mask, num_classes, phi=0, pretrained=False):
super(YoloBodytiny, self).__init__()
self.phi = phi
self.backbone = darknet53_tiny(pretrained)
self.conv_for_P5 = BasicConv(512,256,1)
self.yolo_headP5 = yolo_head([512, len(anchors_mask[0]) * (5 + num_classes)],256)
self.upsample = Upsample(256,128)
self.yolo_headP4 = yolo_head([256, len(anchors_mask[1]) * (5 + num_classes)],384)
if 1 <= self.phi and self.phi <= 4:
self.feat1_att = attention_block[self.phi - 1](256)
self.feat2_att = attention_block[self.phi - 1](512)
self.upsample_att = attention_block[self.phi - 1](128)
def forward(self, x):
#---------------------------------------------------#
# 生成CSPdarknet53_tiny的主干模型
# feat1的shape为26,26,256
# feat2的shape为13,13,512
#---------------------------------------------------#
feat1, feat2 = self.backbone(x)
if 1 <= self.phi and self.phi <= 4:
feat1 = self.feat1_att(feat1)
feat2 = self.feat2_att(feat2)
# 13,13,512 -> 13,13,256
P5 = self.conv_for_P5(feat2)
# 13,13,256 -> 13,13,512 -> 13,13,255
out0 = self.yolo_headP5(P5)
# 13,13,256 -> 13,13,128 -> 26,26,128
P5_Upsample = self.upsample(P5)
# 26,26,256 + 26,26,128 -> 26,26,384
if 1 <= self.phi and self.phi <= 4:
P5_Upsample = self.upsample_att(P5_Upsample)
P4 = torch.cat([P5_Upsample,feat1],axis=1)
# 26,26,384 -> 26,26,256 -> 26,26,255
out1 = self.yolo_headP4(P4)
return out0, out1
|