78 lines
2.9 KiB
Python
78 lines
2.9 KiB
Python
import math
|
|
import torch
|
|
import comfy.utils
|
|
import numpy as np
|
|
import cv2
|
|
|
|
class LoadReplaceImage:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required":
|
|
{"image_t": ("IMAGE", ),
|
|
"mask": ("MASK",),
|
|
"image_bt": ("IMAGE", ),
|
|
"x": ("INT", {"default": 0, "min": 0, "max": 99999, "step": 1}),
|
|
"y": ("INT", {"default": 0, "min": 0, "max": 99999, "step": 1}),
|
|
},
|
|
}
|
|
|
|
CATEGORY = "lam"
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("图片",)
|
|
FUNCTION = "load_replace_image"
|
|
def load_replace_image(self, image_t,mask,image_bt,x,y):
|
|
mask_erosion = mask
|
|
# if image_t.shape != image_bt.shape:
|
|
# image_bt = image_bt.permute(0, 3, 1, 2)
|
|
# image_bt = comfy.utils.common_upscale(image_bt, image_t.shape[2], image_t.shape[1], upscale_method='bicubic', crop='center')
|
|
# image_bt = image_bt.permute(0, 2, 3, 1)
|
|
|
|
replace_image = self.blend_mode(image_t, image_bt, mask_erosion,x,y)
|
|
return (replace_image,)
|
|
|
|
def blend_mode(self, img1, img2, mask,y,x):
|
|
mask_np=mask.numpy()
|
|
if mask_np.ndim==2:
|
|
mask_np=np.expand_dims(mask_np, axis=0)
|
|
|
|
mask_shape=mask_np.shape
|
|
img1_shape=img1.size()
|
|
img2_shape=img2.size()
|
|
print("=====------=======",mask_shape)
|
|
if tuple(mask_shape)!=tuple(img1_shape):
|
|
maskList=[]
|
|
for i in range(mask_shape[0]):
|
|
maskList.append(cv2.resize(mask_np[i],(img1_shape[2],img1_shape[1])))
|
|
mask_np=np.array(maskList)
|
|
mask_shape=mask_np.shape
|
|
mask_np_rgb = np.zeros((mask_shape[0],mask_shape[1],mask_shape[2],3))
|
|
mask_np_rgb[:,:,:,0] = mask_np
|
|
mask_np_rgb[:,:,:,1] = mask_np
|
|
mask_np_rgb[:,:,:,2] = mask_np
|
|
|
|
mask=torch.from_numpy(mask_np_rgb)
|
|
print("-----------",mask.size())
|
|
print("-----------",img1_shape)
|
|
print("-----------",img2_shape)
|
|
#取最少图片数量
|
|
min_size=img2_shape[0] if img1_shape[0]>img2_shape[0] else img1_shape[0]
|
|
#图片大小不一致时，返回被替换内容
|
|
if tuple(mask.size()[1:])!=tuple(img1_shape[1:]) or x+img1_shape[1]>img2_shape[1] or y+img1_shape[2]>img2_shape[2]:
|
|
return img2
|
|
|
|
print("min_size::::",str(min_size))
|
|
img1[:min_size]=img1[:min_size]*mask[:min_size]
|
|
img2[:min_size,x:x+img1_shape[1],y:y+img1_shape[2]]=img2[:min_size,x:x+img1_shape[1],y:y+img1_shape[2]]*(1-mask[:min_size])
|
|
img2[:min_size,x:x+img1_shape[1],y:y+img1_shape[2]]=img1[:min_size]+img2[:min_size,x:x+img1_shape[1],y:y+img1_shape[2]]
|
|
return img2
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"LoadReplaceImage": LoadReplaceImage
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"LoadReplaceImage": "图片局部替换"
|
|
}
|