diff --git a/README.MD b/README.MD index 714f05a..8d59408 100644 --- a/README.MD +++ b/README.MD @@ -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. +### RemBgUltra +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. + +### PixelSpread +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. + + ### MaskGrow 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: diff --git a/README_CN.MD b/README_CN.MD index 266736b..9b343ed 100644 --- a/README_CN.MD +++ b/README_CN.MD @@ -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 +### RemBgUltra +ȥ֮ǰ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: Եɫֵ + +### PixelSpread +ͼֱԵֽԤЧͼϳɵıԵ +![image](image/pixel_spread_example.png) + +ڵѡ˵: +![image](image/pixel_spread_node.png) +* invert_mask: Ƿת֡ + + ### MaskGrow 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 diff --git a/RMBG-1.4/put model here.txt b/RMBG-1.4/put model here.txt new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/RMBG-1.4/put model here.txt @@ -0,0 +1 @@ + diff --git a/image/layermask_nodes.png b/image/layermask_nodes.png index e0e65db..713fe18 100644 Binary files a/image/layermask_nodes.png and b/image/layermask_nodes.png differ diff --git a/image/pixel_spread_example.png b/image/pixel_spread_example.png new file mode 100644 index 0000000..e026134 Binary files /dev/null and b/image/pixel_spread_example.png differ diff --git a/image/pixel_spread_node.png b/image/pixel_spread_node.png new file mode 100644 index 0000000..176c060 Binary files /dev/null and b/image/pixel_spread_node.png differ diff --git a/image/rembg_ultra_example.png b/image/rembg_ultra_example.png new file mode 100644 index 0000000..0189831 Binary files /dev/null and b/image/rembg_ultra_example.png differ diff --git a/image/rembg_ultra_node.png b/image/rembg_ultra_node.png new file mode 100644 index 0000000..df83016 Binary files /dev/null and b/image/rembg_ultra_node.png differ diff --git a/py/briarmbg.py b/py/briarmbg.py new file mode 100644 index 0000000..647bdc0 --- /dev/null +++ b/py/briarmbg.py @@ -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] + diff --git a/py/imagefunc.py b/py/imagefunc.py index 3b11af9..e42e9ac 100644 --- a/py/imagefunc.py +++ b/py/imagefunc.py @@ -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: diff --git a/py/pixel_spread.py b/py/pixel_spread.py new file mode 100644 index 0000000..8784566 --- /dev/null +++ b/py/pixel_spread.py @@ -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" +} \ No newline at end of file diff --git a/py/rembg_ultra.py b/py/rembg_ultra.py new file mode 100644 index 0000000..ac5bde7 --- /dev/null +++ b/py/rembg_ultra.py @@ -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", +} diff --git a/requirements.txt b/requirements.txt index 2bd74b9..91bd24a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,4 +4,6 @@ torch matplotlib Scipy opencv-python -scikit_image \ No newline at end of file +scikit_image +opencv-contrib-python +pymatting \ No newline at end of file diff --git a/workflow/1344x768_beach.png b/workflow/1344x768_beach.png new file mode 100644 index 0000000..cdc1b69 Binary files /dev/null and b/workflow/1344x768_beach.png differ diff --git a/workflow/1344x768_hair.png b/workflow/1344x768_hair.png new file mode 100644 index 0000000..48a9d93 Binary files /dev/null and b/workflow/1344x768_hair.png differ diff --git a/workflow/pixel_spread_example.json b/workflow/pixel_spread_example.json new file mode 100644 index 0000000..8adef58 --- /dev/null +++ b/workflow/pixel_spread_example.json @@ -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 +} \ No newline at end of file diff --git a/workflow/rembg_ultra_example.json b/workflow/rembg_ultra_example.json new file mode 100644 index 0000000..182de1b --- /dev/null +++ b/workflow/rembg_ultra_example.json @@ -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 +} \ No newline at end of file