120 lines
3.4 KiB
Python
120 lines
3.4 KiB
Python
import math
|
|
import torch
|
|
import comfy.utils
|
|
import numpy as np
|
|
import cv2
|
|
from .lama import LamaInpainting
|
|
|
|
def pad64(x):
|
|
return int(np.ceil(float(x) / 64.0) * 64 - x)
|
|
|
|
|
|
def safer_memory(x):
|
|
# Fix many MAC/AMD problems
|
|
return np.ascontiguousarray(x.copy()).copy()
|
|
|
|
def HWC3(x):
|
|
assert x.dtype == np.uint8
|
|
if x.ndim == 2:
|
|
x = x[:, :, None]
|
|
assert x.ndim == 3
|
|
H, W, C = x.shape
|
|
assert C == 1 or C == 3 or C == 4
|
|
if C == 3:
|
|
return x
|
|
if C == 1:
|
|
return np.concatenate([x, x, x], axis=2)
|
|
if C == 4:
|
|
color = x[:, :, 0:3].astype(np.float32)
|
|
alpha = x[:, :, 3:4].astype(np.float32) / 255.0
|
|
y = color * alpha + 255.0 * (1.0 - alpha)
|
|
y = y.clip(0, 255).astype(np.uint8)
|
|
return y
|
|
|
|
|
|
def resize_image_with_pad(input_image, resolution, skip_hwc3=False):
|
|
if skip_hwc3:
|
|
img = input_image
|
|
else:
|
|
img = HWC3(input_image)
|
|
H_raw, W_raw, _ = img.shape
|
|
k = float(resolution) / float(min(H_raw, W_raw))
|
|
interpolation = cv2.INTER_CUBIC if k > 1 else cv2.INTER_AREA
|
|
H_target = int(np.round(float(H_raw) * k))
|
|
W_target = int(np.round(float(W_raw) * k))
|
|
img = cv2.resize(img, (W_target, H_target), interpolation=interpolation)
|
|
H_pad, W_pad = pad64(H_target), pad64(W_target)
|
|
img_padded = np.pad(img, [[0, H_pad], [0, W_pad], [0, 0]], mode='edge')
|
|
|
|
def remove_pad(x):
|
|
return safer_memory(x[:H_target, :W_target])
|
|
|
|
return safer_memory(img_padded), remove_pad
|
|
|
|
class ImageLama:
|
|
|
|
def __init__(self,res=512):
|
|
self.res=res
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required":
|
|
{"images": ("IMAGE", ),
|
|
"mask": ("MASK",),
|
|
},
|
|
}
|
|
|
|
CATEGORY = "lam"
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("图片",)
|
|
FUNCTION = "lama_inpaint"
|
|
def lama_inpaint(self, images,mask):
|
|
mask_np=mask.numpy()
|
|
mask_shape=mask_np.shape
|
|
imaget_np=images.numpy()
|
|
imgt_shape=imaget_np.shape
|
|
prd_images=[]
|
|
for i in range(imgt_shape[0]):
|
|
imaget=np.uint8(imaget_np[i]*255)
|
|
imaget = cv2.cvtColor(imaget, cv2.COLOR_BGR2RGB)
|
|
mask=np.uint8(mask_np*255)
|
|
mask=np.expand_dims(mask, axis=-1)
|
|
H, W, C = imaget.shape
|
|
img=np.zeros((H,W,4))
|
|
img[:,:,0:3]=imaget
|
|
img[:,:,3:4]=mask
|
|
|
|
raw_color = img[:, :, 0:3].copy()
|
|
raw_mask = img[:, :, 3:4].copy()
|
|
|
|
res = 256 # Always use 256 since lama is trained on 256
|
|
|
|
img_res, remove_pad = resize_image_with_pad(img, res, skip_hwc3=True)
|
|
|
|
model_lama = LamaInpainting()
|
|
|
|
# applied auto inversion
|
|
prd_color = model_lama(img_res)
|
|
prd_color = remove_pad(prd_color)
|
|
prd_color = cv2.resize(prd_color, (W, H))
|
|
|
|
alpha = raw_mask.astype(np.float32) / 255.0
|
|
fin_color = prd_color.astype(np.float32) * alpha + raw_color.astype(np.float32) * (1 - alpha)
|
|
fin_color = fin_color.clip(0, 255).astype(np.uint8)
|
|
|
|
fin_color = cv2.cvtColor(fin_color, cv2.COLOR_BGR2RGB)
|
|
prd_images.append(torch.from_numpy(np.array(fin_color).astype(np.float32) / 255.0))
|
|
|
|
return (torch.stack(prd_images),)
|
|
|
|
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"ImageLama": ImageLama
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"ImageLama": "图片局部重绘lama"
|
|
}
|