leonelhs commited on
Commit
617d6e1
1 Parent(s): 67d263e

add composite feature

Browse files
Files changed (2) hide show
  1. app.py +5 -83
  2. requirements.txt +2 -0
app.py CHANGED
@@ -1,13 +1,9 @@
1
  # demo source from: https://github.com/MarcoForte/FBA_Matting
2
-
3
- import cv2
4
  import gradio as gr
5
- import numpy as np
6
- import torch
7
  from huggingface_hub import hf_hub_download
8
 
9
  from networks.models import build_model
10
- from networks.transforms import trimap_transform, normalise_image
11
 
12
  REPO_ID = "leonelhs/FBA-Matting"
13
 
@@ -16,83 +12,8 @@ model = build_model(weights)
16
  model.eval().cpu()
17
 
18
 
19
- def np_to_torch(x, permute=True):
20
- if permute:
21
- return torch.from_numpy(x).permute(2, 0, 1)[None, :, :, :].float().cpu()
22
- else:
23
- return torch.from_numpy(x)[None, :, :, :].float().cpu()
24
-
25
-
26
- def scale_input(x: np.ndarray, scale: float, scale_type) -> np.ndarray:
27
- ''' Scales inputs to multiple of 8. '''
28
- h, w = x.shape[:2]
29
- h1 = int(np.ceil(scale * h / 8) * 8)
30
- w1 = int(np.ceil(scale * w / 8) * 8)
31
- x_scale = cv2.resize(x, (w1, h1), interpolation=scale_type)
32
- return x_scale
33
-
34
-
35
- def inference(image_np: np.ndarray, trimap_np: np.ndarray) -> [np.ndarray]:
36
- ''' Predict alpha, foreground and background.
37
- Parameters:
38
- image_np -- the image in rgb format between 0 and 1. Dimensions: (h, w, 3)
39
- trimap_np -- two channel trimap, first background then foreground. Dimensions: (h, w, 2)
40
- Returns:
41
- fg: foreground image in rgb format between 0 and 1. Dimensions: (h, w, 3)
42
- bg: background image in rgb format between 0 and 1. Dimensions: (h, w, 3)
43
- alpha: alpha matte image between 0 and 1. Dimensions: (h, w)
44
- '''
45
- h, w = trimap_np.shape[:2]
46
- image_scale_np = scale_input(image_np, 1.0, cv2.INTER_LANCZOS4)
47
- trimap_scale_np = scale_input(trimap_np, 1.0, cv2.INTER_LANCZOS4)
48
-
49
- with torch.no_grad():
50
- image_torch = np_to_torch(image_scale_np)
51
- trimap_torch = np_to_torch(trimap_scale_np)
52
-
53
- trimap_transformed_torch = np_to_torch(
54
- trimap_transform(trimap_scale_np), permute=False)
55
- image_transformed_torch = normalise_image(
56
- image_torch.clone())
57
-
58
- output = model(
59
- image_torch,
60
- trimap_torch,
61
- image_transformed_torch,
62
- trimap_transformed_torch)
63
- output = cv2.resize(
64
- output[0].cpu().numpy().transpose(
65
- (1, 2, 0)), (w, h), cv2.INTER_LANCZOS4)
66
-
67
- alpha = output[:, :, 0]
68
- fg = output[:, :, 1:4]
69
- bg = output[:, :, 4:7]
70
-
71
- alpha[trimap_np[:, :, 0] == 1] = 0
72
- alpha[trimap_np[:, :, 1] == 1] = 1
73
- fg[alpha == 1] = image_np[alpha == 1]
74
- bg[alpha == 0] = image_np[alpha == 0]
75
-
76
- return fg, bg, alpha
77
-
78
-
79
- def read_image(name):
80
- return (cv2.imread(name) / 255.0)[:, :, ::-1]
81
-
82
-
83
- def read_trimap(name):
84
- trimap_im = cv2.imread(name, 0) / 255.0
85
- h, w = trimap_im.shape
86
- trimap_np = np.zeros((h, w, 2))
87
- trimap_np[trimap_im == 1, 1] = 1
88
- trimap_np[trimap_im == 0, 0] = 1
89
- return trimap_np
90
-
91
-
92
  def predict(image, trimap):
93
- image_np = read_image(image)
94
- trimap_np = read_trimap(trimap)
95
- return inference(image_np, trimap_np)
96
 
97
 
98
  footer = r"""
@@ -115,8 +36,10 @@ with gr.Blocks(title="FBA Matting") as app:
115
  fg = gr.Image(type="numpy", label="Foreground")
116
  bg = gr.Image(type="numpy", label="Background")
117
  alpha = gr.Image(type="numpy", label="Alpha")
 
 
118
 
119
- run_btn.click(predict, [input_img, input_trimap], [fg, bg, alpha])
120
 
121
  with gr.Row():
122
  blobs = [[
@@ -129,4 +52,3 @@ with gr.Blocks(title="FBA Matting") as app:
129
  gr.HTML(footer)
130
 
131
  app.launch(share=False, debug=True, enable_queue=True, show_error=True)
132
-
 
1
  # demo source from: https://github.com/MarcoForte/FBA_Matting
 
 
2
  import gradio as gr
3
+ from FBA_Matting import inference
 
4
  from huggingface_hub import hf_hub_download
5
 
6
  from networks.models import build_model
 
7
 
8
  REPO_ID = "leonelhs/FBA-Matting"
9
 
 
12
  model.eval().cpu()
13
 
14
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
15
  def predict(image, trimap):
16
+ return inference(model, image, trimap)
 
 
17
 
18
 
19
  footer = r"""
 
36
  fg = gr.Image(type="numpy", label="Foreground")
37
  bg = gr.Image(type="numpy", label="Background")
38
  alpha = gr.Image(type="numpy", label="Alpha")
39
+ composite = gr.Image(type="numpy", label="Composite")
40
+ gr.ClearButton(components=[input_img, input_trimap, fg, bg, alpha, composite], variant="stop")
41
 
42
+ run_btn.click(predict, [input_img, input_trimap], [fg, bg, alpha, composite])
43
 
44
  with gr.Row():
45
  blobs = [[
 
52
  gr.HTML(footer)
53
 
54
  app.launch(share=False, debug=True, enable_queue=True, show_error=True)
 
requirements.txt CHANGED
@@ -1,3 +1,5 @@
1
  torch>=1.4.0
2
  numpy
3
  opencv-python
 
 
 
1
  torch>=1.4.0
2
  numpy
3
  opencv-python
4
+ FBA-Matting
5
+