SuObjectDetection/model.py at main · Core43/SuObjectDetection · GitHub

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
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
import torch
import torch.nn as nn
class Attention(nn.Module):
    def __init__(self, in_features, out_features,res=True):
        super(Attention, self).__init__()
        self.q = nn.Linear(in_features, out_features)
        self.k = nn.Linear(in_features, out_features)
        self.v = nn.Linear(in_features, out_features)
        self.res=res
        if in_features!=out_features:
            self.res=False
        self.softmax = nn.Softmax(dim=-1)
        self._norm_fact = 1 / (out_features ** 0.5)

    def forward(self, x):
        x_ = x.permute(0, 2, 3, 1)
        shape = x_.shape
        x_ = x_.view(shape[0], -1, shape[-1])
        Q = self.q(x_)
        K = self.k(x_)
        V = self.v(x_)

        atten = nn.Softmax(dim=-1)(
            torch.bmm(Q, K.permute(0, 2, 1))) * self._norm_fact

        output = torch.bmm(atten, V)

        output = output.view(*shape[:3], -1)
        output = output.permute(0, 3, 1, 2)
        if self.res:
            output=x+output
        return output


class DownSample(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(DownSample, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=2, padding=1)
        self.norm = nn.BatchNorm2d(out_channels)
        self.relu = nn.LeakyReLU()

    def forward(self, input):
        output = self.conv(input)
        output = self.norm(output)
        output = self.relu(output)
        return output


class CBR(nn.Module):
    def __init__(self, in_channels, out_channels, norm_f=True):
        super(CBR, self).__init__()
        self.norm_f = norm_f
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
        self.norm = nn.BatchNorm2d(out_channels)
        self.relu = nn.LeakyReLU()

    def forward(self, input):
        output = self.conv(input)
        output = self.norm(output)
        output = self.relu(output)
        return output


class PR(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(PR, self).__init__()
        self.d_ = DownSample(in_channels, out_channels)
        self.max_pool = nn.MaxPool2d(2, stride=2)
        self.avg_pool = nn.AvgPool2d(2, stride=2)
        self.cbr = CBR(out_channels + in_channels * 2, out_channels)

    def forward(self, input_):
        output = torch.cat([
            self.d_(input_),
            self.max_pool(input_),
            self.avg_pool(input_)
        ], dim=1)
        output = self.cbr(output)
        return output


class VRes(nn.Module):
    def __init__(self, in_channels, out_channels, mid_scale=12, down_sample=5):
        super(VRes, self).__init__()
        self.down_sample=down_sample
        self.blocks = nn.ModuleList()
        self.r_blocks = nn.ModuleList()
        self.head = nn.Sequential(
            nn.Conv2d(in_channels, mid_scale, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(mid_scale),
            nn.LeakyReLU()
        )
        self.end = nn.Sequential(
            nn.Conv2d(int(mid_scale * (1.5 ** down_sample))*2, int(mid_scale * (1.5 ** down_sample)), kernel_size=3,
                      stride=1, padding=1),
            nn.BatchNorm2d(int(mid_scale * (1.5 ** down_sample))),
            nn.LeakyReLU(),
            nn.Conv2d(int(mid_scale * (1.5 ** down_sample)), out_channels, kernel_size=3, stride=1, padding=1)
        )
        self.softmax = nn.Softmax(dim=1)
        self.sigmoid = nn.Sigmoid()
        for i in range(down_sample):
            self.r_blocks.append(PR(int(mid_scale * (1.5 ** i))*(1 if i==0 else 2), int(mid_scale * (1.5 ** (i+1)))))
            self.blocks.append(nn.Sequential(
                    CBR(int(mid_scale * (1.5 ** i))*(1 if i==0 else 2), int(mid_scale * (1.5 ** i))),
                    Attention(int(mid_scale * (1.5 ** i)), int(mid_scale * (1.5 ** i))) if i >= (
                            down_sample - 2) else nn.Identity(),
                    CBR(int(mid_scale * (1.5 ** i)), int(mid_scale * (1.5 ** i))),
                    Attention(int(mid_scale * (1.5 ** i)), int(mid_scale * (1.5 ** i))) if i == (
                            down_sample - 1) else nn.Identity(),
                    CBR(int(mid_scale * (1.5 ** i)), int(mid_scale * (1.5 ** i))),
                    Attention(int(mid_scale * (1.5 ** i)), int(mid_scale * (1.5 ** i))) if i >= (
                            down_sample - 2) else nn.Identity(),
                    DownSample(int(mid_scale * (1.5 ** i)), int(mid_scale * (1.5 ** (i + 1))))
            ))
    def forward(self,input):
        output=self.head(input)
        for i in range(self.down_sample):
            output=torch.cat([self.r_blocks[i](output),self.blocks[i](output)],dim=1)
        output=self.end(output)
        o1=self.sigmoid(output[:,:5,:,:])
        o2=self.softmax(output[:,5:,:,:])
        output=torch.cat([o1,o2],dim=1)
        return output


if __name__ == '__main__':
    a = torch.rand(10, 3, 512, 512)
    model = VRes(in_channels=3,out_channels=5+4,mid_scale=28)
    print(sum([i.nelement() for i in model.parameters()]))
    print(model(a).shape)