Commit RemBgUltra and PixelSpread nodes

This commit is contained in:
chflame163
2024-02-09 10:21:47 +08:00
parent 8e096802df
commit f7e69e85c3
17 changed files with 1610 additions and 8 deletions
+27 -1
View File
@@ -13,6 +13,7 @@ Nodes are divided into 5 groups according to their functions: LayerStyle, LayerC
[中文说明点这里](./README_CN.MD)
## Update
* Commit [RemBgUltra](#RemBgUltra) and [PixelSpread](#PixelSpread) nodes significantly improved mask quality. *RemBgUltra requires manual model download.
* Commit [TextImage](#TextImage) node, it generate text images and masks.
* Add new types of [blend mode](#blend) between images. now supports up to 19 blend modes. add **color_burn, color_dodge, linear_burn, linear_dodge, overlay, soft_light, hard_light, vivid_light, pin_light, linear_light** and **hard_mix**.
The newly added blend mode is applicable to all nodes that support blend mode.
@@ -561,6 +562,29 @@ Output:
* x: The x-coordinate of the top left corner position.
* y: The y-coordinate of the top left corner position.
### <a id="table1">RemBgUltra</a>
Remove background. compared to the previous RemBg, the cutout of this node has ultra-high edge details.
This node combines the Alpha Matte node of Spacepxl's [ComfyUI-Image-Filters](https://github.com/spacepxl/ComfyUI-Image-Filters) and the functionality of ZO-ZHO-ZHO's [ComfyUI-BRIA_AI-RMBG](https://github.com/ZHO-ZHO-ZHO/ComfyUI-BRIA_AI-RMBG).
*Download the [BRIA Background Removal v1.4](https://huggingface.co/briaai/RMBG-1.4) model file (model. pth) to the customer_modes/ComfyUI-LayerStyle/RMBG-1.4 floder.
This model was developed by BRIA AI and can be used as an open-source model for non-commercial purposes.
![image](image/rembg_ultra_example.png)
Node options:
![image](image/rembg_ultra_node.png)
* detail_range: Edge detail range.
* black_point: Edge black sampling threshold.
* white_point: Edge white sampling threshold.
### <a id="table1">PixelSpread</a>
Pixel expansion preprocessing on the masked edge of an image can effectively improve the edges of image synthesis.
![image](image/pixel_spread_example.png)
Node options:
![image](image/pixel_spread_node.png)
* invert_mask: Whether to reverse the mask.
### <a id="table1">MaskGrow</a>
Grow and shrink edges and blur the mask
@@ -724,7 +748,9 @@ image Some JSON workflow files in the workflow directory, that is example for Co
## How to install
* Open the cmd window in the plugin directory of ComfyUI, like "ComfyUI\custom_nodes\",type```git clone https://github.com/chflame163/ComfyUI_LayerStyle.git```
* Recommended use ComfyUI Manager for installation.
* Or open the cmd window in the plugin directory of ComfyUI, like "ComfyUI\custom_nodes\",type```git clone https://github.com/chflame163/ComfyUI_LayerStyle.git```
or download the zip file and extracted, copy the resulting folder to ComfyUI\custom_ Nodes\
* Install dependency packages, open the cmd window in the WordCloud plugin directory like "ComfyUI\custom_ Nodes\ComfyUI_WordCloud" and enter the following command:
+27 -2
View File
@@ -9,7 +9,8 @@
* [LayerUtility](#LayerUtility)节点组提供图层合成工具和工作流相关的辅助节点。
* [LayerFilter](#LayerFilter)节点组提供图像效果滤镜。
## 更新说明
## 更新说明、
* 添加[RemBgUltra](#RemBgUltra) 和 [PixelSpread](#PixelSpread) 节点,显著提升了遮罩质量。*RemBgUltra需手动下载模型。
* 添加[TextImage](#TextImage) 节点,生成文字图像和遮罩。
* 图像之间的[混合模式](#混合模式)增加新类型,现在支持多达19种混合模式。新增color_burn颜色加深, color_dodge颜色减淡, linear_burn线性加深, linear_dodge线性减淡, overlay叠加, soft_light柔光, hard_light强光, vivid_light亮光, pin_light点光, linear_light线性光, hard_mix实色混合。新增的混合模式适用于所有支持混合模式的节点。
* 添加[ColorMap](#ColorMap) 滤镜节点,用于制作伪彩色热力图效果。
@@ -547,6 +548,29 @@
* x: 左上角位置x坐标输出。
* y: 左上角位置y坐标输出。
### <a id="table1">RemBgUltra</a>
去除背景。与之前的RemBg相比,这个节点的抠像具有超高的边缘细节。
本节点结合了spacepxl的[ComfyUI-Image-Filters](https://github.com/spacepxl/ComfyUI-Image-Filters)的Alpha Matte节点,以及ZHO-ZHO-ZHO的[ComfyUI-BRIA_AI-RMBG](https://github.com/ZHO-ZHO-ZHO/ComfyUI-BRIA_AI-RMBG)的功能。
*将[BRIA Background Removal v1.4](https://huggingface.co/briaai/RMBG-1.4)模型文件(model.pth)下载至/custom_nodes/ComfyUI_LayerStyle/RMBG-1.4。
该模型由 BRIA AI 开发,可作为非商业用途的开源模型。
![image](image/rembg_ultra_example.png)
节点选项说明:
![image](image/rembg_ultra_node.png)
* detail_range: 边缘细节范围。
* black_point: 边缘黑色采样阈值。
* white_point: 边缘黑色采样阈值。
### <a id="table1">PixelSpread</a>
对图像的遮罩边缘部分进行像素扩张预处理,可有效改善图像合成的边缘。
![image](image/pixel_spread_example.png)
节点选项说明:
![image](image/pixel_spread_node.png)
* invert_mask: 是否反转遮罩。
### <a id="table1">MaskGrow</a>
对mask进行扩张收缩边缘和模糊处理
![image](image/mask_grow_example.png)
@@ -700,7 +724,8 @@ mask
## 安装方法
* 在CompyUI插件目录(例如“CompyUI\custom_nodes\”)中打开cmd窗口,键入```git clone https://github.com/chflame163/ComfyUI_LayerStyle.git```安装。或者下载解压zip文件,将得到的文件夹复制到 ComfyUI\custom_nodes\
* 推荐使用 ComfyUI Manager 安装。
* 或者在CompyUI插件目录(例如“CompyUI\custom_nodes\”)中打开cmd窗口,键入```git clone https://github.com/chflame163/ComfyUI_LayerStyle.git```安装。或者下载解压zip文件,将得到的文件夹复制到 ComfyUI\custom_nodes\
* 安装依赖包,在资源管理器ComfyUI\custom_nodes\ComfyUI_WordCloud 插件目录位置打开cmd窗口,输入以下命令:
```..\..\..\python_embeded\python.exe -m pip install -r requirements.txt```
* 重新打开ComfyUI。
+1
View File
@@ -0,0 +1 @@
Binary file not shown.

Before

Width:  |  Height:  |  Size: 313 KiB

After

Width:  |  Height:  |  Size: 382 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.3 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 114 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 3.6 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 140 KiB

+455
View File
@@ -0,0 +1,455 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
class REBNCONV(nn.Module):
def __init__(self,in_ch=3,out_ch=3,dirate=1,stride=1):
super(REBNCONV,self).__init__()
self.conv_s1 = nn.Conv2d(in_ch,out_ch,3,padding=1*dirate,dilation=1*dirate,stride=stride)
self.bn_s1 = nn.BatchNorm2d(out_ch)
self.relu_s1 = nn.ReLU(inplace=True)
def forward(self,x):
hx = x
xout = self.relu_s1(self.bn_s1(self.conv_s1(hx)))
return xout
## upsample tensor 'src' to have the same spatial size with tensor 'tar'
def _upsample_like(src,tar):
src = F.interpolate(src,size=tar.shape[2:],mode='bilinear')
return src
### RSU-7 ###
class RSU7(nn.Module):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3, img_size=512):
super(RSU7,self).__init__()
self.in_ch = in_ch
self.mid_ch = mid_ch
self.out_ch = out_ch
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1) ## 1 -> 1/2
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool3 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool4 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool5 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
self.rebnconv6 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.rebnconv7 = REBNCONV(mid_ch,mid_ch,dirate=2)
self.rebnconv6d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv5d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
def forward(self,x):
b, c, h, w = x.shape
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx = self.pool3(hx3)
hx4 = self.rebnconv4(hx)
hx = self.pool4(hx4)
hx5 = self.rebnconv5(hx)
hx = self.pool5(hx5)
hx6 = self.rebnconv6(hx)
hx7 = self.rebnconv7(hx6)
hx6d = self.rebnconv6d(torch.cat((hx7,hx6),1))
hx6dup = _upsample_like(hx6d,hx5)
hx5d = self.rebnconv5d(torch.cat((hx6dup,hx5),1))
hx5dup = _upsample_like(hx5d,hx4)
hx4d = self.rebnconv4d(torch.cat((hx5dup,hx4),1))
hx4dup = _upsample_like(hx4d,hx3)
hx3d = self.rebnconv3d(torch.cat((hx4dup,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
return hx1d + hxin
### RSU-6 ###
class RSU6(nn.Module):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU6,self).__init__()
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool3 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool4 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.rebnconv6 = REBNCONV(mid_ch,mid_ch,dirate=2)
self.rebnconv5d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
def forward(self,x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx = self.pool3(hx3)
hx4 = self.rebnconv4(hx)
hx = self.pool4(hx4)
hx5 = self.rebnconv5(hx)
hx6 = self.rebnconv6(hx5)
hx5d = self.rebnconv5d(torch.cat((hx6,hx5),1))
hx5dup = _upsample_like(hx5d,hx4)
hx4d = self.rebnconv4d(torch.cat((hx5dup,hx4),1))
hx4dup = _upsample_like(hx4d,hx3)
hx3d = self.rebnconv3d(torch.cat((hx4dup,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
return hx1d + hxin
### RSU-5 ###
class RSU5(nn.Module):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU5,self).__init__()
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool3 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=2)
self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
def forward(self,x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx = self.pool3(hx3)
hx4 = self.rebnconv4(hx)
hx5 = self.rebnconv5(hx4)
hx4d = self.rebnconv4d(torch.cat((hx5,hx4),1))
hx4dup = _upsample_like(hx4d,hx3)
hx3d = self.rebnconv3d(torch.cat((hx4dup,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
return hx1d + hxin
### RSU-4 ###
class RSU4(nn.Module):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU4,self).__init__()
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=2)
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
def forward(self,x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx4 = self.rebnconv4(hx3)
hx3d = self.rebnconv3d(torch.cat((hx4,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
return hx1d + hxin
### RSU-4F ###
class RSU4F(nn.Module):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU4F,self).__init__()
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=2)
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=4)
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=8)
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=4)
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=2)
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
def forward(self,x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx2 = self.rebnconv2(hx1)
hx3 = self.rebnconv3(hx2)
hx4 = self.rebnconv4(hx3)
hx3d = self.rebnconv3d(torch.cat((hx4,hx3),1))
hx2d = self.rebnconv2d(torch.cat((hx3d,hx2),1))
hx1d = self.rebnconv1d(torch.cat((hx2d,hx1),1))
return hx1d + hxin
class myrebnconv(nn.Module):
def __init__(self, in_ch=3,
out_ch=1,
kernel_size=3,
stride=1,
padding=1,
dilation=1,
groups=1):
super(myrebnconv,self).__init__()
self.conv = nn.Conv2d(in_ch,
out_ch,
kernel_size=kernel_size,
stride=stride,
padding=padding,
dilation=dilation,
groups=groups)
self.bn = nn.BatchNorm2d(out_ch)
self.rl = nn.ReLU(inplace=True)
def forward(self,x):
return self.rl(self.bn(self.conv(x)))
class BriaRMBG(nn.Module):
def __init__(self,in_ch=3,out_ch=1):
super(BriaRMBG,self).__init__()
self.conv_in = nn.Conv2d(in_ch,64,3,stride=2,padding=1)
self.pool_in = nn.MaxPool2d(2,stride=2,ceil_mode=True)
self.stage1 = RSU7(64,32,64)
self.pool12 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
self.stage2 = RSU6(64,32,128)
self.pool23 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
self.stage3 = RSU5(128,64,256)
self.pool34 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
self.stage4 = RSU4(256,128,512)
self.pool45 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
self.stage5 = RSU4F(512,256,512)
self.pool56 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
self.stage6 = RSU4F(512,256,512)
# decoder
self.stage5d = RSU4F(1024,256,512)
self.stage4d = RSU4(1024,128,256)
self.stage3d = RSU5(512,64,128)
self.stage2d = RSU6(256,32,64)
self.stage1d = RSU7(128,16,64)
self.side1 = nn.Conv2d(64,out_ch,3,padding=1)
self.side2 = nn.Conv2d(64,out_ch,3,padding=1)
self.side3 = nn.Conv2d(128,out_ch,3,padding=1)
self.side4 = nn.Conv2d(256,out_ch,3,padding=1)
self.side5 = nn.Conv2d(512,out_ch,3,padding=1)
self.side6 = nn.Conv2d(512,out_ch,3,padding=1)
# self.outconv = nn.Conv2d(6*out_ch,out_ch,1)
def forward(self,x):
hx = x
hxin = self.conv_in(hx)
#hx = self.pool_in(hxin)
#stage 1
hx1 = self.stage1(hxin)
hx = self.pool12(hx1)
#stage 2
hx2 = self.stage2(hx)
hx = self.pool23(hx2)
#stage 3
hx3 = self.stage3(hx)
hx = self.pool34(hx3)
#stage 4
hx4 = self.stage4(hx)
hx = self.pool45(hx4)
#stage 5
hx5 = self.stage5(hx)
hx = self.pool56(hx5)
#stage 6
hx6 = self.stage6(hx)
hx6up = _upsample_like(hx6,hx5)
#-------------------- decoder --------------------
hx5d = self.stage5d(torch.cat((hx6up,hx5),1))
hx5dup = _upsample_like(hx5d,hx4)
hx4d = self.stage4d(torch.cat((hx5dup,hx4),1))
hx4dup = _upsample_like(hx4d,hx3)
hx3d = self.stage3d(torch.cat((hx4dup,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.stage2d(torch.cat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.stage1d(torch.cat((hx2dup,hx1),1))
#side output
d1 = self.side1(hx1d)
d1 = _upsample_like(d1,x)
d2 = self.side2(hx2d)
d2 = _upsample_like(d2,x)
d3 = self.side3(hx3d)
d3 = _upsample_like(d3,x)
d4 = self.side4(hx4d)
d4 = _upsample_like(d4,x)
d5 = self.side5(hx5d)
d5 = _upsample_like(d5,x)
d6 = self.side6(hx6)
d6 = _upsample_like(d6,x)
return [F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)],[hx1d,hx2d,hx3d,hx4d,hx5d,hx6]
+16 -4
View File
@@ -199,7 +199,6 @@ def blend_linear_light(background_image:Image, layer_image:Image) -> Image:
img = img * (1 - mask_2) + mask_2
return cv22pil(ski2cv2(img))
def blend_hard_mix(background_image:Image, layer_image:Image) -> Image:
img_1 = cv22ski(pil2cv2(background_image))
img_2 = cv22ski(pil2cv2(layer_image))
@@ -209,7 +208,6 @@ def blend_hard_mix(background_image:Image, layer_image:Image) -> Image:
img = img * mask
return cv22pil(ski2cv2(img))
def shift_image(image:Image, distance_x:int, distance_y:int, background_color:str='#000000', cyclic:bool=False) -> Image:
width = image.width
height = image.height
@@ -273,8 +271,6 @@ def chop_image(background_image:Image, layer_image:Image, blend_mode:str, opacit
ret_image = blend_linear_light(background_image, layer_image)
if blend_mode == 'hard_mix':
ret_image = blend_hard_mix(background_image, layer_image)
# opacity
if opacity == 0:
ret_image = background_image
@@ -624,6 +620,22 @@ def image_beauty(image:Image, level:int=50) -> Image:
ret_image = cv2.cvtColor(img_bit, cv2.COLOR_BGR2RGB)
return cv22pil(ret_image)
def imagebatch2imagelist(image:torch.Tensor) -> torch.Tensor:
images = [image[i:i + 1, ...] for i in range(image.shape[0])]
return images
def imagelist2imagebatch(images:torch.Tensor) -> torch.Tensor:
if len(images) <= 1:
return (images,)
else:
image1 = images[0]
for image2 in images[1:]:
if image1.shape[1:] != image2.shape[1:]:
image2 = comfy.utils.common_upscale(image2.movedim(-1, 1), image1.shape[2], image1.shape[1], "lanczos", "center").movedim(1, -1)
image1 = torch.cat((image1, image2), dim=0)
return image1
'''Mask Functions'''
def expand_mask(mask:torch.Tensor, grow:int, blur:int) -> torch.Tensor:
+61
View File
@@ -0,0 +1,61 @@
import copy
from pymatting import *
from .imagefunc import *
class PixelSpread:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(self):
return {
"required": {
"image": ("IMAGE", ), #
"invert_mask": ("BOOLEAN", {"default": False}), # 反转mask
},
"optional": {
"mask": ("MASK",), #
}
}
RETURN_TYPES = ("IMAGE", )
RETURN_NAMES = ("image", )
FUNCTION = 'pixel_spread'
CATEGORY = '😺dzNodes/LayerMask'
OUTPUT_NODE = True
def pixel_spread(self, image, invert_mask, mask=None):
_image = tensor2pil(image)
if _image.mode == 'RGBA':
_mask = _image.split()[-1]
else:
_mask = Image.new('L', _image.size, 'white')
if mask is not None:
if invert_mask:
mask = 1 - mask
_mask = mask2image(mask).convert('L')
image = pil2tensor(_image.convert('RGB'))
_mask = _mask.convert('RGB')
i_dup = copy.deepcopy(image.cpu().numpy().astype(np.float64))
a_dup = copy.deepcopy(pil2tensor(_mask).cpu().numpy().astype(np.float64))
fg = copy.deepcopy(image.cpu().numpy().astype(np.float64))
for index, image in enumerate(i_dup):
trimap = a_dup[index][:, :, 0] # convert to single channel
trimap = fix_trimap(trimap, 0.01, 0.99)
alpha = estimate_alpha_cf(image, trimap, laplacian_kwargs={"epsilon": 1e-6},
cg_kwargs={"maxiter": 100})
fg[index], _ = estimate_foreground_ml(image, np.array(alpha), return_background=True)
return (torch.from_numpy(fg.astype(np.float32)), # fg
)
NODE_CLASS_MAPPINGS = {
"LayerMask: PixelSpread": PixelSpread
}
NODE_DISPLAY_NAME_MAPPINGS = {
"LayerMask: PixelSpread": "LayerMask: PixelSpread"
}
+89
View File
@@ -0,0 +1,89 @@
import copy
import torch, os
import torch.nn.functional as F
import numpy as np
from PIL import Image
from pymatting import *
# from torchvision.transforms.functional import normalize
import torchvision.transforms.functional as TF
from .briarmbg import BriaRMBG
from .imagefunc import *
current_directory = os.path.dirname(os.path.abspath(__file__))
device = "cuda" if torch.cuda.is_available() else "cpu"
def load_model():
net = BriaRMBG()
model_path = os.path.join(os.path.dirname(current_directory), "RMBG-1.4/model.pth")
net.load_state_dict(torch.load(model_path, map_location=device))
net.to(device)
net.eval()
return net
class RemBgUltra:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
"detail_range": ("INT", {"default": 8, "min": 0, "max": 256, "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}),
},
"optional": {
}
}
RETURN_TYPES = ("IMAGE", "MASK", )
RETURN_NAMES = ("image", "mask", )
FUNCTION = "rembg_ultra"
CATEGORY = '😺dzNodes/LayerMask'
def rembg_ultra(self, image, detail_range, black_point, white_point):
rmbgmodel = load_model()
orig_image = tensor2pil(image).convert('RGB')
w,h = orig_image.size
im_np = np.array(orig_image.resize((1024, 1024), Image.BILINEAR))
im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2,0,1)
im_tensor = torch.unsqueeze(im_tensor,0)
im_tensor = torch.divide(im_tensor,255.0)
# im_tensor = normalize(im_tensor,[0.5,0.5,0.5],[1.0,1.0,1.0])
im_tensor = TF.normalize(im_tensor, [0.5, 0.5, 0.5], [1.0, 1.0, 1.0])
if torch.cuda.is_available():
im_tensor=im_tensor.cuda()
result=rmbgmodel(im_tensor)
result = torch.squeeze(F.interpolate(result[0][0], size=(h,w), mode='bilinear') ,0)
ma = torch.max(result)
mi = torch.min(result)
result = (result-mi)/(ma-mi)
im_array = (result*255).cpu().data.numpy().astype(np.uint8)
_mask = Image.fromarray(np.squeeze(im_array)).convert('L')
# ultra edge process
d = detail_range * 2 + 1
i_dup = copy.deepcopy(image.cpu().numpy().astype(np.float64))
a_dup = copy.deepcopy(pil2tensor(_mask.convert('RGB')).cpu().numpy().astype(np.float64))
for index, img in enumerate(i_dup):
trimap = a_dup[index][:,:,0] # convert to single channel
if detail_range > 0:
trimap = cv2.GaussianBlur(trimap, (d, d), 0)
trimap = fix_trimap(trimap, black_point, white_point)
alpha = estimate_alpha_cf(img, trimap, laplacian_kwargs={"epsilon": 1e-6},
cg_kwargs={"maxiter": 500})
a_dup[index] = np.stack([alpha, alpha, alpha], axis=-1) # convert back to rgb
_mask = tensor2pil(torch.from_numpy(a_dup.astype(np.float32))) # alpha
ret_image = RGB2RGBA(orig_image, _mask.convert('L'))
return (pil2tensor(ret_image), image2mask(_mask),)
NODE_CLASS_MAPPINGS = {
"LayerMask: RemBgUltra": RemBgUltra,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"LayerMask: RemBgUltra": "LayerMask: RemBgUltra",
}
+3 -1
View File
@@ -4,4 +4,6 @@ torch
matplotlib
Scipy
opencv-python
scikit_image
scikit_image
opencv-contrib-python
pymatting
Binary file not shown.

After

Width:  |  Height:  |  Size: 1.1 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.1 MiB

+234
View File
@@ -0,0 +1,234 @@
{
"last_node_id": 46,
"last_link_id": 82,
"nodes": [
{
"id": 39,
"type": "LoadImage",
"pos": [
490,
660
],
"size": {
"0": 330.9508972167969,
"1": 298.0813903808594
},
"flags": {},
"order": 0,
"mode": 0,
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
74,
79
],
"shape": 3,
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": null,
"shape": 3
}
],
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"1344x768_hair (2).png",
"image"
]
},
{
"id": 44,
"type": "LayerMask: RemBgUltra",
"pos": [
869,
784
],
"size": {
"0": 315,
"1": 126
},
"flags": {},
"order": 1,
"mode": 0,
"inputs": [
{
"name": "image",
"type": "IMAGE",
"link": 74
}
],
"outputs": [
{
"name": "image",
"type": "IMAGE",
"links": null,
"shape": 3
},
{
"name": "mask",
"type": "MASK",
"links": [
75,
80
],
"shape": 3,
"slot_index": 1
}
],
"properties": {
"Node name for S&R": "LayerMask: RemBgUltra"
},
"widgets_values": [
10,
0.01,
0.99
]
},
{
"id": 41,
"type": "LayerMask: MaskPreview",
"pos": [
1247,
802
],
"size": {
"0": 309.01806640625,
"1": 246
},
"flags": {},
"order": 2,
"mode": 0,
"inputs": [
{
"name": "mask",
"type": "MASK",
"link": 75
}
],
"properties": {
"Node name for S&R": "LayerMask: MaskPreview"
}
},
{
"id": 42,
"type": "PreviewImage",
"pos": [
1624,
721
],
"size": {
"0": 314.5066223144531,
"1": 246
},
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 81
}
],
"properties": {
"Node name for S&R": "PreviewImage"
}
},
{
"id": 46,
"type": "LayerMask: PixelSpread",
"pos": [
1238,
665
],
"size": {
"0": 315,
"1": 78
},
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"name": "image",
"type": "IMAGE",
"link": 79
},
{
"name": "mask",
"type": "MASK",
"link": 80
}
],
"outputs": [
{
"name": "image",
"type": "IMAGE",
"links": [
81
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "LayerMask: PixelSpread"
},
"widgets_values": [
false
]
}
],
"links": [
[
74,
39,
0,
44,
0,
"IMAGE"
],
[
75,
44,
1,
41,
0,
"MASK"
],
[
79,
39,
0,
46,
0,
"IMAGE"
],
[
80,
44,
1,
46,
1,
"MASK"
],
[
81,
46,
0,
42,
0,
"IMAGE"
]
],
"groups": [],
"config": {},
"extra": {},
"version": 0.4
}
+697
View File
@@ -0,0 +1,697 @@
{
"last_node_id": 37,
"last_link_id": 70,
"nodes": [
{
"id": 28,
"type": "PreviewImage",
"pos": [
1863,
741
],
"size": {
"0": 629.1258544921875,
"1": 378.28240966796875
},
"flags": {},
"order": 11,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 47
}
],
"properties": {
"Node name for S&R": "PreviewImage"
}
},
{
"id": 30,
"type": "LayerMask: MaskPreview",
"pos": [
2512,
737
],
"size": {
"0": 629.4749145507812,
"1": 381.56378173828125
},
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"name": "mask",
"type": "MASK",
"link": 52
}
],
"properties": {
"Node name for S&R": "LayerMask: MaskPreview"
}
},
{
"id": 33,
"type": "PreviewImage",
"pos": [
1865.9044634155273,
1310.3026451171872
],
"size": {
"0": 629.1258544921875,
"1": 378.28240966796875
},
"flags": {},
"order": 12,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 60
}
],
"properties": {
"Node name for S&R": "PreviewImage"
}
},
{
"id": 34,
"type": "LayerMask: MaskPreview",
"pos": [
2520.904463415527,
1310.3026451171872
],
"size": {
"0": 629.4749145507812,
"1": 381.56378173828125
},
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "mask",
"type": "MASK",
"link": 61
}
],
"properties": {
"Node name for S&R": "LayerMask: MaskPreview"
}
},
{
"id": 27,
"type": "LayerUtility: ImageBlendAdvance",
"pos": [
1490,
790
],
"size": {
"0": 315,
"1": 338
},
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "background_image",
"type": "IMAGE",
"link": 57
},
{
"name": "layer_image",
"type": "IMAGE",
"link": 67
},
{
"name": "layer_mask",
"type": "MASK",
"link": 45
}
],
"outputs": [
{
"name": "image",
"type": "IMAGE",
"links": [
47
],
"shape": 3,
"slot_index": 0
},
{
"name": "mask",
"type": "MASK",
"links": null,
"shape": 3
}
],
"properties": {
"Node name for S&R": "LayerUtility: ImageBlendAdvance"
},
"widgets_values": [
false,
"normal",
100,
50,
50,
"None",
1,
1,
0,
"lanczos",
0
]
},
{
"id": 15,
"type": "LayerUtility: ImageBlendAdvance",
"pos": [
1490,
1280
],
"size": {
"0": 315,
"1": 338
},
"flags": {},
"order": 10,
"mode": 0,
"inputs": [
{
"name": "background_image",
"type": "IMAGE",
"link": 58
},
{
"name": "layer_image",
"type": "IMAGE",
"link": 70
},
{
"name": "layer_mask",
"type": "MASK",
"link": 50
}
],
"outputs": [
{
"name": "image",
"type": "IMAGE",
"links": [
60
],
"shape": 3,
"slot_index": 0
},
{
"name": "mask",
"type": "MASK",
"links": null,
"shape": 3
}
],
"properties": {
"Node name for S&R": "LayerUtility: ImageBlendAdvance"
},
"widgets_values": [
false,
"normal",
100,
50,
50,
"None",
1,
1,
0,
"lanczos",
0
]
},
{
"id": 36,
"type": "Reroute",
"pos": [
1039,
1096
],
"size": [
75,
26
],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"name": "",
"type": "*",
"link": 65
}
],
"outputs": [
{
"name": "",
"type": "IMAGE",
"links": [
66,
67
],
"slot_index": 0
}
],
"properties": {
"showOutputText": false,
"horizontal": false
}
},
{
"id": 26,
"type": "Image Remove Background Rembg (mtb)",
"pos": [
1139,
1067
],
"size": {
"0": 294,
"1": 230
},
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"name": "image",
"type": "IMAGE",
"link": 66
}
],
"outputs": [
{
"name": "Image (rgba)",
"type": "IMAGE",
"links": null,
"shape": 3
},
{
"name": "Mask",
"type": "MASK",
"links": [
45,
52
],
"shape": 3,
"slot_index": 1
},
{
"name": "Image",
"type": "IMAGE",
"links": null,
"shape": 3
}
],
"properties": {
"Node name for S&R": "Image Remove Background Rembg (mtb)"
},
"widgets_values": [
false,
240,
10,
10,
false,
"#000000"
]
},
{
"id": 2,
"type": "LoadImage",
"pos": [
957,
677
],
"size": {
"0": 224.95281982421875,
"1": 314
},
"flags": {},
"order": 0,
"mode": 0,
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
62,
65
],
"shape": 3,
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": null,
"shape": 3
}
],
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"1344x768_hair (4).png",
"image"
]
},
{
"id": 16,
"type": "LoadImage",
"pos": [
1220,
676
],
"size": {
"0": 235.49658203125,
"1": 314
},
"flags": {},
"order": 1,
"mode": 0,
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
57,
58
],
"shape": 3,
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": null,
"shape": 3
}
],
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"1344x768_beach (2).png",
"image"
]
},
{
"id": 35,
"type": "Reroute",
"pos": [
1015,
1343
],
"size": [
75,
26
],
"flags": {},
"order": 2,
"mode": 0,
"inputs": [
{
"name": "",
"type": "*",
"link": 62
}
],
"outputs": [
{
"name": "",
"type": "IMAGE",
"links": [
64,
68
],
"slot_index": 0
}
],
"properties": {
"showOutputText": false,
"horizontal": false
}
},
{
"id": 37,
"type": "LayerMask: PixelSpread",
"pos": [
1126,
1352
],
"size": {
"0": 315,
"1": 78
},
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"name": "image",
"type": "IMAGE",
"link": 68
},
{
"name": "mask",
"type": "MASK",
"link": 69
}
],
"outputs": [
{
"name": "image",
"type": "IMAGE",
"links": [
70
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "LayerMask: PixelSpread"
},
"widgets_values": [
false
]
},
{
"id": 29,
"type": "LayerMask: RemBgUltra",
"pos": [
1124,
1484
],
"size": {
"0": 315,
"1": 126
},
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{
"name": "image",
"type": "IMAGE",
"link": 64
}
],
"outputs": [
{
"name": "image",
"type": "IMAGE",
"links": null,
"shape": 3
},
{
"name": "mask",
"type": "MASK",
"links": [
50,
61,
69
],
"shape": 3,
"slot_index": 1
}
],
"properties": {
"Node name for S&R": "LayerMask: RemBgUltra"
},
"widgets_values": [
10,
0.01,
0.99
]
}
],
"links": [
[
45,
26,
1,
27,
2,
"MASK"
],
[
47,
27,
0,
28,
0,
"IMAGE"
],
[
50,
29,
1,
15,
2,
"MASK"
],
[
52,
26,
1,
30,
0,
"MASK"
],
[
57,
16,
0,
27,
0,
"IMAGE"
],
[
58,
16,
0,
15,
0,
"IMAGE"
],
[
60,
15,
0,
33,
0,
"IMAGE"
],
[
61,
29,
1,
34,
0,
"MASK"
],
[
62,
2,
0,
35,
0,
"*"
],
[
64,
35,
0,
29,
0,
"IMAGE"
],
[
65,
2,
0,
36,
0,
"*"
],
[
66,
36,
0,
26,
0,
"IMAGE"
],
[
67,
36,
0,
27,
1,
"IMAGE"
],
[
68,
35,
0,
37,
0,
"IMAGE"
],
[
69,
29,
1,
37,
1,
"MASK"
],
[
70,
37,
0,
15,
1,
"IMAGE"
]
],
"groups": [
{
"title": "Previous RemBg",
"bounding": [
1848,
642,
1306,
498
],
"color": "#3f789e",
"font_size": 30
},
{
"title": "RemBgUltra + PixelSpread",
"bounding": [
1852,
1215,
1312,
497
],
"color": "#b06634",
"font_size": 30
}
],
"config": {},
"extra": {},
"version": 0.4
}