Spaces:
Running
Running
| 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() |