Create app.py

app.py

@@ -0,0 +1,112 @@
+from PIL import Image
+import torch
+import torch.nn.functional as F
+from torchvision.transforms.functional import normalize
+import numpy as np
+from briarmbg import BriaRMBG
+import gradio as gr
+import os
+
+
+def preprocess_image(im: np.ndarray, model_input_size: list) -> torch.Tensor:
+    if len(im.shape) < 3:
+        im = im[:, :, np.newaxis]
+    im_tensor = torch.tensor(im, dtype=torch.float32).permute(2, 0, 1)
+    im_tensor = F.interpolate(torch.unsqueeze(im_tensor, 0), size=model_input_size, mode='bilinear').type(torch.uint8)
+    image = torch.divide(im_tensor, 255.0)
+    image = normalize(image, [0.5, 0.5, 0.5], [1.0, 1.0, 1.0])
+    return image
+
+
+def postprocess_image(result: torch.Tensor, im_size: list) -> np.ndarray:
+    result = torch.squeeze(F.interpolate(result, size=im_size, mode='bilinear'), 0)
+    ma = torch.max(result)
+    mi = torch.min(result)
+    result = (result - mi) / (ma - mi)
+    im_array = (result * 255).permute(1, 2, 0).cpu().data.numpy().astype(np.uint8)
+    im_array = np.squeeze(im_array)
+    return im_array
+
+
+def example_inference(image):
+    orig_im = image.copy()
+    orig_image = image.copy()
+
+    model_path = "model.pth"
+
+    net = BriaRMBG()
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    net.load_state_dict(torch.load(model_path, map_location=device))
+    net.to(device)
+    net.eval()
+
+    # prepare input
+    model_input_size = [1024, 1024]
+    orig_im_size = orig_im.size
+    orig_im_size = (orig_im_size[1], orig_im_size[0])
+    orig_im = np.array(orig_im)
+    image = preprocess_image(orig_im, model_input_size).to(device)
+
+    # inference
+    result = net(image)
+
+    # post process
+    result_image = postprocess_image(result[0][0], orig_im_size)
+
+    # save result
+    pil_im = Image.fromarray(result_image)
+    no_bg_image = Image.new("RGBA", pil_im.size, (0, 0, 0, 0))
+    no_bg_image.paste(orig_image, mask=pil_im)
+
+    return no_bg_image
+
+
+original_image, binary_image = None, None
+paths = [os.path.join("bg_images", file) for file in os.listdir("bg_images")]
+images = [Image.open(path) for path in paths]
+
+with gr.Blocks(
+        theme=gr.themes.Default(primary_hue=gr.themes.colors.red, secondary_hue=gr.themes.colors.indigo)) as demo:
+    with gr.Row():
+        input_img = gr.Image(label="Input", interactive=True, type='pil')
+        hidden_img = gr.Image(label="Chosen Background", visible=False)
+        output_img = gr.Image(label="Output", interactive=False, type='pil')
+
+
+    def clearFunc():
+        global original_image
+        global binary_image
+
+
+    def update_visibility():
+        return gr.Image(visible=True)
+
+        torch.cuda.empty_cache()
+        gc.collect()
+        return gr.Image(visible=False, value=None)
+
+
+    with gr.Row():
+        examples = gr.Examples(examples=images, inputs=[hidden_img])
+
+    with gr.Row():
+        submit = gr.Button("Submit")
+        clear = gr.ClearButton(components=[input_img, output_img, hidden_img], value="Reset", variant="stop")
+
+
+    def generate_img(image, background):
+        orig_img = example_inference(image)
+        height, width = orig_img.size
+
+        background = Image.fromarray(background).resize((width, height))
+        orig_img = Image.fromarray(np.array(orig_img)).resize((width, height))
+        background.paste(orig_img, (0, 0), mask=orig_img)
+        return background
+
+
+    hidden_img.change(fn=update_visibility, inputs=[], outputs=[hidden_img])
+
+    submit.click(generate_img, inputs=[input_img, hidden_img], outputs=[output_img])
+    clear.click(fn=clearFunc, inputs=[], outputs=[hidden_img])
+
+demo.launch(share=True, debug=True)