-
Notifications
You must be signed in to change notification settings - Fork 3
/
ours_model.py
122 lines (95 loc) · 3.97 KB
/
ours_model.py
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
import torch
import math
import numbers
import numpy as np
from torch import nn
from torch.nn import functional as F
import ours_extractors
class PSPModule(nn.Module):
def __init__(self, features, out_features=1024, sizes=(1, 2, 3, 6)):
super().__init__()
self.stages = []
self.stages = nn.ModuleList([self._make_stage(features, size) for size in sizes])
self.bottleneck = nn.Conv2d(features * (len(sizes) + 1), out_features, kernel_size=1)
self.relu = nn.ReLU()
def _make_stage(self, features, size):
prior = nn.AdaptiveAvgPool2d(output_size=(size, size))
conv = nn.Conv2d(features, features, kernel_size=1, bias=False)
return nn.Sequential(prior, conv)
def forward(self, feats):
h, w = feats.size(2), feats.size(3)
priors = [F.interpolate(input=stage(feats), size=(h, w), mode='bilinear', align_corners=False) for stage in self.stages] + [feats]
bottle = self.bottleneck(torch.cat(priors, 1))
return self.relu(bottle)
class PSPUpsample(nn.Module):
def __init__(self, in_channels, out_channels):
super().__init__()
self.conv = nn.Sequential(
nn.Conv2d(in_channels, out_channels, 3, padding=1),
nn.BatchNorm2d(out_channels),
nn.PReLU()
)
def forward(self, x):
h, w = 2 * x.size(2), 2 * x.size(3)
p = F.interpolate(input=x, size=(h, w), mode='bilinear', align_corners=False)
return self.conv(p)
class PSPNet(nn.Module):
def __init__(self, n_classes=3, sizes=(1, 2, 3, 6), psp_size=2048, backend='resnet50',
pretrained=True):
super().__init__()
self.feats = getattr(ours_extractors, backend)(pretrained)
self.convfore = getattr(ours_extractors, 'resnet18half')(pretrained)
self.psp = PSPModule(psp_size, 1024, sizes)
self.drop_1 = nn.Dropout2d(p=0.3)
self.up_1 = PSPUpsample(1024, 256)
self.up_2 = PSPUpsample(256, 64)
self.up_3 = PSPUpsample(64, 64)
self.drop_2 = nn.Dropout2d(p=0.15)
self.final = nn.Sequential(
nn.Conv2d(64, n_classes, kernel_size=1),
)
def forward(self, x, training, device, gt_fore, use_gt_fore):
fore_sigmoid = self.convfore(x)
if use_gt_fore:
f = self.feats(x, gt_fore, training, device)
else:
f = self.feats(x, fore_sigmoid, training, device)
p = self.psp(f)
p = self.drop_1(p)
p = self.up_1(p)
p = self.drop_2(p)
p = self.up_2(p)
p = self.drop_2(p)
p = self.up_3(p)
p = self.drop_2(p)
return self.final(p), fore_sigmoid
# x_final = self.final(p)
# x_patch = ours_extractors.MaxpoolingAsInpainting(x_final, fore_sigmoid)
# x_final = x_patch * (F.interpolate(fore_sigmoid, size=x.size()[2:4]) > 0.5).float() +\
# x_final * (F.interpolate(fore_sigmoid, size=x.size()[2:4]) < 0.5).float()
# return x_final, fore_sigmoid
class PSPNetShareEarlyLayer(nn.Module):
def __init__(self, n_classes=3, sizes=(1, 2, 3, 6), psp_size=2048, backend='resnet50shareearlylayer',
pretrained=True):
super().__init__()
self.feats = getattr(ours_extractors, backend)(pretrained)
self.psp = PSPModule(psp_size, 1024, sizes)
self.drop_1 = nn.Dropout2d(p=0.3)
self.up_1 = PSPUpsample(1024, 256)
self.up_2 = PSPUpsample(256, 64)
self.up_3 = PSPUpsample(64, 64)
self.drop_2 = nn.Dropout2d(p=0.15)
self.final = nn.Sequential(
nn.Conv2d(64, n_classes, kernel_size=1),
)
def forward(self, x, training, device, gt_fore, use_gt_fore):
f, fore_sigmoid = self.feats(x, training, device)
p = self.psp(f)
p = self.drop_1(p)
p = self.up_1(p)
p = self.drop_2(p)
p = self.up_2(p)
p = self.drop_2(p)
p = self.up_3(p)
p = self.drop_2(p)
return self.final(p), fore_sigmoid