# The implementation here is modified based on F3Net, # originally Apache 2.0 License and publicly available at https://github.com/weijun88/F3Net import torch import torch.nn as nn import torch.nn.functional as F class Bottleneck(nn.Module): def __init__(self, inplanes, planes, stride=1, downsample=None, dilation=1): super(Bottleneck, self).__init__() self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) self.bn1 = nn.BatchNorm2d(planes) self.conv2 = nn.Conv2d( planes, planes, kernel_size=3, stride=stride, padding=(3 * dilation - 1) // 2, bias=False, dilation=dilation) self.bn2 = nn.BatchNorm2d(planes) self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False) self.bn3 = nn.BatchNorm2d(planes * 4) self.downsample = downsample def forward(self, x): out = F.relu(self.bn1(self.conv1(x)), inplace=True) out = F.relu(self.bn2(self.conv2(out)), inplace=True) out = self.bn3(self.conv3(out)) if self.downsample is not None: x = self.downsample(x) return F.relu(out + x, inplace=True) class ResNet(nn.Module): def __init__(self): super(ResNet, self).__init__() self.inplanes = 64 self.conv1 = nn.Conv2d( 3, 64, kernel_size=7, stride=2, padding=3, bias=False) self.bn1 = nn.BatchNorm2d(64) self.layer1 = self.make_layer(64, 3, stride=1, dilation=1) self.layer2 = self.make_layer(128, 4, stride=2, dilation=1) self.layer3 = self.make_layer(256, 6, stride=2, dilation=1) self.layer4 = self.make_layer(512, 3, stride=2, dilation=1) def make_layer(self, planes, blocks, stride, dilation): downsample = nn.Sequential( nn.Conv2d( self.inplanes, planes * 4, kernel_size=1, stride=stride, bias=False), nn.BatchNorm2d(planes * 4)) layers = [ Bottleneck( self.inplanes, planes, stride, downsample, dilation=dilation) ] self.inplanes = planes * 4 for _ in range(1, blocks): layers.append(Bottleneck(self.inplanes, planes, dilation=dilation)) return nn.Sequential(*layers) def forward(self, x): x = x.reshape(1, 3, 448, 448) out1 = F.relu(self.bn1(self.conv1(x)), inplace=True) out1 = F.max_pool2d(out1, kernel_size=3, stride=2, padding=1) out2 = self.layer1(out1) out3 = self.layer2(out2) out4 = self.layer3(out3) out5 = self.layer4(out4) return out2, out3, out4, out5 class CFM(nn.Module): def __init__(self): super(CFM, self).__init__() self.conv1h = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1) self.bn1h = nn.BatchNorm2d(64) self.conv2h = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1) self.bn2h = nn.BatchNorm2d(64) self.conv3h = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1) self.bn3h = nn.BatchNorm2d(64) self.conv4h = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1) self.bn4h = nn.BatchNorm2d(64) self.conv1v = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1) self.bn1v = nn.BatchNorm2d(64) self.conv2v = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1) self.bn2v = nn.BatchNorm2d(64) self.conv3v = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1) self.bn3v = nn.BatchNorm2d(64) self.conv4v = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1) self.bn4v = nn.BatchNorm2d(64) def forward(self, left, down): if down.size()[2:] != left.size()[2:]: down = F.interpolate(down, size=left.size()[2:], mode='bilinear') out1h = F.relu(self.bn1h(self.conv1h(left)), inplace=True) out2h = F.relu(self.bn2h(self.conv2h(out1h)), inplace=True) out1v = F.relu(self.bn1v(self.conv1v(down)), inplace=True) out2v = F.relu(self.bn2v(self.conv2v(out1v)), inplace=True) fuse = out2h * out2v out3h = F.relu(self.bn3h(self.conv3h(fuse)), inplace=True) + out1h out4h = F.relu(self.bn4h(self.conv4h(out3h)), inplace=True) out3v = F.relu(self.bn3v(self.conv3v(fuse)), inplace=True) + out1v out4v = F.relu(self.bn4v(self.conv4v(out3v)), inplace=True) return out4h, out4v class Decoder(nn.Module): def __init__(self): super(Decoder, self).__init__() self.cfm45 = CFM() self.cfm34 = CFM() self.cfm23 = CFM() def forward(self, out2h, out3h, out4h, out5v, fback=None): if fback is not None: refine5 = F.interpolate( fback, size=out5v.size()[2:], mode='bilinear') refine4 = F.interpolate( fback, size=out4h.size()[2:], mode='bilinear') refine3 = F.interpolate( fback, size=out3h.size()[2:], mode='bilinear') refine2 = F.interpolate( fback, size=out2h.size()[2:], mode='bilinear') out5v = out5v + refine5 out4h, out4v = self.cfm45(out4h + refine4, out5v) out3h, out3v = self.cfm34(out3h + refine3, out4v) out2h, pred = self.cfm23(out2h + refine2, out3v) else: out4h, out4v = self.cfm45(out4h, out5v) out3h, out3v = self.cfm34(out3h, out4v) out2h, pred = self.cfm23(out2h, out3v) return out2h, out3h, out4h, out5v, pred class F3Net(nn.Module): def __init__(self): super(F3Net, self).__init__() self.bkbone = ResNet() self.squeeze5 = nn.Sequential( nn.Conv2d(2048, 64, 1), nn.BatchNorm2d(64), nn.ReLU(inplace=True)) self.squeeze4 = nn.Sequential( nn.Conv2d(1024, 64, 1), nn.BatchNorm2d(64), nn.ReLU(inplace=True)) self.squeeze3 = nn.Sequential( nn.Conv2d(512, 64, 1), nn.BatchNorm2d(64), nn.ReLU(inplace=True)) self.squeeze2 = nn.Sequential( nn.Conv2d(256, 64, 1), nn.BatchNorm2d(64), nn.ReLU(inplace=True)) self.decoder1 = Decoder() self.decoder2 = Decoder() self.linearp1 = nn.Conv2d(64, 1, kernel_size=3, stride=1, padding=1) self.linearp2 = nn.Conv2d(64, 1, kernel_size=3, stride=1, padding=1) self.linearr2 = nn.Conv2d(64, 1, kernel_size=3, stride=1, padding=1) self.linearr3 = nn.Conv2d(64, 1, kernel_size=3, stride=1, padding=1) self.linearr4 = nn.Conv2d(64, 1, kernel_size=3, stride=1, padding=1) self.linearr5 = nn.Conv2d(64, 1, kernel_size=3, stride=1, padding=1) def forward(self, x, shape=None): x = x.reshape(1, 3, 448, 448) out2h, out3h, out4h, out5v = self.bkbone(x) out2h, out3h, out4h, out5v = self.squeeze2(out2h), self.squeeze3( out3h), self.squeeze4(out4h), self.squeeze5(out5v) out2h, out3h, out4h, out5v, pred1 = self.decoder1( out2h, out3h, out4h, out5v) out2h, out3h, out4h, out5v, pred2 = self.decoder2( out2h, out3h, out4h, out5v, pred1) shape = x.size()[2:] if shape is None else shape pred1 = F.interpolate( self.linearp1(pred1), size=shape, mode='bilinear') pred2 = F.interpolate( self.linearp2(pred2), size=shape, mode='bilinear') out2h = F.interpolate( self.linearr2(out2h), size=shape, mode='bilinear') out3h = F.interpolate( self.linearr3(out3h), size=shape, mode='bilinear') out4h = F.interpolate( self.linearr4(out4h), size=shape, mode='bilinear') out5h = F.interpolate( self.linearr5(out5v), size=shape, mode='bilinear') return pred1, pred2, out2h, out3h, out4h, out5h