import torch import torch.nn as nn from FBA_Matting.networks import ResNet import FBA_Matting.networks.layers_WS as L def build_model(weights): net_encoder = fba_encoder() net_decoder = fba_decoder() model = MattingModule(net_encoder, net_decoder) if weights != 'default': device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') sd = torch.load(weights, map_location=device) model.load_state_dict(sd, strict=True) return model class MattingModule(nn.Module): def __init__(self, net_enc, net_dec): super(MattingModule, self).__init__() self.encoder = net_enc self.decoder = net_dec def forward(self, image, two_chan_trimap, image_n, trimap_transformed): resnet_input = torch.cat((image_n, trimap_transformed, two_chan_trimap), 1) conv_out, indices = self.encoder(resnet_input, return_feature_maps=True) return self.decoder(conv_out, image, indices, two_chan_trimap) def fba_encoder(): orig_resnet = ResNet() net_encoder = ResnetDilated(orig_resnet, dilate_scale=8) num_channels = 3 + 6 + 2 print(f'modifying input layer to accept {num_channels} channels') net_encoder_sd = net_encoder.state_dict() conv1_weights = net_encoder_sd['conv1.weight'] c_out, c_in, h, w = conv1_weights.size() conv1_mod = torch.zeros(c_out, num_channels, h, w) conv1_mod[:, :3, :, :] = conv1_weights conv1 = net_encoder.conv1 conv1.in_channels = num_channels conv1.weight = torch.nn.Parameter(conv1_mod) net_encoder.conv1 = conv1 net_encoder_sd['conv1.weight'] = conv1_mod net_encoder.load_state_dict(net_encoder_sd) return net_encoder class ResnetDilated(nn.Module): def __init__(self, orig_resnet, dilate_scale=8): super(ResnetDilated, self).__init__() from functools import partial if dilate_scale == 8: orig_resnet.layer3.apply( partial(self._nostride_dilate, dilate=2)) orig_resnet.layer4.apply( partial(self._nostride_dilate, dilate=4)) elif dilate_scale == 16: orig_resnet.layer4.apply( partial(self._nostride_dilate, dilate=2)) # take pretrained resnet, except AvgPool and FC self.conv1 = orig_resnet.conv1 self.bn1 = orig_resnet.bn1 self.relu = orig_resnet.relu self.maxpool = orig_resnet.maxpool self.layer1 = orig_resnet.layer1 self.layer2 = orig_resnet.layer2 self.layer3 = orig_resnet.layer3 self.layer4 = orig_resnet.layer4 def _nostride_dilate(self, m, dilate): classname = m.__class__.__name__ if classname.find('Conv') != -1: # the convolution with stride if m.stride == (2, 2): m.stride = (1, 1) if m.kernel_size == (3, 3): m.dilation = (dilate // 2, dilate // 2) m.padding = (dilate // 2, dilate // 2) # other convoluions else: if m.kernel_size == (3, 3): m.dilation = (dilate, dilate) m.padding = (dilate, dilate) def forward(self, x, return_feature_maps=False): conv_out = [x] x = self.relu(self.bn1(self.conv1(x))) conv_out.append(x) x, indices = self.maxpool(x) x = self.layer1(x) conv_out.append(x) x = self.layer2(x) conv_out.append(x) x = self.layer3(x) conv_out.append(x) x = self.layer4(x) conv_out.append(x) if return_feature_maps: return conv_out, indices return [x] def fba_fusion(alpha, img, F, B): F = (alpha * img + (1 - alpha ** 2) * F - alpha * (1 - alpha) * B) B = ((1 - alpha) * img + (2 * alpha - alpha ** 2) * B - alpha * (1 - alpha) * F) F = torch.clamp(F, 0, 1) B = torch.clamp(B, 0, 1) la = 0.1 alpha = (alpha * la + torch.sum((img - B) * (F - B), 1, keepdim=True)) / ( torch.sum((F - B) * (F - B), 1, keepdim=True) + la) alpha = torch.clamp(alpha, 0, 1) return alpha, F, B class fba_decoder(nn.Module): def __init__(self): super(fba_decoder, self).__init__() pool_scales = (1, 2, 3, 6) self.ppm = [] for scale in pool_scales: self.ppm.append(nn.Sequential( nn.AdaptiveAvgPool2d(scale), L.Conv2d(2048, 256, kernel_size=1, bias=True), L.norm(256), nn.LeakyReLU() )) self.ppm = nn.ModuleList(self.ppm) self.conv_up1 = nn.Sequential( L.Conv2d(2048 + len(pool_scales) * 256, 256, kernel_size=3, padding=1, bias=True), L.norm(256), nn.LeakyReLU(), L.Conv2d(256, 256, kernel_size=3, padding=1), L.norm(256), nn.LeakyReLU() ) self.conv_up2 = nn.Sequential( L.Conv2d(256 + 256, 256, kernel_size=3, padding=1, bias=True), L.norm(256), nn.LeakyReLU() ) self.conv_up3 = nn.Sequential( L.Conv2d(256 + 64, 64, kernel_size=3, padding=1, bias=True), L.norm(64), nn.LeakyReLU() ) self.unpool = nn.MaxUnpool2d(2, stride=2) self.conv_up4 = nn.Sequential( nn.Conv2d(64 + 3 + 3 + 2, 32, kernel_size=3, padding=1, bias=True), nn.LeakyReLU(), nn.Conv2d(32, 16, kernel_size=3, padding=1, bias=True), nn.LeakyReLU(), nn.Conv2d(16, 7, kernel_size=1, padding=0, bias=True) ) def forward(self, conv_out, img, indices, two_chan_trimap): conv5 = conv_out[-1] input_size = conv5.size() ppm_out = [conv5] for pool_scale in self.ppm: ppm_out.append(nn.functional.interpolate( pool_scale(conv5), (input_size[2], input_size[3]), mode='bilinear', align_corners=False)) ppm_out = torch.cat(ppm_out, 1) x = self.conv_up1(ppm_out) x = torch.nn.functional.interpolate(x, scale_factor=2, mode='bilinear', align_corners=False) x = torch.cat((x, conv_out[-4]), 1) x = self.conv_up2(x) x = torch.nn.functional.interpolate(x, scale_factor=2, mode='bilinear', align_corners=False) x = torch.cat((x, conv_out[-5]), 1) x = self.conv_up3(x) x = torch.nn.functional.interpolate(x, scale_factor=2, mode='bilinear', align_corners=False) x = torch.cat((x, conv_out[-6][:, :3], img, two_chan_trimap), 1) output = self.conv_up4(x) alpha = torch.clamp(output[:, 0][:, None], 0, 1) F = torch.sigmoid(output[:, 1:4]) B = torch.sigmoid(output[:, 4:7]) # FBA Fusion alpha, F, B = fba_fusion(alpha, img, F, B) output = torch.cat((alpha, F, B), 1) return output