Spaces:
Sleeping
Sleeping
File size: 2,312 Bytes
d7713d2 e3214c0 d7713d2 e3214c0 d7713d2 e3214c0 d7713d2 e3214c0 d7713d2 e3214c0 d7713d2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 |
import gradio as gr
from maskformer import Mask2FormerSegmenter
from inpainter import Inpainter
import os
from tqdm import tqdm
from PIL import Image
def process_image(image1, image2, refine='False'):
""" Process a pair of images to swap their backgrounds.
image1: PIL image of the first image.
image2: PIL image of the second image.
refine: if True, use an additional network in the inpainter to refine the background.\
useful in high-resolution images and large inpainting masks. It takes longer to run."""
segmenter = Mask2FormerSegmenter()
segmenter.load_models(checkpoint_name = "facebook/mask2former-swin-large-ade-semantic")
inpainter = Inpainter({'scale_factor': None, 'pad_out_to_modulo': 8, 'predict': {'out_key': 'inpainted'}})
inpainter.load_model_from_checkpoint('big-lama', 'best.ckpt')
fg_img1, mask_img1 = segmenter.retrieve_fg_image_and_mask(image1, verbose=False)
new_bg_img1 = inpainter.inpaint_img(image1, mask_img1, refine=refine)
fg_img2, mask_img2 = segmenter.retrieve_fg_image_and_mask(image2, verbose=False)
new_bg_img2 = inpainter.inpaint_img(image2, mask_img2, refine=refine)
image_a = Image.alpha_composite(new_bg_img2.convert('RGBA'), fg_img1)
image_b = Image.alpha_composite(new_bg_img1.convert('RGBA'), fg_img2)
return image_a, image_b
def run_gradio_fn(image1, image2, refine=False):
""" Run the background swapper app on a pair of images."""
img1 = Image.fromarray(image1.astype('uint8'), 'RGB')
img2 = Image.fromarray(image2.astype('uint8'), 'RGB')
return process_image(img1, img2, refine)
iface = gr.Interface(
fn=run_gradio_fn,
inputs=["image", "image", gr.inputs.Checkbox(label="Use Refiner on background")],
outputs=["image", "image"],
title="Background Swapper App",
description="Upload two images to see their backgrounds swapped. \
The first image will have the background of the second image, and vice versa. \
Refiner option: Uses an additional network in the inpainter to refine the background.\
Useful in high-resolution images and large inpainting masks. \
It takes about 20s longer to run on 1024x1024",
)
iface.launch() |