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