commit SegformerB2ClothesUltra node

This commit is contained in:
chflame
2024-05-24 12:41:03 +08:00
parent 78888f0b94
commit f1dd56b6e6
6 changed files with 234 additions and 0 deletions
+34
View File
@@ -70,6 +70,7 @@ When this error has occurred, please check the network environment.
## Update
<font size="4">**If the dependency package error after updating, please reinstall the relevant dependency packages. </font><br />
* Commit [SegformerB2ClothesUltra](#SegformerB2ClothesUltra) node, it used to segment character clothing. The model segmentation code is from[StartHua](https://github.com/StartHua/Comfyui_segformer_b2_clothes), thanks to the original author.
* [SaveImagePlus](#SaveImagePlus) node adds the output workflow to the json function, supports ```%date``` and ```%time``` to embeddint date or time to path and filename, and adds the preview switch.
* Commit [SaveImagePlus](#SaveImagePlus) node,It can customize the directory where the picture is saved, add a timestamp to the file name, select the save format, set the image compression rate, set whether to save the workflow, and optionally add invisible watermarks to the picture.
* Commit [AddBlindWaterMark](#AddBlindWaterMark), [ShowBlindWaterMark](#ShowBlindWaterMark) nodes, Add invisible watermark and decoded watermark to the picture. Commit [CreateQRCode](#CreateQRCode), [DecodeQRCode](#DecodeQRCode) nodes, It can generate two-dimensional code pictures and decode two-dimensional codes.
@@ -1304,6 +1305,39 @@ On the basis of PersonMaskUltra, the following changes have been made:
* detail_erode: Mask the erosion range inward from the edge. the larger the value, the larger the range of inward repair.
* detail_dilate: The edge of the mask expands outward. the larger the value, the wider the range of outward repair.
### <a id="table1">SegformerB2ClothesUltra</a>
![image](image/segformer_ultra_example.jpg)
Generate masks for characters' faces, hair, arms, legs, and clothing, mainly used for segmenting clothing.
The model segmentation code is from[StartHua](https://github.com/StartHua/Comfyui_segformer_b2_clothes),thanks to the original author.
Compared to the comfyui_segformer_b2_clothes, this node has ultra-high edge details. (Note: Generating images with edges exceeding 2K in size using the VITMatte method will consume a lot of memory)
*Download all model files from [https://huggingface.co/mattmdjaga/segformer_b2_clothes](https://huggingface.co/mattmdjaga/segformer_b2_clothes) to ```ComfyUI/models/segformer_b2_clothes``` folder.
Node Options:
![image](image/segformer_ultra_node.jpg)
* face: Facial recognition switch.
* hair: Hair recognition switch.
* hat: Hat recognition switch.
* sunglass: Sunglass recognition switch.
* left_arm: Left arm recognition switch.
* right_arm: Right arm recognition switch.
* left_leg: Left leg recognition switch.
* right_leg: Right leg recognition switch.
* skirt: Skirt recognition switch.
* pants: Pants recognition switch.
* dress: Dress recognition switch.
* belt: Belt recognition switch.
* shoe: Shoes recognition switch.
* bag: Bag recognition switch.
* scarf: Scarf recognition switch.
* detail_method: Edge processing methods. provides VITMatte, VITMatte(local), PyMatting, GuidedFilter. If the model has been downloaded after the first use of VITMatte, you can use VITMatte (local) afterwards.
* detail_erode: Mask the erosion range inward from the edge. the larger the value, the larger the range of inward repair.
* detail_dilate: The edge of the mask expands outward. the larger the value, the wider the range of outward repair.
* black_point: Edge black sampling threshold.
* white_point: Edge white sampling threshold.
* process_detail: Set to false here will skip edge processing to save runtime.
### <a id="table1">YoloV8Detect</a>
Use the YoloV8 model to detect faces, hand box areas, or character segmentation. Supports the output of the selected number of channels.
Download the model files from [GoogleDrive](https://drive.google.com/drive/folders/1I5TISO2G1ArSkKJu1O9b4Uvj3DVgn5d2) or [BaiduNetdisk](https://pan.baidu.com/s/1ImoJrzL1zDgaCqaSzrNEtw?pwd=5xgk) to ```ComfyUI/models/yolo``` folder.
+33
View File
@@ -70,6 +70,7 @@ git clone https://github.com/chflame163/ComfyUI_LayerStyle.git
## 更新说明
<font size="4">**如果本插件更新后出现依赖包错误,请重新安装相关依赖包。
* 添加 [SegformerB2ClothesUltra](#SegformerB2ClothesUltra)节点,用于分割人物服装。模型分割代码来自[StartHua](https://github.com/StartHua/Comfyui_segformer_b2_clothes),感谢原作者。
* [SaveImagePlus](#SaveImagePlus)节点增加输出工作流为json功能,支持使用```%date```和```%time```在路径和文件名嵌入时间,增加预览开关。
* 添加 [SaveImagePlus](#SaveImagePlus)节点,可自定义保存图片的目录,文件名增加时间戳,选择保存格式,设置图片压缩率,设置是否保存工作流,以及可选给图片添加隐形水印。
* 添加 [AddBlindWaterMark](#AddBlindWaterMark), [ShowBlindWaterMark](#ShowBlindWaterMark)节点,为图片增加隐形水印和解码水印。添加 [CreateQRCode](#CreateQRCode), [DecodeQRCode](#DecodeQRCode)节点,可生成二维码图片和解码二维码。
@@ -1291,6 +1292,38 @@ PersonMaskUltra的V2升级版,增加了VITMatte边缘处理方法。(注意:
* detail_dilate: 遮罩边缘向外扩张范围。数值越大,向外修复的范围越大。
### <a id="table1">SegformerB2ClothesUltra</a>
![image](image/segformer_ultra_example.jpg)
为人物生成脸、头发、手臂、腿以及服饰的遮罩,主要用于分割服装。模型分割代码来自[StartHua](https://github.com/StartHua/Comfyui_segformer_b2_clothes),感谢原作者。
与comfyui_segformer_b2_clothes节点相比,这个节点具有超高的边缘细节。(注意:生成边缘超过2K尺寸的图片使用VITMatte方法将占用大量内存)
*从[https://huggingface.co/mattmdjaga/segformer_b2_clothes](https://huggingface.co/mattmdjaga/segformer_b2_clothes)下载全部文件至```ComfyUI/models/segformer_b2_clothes```文件夹。
节点选项说明:
![image](image/segformer_ultra_node.jpg)
* face: 脸部识别。
* hair: 头发识别。
* hat: 帽子识别。
* sunglass: 墨镜识别。
* left_arm:左手臂识别。
* right_arm:右手臂识别。
* left_leg:左腿识别。
* right_leg:右腿识别。
* skirt:短裙识别。
* pants:裤子识别。
* dress:连衣裙识别。
* belt:腰带识别。
* shoe:鞋子识别。
* bag:背包识别。
* scarf:围巾识别。
* detail_method: 边缘处理方法。提供了VITMatte, VITMatte(local), PyMatting, GuidedFilter。如果首次使用VITMatte后模型已经下载,之后可以使用VITMatte(local)。
* detail_erode: 遮罩边缘向内侵蚀范围。数值越大,向内修复的范围越大。
* detail_dilate: 遮罩边缘向外扩张范围。数值越大,向外修复的范围越大。
* black_point: 边缘黑色采样阈值。
* white_point: 边缘黑色采样阈值。
* process_detail: 此处设为False将跳过边缘处理以节省运行时间。
### <a id="table1">YoloV8Detect</a>
使用YoloV8模型检测人脸、手部box区域,或者人物分割。支持输出所选择数量的通道。
请在 [GoogleDrive](https://drive.google.com/drive/folders/1I5TISO2G1ArSkKJu1O9b4Uvj3DVgn5d2) 或者 [百度网盘](https://pan.baidu.com/s/1ImoJrzL1zDgaCqaSzrNEtw?pwd=5xgk) 下载模型文件并放到 ```ComfyUI/models/yolo``` 文件夹。
Binary file not shown.

After

Width:  |  Height:  |  Size: 170 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 102 KiB

+2
View File
@@ -1272,6 +1272,8 @@ def load_RMBG_model():
net.eval()
return net
def RMBG(image:Image) -> Image:
rmbgmodel = load_RMBG_model()
w, h = image.size
+165
View File
@@ -0,0 +1,165 @@
'''
原始代码来自 https://github.com/StartHua/Comfyui_segformer_b2_clothes
'''
from transformers import SegformerImageProcessor, AutoModelForSemanticSegmentation
import torch.nn as nn
from .imagefunc import *
NODE_NAME = 'SegformerB2ClothesUltra'
# 切割服装
def get_segmentation(tensor_image):
cloth = tensor2pil(tensor_image)
model_folder_path = os.path.join(folder_paths.models_dir, "segformer_b2_clothes")
try:
model_folder_path = os.path.normpath(folder_paths.folder_names_and_paths['segformer_b2_clothes'][0][0])
except:
pass
processor = SegformerImageProcessor.from_pretrained(model_folder_path)
model = AutoModelForSemanticSegmentation.from_pretrained(model_folder_path)
# 预处理和预测
inputs = processor(images=cloth, return_tensors="pt")
outputs = model(**inputs)
logits = outputs.logits.cpu()
upsampled_logits = nn.functional.interpolate(logits, size=cloth.size[::-1], mode="bilinear", align_corners=False)
pred_seg = upsampled_logits.argmax(dim=1)[0].numpy()
return pred_seg,cloth
class Segformer_B2_Clothes:
def __init__(self):
pass
# Labels: 0: "Background", 1: "Hat", 2: "Hair", 3: "Sunglasses", 4: "Upper-clothes", 5: "Skirt",
# 6: "Pants", 7: "Dress", 8: "Belt", 9: "Left-shoe", 10: "Right-shoe", 11: "Face",
# 12: "Left-leg", 13: "Right-leg", 14: "Left-arm", 15: "Right-arm", 16: "Bag", 17: "Scarf"
@classmethod
def INPUT_TYPES(cls):
method_list = ['VITMatte', 'VITMatte(local)', 'PyMatting', 'GuidedFilter', ]
return {"required":
{
"image":("IMAGE",),
"face": ("BOOLEAN", {"default": True}),
"hair": ("BOOLEAN", {"default": True}),
"hat": ("BOOLEAN", {"default": True}),
"sunglass": ("BOOLEAN", {"default": True}),
"left_arm": ("BOOLEAN", {"default": True}),
"right_arm": ("BOOLEAN", {"default": True}),
"left_leg": ("BOOLEAN", {"default": True}),
"right_leg": ("BOOLEAN", {"default": True}),
"upper_clothes": ("BOOLEAN", {"default": True}),
"skirt": ("BOOLEAN", {"default": True}),
"pants": ("BOOLEAN", {"default": True}),
"dress": ("BOOLEAN", {"default": True}),
"belt": ("BOOLEAN", {"default": True}),
"shoe": ("BOOLEAN", {"default": True}),
"bag": ("BOOLEAN", {"default": True}),
"scarf": ("BOOLEAN", {"default": True}),
"detail_method": (method_list,),
"detail_erode": ("INT", {"default": 12, "min": 1, "max": 255, "step": 1}),
"detail_dilate": ("INT", {"default": 6, "min": 1, "max": 255, "step": 1}),
"black_point": ("FLOAT", {"default": 0.01, "min": 0.01, "max": 0.98, "step": 0.01}),
"white_point": ("FLOAT", {"default": 0.99, "min": 0.02, "max": 0.99, "step": 0.01}),
"process_detail": ("BOOLEAN", {"default": True}),
}
}
RETURN_TYPES = ("IMAGE", "MASK", )
RETURN_NAMES = ("image", "mask", )
FUNCTION = "segformer_ultra"
CATEGORY = '😺dzNodes/LayerMask'
def segformer_ultra(self, image,
face, hat, hair, sunglass, upper_clothes, skirt, pants, dress, belt, shoe,
left_leg, right_leg, left_arm, right_arm, bag, scarf, detail_method,
detail_erode, detail_dilate, black_point, white_point, process_detail
):
ret_images = []
ret_masks = []
if detail_method == 'VITMatte(local)':
local_files_only = True
else:
local_files_only = False
for i in image:
pred_seg, cloth = get_segmentation(i)
i = torch.unsqueeze(i, 0)
i = pil2tensor(tensor2pil(i).convert('RGB'))
orig_image = tensor2pil(i).convert('RGB')
labels_to_keep = [0]
if not hat:
labels_to_keep.append(1)
if not hair:
labels_to_keep.append(2)
if not sunglass:
labels_to_keep.append(3)
if not upper_clothes:
labels_to_keep.append(4)
if not skirt:
labels_to_keep.append(5)
if not pants:
labels_to_keep.append(6)
if not dress:
labels_to_keep.append(7)
if not belt:
labels_to_keep.append(8)
if not shoe:
labels_to_keep.append(9)
labels_to_keep.append(10)
if not face:
labels_to_keep.append(11)
if not left_leg:
labels_to_keep.append(12)
if not right_leg:
labels_to_keep.append(13)
if not left_arm:
labels_to_keep.append(14)
if not right_arm:
labels_to_keep.append(15)
if not bag:
labels_to_keep.append(16)
if not scarf:
labels_to_keep.append(17)
mask = np.isin(pred_seg, labels_to_keep).astype(np.uint8)
# 创建agnostic-mask图像
mask_image = Image.fromarray((1 - mask) * 255)
mask_image = mask_image.convert("L")
_mask = pil2tensor(mask_image)
detail_range = detail_erode + detail_dilate
if process_detail:
if detail_method == 'GuidedFilter':
_mask = guided_filter_alpha(i, _mask, detail_range // 6 + 1)
_mask = tensor2pil(histogram_remap(_mask, black_point, white_point))
elif detail_method == 'PyMatting':
_mask = tensor2pil(mask_edge_detail(i, _mask, detail_range // 8 + 1, black_point, white_point))
else:
_trimap = generate_VITMatte_trimap(_mask, detail_erode, detail_dilate)
_mask = generate_VITMatte(orig_image, _trimap, local_files_only=local_files_only)
_mask = tensor2pil(histogram_remap(pil2tensor(_mask), black_point, white_point))
else:
_mask = mask2image(_mask)
ret_image = RGB2RGBA(orig_image, _mask.convert('L'))
ret_images.append(pil2tensor(ret_image))
ret_masks.append(image2mask(_mask))
log(f"{NODE_NAME} Processed {len(ret_images)} image(s).", message_type='finish')
return (torch.cat(ret_images, dim=0), torch.cat(ret_masks, dim=0),)
NODE_CLASS_MAPPINGS = {
"LayerMask: SegformerB2ClothesUltra": Segformer_B2_Clothes
}
NODE_DISPLAY_NAME_MAPPINGS = {
"LayerMask: SegformerB2ClothesUltra": "LayerMask: Segformer B2 Clothes Ultra"
}