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()