from maskformer import Mask2FormerSegmenter
from inpainter import Inpainter

import os
from tqdm import tqdm
from PIL import Image

if __name__ == "__main__":
    dirpath = "ocus_images/final_images"
    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')    
    
    # List image files in the input directory
    image_files = [file for file in os.listdir(dirpath) if file.lower().endswith(('.jpg', '.jpeg', '.png'))]

    #for file in tqdm(image_files, desc="Processing images"):
    for i in tqdm(range(1, len(image_files), 2), desc="Processing image pairs"):
        filepath1 = os.path.join(dirpath, image_files[i-1])
        filepath2 = os.path.join(dirpath, image_files[i])
        image1 = Image.open(filepath1).convert('RGB')
        image2 = Image.open(filepath2).convert('RGB')
        
        fg_img1, mask_img1 = segmenter.retrieve_fg_image_and_mask(image1, verbose=False)
        new_bg_img1 = inpainter.inpaint_img(image1, mask_img1, refine=False)
        fg_img2, mask_img2 = segmenter.retrieve_fg_image_and_mask(image2, verbose=False)
        new_bg_img2 = inpainter.inpaint_img(image2, mask_img2, refine=False)
        
        image_a = Image.alpha_composite(new_bg_img2.convert('RGBA'), fg_img1)
        image_b = Image.alpha_composite(new_bg_img1.convert('RGBA'), fg_img2)
        image_a.save(f"results/joint/{os.path.basename(filepath1).split('.')[0]}_swapped.png")
        image_b.save(f"results/joint/{os.path.basename(filepath2).split('.')[0]}_swapped.png")