| from torch import nn | |
| from torch import nn | |
| import torch.nn.functional as F | |
| import torch | |
| # from sync_batchnorm import SynchronizedBatchNorm2d as BatchNorm2d | |
| # class ResBlock2d(nn.Module): | |
| # def __init__(self, in_features, kernel_size, padding): | |
| # super(ResBlock2d, self).__init__() | |
| # self.conv1 = nn.Conv2d(in_channels=in_features, out_channels=in_features, kernel_size=kernel_size, | |
| # padding=padding) | |
| # self.conv2 = nn.Conv2d(in_channels=in_features, out_channels=in_features, kernel_size=kernel_size, | |
| # padding=padding) | |
| # self.norm1 = BatchNorm2d(in_features) | |
| # self.norm2 = BatchNorm2d(in_features) | |
| # self.relu = nn.ReLU() | |
| # def forward(self, x): | |
| # out = self.norm1(x) | |
| # out = self.relu(out) | |
| # out = self.conv1(out) | |
| # out = self.norm2(out) | |
| # out = self.relu(out) | |
| # out = self.conv2(out) | |
| # out += x | |
| # return out | |
| class UpBlock2d(nn.Module): | |
| def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1): | |
| super(UpBlock2d, self).__init__() | |
| self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, | |
| padding=padding, groups=groups) | |
| # self.norm = BatchNorm2d(out_features) | |
| self.relu = nn.ReLU() | |
| def forward(self, x): | |
| out = x | |
| # out = F.interpolate(x, scale_factor=2) | |
| out = self.conv(out) | |
| # out = self.norm(out) | |
| out = F.relu(out) | |
| return out | |
| class DownBlock2d(nn.Module): | |
| def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1): | |
| super(DownBlock2d, self).__init__() | |
| self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, | |
| padding=padding, groups=groups) | |
| # self.norm = BatchNorm2d(out_features) | |
| # self.pool = nn.AvgPool2d(kernel_size=(2, 2)) | |
| self.relu = nn.ReLU() | |
| def forward(self, x): | |
| out = self.conv(x) | |
| # out = self.norm(out) | |
| out = self.relu(out) | |
| # out = self.pool(out) | |
| return out | |
| class SameBlock2d(nn.Module): | |
| def __init__(self, in_features, out_features, groups=1, kernel_size=3, padding=1): | |
| super(SameBlock2d, self).__init__() | |
| self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features, | |
| kernel_size=kernel_size, padding=padding, groups=groups) | |
| # self.norm = BatchNorm2d(out_features) | |
| self.relu = nn.ReLU() | |
| def forward(self, x): | |
| out = self.conv(x) | |
| # out = self.norm(out) | |
| out = self.relu(out) | |
| return out | |
| class HourglassEncoder(nn.Module): | |
| def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256): | |
| super(HourglassEncoder, self).__init__() | |
| down_blocks = [] | |
| for i in range(num_blocks): | |
| down_blocks.append(DownBlock2d(in_features if i == 0 else min(max_features, block_expansion * (2 ** i)), | |
| min(max_features, block_expansion * (2 ** (i + 1))), | |
| kernel_size=3, padding=1)) | |
| self.down_blocks = nn.ModuleList(down_blocks) | |
| def forward(self, x): | |
| outs = [x] | |
| for down_block in self.down_blocks: | |
| outs.append(down_block(outs[-1])) | |
| outs = outs[1:] | |
| return outs | |
| class HourglassDecoder(nn.Module): | |
| def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256): | |
| super(HourglassDecoder, self).__init__() | |
| up_blocks = [] | |
| for i in range(num_blocks)[::-1]: | |
| in_filters = (1 if i == num_blocks - 1 else 2) * min(max_features, block_expansion * (2 ** (i + 1))) | |
| out_filters = min(max_features, block_expansion * (2 ** i)) | |
| up_blocks.append(UpBlock2d(in_filters, out_filters, kernel_size=3, padding=1)) | |
| self.up_blocks = nn.ModuleList(up_blocks) | |
| self.out_filters = block_expansion | |
| def forward(self, x): | |
| new_out = None | |
| for up_block in self.up_blocks: | |
| out = x.pop() | |
| if new_out is not None: | |
| out = torch.cat([out, new_out], dim=1) | |
| new_out = up_block(out) | |
| return new_out | |
| class Hourglass(nn.Module): | |
| def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256): | |
| super(Hourglass, self).__init__() | |
| self.encoder = HourglassEncoder(block_expansion, in_features, num_blocks, max_features) | |
| self.decoder = HourglassDecoder(block_expansion, in_features, num_blocks, max_features) | |
| self.out_filters = self.decoder.out_filters | |
| def forward(self, x): | |
| return self.decoder(self.encoder(x)) | |
| # class AntiAliasInterpolation2d(nn.Module): | |
| # """ | |
| # Band-limited downsampling, for better preservation of the input signal. | |
| # """ | |
| # def __init__(self, channels, scale): | |
| # super(AntiAliasInterpolation2d, self).__init__() | |
| # sigma = (1 / scale - 1) / 2 | |
| # kernel_size = 2 * round(sigma * 4) + 1 | |
| # self.ka = kernel_size // 2 | |
| # self.kb = self.ka - 1 if kernel_size % 2 == 0 else self.ka | |
| # kernel_size = [kernel_size, kernel_size] | |
| # sigma = [sigma, sigma] | |
| # # The gaussian kernel is the product of the | |
| # # gaussian function of each dimension. | |
| # kernel = 1 | |
| # meshgrids = torch.meshgrid( | |
| # [ | |
| # torch.arange(size, dtype=torch.float32) | |
| # for size in kernel_size | |
| # ] | |
| # ) | |
| # for size, std, mgrid in zip(kernel_size, sigma, meshgrids): | |
| # mean = (size - 1) / 2 | |
| # kernel *= torch.exp(-(mgrid - mean) ** 2 / (2 * std ** 2)) | |
| # # Make sure sum of values in gaussian kernel equals 1. | |
| # kernel = kernel / torch.sum(kernel) | |
| # # Reshape to depthwise convolutional weight | |
| # kernel = kernel.view(1, 1, *kernel.size()) | |
| # kernel = kernel.repeat(channels, *[1] * (kernel.dim() - 1)) | |
| # self.register_buffer('weight', kernel) | |
| # self.groups = channels | |
| # self.scale = scale | |
| # def forward(self, input): | |
| # if self.scale == 1.0: | |
| # return input | |
| # out = F.pad(input, (self.ka, self.kb, self.ka, self.kb)) | |
| # out = F.conv2d(out, weight=self.weight, groups=self.groups) | |
| # out = F.interpolate(out, scale_factor=(self.scale, self.scale)) | |
| # return out | |
| # class Encoder(nn.Module): | |
| # def __init__(self, num_channels, num_down_blocks=3, block_expansion=64, max_features=512, | |
| # ): | |
| # super(Encoder, self).__init__() | |
| # self.in_conv = SameBlock2d(num_channels, block_expansion, kernel_size=(7, 7), padding=(3, 3)) | |
| # down_blocks = [] | |
| # for i in range(num_down_blocks): | |
| # in_features = min(max_features, block_expansion * (2 ** i)) | |
| # out_features = min(max_features, block_expansion * (2 ** (i + 1))) | |
| # down_blocks.append(DownBlock2d(in_features, out_features, kernel_size=(3, 3), padding=(1, 1))) | |
| # self.down_blocks = nn.Sequential(*down_blocks) | |
| # def forward(self, image): | |
| # out = self.in_conv(image) | |
| # out = self.down_blocks(out) | |
| # return out | |
| # class Bottleneck(nn.Module): | |
| # def __init__(self, num_bottleneck_blocks,num_down_blocks=3, block_expansion=64, max_features=512): | |
| # super(Bottleneck, self).__init__() | |
| # bottleneck = [] | |
| # in_features = min(max_features, block_expansion * (2 ** num_down_blocks)) | |
| # for i in range(num_bottleneck_blocks): | |
| # bottleneck.append(ResBlock2d(in_features, kernel_size=(3, 3), padding=(1, 1))) | |
| # self.bottleneck = nn.Sequential(*bottleneck) | |
| # def forward(self, feature_map): | |
| # out = self.bottleneck(feature_map) | |
| # return out | |
| class Decoder(nn.Module): | |
| def __init__(self,num_channels, num_down_blocks=3, block_expansion=64, max_features=512): | |
| super(Decoder, self).__init__() | |
| up_blocks = [] | |
| for i in range(num_down_blocks): | |
| in_features = min(max_features, block_expansion * (2 ** (num_down_blocks - i))) | |
| out_features = min(max_features, block_expansion * (2 ** (num_down_blocks - i - 1))) | |
| up_blocks.append(UpBlock2d(in_features, out_features, kernel_size=(3, 3), padding=(1, 1))) | |
| self.up_blocks = nn.Sequential(*up_blocks) | |
| self.out_conv = nn.Conv2d(block_expansion, num_channels, kernel_size=(7, 7), padding=(3, 3)) | |
| self.sigmoid = nn.Sigmoid() | |
| def forward(self, feature_map): | |
| out = self.up_blocks(feature_map) | |
| out = self.out_conv(out) | |
| out = self.sigmoid(out) | |
| return out | |
| # def warp_image(image, motion_flow): | |
| # _, h_old, w_old, _ = motion_flow.shape | |
| # _, _, h, w = image.shape | |
| # if h_old != h or w_old != w: | |
| # motion_flow = motion_flow.permute(0, 3, 1, 2) | |
| # motion_flow = F.interpolate(motion_flow, size=(h, w), mode='bilinear') | |
| # motion_flow = motion_flow.permute(0, 2, 3, 1) | |
| # return F.grid_sample(image, motion_flow) | |
| # def make_coordinate_grid(spatial_size, type): | |
| # h, w = spatial_size | |
| # x = torch.arange(w).type(type) | |
| # y = torch.arange(h).type(type) | |
| # x = (2 * (x / (w - 1)) - 1) | |
| # y = (2 * (y / (h - 1)) - 1) | |
| # yy = y.view(-1, 1).repeat(1, w) | |
| # xx = x.view(1, -1).repeat(h, 1) | |
| # meshed = torch.cat([xx.unsqueeze_(2), yy.unsqueeze_(2)], 2) | |
| # return meshed | |
| class ForegroundMatting(nn.Module): | |
| def __init__(self, num_channels, num_blocks=3, block_expansion=64, max_features=512): | |
| super(ForegroundMatting, self).__init__() | |
| # self.down_sample_image = AntiAliasInterpolation2d(num_channels, scale_factor) | |
| # self.down_sample_flow = AntiAliasInterpolation2d(2, scale_factor) | |
| self.hourglass = Hourglass( | |
| block_expansion=block_expansion, | |
| in_features=num_channels * 2 + 2, | |
| max_features=max_features, | |
| num_blocks=num_blocks | |
| ) | |
| # self.foreground_mask = nn.Conv2d(self.hourglass.out_filters, 1, kernel_size=(7, 7), padding=(3, 3)) | |
| self.matting_mask = nn.Conv2d(self.hourglass.out_filters, 1, kernel_size=(7, 7), padding=(3, 3)) | |
| self.matting = nn.Conv2d(self.hourglass.out_filters, num_channels, kernel_size=(7, 7), padding=(3, 3)) | |
| # self.scale_factor = scale_factor | |
| self.sigmoid = nn.Sigmoid() | |
| def forward(self, reference_image, dense_flow, warped_image): | |
| ''' | |
| source_image : b, c, h, w | |
| dense_tensor: b, 2, h, w | |
| warped_image: b, c, h, w | |
| ''' | |
| # res_out = {} | |
| # batch, _, h, w = reference_image.shape | |
| # warped_image = warp_image(reference_image, dense_flow)#warp the image with dense flow | |
| # res_out['warped_image'] = warped_image | |
| hourglass_input = torch.cat([reference_image, dense_flow, warped_image], dim=1) | |
| hourglass_out = self.hourglass(hourglass_input) | |
| # foreground_mask = self.foreground_mask(hourglass_out) # compute foreground mask | |
| # foreground_mask = self.sigmoid(foreground_mask).permute(0,2,3,1) | |
| # res_out['foreground_mask'] = foreground_mask | |
| # grid_flow = make_coordinate_grid((h, w), dense_flow.type()) | |
| # dense_flow_foreground = dense_flow * foreground_mask + (1-foreground_mask) * grid_flow.unsqueeze(0) ## revise the dense flow | |
| # res_out['dense_flow_foreground'] = dense_flow_foreground | |
| # res_out['dense_flow_foreground_vis'] = dense_flow * foreground_mask | |
| matting_mask = self.matting_mask(hourglass_out) # compute matting mask | |
| matting_mask = self.sigmoid(matting_mask) | |
| # res_out['matting_mask'] = matting_mask | |
| matting_image = self.matting(hourglass_out) # computing matting image | |
| # res_out['matting_image'] = matting_image | |
| out = warped_image * matting_mask + matting_image * (1 - matting_mask) | |
| return out, matting_mask | |
| if __name__ == '__main__': | |
| device = 'cuda' | |
| b, c, h, w = 2, 1280, 40, 40 | |
| m = ForegroundMatting(c).to(device) | |
| print(m) | |
| reference_image = torch.randn(b, c, h, w).to(device) | |
| dense_flow = torch.randn(b, 2, h, w).to(device) | |
| warped_image = torch.randn(b, c, h, w).to(device) | |
| o = m(reference_image, dense_flow, warped_image) |