import gradio as gr
from maskformer import Mask2FormerSegmenter
from inpainter import Inpainter

import os
from tqdm import tqdm
from PIL import Image

def main(image1, image2, refine='False'):
       
    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 process_image(image1, image2, refine=False):
    img1 = Image.fromarray(image1.astype('uint8'), 'RGB')
    img2 = Image.fromarray(image2.astype('uint8'), 'RGB')
    return main(img1, img2, refine)


iface = gr.Interface(
    fn=process_image,
    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.",
)
iface.launch()