Spaces:
Building
Building
| # https://github.com/MarcoForte/FBA_Matting | |
| import cv2 | |
| import gradio as gr | |
| import numpy as np | |
| import torch | |
| 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" | |
| weights = hf_hub_download(repo_id=REPO_ID, filename="FBA.pth") | |
| 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) | |
| footer = r""" | |
| <center> | |
| <b> | |
| Demo for <a href='https://github.com/MarcoForte/FBA_Matting'>FBA Matting</a> | |
| </b> | |
| </center> | |
| """ | |
| with gr.Blocks(title="FBA Matting") as app: | |
| gr.HTML("<center><h1>FBA Matting</h1></center>") | |
| gr.HTML("<center><h3>Foreground, Background, Alpha Matting Generator.</h3></center>") | |
| with gr.Row().style(equal_height=False): | |
| with gr.Column(): | |
| input_img = gr.Image(type="filepath", label="Input image") | |
| input_trimap = gr.Image(type="filepath", label="Trimap image") | |
| run_btn = gr.Button(variant="primary") | |
| with gr.Column(): | |
| fg = gr.Image(type="numpy", label="Foreground") | |
| bg = gr.Image(type="numpy", label="Background") | |
| alpha = gr.Image(type="numpy", label="Alpha") | |
| run_btn.click(predict, [input_img, input_trimap], [fg, bg, alpha]) | |
| with gr.Row(): | |
| examples_data = [[f"examples/{x:02d}.jpg"] for x in range(1, 4)] | |
| examples = gr.Dataset(components=[input_img], samples=examples_data) | |
| examples.click(lambda x: x[0], [examples], [input_img]) | |
| with gr.Row(): | |
| gr.HTML(footer) | |
| app.launch(share=False, debug=True, enable_queue=True, show_error=True) | |