174 lines
5.4 KiB
Python
174 lines
5.4 KiB
Python
"""
|
||
InpaintRegionEditor Node
|
||
|
||
增强版遮罩编辑器，支持 Photopea 集成和可配置选区
|
||
|
||
概念说明：
|
||
- 遮罩 (Mask): 定义重绘区域，重绘完成后只在遮罩内合成新内容
|
||
- 选区 (Region): 定义参考区域，大模型重绘时参考这个区域内的原图内容
|
||
|
||
约束：选区必须 >= 遮罩
|
||
"""
|
||
|
||
import torch
|
||
import numpy as np
|
||
import folder_paths
|
||
from PIL import Image
|
||
import os
|
||
import json
|
||
|
||
|
||
class InpaintRegionEditor:
|
||
"""
|
||
增强版遮罩编辑器，支持 Photopea 集成和可配置选区
|
||
"""
|
||
|
||
# 预设尺寸映射
|
||
PRESET_SIZES = {
|
||
"512×512 (SD)": (512, 512),
|
||
"768×768": (768, 768),
|
||
"1024×1024 (SDXL)": (1024, 1024),
|
||
"1280×1280": (1280, 1280),
|
||
"1536×1536": (1536, 1536),
|
||
}
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(s):
|
||
input_dir = folder_paths.get_input_directory()
|
||
files = [f for f in os.listdir(input_dir)
|
||
if os.path.isfile(os.path.join(input_dir, f))]
|
||
files = folder_paths.filter_files_content_types(files, ["image"])
|
||
|
||
return {
|
||
"required": {
|
||
# 图像输入
|
||
"image": (sorted(files), {"image_upload": True}),
|
||
|
||
# 选区尺寸预设
|
||
"region_size": (
|
||
list(s.PRESET_SIZES.keys()) + ["Custom (自定义)"],
|
||
{"default": "512×512 (SD)"}
|
||
),
|
||
|
||
# 自定义选区尺寸
|
||
"region_width": ("INT", {
|
||
"default": 512,
|
||
"min": 64,
|
||
"max": 4096,
|
||
"step": 64
|
||
}),
|
||
"region_height": ("INT", {
|
||
"default": 512,
|
||
"min": 64,
|
||
"max": 4096,
|
||
"step": 64
|
||
}),
|
||
},
|
||
"hidden": {
|
||
# 遮罩边界信息（前端传递）
|
||
"mask_bounds": "STRING",
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE", "MASK", "INT", "INT", "INT", "INT")
|
||
RETURN_NAMES = ("image", "mask", "region_top", "region_left", "region_width", "region_height")
|
||
FUNCTION = "process"
|
||
CATEGORY = "image/inpaint"
|
||
|
||
def process(self, image, region_size, region_width, region_height,
|
||
mask_bounds=None):
|
||
"""
|
||
处理图像和遮罩，计算选区位置
|
||
|
||
Args:
|
||
image: 图像文件名
|
||
region_size: 选区预设名称
|
||
region_width: 自定义选区宽度
|
||
region_height: 自定义选区高度
|
||
mask_bounds: 遮罩边界 JSON 字符串
|
||
"""
|
||
# 加载图像
|
||
image_path = folder_paths.get_annotated_filepath(image)
|
||
img = Image.open(image_path)
|
||
|
||
img_width, img_height = img.size
|
||
|
||
# 提取或创建遮罩
|
||
if 'A' in img.getbands():
|
||
mask = np.array(img.getchannel('A')).astype(np.float32) / 255.0
|
||
mask = 1.0 - mask # ComfyUI 遮罩约定
|
||
else:
|
||
mask = np.zeros((img_height, img_width), dtype=np.float32)
|
||
|
||
# 解析遮罩边界
|
||
mask_bounds_dict = None
|
||
if mask_bounds:
|
||
try:
|
||
mask_bounds_dict = json.loads(mask_bounds)
|
||
except:
|
||
pass
|
||
|
||
# 如果没有前端传递的遮罩边界，自动计算
|
||
if not mask_bounds_dict:
|
||
mask_bounds_dict = self._calculate_mask_bounds(mask)
|
||
|
||
# 确定选区尺寸
|
||
if region_size in self.PRESET_SIZES:
|
||
r_width, r_height = self.PRESET_SIZES[region_size]
|
||
else:
|
||
r_width, r_height = region_width, region_height
|
||
|
||
# 约束验证：选区必须 >= 遮罩
|
||
mask_w = mask_bounds_dict.get('width', 0)
|
||
mask_h = mask_bounds_dict.get('height', 0)
|
||
|
||
if r_width < mask_w:
|
||
r_width = int(mask_w)
|
||
if r_height < mask_h:
|
||
r_height = int(mask_h)
|
||
|
||
# 计算选区位置（居中于遮罩）
|
||
mask_x = mask_bounds_dict.get('x', 0)
|
||
mask_y = mask_bounds_dict.get('y', 0)
|
||
|
||
region_x = mask_x + mask_w / 2 - r_width / 2
|
||
region_y = mask_y + mask_h / 2 - r_height / 2
|
||
|
||
# 确保选区在图像范围内
|
||
region_x = max(0, min(region_x, img_width - r_width))
|
||
region_y = max(0, min(region_y, img_height - r_height))
|
||
|
||
# 转换为 tensor
|
||
image_tensor = torch.from_numpy(
|
||
np.array(img.convert("RGB")).astype(np.float32) / 255.0
|
||
).unsqueeze(0)
|
||
|
||
mask_tensor = torch.from_numpy(mask).unsqueeze(0)
|
||
|
||
return (
|
||
image_tensor,
|
||
mask_tensor,
|
||
int(region_y), # top
|
||
int(region_x), # left
|
||
int(r_width),
|
||
int(r_height)
|
||
)
|
||
|
||
def _calculate_mask_bounds(self, mask):
|
||
"""计算遮罩的边界框"""
|
||
rows = np.any(mask > 0.5, axis=1)
|
||
cols = np.any(mask > 0.5, axis=0)
|
||
|
||
if not np.any(rows) or not np.any(cols):
|
||
return {'x': 0, 'y': 0, 'width': 0, 'height': 0}
|
||
|
||
rmin, rmax = np.where(rows)[0][[0, -1]]
|
||
cmin, cmax = np.where(cols)[0][[0, -1]]
|
||
|
||
return {
|
||
'x': int(cmin),
|
||
'y': int(rmin),
|
||
'width': int(cmax - cmin + 1),
|
||
'height': int(rmax - rmin + 1)
|
||
}
|