import os import sys import cv2 import numpy as np import torch from PIL import Image from .utils import gen_new_name, prompts import torch from omegaconf import OmegaConf import numpy as np from .inpainting_src.ldm_inpainting.ldm.models.diffusion.ddim import DDIMSampler from .inpainting_src.ldm_inpainting.ldm.util import instantiate_from_config from .utils import cal_dilate_factor, dilate_mask def make_batch(image, mask, device): image = image.astype(np.float32) / 255.0 image = image[None].transpose(0, 3, 1, 2) image = torch.from_numpy(image) mask = mask.astype(np.float32) / 255.0 mask = mask[None, None] mask[mask < 0.5] = 0 mask[mask >= 0.5] = 1 mask = torch.from_numpy(mask) masked_image = (1 - mask) * image batch = {"image": image, "mask": mask, "masked_image": masked_image} for k in batch: batch[k] = batch[k].to(device=device) batch[k] = batch[k] * 2.0 - 1.0 return batch class LDMInpainting: def __init__(self, device): ckpt_path = 'model_zoo/ldm_inpainting_big.ckpt' config = './iGPT/models/inpainting_src/ldm_inpainting/config.yaml' self.ddim_steps = 50 self.device = device config = OmegaConf.load(config) model = instantiate_from_config(config.model) model.load_state_dict(torch.load(ckpt_path)["state_dict"], strict=False) self.model = model.to(device=device) self.sampler = DDIMSampler(model) @prompts(name="Remove the Masked Object", description="useful when you want to remove an object by masking the region in the image. " "like: remove masked object or inpaint the masked region.. " "The input to this tool should be a comma separated string of two, " "representing the image_path and mask_path") @torch.no_grad() def inference(self, inputs): print(f'inputs: {inputs}') # image, mask, device img_path, mask_path = inputs.split(',')[0], inputs.split(',')[1] img_path = img_path.strip() mask_path = mask_path.strip() image = Image.open(img_path) mask = Image.open(mask_path).convert('L') w, h = image.size image = image.resize((512, 512)) mask = mask.resize((512, 512)) image = np.array(image) mask = np.array(mask) dilate_factor = cal_dilate_factor(mask.astype(np.uint8)) mask = dilate_mask(mask, dilate_factor) with self.model.ema_scope(): batch = make_batch(image, mask, device=self.device) # encode masked image and concat downsampled mask c = self.model.cond_stage_model.encode(batch["masked_image"]) cc = torch.nn.functional.interpolate(batch["mask"], size=c.shape[-2:]) c = torch.cat((c, cc), dim=1) shape = (c.shape[1] - 1,) + c.shape[2:] samples_ddim, _ = self.sampler.sample(S=self.ddim_steps, conditioning=c, batch_size=c.shape[0], shape=shape, verbose=False) x_samples_ddim = self.model.decode_first_stage(samples_ddim) image = torch.clamp((batch["image"] + 1.0) / 2.0, min=0.0, max=1.0) mask = torch.clamp((batch["mask"] + 1.0) / 2.0, min=0.0, max=1.0) predicted_image = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0) inpainted = (1 - mask) * image + mask * predicted_image inpainted = inpainted.cpu().numpy().transpose(0, 2, 3, 1)[0] * 255 # print(type(inpainted)) inpainted = inpainted.astype(np.uint8) new_img_name = gen_new_name(img_path, 'LDMInpainter') new_img = Image.fromarray(inpainted) new_img = new_img.resize((w, h)) new_img.save(new_img_name) print( f"\nProcessed LDMInpainting, Inputs: {inputs}, " f"Output Image: {new_img_name}") return new_img_name # return inpainted ''' if __name__ == '__main__': painting = LDMInpainting('cuda:0') res = painting.inference(f'image/82e612_fe54ca_raw.png,image/04a785_fe54ca_mask.png.') print(res) '''