Spaces:
Runtime error
Runtime error
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from transformers import Mask2FormerForUniversalSegmentation | |
| from transformers.models.mask2former.configuration_mask2former import Mask2FormerConfig | |
| class StyleMatte(nn.Module): | |
| def __init__(self): | |
| super(StyleMatte, self).__init__() | |
| self.fpn = FPN_fuse(feature_channels=[256, 256, 256, 256], fpn_out=256) | |
| config = Mask2FormerConfig.from_json_file('./configs/stylematte_config.json') | |
| self.pixel_decoder = Mask2FormerForUniversalSegmentation(config).base_model.pixel_level_module | |
| self.fgf = FastGuidedFilter(eps=1e-4) | |
| self.conv = nn.Conv2d(256, 1, kernel_size=3, padding=1) | |
| def forward(self, image, normalize=False): | |
| decoder_out = self.pixel_decoder(image) | |
| decoder_states = list(decoder_out.decoder_hidden_states) | |
| decoder_states.append(decoder_out.decoder_last_hidden_state) | |
| out_pure = self.fpn(decoder_states) | |
| image_lr = nn.functional.interpolate(image.mean(1, keepdim=True), | |
| scale_factor=0.25, | |
| mode='bicubic', | |
| align_corners=True | |
| ) | |
| out = self.conv(out_pure) | |
| out = self.fgf(image_lr, out, image.mean(1, keepdim=True)) | |
| return torch.sigmoid(out) | |
| def get_training_params(self): | |
| return list(self.fpn.parameters())+list(self.conv.parameters()) | |
| def conv2d_relu(input_filters, output_filters, kernel_size=3, bias=True): | |
| return nn.Sequential( | |
| nn.Conv2d(input_filters, output_filters, | |
| kernel_size=kernel_size, padding=kernel_size//2, bias=bias), | |
| nn.LeakyReLU(0.2, inplace=True), | |
| nn.BatchNorm2d(output_filters) | |
| ) | |
| def up_and_add(x, y): | |
| return F.interpolate(x, size=(y.size(2), y.size(3)), mode='bilinear', align_corners=True) + y | |
| class FPN_fuse(nn.Module): | |
| def __init__(self, feature_channels=[256, 512, 1024, 2048], fpn_out=256): | |
| super(FPN_fuse, self).__init__() | |
| assert feature_channels[0] == fpn_out | |
| self.conv1x1 = nn.ModuleList([nn.Conv2d(ft_size, fpn_out, kernel_size=1) | |
| for ft_size in feature_channels[1:]]) | |
| self.smooth_conv = nn.ModuleList([nn.Conv2d(fpn_out, fpn_out, kernel_size=3, padding=1)] | |
| * (len(feature_channels)-1)) | |
| self.conv_fusion = nn.Sequential( | |
| nn.Conv2d(2*fpn_out, fpn_out, kernel_size=3, | |
| padding=1, bias=False), | |
| nn.BatchNorm2d(fpn_out), | |
| nn.ReLU(inplace=True), | |
| ) | |
| def forward(self, features): | |
| features[:-1] = [conv1x1(feature) for feature, | |
| conv1x1 in zip(features[:-1], self.conv1x1)] | |
| feature = up_and_add(self.smooth_conv[0](features[0]), features[1]) | |
| feature = up_and_add(self.smooth_conv[1](feature), features[2]) | |
| feature = up_and_add(self.smooth_conv[2](feature), features[3]) | |
| H, W = features[-1].size(2), features[-1].size(3) | |
| x = [feature, features[-1]] | |
| x = [F.interpolate(x_el, size=(H, W), mode='bilinear', | |
| align_corners=True) for x_el in x] | |
| x = self.conv_fusion(torch.cat(x, dim=1)) | |
| return x | |
| class PSPModule(nn.Module): | |
| # In the original inmplementation they use precise RoI pooling | |
| # Instead of using adaptative average pooling | |
| def __init__(self, in_channels, bin_sizes=[1, 2, 4, 6]): | |
| super(PSPModule, self).__init__() | |
| out_channels = in_channels // len(bin_sizes) | |
| self.stages = nn.ModuleList([self._make_stages(in_channels, out_channels, b_s) | |
| for b_s in bin_sizes]) | |
| self.bottleneck = nn.Sequential( | |
| nn.Conv2d(in_channels+(out_channels * len(bin_sizes)), in_channels, | |
| kernel_size=3, padding=1, bias=False), | |
| nn.BatchNorm2d(in_channels), | |
| nn.ReLU(inplace=True), | |
| nn.Dropout2d(0.1) | |
| ) | |
| def _make_stages(self, in_channels, out_channels, bin_sz): | |
| prior = nn.AdaptiveAvgPool2d(output_size=bin_sz) | |
| conv = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False) | |
| bn = nn.BatchNorm2d(out_channels) | |
| relu = nn.ReLU(inplace=True) | |
| return nn.Sequential(prior, conv, bn, relu) | |
| def forward(self, features): | |
| h, w = features.size()[2], features.size()[3] | |
| pyramids = [features] | |
| pyramids.extend([F.interpolate(stage(features), size=(h, w), mode='bilinear', | |
| align_corners=True) for stage in self.stages]) | |
| output = self.bottleneck(torch.cat(pyramids, dim=1)) | |
| return output | |
| class GuidedFilter(nn.Module): | |
| def __init__(self, r, eps=1e-8): | |
| super(GuidedFilter, self).__init__() | |
| self.r = r | |
| self.eps = eps | |
| self.boxfilter = BoxFilter(r) | |
| def forward(self, x, y): | |
| n_x, c_x, h_x, w_x = x.size() | |
| n_y, c_y, h_y, w_y = y.size() | |
| assert n_x == n_y | |
| assert c_x == 1 or c_x == c_y | |
| assert h_x == h_y and w_x == w_y | |
| assert h_x > 2 * self.r + 1 and w_x > 2 * self.r + 1 | |
| # N | |
| N = self.boxfilter((x.data.new().resize_((1, 1, h_x, w_x)).fill_(1.0))) | |
| # mean_x | |
| mean_x = self.boxfilter(x) / N | |
| # mean_y | |
| mean_y = self.boxfilter(y) / N | |
| # cov_xy | |
| cov_xy = self.boxfilter(x * y) / N - mean_x * mean_y | |
| # var_x | |
| var_x = self.boxfilter(x * x) / N - mean_x * mean_x | |
| # A | |
| A = cov_xy / (var_x + self.eps) | |
| # b | |
| b = mean_y - A * mean_x | |
| # mean_A; mean_b | |
| mean_A = self.boxfilter(A) / N | |
| mean_b = self.boxfilter(b) / N | |
| return mean_A * x + mean_b | |
| class FastGuidedFilter(nn.Module): | |
| def __init__(self, r=1, eps=1e-8): | |
| super(FastGuidedFilter, self).__init__() | |
| self.r = r | |
| self.eps = eps | |
| self.boxfilter = BoxFilter(r) | |
| def forward(self, lr_x, lr_y, hr_x): | |
| n_lrx, c_lrx, h_lrx, w_lrx = lr_x.size() | |
| n_lry, c_lry, h_lry, w_lry = lr_y.size() | |
| n_hrx, c_hrx, h_hrx, w_hrx = hr_x.size() | |
| assert n_lrx == n_lry and n_lry == n_hrx | |
| assert c_lrx == c_hrx and (c_lrx == 1 or c_lrx == c_lry) | |
| assert h_lrx == h_lry and w_lrx == w_lry | |
| assert h_lrx > 2*self.r+1 and w_lrx > 2*self.r+1 | |
| # N | |
| N = self.boxfilter(lr_x.new().resize_((1, 1, h_lrx, w_lrx)).fill_(1.0)) | |
| # mean_x | |
| mean_x = self.boxfilter(lr_x) / N | |
| # mean_y | |
| mean_y = self.boxfilter(lr_y) / N | |
| # cov_xy | |
| cov_xy = self.boxfilter(lr_x * lr_y) / N - mean_x * mean_y | |
| # var_x | |
| var_x = self.boxfilter(lr_x * lr_x) / N - mean_x * mean_x | |
| # A | |
| A = cov_xy / (var_x + self.eps) | |
| # b | |
| b = mean_y - A * mean_x | |
| # mean_A; mean_b | |
| mean_A = F.interpolate( | |
| A, (h_hrx, w_hrx), mode='bilinear', align_corners=True) | |
| mean_b = F.interpolate( | |
| b, (h_hrx, w_hrx), mode='bilinear', align_corners=True) | |
| return mean_A*hr_x+mean_b | |
| class DeepGuidedFilterRefiner(nn.Module): | |
| def __init__(self, hid_channels=16): | |
| super().__init__() | |
| self.box_filter = nn.Conv2d( | |
| 4, 4, kernel_size=3, padding=1, bias=False, groups=4) | |
| self.box_filter.weight.data[...] = 1 / 9 | |
| self.conv = nn.Sequential( | |
| nn.Conv2d(4 * 2 + hid_channels, hid_channels, | |
| kernel_size=1, bias=False), | |
| nn.BatchNorm2d(hid_channels), | |
| nn.ReLU(True), | |
| nn.Conv2d(hid_channels, hid_channels, kernel_size=1, bias=False), | |
| nn.BatchNorm2d(hid_channels), | |
| nn.ReLU(True), | |
| nn.Conv2d(hid_channels, 4, kernel_size=1, bias=True) | |
| ) | |
| def forward(self, fine_src, base_src, base_fgr, base_pha, base_hid): | |
| fine_x = torch.cat([fine_src, fine_src.mean(1, keepdim=True)], dim=1) | |
| base_x = torch.cat([base_src, base_src.mean(1, keepdim=True)], dim=1) | |
| base_y = torch.cat([base_fgr, base_pha], dim=1) | |
| mean_x = self.box_filter(base_x) | |
| mean_y = self.box_filter(base_y) | |
| cov_xy = self.box_filter(base_x * base_y) - mean_x * mean_y | |
| var_x = self.box_filter(base_x * base_x) - mean_x * mean_x | |
| A = self.conv(torch.cat([cov_xy, var_x, base_hid], dim=1)) | |
| b = mean_y - A * mean_x | |
| H, W = fine_src.shape[2:] | |
| A = F.interpolate(A, (H, W), mode='bilinear', align_corners=False) | |
| b = F.interpolate(b, (H, W), mode='bilinear', align_corners=False) | |
| out = A * fine_x + b | |
| fgr, pha = out.split([3, 1], dim=1) | |
| return fgr, pha | |
| def diff_x(input, r): | |
| assert input.dim() == 4 | |
| left = input[:, :, r:2 * r + 1] | |
| middle = input[:, :, 2 * r + 1:] - input[:, :, :-2 * r - 1] | |
| right = input[:, :, -1:] - input[:, :, -2 * r - 1: -r - 1] | |
| output = torch.cat([left, middle, right], dim=2) | |
| return output | |
| def diff_y(input, r): | |
| assert input.dim() == 4 | |
| left = input[:, :, :, r:2 * r + 1] | |
| middle = input[:, :, :, 2 * r + 1:] - input[:, :, :, :-2 * r - 1] | |
| right = input[:, :, :, -1:] - input[:, :, :, -2 * r - 1: -r - 1] | |
| output = torch.cat([left, middle, right], dim=3) | |
| return output | |
| class BoxFilter(nn.Module): | |
| def __init__(self, r): | |
| super(BoxFilter, self).__init__() | |
| self.r = r | |
| def forward(self, x): | |
| assert x.dim() == 4 | |
| return diff_y(diff_x(x.cumsum(dim=2), self.r).cumsum(dim=3), self.r) | |