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