Spaces:
Running
Running
File size: 1,791 Bytes
d7713d2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 |
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")
|