Files
rui40000-RUI-Nodes/mask_selector_node.py
rui40000andClaude Opus 4.8 5d3834503e fix: 八方向拆分不再切断角色;全仓库参数补齐中文 tooltip
【修复】角色被格线切断(脚、手杖、飘起的斗篷被削掉)
新增 expand_beyond_cell(默认开启):格子只用来判定「这是哪个方向」,
角色的实际范围由它自身的连通区域决定,按质心归属确保邻居不混入。
实测 8/8 方向的裁剪框边缘 alpha 从 1.00(内容顶到边界=被切断)
降到 0.00,S 方向高度 326→356、E 方向宽度 150→188 把缺的部分找了回来。
代价是需要两遍扫描(先求全序列并集框再提取),耗时 4.7s→14.9s。

【规则】每个参数都必须有中文 tooltip,作为以后的统一约定
全仓库 26 个节点 169 个参数,此前缺 115 个,现已 100% 覆盖。
tooltip 写「怎么调」而不只是「是什么」:给取值区间的实际影响、
推荐值与踩坑提示(如 OpenAI/ZenMux 的地址栏不能带 :// ,
素材拆分节点用于动画序列时顺序会漂移等)。

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-29 19:07:01 +08:00

184 lines
7.9 KiB
Python

import torch
import json
class MaskSelector:
"""
遮罩筛选节点:
- 对输入遮罩进行排序(按面积、从左到右、从上到下)
- 输出指定编号的遮罩、剩余遮罩的合并结果,以及 JSON 信息
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"masks": ("MASK", {
"tooltip": "待筛选的遮罩批次(N,H,W)。单张遮罩也可以,会自动补成批次。"
}),
"sort_method": (["按面积排序 / By Area", "从左到右 / Left to Right", "从上到下 / Top to Bottom"], {
"default": "按面积排序 / By Area",
"tooltip": "先按此规则排序,再用下面的编号取第几个。\n"
"按面积:从大到小,取主体用它最稳\n"
"从左到右 / 从上到下:按遮罩质心的坐标排\n"
"空遮罩一律排到末尾。"
}),
"index": ("INT", {
"default": 1,
"min": 1,
"step": 1,
"tooltip": "取排序后的第几个遮罩,从 1 开始。\n"
"超出总数会自动夹到最后一个,并在 info 里标记 index_clamped。"
}),
}
}
RETURN_TYPES = ("MASK", "MASK", "STRING")
RETURN_NAMES = ("选中遮罩 / Selected", "剩余遮罩 / Remaining", "信息 / Info")
FUNCTION = "select_mask"
CATEGORY = "Rui-Node🐶/遮罩处理🎭"
def select_mask(self, masks, sort_method, index):
"""
主处理流程:
- 标准化输入形状为 (N, H, W)
- 计算每个遮罩的面积与质心
- 按指定规则排序
- 选择指定编号并返回三个输出
"""
# 标准化输入形状
if masks.dim() == 2:
# (H, W) -> (1, H, W)
masks = masks.unsqueeze(0)
elif masks.dim() == 3:
pass # (N, H, W)
else:
# 非法形状,返回空结果
H = W = 0
empty = torch.zeros((1, H, W), dtype=torch.float32)
info = json.dumps({"error": "invalid mask shape"}, ensure_ascii=False)
return (empty, empty, info)
N, H, W = masks.shape
# 计算面积与质心
# 使用阈值 0.5 判定非零像素
bin_masks = (masks > 0.5)
areas = []
centers = []
all_empty = True
for i in range(N):
nz = torch.nonzero(bin_masks[i], as_tuple=False)
area = nz.shape[0]
areas.append(int(area))
if area > 0:
all_empty = False
# 质心坐标:x 为列索引平均,y 为行索引平均
y_mean = float(nz[:, 0].float().mean().item())
x_mean = float(nz[:, 1].float().mean().item())
centers.append((x_mean, y_mean))
else:
# 空遮罩:用占位中心使其排序在末尾
centers.append((float('inf'), float('inf')))
# 映射中英双语排序选项到内部代码
method_map = {
"按面积排序 / By Area": "by_area",
"从左到右 / Left to Right": "left_to_right",
"从上到下 / Top to Bottom": "top_to_bottom",
# 兼容老值
"by_area": "by_area",
"left_to_right": "left_to_right",
"top_to_bottom": "top_to_bottom",
}
internal_method = method_map.get(sort_method, "by_area")
# 排序索引
if internal_method == "by_area":
# 面积从大到小
sort_key = [(areas[i], -i) for i in range(N)]
order = sorted(range(N), key=lambda i: sort_key[i], reverse=True)
elif internal_method == "left_to_right":
# 按质心 X 升序
sort_key = [centers[i][0] for i in range(N)]
order = sorted(range(N), key=lambda i: sort_key[i])
elif internal_method == "top_to_bottom":
# 按质心 Y 升序
sort_key = [centers[i][1] for i in range(N)]
order = sorted(range(N), key=lambda i: sort_key[i])
else:
# 未知排序方式,默认按面积
sort_key = [(areas[i], -i) for i in range(N)]
order = sorted(range(N), key=lambda i: sort_key[i], reverse=True)
# 1-based 索引选择与夹取
index_clamped = False
target_idx = index - 1
if target_idx < 0:
target_idx = 0
index_clamped = True
if target_idx >= N:
target_idx = N - 1
index_clamped = True
selected_i = order[target_idx]
selected = masks[selected_i].unsqueeze(0) # (1, H, W)
# 合并剩余遮罩:逻辑或
remaining_indices = [i for i in order if i != selected_i]
if len(remaining_indices) == 0:
remaining = torch.zeros((1, H, W), dtype=masks.dtype, device=masks.device)
else:
remaining_stack = masks[remaining_indices] # (M, H, W)
# 使用 max 合并为逻辑或
remaining = torch.max(remaining_stack, dim=0).values.unsqueeze(0)
# JSON 信息(包含中文注释的字符串)
selected_area = areas[selected_i]
sel_center = centers[selected_i]
# 对空遮罩设定中心为 [0, 0]
if sel_center[0] == float('inf') or sel_center[1] == float('inf'):
sel_center_out = [0, 0]
else:
sel_center_out = [int(round(sel_center[0])), int(round(sel_center[1]))]
info_obj = {
"total_masks": N, # 遮罩总数:检测到的遮罩总数量
"selected_index": index, # 选中编号:用户指定输出的遮罩编号(原始输入)
"sort_method": internal_method, # 排序方式:当前使用的排序规则代码
"selected_area": selected_area, # 选中面积:选中遮罩的像素面积
"selected_center": sel_center_out, # 选中中心点:选中遮罩的质心坐标 [x, y]
"index_clamped": index_clamped # 编号越界:若编号超出范围被自动修正则为 true
}
if all_empty:
info_obj["warning"] = "所有遮罩为空 / All masks empty"
# 生成带注释的 JSON 字符串
# 注意:标准 JSON 不支持注释,这里按照需求输出 JSON 风格字符串并附带注释
# 若后续需严格 JSON,可移除注释并使用 json.dumps(info_obj, ensure_ascii=False)
info_lines = [
"{",
f' "total_masks": {info_obj["total_masks"]}, // 遮罩总数:检测到的遮罩总数量',
f' "selected_index": {info_obj["selected_index"]}, // 选中编号:用户指定输出的遮罩编号',
f' "sort_method": "{info_obj["sort_method"]}", // 排序方式:当前使用的排序规则代码',
f' "selected_area": {info_obj["selected_area"]}, // 选中面积:选中遮罩的像素面积',
f' "selected_center": {json.dumps(info_obj["selected_center"], ensure_ascii=False)}, // 选中中心点:选中遮罩的质心坐标 [x, y]',
f' "index_clamped": {"true" if info_obj["index_clamped"] else "false"} // 编号越界:如果用户输入的编号超出范围被自动修正则为 true'
]
if "warning" in info_obj:
info_lines.append(f' ,"warning": "{info_obj["warning"]}" // 警告:所有遮罩为空')
info_lines.append("}")
info_str = "\n".join(info_lines)
return (selected, remaining, info_str)
# 节点映射字典,用于 ComfyUI 注册节点
NODE_CLASS_MAPPINGS = {
"MaskSelector": MaskSelector
}
# 节点显示名称映射,用于在 UI 中显示友好名称
NODE_DISPLAY_NAME_MAPPINGS = {
"MaskSelector": "遮罩筛选 / Mask Selector"
}