Merge pull request #480 from PVSemk/fba-matting

FBA Matting

Author: vikasgupta-github

FBAMatting/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.

FBAMatting/README.md

@@ -0,0 +1,45 @@
+import os
+
+import cv2
+import numpy as np
+from torch.utils.data import Dataset
+
+
+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 os.path.isfile(os.path.join(self.img_dir, x))
+        ]
+
+    def __len__(self):
+        return len(self.img_names)
+
+    def __getitem__(self, idx):
+        img_name = self.img_names[idx]
+        trimap_name = img_name[:-3] + "png"
+
+        image = read_image(os.path.join(self.img_dir, img_name))
+        trimap = read_trimap(os.path.join(self.trimap_dir, trimap_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

FBAMatting/dataloader.py

@@ -0,0 +1,136 @@
+import argparse
+import os
+
+import cv2
+import numpy as np
+import torch
+from dataloader import PredDataset
+from networks.models import build_model
+from networks.transforms import (
+    groupnorm_normalise_image,
+    trimap_transform,
+)
+from tqdm import tqdm
+
+
+def np_to_torch(x):
+    return torch.from_numpy(x).permute(2, 0, 1)[None, :, :, :].float()
+
+
+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 swap_bg(image, alpha):
+    green_bg = np.zeros_like(image).astype(np.float32)
+    green_bg[:, :, 1] = 255
+
+    alpha = alpha[:, :, np.newaxis]
+    result = alpha * image.astype(np.float32) + (1 - alpha) * green_bg
+    result = np.clip(result, 0, 255).astype(np.uint8)
+
+    return result
+
+
+def predict_fba_folder(model, args):
+    save_dir = args.output_dir
+    os.makedirs(save_dir, exist_ok=True)
+
+    dataset_test = PredDataset(args.image_dir, args.trimap_dir)
+
+    gen = iter(dataset_test)
+    for item_dict in tqdm(gen):
+        image_np = item_dict["image"]
+        trimap_np = item_dict["trimap"]
+
+        fg, bg, alpha = pred(image_np, trimap_np, model, args)
+
+        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,
+        )
+
+        example_swap_bg = swap_bg(fg[:, :, ::-1] * 255, alpha)
+        cv2.imwrite(
+            os.path.join(save_dir, item_dict["name"][:-4] + "_swapped_bg.png"), example_swap_bg,
+        )
+
+
+def pred(image_np: np.ndarray, trimap_np: np.ndarray, model, args) -> 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).to(args.device)
+        trimap_torch = np_to_torch(trimap_scale_np).to(args.device)
+
+        trimap_transformed_torch = np_to_torch(trimap_transform(trimap_scale_np)).to(
+            args.device,
+        )
+        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="")
+    parser.add_argument("--device", default="cpu", help="Device for inference on")
+
+    args = parser.parse_args()
+    model = build_model(args).to(args.device)
+    model.eval()
+    predict_fba_folder(model, args)