Spaces:
Running
Running
| import gradio as gr | |
| from tools import Inference, Matting, log | |
| from omegaconf import OmegaConf | |
| import os | |
| import sys | |
| import numpy as np | |
| import torchvision.transforms.functional as tf | |
| from PIL import Image | |
| args = OmegaConf.load(os.path.join(f"./config/test.yaml")) | |
| global_comp = None | |
| global_mask = None | |
| log("Model loading") | |
| phnet = Inference(**args) | |
| stylematte = Matting(**args) | |
| log("Model loaded") | |
| def harmonize(comp, mask): | |
| log("Inference started") | |
| if comp is None or mask is None: | |
| log("Empty source") | |
| return np.zeros((16, 16, 3)) | |
| comp = comp.convert('RGB') | |
| mask = mask.convert('1') | |
| in_shape = comp.size[::-1] | |
| comp = tf.resize(comp, [args.image_size, args.image_size]) | |
| mask = tf.resize(mask, [args.image_size, args.image_size]) | |
| compt = tf.to_tensor(comp) | |
| maskt = tf.to_tensor(mask) | |
| res = phnet.harmonize(compt, maskt) | |
| res = tf.resize(res, in_shape) | |
| log("Inference finished") | |
| return np.uint8((res*255)[0].permute(1, 2, 0).numpy()) | |
| def extract_matte(img, back): | |
| mask, fg = stylematte.extract(img) | |
| fg_pil = Image.fromarray(np.uint8(fg)) | |
| composite = fg + (1 - mask[:, :, None]) * \ | |
| np.array(back.resize(mask.shape[::-1])) | |
| composite_pil = Image.fromarray(np.uint8(composite)) | |
| global_comp = composite_pil | |
| global_mask = mask | |
| return [composite_pil, mask, fg_pil] | |
| def css(height=3, scale=2): | |
| return f".output_image {{height: {height}rem !important; width: {scale}rem !important;}}" | |
| with gr.Blocks() as demo: | |
| gr.Markdown( | |
| """ | |
| # Welcome to portrait transfer demo app! | |
| Select source portrait image and new background. | |
| """) | |
| btn_compose = gr.Button(value="Compose") | |
| with gr.Row(): | |
| input_ui = gr.Image( | |
| type="numpy", label='Source image to extract foreground') | |
| back_ui = gr.Image(type="pil", label='The new background') | |
| gr.Examples( | |
| examples=[["./assets/comp.jpg", "./assets/back.jpg"]], | |
| inputs=[input_ui, back_ui], | |
| ) | |
| gr.Markdown( | |
| """ | |
| ## Resulting alpha matte and extracted foreground. | |
| """) | |
| with gr.Row(): | |
| matte_ui = gr.Image(type="pil", label='Alpha matte') | |
| fg_ui = gr.Image(type="pil", image_mode='RGBA', | |
| label='Extracted foreground') | |
| gr.Markdown( | |
| """ | |
| ## Click the button and compare the composite with the harmonized version. | |
| """) | |
| btn_harmonize = gr.Button(value="Harmonize composite") | |
| with gr.Row(): | |
| composite_ui = gr.Image(type="pil", label='Composite') | |
| harmonized_ui = gr.Image( | |
| type="pil", label='Harmonized composite', css=css(3, 3)) | |
| btn_compose.click(extract_matte, inputs=[input_ui, back_ui], outputs=[ | |
| composite_ui, matte_ui, fg_ui]) | |
| btn_harmonize.click(harmonize, inputs=[ | |
| composite_ui, matte_ui], outputs=[harmonized_ui]) | |
| log("Interface created") | |
| demo.launch(share=True) | |