diff --git a/README.MD b/README.MD index 334b335..0679b8a 100644 --- a/README.MD +++ b/README.MD @@ -70,6 +70,7 @@ When this error has occurred, please check the network environment. ## Update **If the dependency package error after updating, please reinstall the relevant dependency packages.
+* 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. +### SegformerB2ClothesUltra +![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. + + ### YoloV8Detect 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. diff --git a/README_CN.MD b/README_CN.MD index b4fa7c8..d57fba7 100644 --- a/README_CN.MD +++ b/README_CN.MD @@ -70,6 +70,7 @@ git clone https://github.com/chflame163/ComfyUI_LayerStyle.git ## 更新说明 **如果本插件更新后出现依赖包错误，请重新安装相关依赖包。 +* 添加 [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: 遮罩边缘向外扩张范围。数值越大，向外修复的范围越大。 +### SegformerB2ClothesUltra +![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将跳过边缘处理以节省运行时间。 + + ### YoloV8Detect 使用YoloV8模型检测人脸、手部box区域，或者人物分割。支持输出所选择数量的通道。 请在 [GoogleDrive](https://drive.google.com/drive/folders/1I5TISO2G1ArSkKJu1O9b4Uvj3DVgn5d2) 或者 [百度网盘](https://pan.baidu.com/s/1ImoJrzL1zDgaCqaSzrNEtw?pwd=5xgk) 下载模型文件并放到 ```ComfyUI/models/yolo``` 文件夹。 diff --git a/image/segformer_ultra_example.jpg b/image/segformer_ultra_example.jpg new file mode 100644 index 0000000..f2cb386 Binary files /dev/null and b/image/segformer_ultra_example.jpg differ diff --git a/image/segformer_ultra_node.jpg b/image/segformer_ultra_node.jpg new file mode 100644 index 0000000..387a883 Binary files /dev/null and b/image/segformer_ultra_node.jpg differ diff --git a/py/imagefunc.py b/py/imagefunc.py index 6202217..3912121 100644 --- a/py/imagefunc.py +++ b/py/imagefunc.py @@ -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 diff --git a/py/segformer_ultra.py b/py/segformer_ultra.py new file mode 100644 index 0000000..0034770 --- /dev/null +++ b/py/segformer_ultra.py @@ -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" +}