add composite feature

app.py

@@ -1,13 +1,9 @@
 # demo source from: https://github.com/MarcoForte/FBA_Matting
-
-import cv2
 import gradio as gr
-import numpy as np
-import torch
+from FBA_Matting import inference
 from huggingface_hub import hf_hub_download
 
 from networks.models import build_model
-from networks.transforms import trimap_transform, normalise_image
 
 REPO_ID = "leonelhs/FBA-Matting"
 
@@ -16,83 +12,8 @@ model = build_model(weights)
 model.eval().cpu()
 
 
-def np_to_torch(x, permute=True):
-    if permute:
-        return torch.from_numpy(x).permute(2, 0, 1)[None, :, :, :].float().cpu()
-    else:
-        return torch.from_numpy(x)[None, :, :, :].float().cpu()
-
-
-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 inference(image_np: np.ndarray, trimap_np: np.ndarray) -> [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), permute=False)
-        image_transformed_torch = normalise_image(
-            image_torch.clone())
-
-        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
-
-
-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 = np.zeros((h, w, 2))
-    trimap_np[trimap_im == 1, 1] = 1
-    trimap_np[trimap_im == 0, 0] = 1
-    return trimap_np
-
-
 def predict(image, trimap):
-    image_np = read_image(image)
-    trimap_np = read_trimap(trimap)
-    return inference(image_np, trimap_np)
+    return inference(model, image, trimap)
 
 
 footer = r"""
@@ -115,8 +36,10 @@ with gr.Blocks(title="FBA Matting") as app:
             fg = gr.Image(type="numpy", label="Foreground")
             bg = gr.Image(type="numpy", label="Background")
             alpha = gr.Image(type="numpy", label="Alpha")
+            composite = gr.Image(type="numpy", label="Composite")
+            gr.ClearButton(components=[input_img, input_trimap, fg, bg, alpha, composite], variant="stop")
 
-    run_btn.click(predict, [input_img, input_trimap], [fg, bg, alpha])
+    run_btn.click(predict, [input_img, input_trimap], [fg, bg, alpha, composite])
 
     with gr.Row():
         blobs = [[
@@ -129,4 +52,3 @@ with gr.Blocks(title="FBA Matting") as app:
         gr.HTML(footer)
 
 app.launch(share=False, debug=True, enable_queue=True, show_error=True)
-

requirements.txt

@@ -1,3 +1,5 @@
 torch>=1.4.0
 numpy
 opencv-python
+FBA-Matting
+