Add files via upload

FBA_Matting/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2020 Marco Forte
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

FBA_Matting/dataloader.py

@@ -0,0 +1,39 @@
+from torch.utils.data import Dataset
+import numpy as np
+import cv2
+import os
+
+
+class PredDataset(Dataset):
+    ''' Reads image and trimap pairs from folder.
+
+    '''
+
+    def __init__(self, img_dir, trimap_dir):
+        self.img_dir, self.trimap_dir = img_dir, trimap_dir
+        self.img_names = [x for x in os.listdir(self.img_dir) if 'png' in x]
+
+    def __len__(self):
+        return len(self.img_names)
+
+    def __getitem__(self, idx):
+        img_name = self.img_names[idx]
+
+        image = read_image(os.path.join(self.img_dir, img_name))
+        trimap = read_trimap(os.path.join(self.trimap_dir, img_name))
+        pred_dict = {'image': image, 'trimap': trimap, 'name': img_name}
+
+        return pred_dict
+
+
+def read_image(name):
+    return (cv2.imread(name) / 255.0)[:, :, ::-1]
+
+
+def read_trimap(name):
+    trimap_im = cv2.imread(name, 0) / 255.0
+    h, w = trimap_im.shape
+    trimap = np.zeros((h, w, 2))
+    trimap[trimap_im == 1, 1] = 1
+    trimap[trimap_im == 0, 0] = 1
+    return trimap

FBA_Matting/demo.py

@@ -0,0 +1,99 @@
+# Our libs
+from networks.transforms import trimap_transform, groupnorm_normalise_image
+from networks.models import build_model
+from dataloader import PredDataset
+
+# System libs
+import os
+import argparse
+
+# External libs
+import cv2
+import numpy as np
+import torch
+
+CUDA_DEVICE = 'gpu' if torch.cuda.is_available() else 'cpu'
+def np_to_torch(x):
+    val = torch.from_numpy(x).permute(2, 0, 1)[None, :, :, :].float()
+    if CUDA_DEVICE == 'gpu':
+        val = val.cuda()
+    return val
+
+def scale_input(x: np.ndarray, scale: float, scale_type) -> np.ndarray:
+    ''' Scales inputs to multiple of 8. '''
+    h, w = x.shape[:2]
+    h1 = int(np.ceil(scale * h / 8) * 8)
+    w1 = int(np.ceil(scale * w / 8) * 8)
+    x_scale = cv2.resize(x, (w1, h1), interpolation=scale_type)
+    return x_scale
+
+
+def predict_fba_folder(model, args):
+    save_dir = args.output_dir
+
+    dataset_test = PredDataset(args.image_dir, args.trimap_dir)
+
+    gen = iter(dataset_test)
+    for item_dict in gen:
+        image_np = item_dict['image']
+        trimap_np = item_dict['trimap']
+
+        fg, bg, alpha = pred(image_np, trimap_np, model)
+
+        cv2.imwrite(os.path.join(save_dir, item_dict['name'][:-4] + '_fg.png'), fg[:, :, ::-1] * 255)
+        cv2.imwrite(os.path.join(save_dir, item_dict['name'][:-4] + '_bg.png'), bg[:, :, ::-1] * 255)
+        cv2.imwrite(os.path.join(save_dir, item_dict['name'][:-4] + '_alpha.png'), alpha * 255)
+
+
+def pred(image_np: np.ndarray, trimap_np: np.ndarray, model) -> np.ndarray:
+    ''' Predict alpha, foreground and background.
+        Parameters:
+        image_np -- the image in rgb format between 0 and 1. Dimensions: (h, w, 3)
+        trimap_np -- two channel trimap, first background then foreground. Dimensions: (h, w, 2)
+        Returns:
+        fg: foreground image in rgb format between 0 and 1. Dimensions: (h, w, 3)
+        bg: background image in rgb format between 0 and 1. Dimensions: (h, w, 3)
+        alpha: alpha matte image between 0 and 1. Dimensions: (h, w)
+    '''
+    h, w = trimap_np.shape[:2]
+
+    image_scale_np = scale_input(image_np, 1.0, cv2.INTER_LANCZOS4)
+    trimap_scale_np = scale_input(trimap_np, 1.0, cv2.INTER_LANCZOS4)
+
+    with torch.no_grad():
+
+        image_torch = np_to_torch(image_scale_np)
+        trimap_torch = np_to_torch(trimap_scale_np)
+
+        trimap_transformed_torch = np_to_torch(trimap_transform(trimap_scale_np))
+        image_transformed_torch = groupnorm_normalise_image(image_torch.clone(), format='nchw')
+
+        output = model(image_torch, trimap_torch, image_transformed_torch, trimap_transformed_torch)
+
+        output = cv2.resize(output[0].cpu().numpy().transpose((1, 2, 0)), (w, h), cv2.INTER_LANCZOS4)
+    alpha = output[:, :, 0]
+    fg = output[:, :, 1:4]
+    bg = output[:, :, 4:7]
+
+    alpha[trimap_np[:, :, 0] == 1] = 0
+    alpha[trimap_np[:, :, 1] == 1] = 1
+    fg[alpha == 1] = image_np[alpha == 1]
+    bg[alpha == 0] = image_np[alpha == 0]
+    return fg, bg, alpha
+
+
+if __name__ == '__main__':
+
+    parser = argparse.ArgumentParser()
+    # Model related arguments
+    parser.add_argument('--encoder', default='resnet50_GN_WS', help="encoder model")
+    parser.add_argument('--decoder', default='fba_decoder', help="Decoder model")
+    parser.add_argument('--weights', default='FBA.pth')
+    parser.add_argument('--image_dir', default='./examples/images', help="")
+    parser.add_argument('--trimap_dir', default='./examples/trimaps', help="")
+    parser.add_argument('--output_dir', default='./examples/predictions', help="")
+
+    args = parser.parse_args()
+    model = build_model(args)
+    model.eval()
+    predict_fba_folder(model, args)

FBA_Matting/networks/layers_WS.py

@@ -0,0 +1,27 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+
+class Conv2d(nn.Conv2d):
+
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
+                 padding=0, dilation=1, groups=1, bias=True):
+        super(Conv2d, self).__init__(in_channels, out_channels, kernel_size, stride,
+                                     padding, dilation, groups, bias)
+
+    def forward(self, x):
+        # return super(Conv2d, self).forward(x)
+        weight = self.weight
+        weight_mean = weight.mean(dim=1, keepdim=True).mean(dim=2,
+                                                            keepdim=True).mean(dim=3, keepdim=True)
+        weight = weight - weight_mean
+        # std = (weight).view(weight.size(0), -1).std(dim=1).view(-1, 1, 1, 1) + 1e-5
+        std = torch.sqrt(torch.var(weight.view(weight.size(0), -1), dim=1) + 1e-12).view(-1, 1, 1, 1) + 1e-5
+        weight = weight / std.expand_as(weight)
+        return F.conv2d(x, weight, self.bias, self.stride,
+                        self.padding, self.dilation, self.groups)
+
+
+def BatchNorm2d(num_features):
+    return nn.GroupNorm(num_channels=num_features, num_groups=32)