40 changed files with 384 additions and 12802 deletions
-3
View File
@@ -1,3 +0,0 @@
CLIPDROP_API_KEY=
COMFY_PYTHON_PATH=/home/ubuntu/.conda/envs/comfy/bin/python
PHOTOROOM_API_KEY=3603b83dfa1846bc3c7270ead7876
+1 -7
View File
@@ -5,11 +5,5 @@ venv
.DS_Store
checkpoints/
checkpoint/
.env
.pth
cloth-segmentation/model/cloth_segm.pth
dwpose/keypoints/
huggingface/
safetychecker/model.safetensors
dwpose/keypoints/
-1244
View File
File diff suppressed because it is too large Load Diff
+377 -2130
View File
File diff suppressed because it is too large Load Diff
+2 -22
View File
@@ -88,29 +88,9 @@ def get_palette(num_cls):
return palette
def download_model_restore(model_restore_path):
import os
import gdown
from pathlib import Path
# Ensure the directory for the model path exists
os.makedirs(os.path.dirname(model_restore_path), exist_ok=True)
# Check if the model file already exists
if not Path(model_restore_path).is_file():
print("Model file does not exist, downloading...")
# Google Drive ID for the file
# file_id = '1ruJg4lqR_jgQPj-9K0PP-L2vJERYOxLP'
file_id="1AVVLm1LxOs3W1Fp_GLefIz6fdWfEYg88"
gdown.download(id=file_id, output=model_restore_path, quiet=False)
print("Download complete.")
else:
print("Model file already exists.")
def main():
args = get_arguments()
gpus = [int(i) for i in args.gpu.split(',')]
assert len(gpus) == 1
if not args.gpu == 'None':
@@ -122,7 +102,7 @@ def main():
print("Evaluating total class number {} with {}".format(num_classes, label))
model = networks.init_model('resnet101', num_classes=num_classes, pretrained=None)
download_model_restore(args.model_restore)
state_dict = torch.load(args.model_restore)['state_dict']
from collections import OrderedDict
new_state_dict = OrderedDict()
-47
View File
@@ -1,47 +0,0 @@
import PIL
import torch
import os
from process import load_seg_model, get_palette, generate_mask
device = 'cuda'
def initialize_and_load_models():
checkpoint_path = 'model/cloth_segm.pth'
net = load_seg_model(checkpoint_path, device=device)
return net
net = initialize_and_load_models()
def run(img, image_id, output_dir):
palette = get_palette(4)
mask0, mask1, mask2, cloth_seg = generate_mask(img, net=net, device=device, image_id=image_id, output_dir=output_dir)
return mask0, mask1, mask2, cloth_seg
INPUT_PATH = "./input/"
OUTPUT_PATH = "./output/"
import os
mask0_paths = []
mask1_paths = []
mask2_paths = []
cloth_paths = []
for idx, cur_image in enumerate(os.listdir(INPUT_PATH)):
img = PIL.Image.open(INPUT_PATH + cur_image)
mask0, mask1, mask2, cloth_seg = run(img, image_id=idx, output_dir=OUTPUT_PATH)
# Save masks and cloth_seg with unique names (already saved in generate_mask)
mask0_path = os.path.join(OUTPUT_PATH, f"{idx}__mask0.png")
mask1_path = os.path.join(OUTPUT_PATH, f"{idx}__mask1.png")
mask2_path = os.path.join(OUTPUT_PATH, f"{idx}__mask2.png")
cloth_path = os.path.join(OUTPUT_PATH, f"{idx}__extracted_garment.png")
mask0_paths.append(mask0_path)
mask1_paths.append(mask1_path)
mask2_paths.append(mask2_path)
cloth_paths.append(cloth_path)
print("Mask0 batch:", mask0_paths)
print("Mask1 batch:", mask1_paths)
print("Mask2 batch:", mask2_paths)
print("Garment batch:", cloth_paths)
-1
View File
@@ -1 +0,0 @@
/*upload model */
-560
View File
@@ -1,560 +0,0 @@
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):
super(REBNCONV, self).__init__()
self.conv_s1 = nn.Conv2d(
in_ch, out_ch, 3, padding=1 * dirate, dilation=1 * dirate
)
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.upsample(src, size=tar.shape[2:], mode="bilinear")
return src
### RSU-7 ###
class RSU7(nn.Module): # UNet07DRES(nn.Module):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU7, 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.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):
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))
"""
del hx1, hx2, hx3, hx4, hx5, hx6, hx7
del hx6d, hx5d, hx3d, hx2d
del hx2dup, hx3dup, hx4dup, hx5dup, hx6dup
"""
return hx1d + hxin
### RSU-6 ###
class RSU6(nn.Module): # UNet06DRES(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))
"""
del hx1, hx2, hx3, hx4, hx5, hx6
del hx5d, hx4d, hx3d, hx2d
del hx2dup, hx3dup, hx4dup, hx5dup
"""
return hx1d + hxin
### RSU-5 ###
class RSU5(nn.Module): # UNet05DRES(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))
"""
del hx1, hx2, hx3, hx4, hx5
del hx4d, hx3d, hx2d
del hx2dup, hx3dup, hx4dup
"""
return hx1d + hxin
### RSU-4 ###
class RSU4(nn.Module): # UNet04DRES(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))
"""
del hx1, hx2, hx3, hx4
del hx3d, hx2d
del hx2dup, hx3dup
"""
return hx1d + hxin
### RSU-4F ###
class RSU4F(nn.Module): # UNet04FRES(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))
"""
del hx1, hx2, hx3, hx4
del hx3d, hx2d
"""
return hx1d + hxin
##### U^2-Net ####
class U2NET(nn.Module):
def __init__(self, in_ch=3, out_ch=1):
super(U2NET, self).__init__()
self.stage1 = RSU7(in_ch, 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
# stage 1
hx1 = self.stage1(hx)
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)
d2 = self.side2(hx2d)
d2 = _upsample_like(d2, d1)
d3 = self.side3(hx3d)
d3 = _upsample_like(d3, d1)
d4 = self.side4(hx4d)
d4 = _upsample_like(d4, d1)
d5 = self.side5(hx5d)
d5 = _upsample_like(d5, d1)
d6 = self.side6(hx6)
d6 = _upsample_like(d6, d1)
d0 = self.outconv(torch.cat((d1, d2, d3, d4, d5, d6), 1))
"""
del hx1, hx2, hx3, hx4, hx5, hx6
del hx5d, hx4d, hx3d, hx2d, hx1d
del hx6up, hx5dup, hx4dup, hx3dup, hx2dup
"""
return d0, d1, d2, d3, d4, d5, d6
### U^2-Net small ###
class U2NETP(nn.Module):
def __init__(self, in_ch=3, out_ch=1):
super(U2NETP, self).__init__()
self.stage1 = RSU7(in_ch, 16, 64)
self.pool12 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
self.stage2 = RSU6(64, 16, 64)
self.pool23 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
self.stage3 = RSU5(64, 16, 64)
self.pool34 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
self.stage4 = RSU4(64, 16, 64)
self.pool45 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
self.stage5 = RSU4F(64, 16, 64)
self.pool56 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
self.stage6 = RSU4F(64, 16, 64)
# decoder
self.stage5d = RSU4F(128, 16, 64)
self.stage4d = RSU4(128, 16, 64)
self.stage3d = RSU5(128, 16, 64)
self.stage2d = RSU6(128, 16, 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(64, out_ch, 3, padding=1)
self.side4 = nn.Conv2d(64, out_ch, 3, padding=1)
self.side5 = nn.Conv2d(64, out_ch, 3, padding=1)
self.side6 = nn.Conv2d(64, out_ch, 3, padding=1)
self.outconv = nn.Conv2d(6 * out_ch, out_ch, 1)
def forward(self, x):
hx = x
# stage 1
hx1 = self.stage1(hx)
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)
d2 = self.side2(hx2d)
d2 = _upsample_like(d2, d1)
d3 = self.side3(hx3d)
d3 = _upsample_like(d3, d1)
d4 = self.side4(hx4d)
d4 = _upsample_like(d4, d1)
d5 = self.side5(hx5d)
d5 = _upsample_like(d5, d1)
d6 = self.side6(hx6)
d6 = _upsample_like(d6, d1)
d0 = self.outconv(torch.cat((d1, d2, d3, d4, d5, d6), 1))
return d0, d1, d2, d3, d4, d5, d6
-12
View File
@@ -1,12 +0,0 @@
import os.path as osp
import os
class parser(object):
def __init__(self):
self.output = "./output" # output image folder path
self.logs_dir = './logs'
self.device = 'cuda:0'
opt = parser()
-257
View File
@@ -1,257 +0,0 @@
from network import U2NET
import os
from PIL import Image
import cv2
import gdown
import argparse
import numpy as np
import torch
import torch.nn.functional as F
import torchvision.transforms as transforms
from collections import OrderedDict
from options import opt
def load_checkpoint(model, checkpoint_path):
if not os.path.exists(checkpoint_path):
print("----No checkpoints at given path----")
return
model_state_dict = torch.load(checkpoint_path, map_location=torch.device("cpu"))
new_state_dict = OrderedDict()
for k, v in model_state_dict.items():
name = k[7:] # remove `module.`
new_state_dict[name] = v
model.load_state_dict(new_state_dict)
print("----checkpoints loaded from path: {}----".format(checkpoint_path))
return model
def get_palette(num_cls):
""" Returns the color map for visualizing the segmentation mask.
Args:
num_cls: Number of classes
Returns:
The color map
"""
n = num_cls
palette = [0] * (n * 3)
for j in range(0, n):
lab = j
palette[j * 3 + 0] = 0
palette[j * 3 + 1] = 0
palette[j * 3 + 2] = 0
i = 0
while lab:
palette[j * 3 + 0] |= (((lab >> 0) & 1) << (7 - i))
palette[j * 3 + 1] |= (((lab >> 1) & 1) << (7 - i))
palette[j * 3 + 2] |= (((lab >> 2) & 1) << (7 - i))
i += 1
lab >>= 3
return palette
class Normalize_image(object):
"""Normalize given tensor into given mean and standard dev
Args:
mean (float): Desired mean to substract from tensors
std (float): Desired std to divide from tensors
"""
def __init__(self, mean, std):
assert isinstance(mean, (float))
if isinstance(mean, float):
self.mean = mean
if isinstance(std, float):
self.std = std
self.normalize_1 = transforms.Normalize(self.mean, self.std)
self.normalize_3 = transforms.Normalize([self.mean] * 3, [self.std] * 3)
self.normalize_18 = transforms.Normalize([self.mean] * 18, [self.std] * 18)
def __call__(self, image_tensor):
if image_tensor.shape[0] == 1:
return self.normalize_1(image_tensor)
elif image_tensor.shape[0] == 3:
return self.normalize_3(image_tensor)
elif image_tensor.shape[0] == 18:
return self.normalize_18(image_tensor)
else:
assert "Please set proper channels! Normlization implemented only for 1, 3 and 18"
def apply_transform(img):
transforms_list = []
transforms_list += [transforms.ToTensor()]
transforms_list += [Normalize_image(0.5, 0.5)]
transform_rgb = transforms.Compose(transforms_list)
return transform_rgb(img)
from PIL import Image
def generate_mask(input_image, net, device='cpu', image_id=None, output_dir=None):
img = input_image
img_size = img.size
img = img.resize((768, 768), Image.BICUBIC)
image_tensor = apply_transform(img)
image_tensor = torch.unsqueeze(image_tensor, 0)
# Allow output_dir override for batch processing
if output_dir is None:
output_dir = os.path.join(opt.output, 'extracted_garment')
os.makedirs(output_dir, exist_ok=True)
with torch.no_grad():
output_tensor = net(image_tensor.to(device))
output_tensor = F.log_softmax(output_tensor[0], dim=1)
output_tensor = torch.max(output_tensor, dim=1, keepdim=True)[1]
output_tensor = torch.squeeze(output_tensor, dim=0)
output_arr = output_tensor.cpu().numpy()
# Create and save individual masks for classes 1, 2, 3
classes_of_interest = [1, 2, 3]
mask_imgs = []
for idx, cls in enumerate(classes_of_interest):
mask = np.zeros_like(output_arr, dtype=np.uint8)
mask[output_arr == cls] = 255
if mask.ndim > 2:
mask = mask.squeeze()
if mask.ndim != 2:
raise ValueError(f"mask{idx} must be a 2-dimensional array")
mask_img = Image.fromarray(mask, mode='L').resize(img_size, Image.BICUBIC)
# Save with unique name if image_id is provided
if image_id is not None:
mask_path = os.path.join(output_dir, f'{image_id}__mask{idx}.png')
else:
mask_path = os.path.join(output_dir, f'mask{idx}.png')
mask_img.save(mask_path, format="PNG")
print(f"Saved mask{idx} at: {mask_path}")
mask_imgs.append(mask_img)
# Create a binary mask where selected classes are 1, others are 0
binary_mask = np.zeros_like(output_arr, dtype=np.uint8)
for cls in classes_of_interest:
binary_mask[output_arr == cls] = 255
if binary_mask.ndim > 2:
binary_mask = binary_mask.squeeze()
if binary_mask.ndim != 2:
raise ValueError("binary_mask must be a 2-dimensional array")
binary_mask_img = Image.fromarray(binary_mask, mode='L').resize(img_size, Image.BICUBIC)
# Create an RGBA image for the output
extracted_garment = Image.new("RGBA", img_size)
original_img = img.resize(img_size)
extracted_garment.paste(original_img, mask=binary_mask_img)
# Save the garment image with transparency
if image_id is not None:
garment_path = os.path.join(output_dir, f'{image_id}__extracted_garment.png')
else:
garment_path = os.path.join(output_dir, 'extracted_garment.png')
extracted_garment.save(garment_path, format="PNG")
print(f"Saved extracted garment at: {garment_path}")
return (*mask_imgs, extracted_garment)
# def generate_mask(input_image, net, device='cpu'):
# img = input_image
# img_size = img.size
# img = img.resize((768, 768), Image.BICUBIC)
# image_tensor = apply_transform(img)
# image_tensor = torch.unsqueeze(image_tensor, 0)
# output_dir = os.path.join(opt.output, 'extracted_garment')
# os.makedirs(output_dir, exist_ok=True)
# with torch.no_grad():
# output_tensor = net(image_tensor.to(device))
# output_tensor = F.log_softmax(output_tensor[0], dim=1)
# output_tensor = torch.max(output_tensor, dim=1, keepdim=True)[1]
# output_tensor = torch.squeeze(output_tensor, dim=0)
# output_arr = output_tensor.cpu().numpy()
# # Create a binary mask where selected classes are 1, others are 0
# binary_mask = np.zeros_like(output_arr, dtype=np.uint8)
# classes_of_interest = [1, 2, 3] # Modify this list according to your classes of interest
# for cls in classes_of_interest:
# binary_mask[output_arr == cls] = 255
# # Convert binary mask to a 3-channel image to use as a mask
# # Ensure binary_mask is 2D
# if binary_mask.ndim > 2:
# binary_mask = binary_mask.squeeze() # Removes single-dimensional entries from the shape
# if binary_mask.ndim != 2:
# raise ValueError("binary_mask must be a 2-dimensional array")
# binary_mask_img = Image.fromarray(binary_mask, mode='L').resize(img_size, Image.BICUBIC)
# binary_mask_3ch = binary_mask_img.convert('RGB') # Convert to RGB
# # Apply mask to the original image
# original_img = img.resize(img_size) # Resize the processed image back to original size
# extracted_garment = Image.new("RGB", original_img.size)
# extracted_garment.paste(original_img, mask=binary_mask_img)
# # Save the garment image
# garment_path = os.path.join(output_dir, 'extracted_garment.png')
# extracted_garment.save(garment_path)
# return extracted_garment
def check_or_download_model(file_path):
if not os.path.exists(file_path):
os.makedirs(os.path.dirname(file_path), exist_ok=True)
url = "https://drive.google.com/uc?export=download&id=1qVv720hAd11JSCuIVJuqfjCGolwb1H8o"
gdown.download(url, file_path, quiet=False)
print("Model downloaded successfully.")
else:
print("Model already exists.")
def load_seg_model(checkpoint_path, device='cpu'):
net = U2NET(in_ch=3, out_ch=4)
check_or_download_model(checkpoint_path)
net = load_checkpoint(net, checkpoint_path)
net = net.to(device)
net = net.eval()
return net
def main(args):
device = 'cuda:0' if args.cuda else 'cpu'
# Create an instance of your model
model = load_seg_model(args.checkpoint_path, device=device)
palette = get_palette(4)
img = Image.open(args.image).convert('RGB')
cloth_seg = generate_mask(img, net=model, palette=palette, device=device)
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Help to set arguments for Cloth Segmentation.')
parser.add_argument('--image', type=str, help='Path to the input image')
parser.add_argument('--cuda', action='store_true', help='Enable CUDA (default: False)')
parser.add_argument('--checkpoint_path', type=str, default='model/cloth_segm.pth', help='Path to the checkpoint file')
args = parser.parse_args()
main(args)
-189
View File
@@ -1,189 +0,0 @@
import cv2
import os
import torch
import numpy as np
def from_torch_image(image):
image = image.cpu().numpy() * 255.0
image = np.clip(image, 0, 255).astype(np.uint8)
return image
def to_torch_image(image):
image = image.astype(dtype=np.float32)
image /= 255.0
image = torch.from_numpy(image)
return image
def get_histogram(array):
array = array.flatten().astype(dtype=np.float64)
hist = np.histogram(array, bins=256, range=(0, 256))
array = hist[0].astype(dtype=np.float64)
array /= len(array)
return array
def get_limits(array, threshold_fraction):
array = get_histogram(array)
left_sum = 0
right_sum = 0
left_start = 0
right_start = len(array) - 1
for i in range(len(array)):
left_index = i
right_index = len(array) - i - 1
left_sum += array[left_index]
right_sum += array[right_index]
if left_sum < threshold_fraction:
left_start = left_index
if right_sum < threshold_fraction:
right_start = right_index
if (left_sum > threshold_fraction) and (right_sum
> threshold_fraction):
return (left_start, right_start)
def do_rescale(x, y1, y2, x1, x2):
x = x.astype(dtype=np.float64)
if x1 > x2:
x1, x2 = x2, x1
if y1 > y2:
y1, y2 = y2, y1
epsilon = 0.0001
y = (x - x1)
y /= (x2 - x1 + epsilon)
y *= (y2 - y1)
y += y1
y = np.clip(y, y1, y2)
for iy in range(y.shape[0]):
for ix in range(y.shape[1]):
if y[iy, ix] > 255:
print(iy, ix)
y = y.astype(dtype=np.uint8)
return y
class get_histogram_limits:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"luminosity_as_mask": ("MASK", ),
"threshold_fraction": ("FLOAT", {
"default": 0.001,
"min": 0.0,
"max": 0.5,
"step": 0.00001,
"round": 0.000001,
"display": "number"
}),
},
}
RETURN_TYPES = ("INT", "INT")
RETURN_NAMES = ("histogram lower limit (x1) as INT",
"histogram upper limit (x2) as INT")
FUNCTION = "test"
#OUTPUT_NODE = False
CATEGORY = "TRI3D"
def test(self, luminosity_as_mask, threshold_fraction):
luminosity_as_mask = from_torch_image(image=luminosity_as_mask)
(left_start,
right_start) = get_limits(array=luminosity_as_mask[0],
threshold_fraction=threshold_fraction)
return (left_start, right_start)
class simple_rescale_histogram:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"layer_as_mask": ("MASK", ),
"y1": ("INT", {
"default": 100,
"min": 0,
"max": 255,
"step": 1,
"display": "number"
}),
"y2": ("INT", {
"default": 200,
"min": 0,
"max": 255,
"step": 1,
"display": "number"
}),
"x1": ("INT", {
"default": 100,
"min": 0,
"max": 255,
"step": 1,
"display": "number"
}),
"x2": ("INT", {
"default": 200,
"min": 0,
"max": 255,
"step": 1,
"display": "number"
})
},
}
RETURN_TYPES = ("MASK", )
RETURN_NAMES = ("rescaled layer as MASK", )
FUNCTION = "test"
CATEGORY = "TRI3D"
def test(self, layer_as_mask, y1, y2, x1, x2):
layer_as_mask = from_torch_image(image=layer_as_mask[0])
layer_as_mask = do_rescale(x=layer_as_mask, y1=y1, y2=y2, x1=x1, x2=x2)
layer_as_mask = to_torch_image(image=layer_as_mask)
layer_as_mask = layer_as_mask.unsqueeze(0)
return (layer_as_mask, )
NODE_CLASS_MAPPINGS = {
"get_histogram_limits": get_histogram_limits,
'simple_rescale_histogram': simple_rescale_histogram
}
NODE_DISPLAY_NAME_MAPPINGS = {
"get_histogram_limits": "get_histogram_limits",
"simple_rescale_histogram": "simple_rescale_histogram"
}
+1 -3
View File
@@ -274,6 +274,4 @@ def switch_to_backpose(input_keypoints, input_width):
x,y = input_keypoints[i]
input_keypoints[i] = [input_width - x, y]
return input_keypoints
return input_keypoints
-203
View File
@@ -1,203 +0,0 @@
#!/usr/bin/python3
from PIL import Image, ImageOps, ImageSequence, ImageFile
from PIL.PngImagePlugin import PngInfo
import cv2
import hashlib
import json
import logging
import math
import numpy as np
import os
import random
import safetensors.torch
import sys
import time
import torch
import traceback
def load_image(path):
return torch.from_numpy(cv2.imread(
path, cv2.IMREAD_COLOR)).to(dtype=torch.float32) / 255.0
def do_stack(img1, img2):
dim = max(max(img1.shape[0], img2.shape[0]), img1.shape[1] + img2.shape[1])
out = torch.zeros((dim, dim, 3), dtype=img1.dtype, device=img1.device) + 1
diff1 = (out.shape[0] - img1.shape[0]) // 2
diff2 = (out.shape[0] - img2.shape[0]) // 2
part0 = 0
part1 = img1.shape[1]
part2 = img2.shape[1] + img1.shape[1]
out[diff1:diff1 + img1.shape[0], part0:part1, :] = img1
out[diff2:diff2 + img2.shape[0], part1:part2, :] = img2
return out
def save_image(image, outpath):
cv2.imwrite(outpath,
(image * 255).to(dtype=torch.uint8).detach().cpu().numpy())
class H_Stack_Images:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image_L": ("IMAGE", ),
"image_R": ("IMAGE", ),
},
}
RETURN_TYPES = ("IMAGE", )
FUNCTION = "test"
CATEGORY = "TRI3D"
def test(self, image_L, image_R):
return (do_stack(img1=image_L[0], img2=image_R[0]).unsqueeze(0), )
class SaveImage_absolute:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"images": ("IMAGE", {
"tooltip": "The images to save."
}),
"absolute_filename": ("STRING", {
"default":
"image.png",
"tooltip":
"The absolute path to the file to save."
})
},
}
RETURN_TYPES = ("STRING", )
RETURN_NAMES = ("text to control order", )
FUNCTION = "save_images"
OUTPUT_NODE = True
CATEGORY = "image"
DESCRIPTION = "Saves the input images to an absolute path."
def save_images(self, images, absolute_filename):
i = 255.0 * images[0].cpu().numpy()
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
img.save(absolute_filename)
return (absolute_filename, )
class SaveText_absolute:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"text": ("STRING", {
"multiline": True,
"dynamicPrompts": True,
"tooltip": "Text to be saved to the file."
}),
"absolute_filename": ("STRING", {
"default":
"image.txt",
"tooltip":
"The absolute path to the file to save."
})
},
"optional": {
"text_opt": ("STRING", {
"multiline":
True,
"dynamicPrompts":
True,
"tooltip":
"Text to provide order when necessary (to create work files after txt files)."
}),
}
}
RETURN_TYPES = ("STRING", )
RETURN_NAMES = ("same text as input", )
FUNCTION = "save_text"
OUTPUT_NODE = True
CATEGORY = "text"
DESCRIPTION = "Saves the input text to an absolute path."
def save_text(self, text, absolute_filename, text_opt=''):
open(absolute_filename, "w").write(text)
return (text, )
class Wait_And_Read_File:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"absolute_filename": ("STRING", {
"default":
"image.txt",
"tooltip":
"The absolute path to the file to read."
})
},
"optional": {
"text": ("STRING", {
"multiline":
True,
"dynamicPrompts":
True,
"tooltip":
"Text to provide order when necessary (to wait on done file)."
}),
}
}
RETURN_TYPES = ("STRING", )
RETURN_NAMES = ("text from file", )
FUNCTION = "read_text"
OUTPUT_NODE = True
CATEGORY = "text"
DESCRIPTION = "Saves the input text to an absolute path."
def read_text(self, absolute_filename, text=''):
while not os.path.exists(absolute_filename):
time.sleep(0.1)
res = open(absolute_filename, "r").read()
os.unlink(absolute_filename)
return (res, )
-202
View File
@@ -1,202 +0,0 @@
import numpy as np
import cv2
import math
import torch
def from_torch_image(image):
image = image.cpu().numpy() * 255.0
image = np.clip(image, 0, 255).astype(np.uint8)
return image
def to_torch_image(image):
image = image.astype(dtype=np.float32)
image /= 255.0
image = torch.from_numpy(image)
return image
def smooth_step_plain(x):
if x < -1:
return -1
elif x <= 1:
return math.sin(x * np.pi / 2.0)
else:
return 1
def smooth_step_np(x):
truths = np.logical_and(-1 < x, x < 1).astype(np.float32)
x1 = np.clip(x, -1, 1)
x2 = np.sin(x * np.pi / 2.0)
ret = (truths * x2) + ((1 - truths) * x1)
return ret
def smooth_step_stretch(x, a, b):
if b < a:
tmp = b
b = a
a = tmp
if a < 0:
a = 0
if b > 1:
b = 1
if a == b:
a = 0
b = 1
return smooth_step_np((2 * (x - a) / (b - a)) - 1)
def get_light_layer(image,
ref_r=255,
ref_g=255,
ref_b=255,
do_scale=True,
scale_a=0.0,
scale_b=1.0):
sqmax = 3 * 255 * 255
scalemax = math.sqrt(sqmax)
b = image[:, :, 0].astype(dtype=np.float32)
g = image[:, :, 1].astype(dtype=np.float32)
r = image[:, :, 2].astype(dtype=np.float32)
b2 = b * b
g2 = g * g
r2 = r * r
d2 = np.zeros(b2.shape, dtype=np.float32)
d2 += sqmax - b2 - g2 - r2
d = np.sqrt(d2)
ref_r2 = ref_r * ref_r
ref_g2 = ref_g * ref_g
ref_b2 = ref_b * ref_b
ref_d2 = sqmax - ref_r2 - ref_g2 - ref_b2
ref_d = math.sqrt(ref_d2)
dot = (b * ref_b) + (g * ref_g) + (r * ref_r) + (d * ref_d)
dot /= sqmax
if do_scale:
dot = smooth_step_stretch(x=dot, a=scale_a, b=scale_b)
dot *= 255
dot = dot.astype(np.uint8)
return dot
class main_light_layer():
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE", ),
"ref_r": ("INT", {
"default": 255,
"min": 0,
"max": 255,
"step": 1
}),
"ref_g": ("INT", {
"default": 255,
"min": 0,
"max": 255,
"step": 1
}),
"ref_b": ("INT", {
"default": 255,
"min": 0,
"max": 255,
"step": 1
}),
"do_scale": (["enable", "disable"], ),
"thresh_low": ("FLOAT", {
"default": 0.6,
"min": 0.0,
"max": 1.0,
"step": 0.01
}),
"thresh_high": ("FLOAT", {
"default": 1.0,
"min": 0.0,
"max": 1.0,
"step": 0.01
}),
},
}
FUNCTION = "run"
RETURN_TYPES = ("MASK", )
CATEGORY = "HackNode"
def run(
self,
image,
ref_r,
ref_g,
ref_b,
do_scale,
thresh_low,
thresh_high,
):
do_scale = (do_scale == "enable")
print('do_scale', do_scale)
image = from_torch_image(image)
print('image.shape', image.shape)
batch_size = image.shape[0]
print('batch_size', batch_size)
mask = []
for i in range(batch_size):
tmp_img = image[i]
print('tmp_img.shape', tmp_img.shape)
tmp_mask = get_light_layer(
tmp_img,
ref_b,
ref_g,
ref_r,
do_scale,
scale_a=thresh_low,
scale_b=thresh_high,
)
print('tmp_mask.shape', tmp_mask.shape)
mask.append(tmp_mask)
mask = np.array(mask)
mask = to_torch_image(mask)
print(mask.shape)
return (mask, )
NODE_CLASS_MAPPINGS = {
'main_light_layer': main_light_layer,
}
NODE_DISPLAY_NAME_MAPPINGS = {
'main_light_layer': 'main_light_layer',
}
-166
View File
@@ -1,166 +0,0 @@
from __future__ import annotations
from weakref import ref as WeakRef
from pathlib import Path
from tqdm import tqdm
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from transformers import CLIPImageProcessor, CLIPConfig, CLIPVisionModel, PreTrainedModel
from kornia.filters import box_blur
def cosine_similarity(image_embeds: Tensor, text_embeds: Tensor):
if image_embeds.dim() == 2 and text_embeds.dim() == 2:
image_embeds = image_embeds.unsqueeze(1)
return F.cosine_similarity(image_embeds, text_embeds, dim=-1)
class CLIPSafetyChecker(PreTrainedModel):
# https://huggingface.co/CompVis/stable-diffusion-safety-checker
# Adapted from:
# https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/safety_checker.py
config_class = CLIPConfig
_no_split_modules = ["CLIPEncoderLayer"]
def __init__(self, config: CLIPConfig):
super().__init__(config)
projdim = config.projection_dim
self.vision_model = CLIPVisionModel(config.vision_config)
self.visual_projection = nn.Linear(config.vision_config.hidden_size, projdim, bias=False)
self.concept_embeds = nn.Parameter(torch.ones(17, projdim), requires_grad=False)
self.special_care_embeds = nn.Parameter(torch.ones(3, projdim), requires_grad=False)
self.concept_embeds_weights = nn.Parameter(torch.ones(17), requires_grad=False)
self.special_care_embeds_weights = nn.Parameter(torch.ones(3), requires_grad=False)
def forward(self, clip_input, images: Tensor, sensitivity: float, alternate_image: Tensor):
with torch.no_grad():
image_batch = self.vision_model(clip_input)[1]
image_embeds = self.visual_projection(image_batch)
sensitivity = -0.1 + 0.14 * sensitivity
special_cos_dist = cosine_similarity(image_embeds, self.special_care_embeds)
special_scores_threshold = self.special_care_embeds_weights.unsqueeze(0)
special_scores = special_cos_dist - special_scores_threshold + sensitivity
if torch.any(special_scores > 0):
sensitivity = sensitivity + 0.01
cos_dist = cosine_similarity(image_embeds, self.concept_embeds)
concept_threshold = self.concept_embeds_weights.unsqueeze(0)
concept_scores = cos_dist - concept_threshold + sensitivity
is_nsfw = [torch.any(concept_scores[i] > 0) for i in range(concept_scores.shape[0])]
is_nsfw = [x.item() for x in is_nsfw]
return self.filter_images(images, alternate_image, is_nsfw)
def filter_images(self, images: Tensor, alternate_image: Tensor, is_nsfw: list[bool]):
if not any(is_nsfw):
return images
images = images.clone()
for idx, nsfw in enumerate(is_nsfw):
if nsfw:
# Resize alternate image to match original image dimensions
resized_alternate = F.interpolate(
alternate_image[idx:idx+1], # Add batch dimension
size=(images[idx].shape[1], images[idx].shape[2]), # Target height, width
mode='bilinear',
align_corners=False
)
images[idx] = resized_alternate.squeeze(0) # Remove batch dimension
return images
class CachedModels:
_instance: WeakRef | None = None
def __init__(self):
model_dir = Path(__file__).parent / "safetychecker"
model_file = model_dir / "model.safetensors"
if not model_file.exists():
self.download(
"https://huggingface.co/CompVis/stable-diffusion-safety-checker/resolve/refs%2Fpr%2F41/model.safetensors",
target=model_file,
)
self.feature_extractor = CLIPImageProcessor.from_pretrained(model_dir)
self.safety_checker = CLIPSafetyChecker.from_pretrained(model_dir)
@classmethod
def load(cls):
models = cls._instance and cls._instance()
if models is None:
models = cls()
cls._instance = WeakRef(models)
return models
def download(self, url: str, target: Path):
import requests
try:
target_temp = target.with_suffix(".download")
with requests.get(url, stream=True) as response:
text = "NSFWFilter model download"
total = int(response.headers.get("content-length", 0))
pbar = tqdm(None, total=total, unit="b", unit_scale=True, desc=text)
with open(target_temp, "wb") as f:
for chunk in response.iter_content(chunk_size=8192):
f.write(chunk)
pbar.update(len(chunk))
pbar.close()
target_temp.rename(target)
except Exception as e:
raise RuntimeError(
f"NSFWFilter: Failed to download safety-checker model from {url} to target location {target}: {e}"
) from e
def to_bchw(image: torch.Tensor):
if image.ndim == 3:
image = image.unsqueeze(0)
return image.movedim(-1, 1)
def to_bhwc(image: torch.Tensor):
return image.movedim(1, -1)
def mask_batch(mask: torch.Tensor):
if mask.ndim == 2:
mask = mask.unsqueeze(0)
return mask
class TRI3DNSFWFilter:
models: CachedModels
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
"alternate_image": ("IMAGE",),
"sensitivity": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.10}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "check"
CATEGORY = "TRI3D NSFW"
def __init__(self):
self.models = CachedModels.load()
def check(self, image, alternate_image,sensitivity):
image = to_bchw(image)
alternate_image = to_bchw(alternate_image)
input = self.models.feature_extractor(image, do_rescale=False, return_tensors="pt")
filtered = self.models.safety_checker(
images=image, clip_input=input.pixel_values, sensitivity=sensitivity, alternate_image=alternate_image
)
return (to_bhwc(filtered),)
-131
View File
@@ -1,131 +0,0 @@
import http.client
import mimetypes
import os
import uuid
import requests
import numpy as np
import torch
import cv2
from PIL import Image
import io
class TRI3D_photoroom_bgremove_api:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"images": ("IMAGE", ),
},
}
FUNCTION = "run"
RETURN_TYPES = ("IMAGE", )
CATEGORY = "TRI3D"
def run(self, images):
import http.client
import mimetypes
import os
import uuid
import dotenv
dotenv.load_dotenv()
# Read the API key from the environment variable
PHOTOROOM_API_KEY = os.getenv('PHOTOROOM_API_KEY','EMPTY')
if PHOTOROOM_API_KEY == 'EMPTY':
return images
def tensor_to_cv2_img(tensor, remove_alpha=False):
i = 255. * tensor.cpu().numpy() # This will give us (H, W, C)
img = np.clip(i, 0, 255).astype(np.uint8)
return img
def cv2_img_to_tensor(img):
img = img.astype(np.float32) / 255.0
img = torch.from_numpy(img)[
None,
]
return img
# Please replace with your own apiKey
def remove_background(input_image_path, output_image_path,apiKey):
# Define multipart boundary
boundary = '----------{}'.format(uuid.uuid4().hex)
# Get mimetype of image
content_type, _ = mimetypes.guess_type(input_image_path)
if content_type is None:
content_type = 'application/octet-stream' # Default type if guessing fails
# Prepare the POST data
with open(input_image_path, 'rb') as f:
image_data = f.read()
filename = os.path.basename(input_image_path)
body = (
f"--{boundary}\r\n"
f"Content-Disposition: form-data; name=\"image_file\"; filename=\"{filename}\"\r\n"
f"Content-Type: {content_type}\r\n\r\n"
).encode('utf-8') + image_data + f"\r\n--{boundary}--\r\n".encode('utf-8')
# Set up the HTTP connection and headers
conn = http.client.HTTPSConnection('sdk.photoroom.com')
headers = {
'Content-Type': f'multipart/form-data; boundary={boundary}',
'x-api-key': apiKey
}
# Make the POST request
conn.request('POST', '/v1/segment', body=body, headers=headers)
response = conn.getresponse()
# Handle the response
if response.status == 200:
response_data = response.read()
with open(output_image_path, 'wb') as out_f:
out_f.write(response_data)
print("Image saved to", output_image_path)
else:
print(f"Error: {response.status} - {response.reason}")
print(response.read())
# Close the connection
conn.close()
OUTPUT_FOLDER = "output/"
batch_results = []
for i in range(images.shape[0]):
image = images[i]
cv2_image = tensor_to_cv2_img(image)
cv2_image = cv2.cvtColor(cv2_image, cv2.COLOR_BGR2RGB)
import random
random_number = random.randint(0, 100000)
output_path = OUTPUT_FOLDER + f"output{i}_{random_number}.png"
input_path = OUTPUT_FOLDER + f"input{i}_{random_number}.png"
cv2.imwrite(input_path, cv2_image)
remove_background(input_path, output_path, PHOTOROOM_API_KEY)
print(input_path, output_path)
cv2_segm = cv2.imread(output_path, cv2.IMREAD_UNCHANGED)
cv2_segm = cv2.cvtColor(cv2_segm, cv2.COLOR_BGRA2RGBA)
b_tensor_img = cv2_img_to_tensor(cv2_segm)
batch_results.append(b_tensor_img.squeeze(0))
batch_results = torch.stack(batch_results)
return (batch_results,)
+1 -6
View File
@@ -3,9 +3,4 @@ ninja
pillow
torch
torchvision
transparent-background
wget
gdown
matplotlib
python-dotenv
git+https://github.com/FacePerceiver/facer.git@main
gdown
-171
View File
@@ -1,171 +0,0 @@
{
"_name_or_path": "clip-vit-large-patch14/",
"architectures": [
"SafetyChecker"
],
"initializer_factor": 1.0,
"logit_scale_init_value": 2.6592,
"model_type": "clip",
"projection_dim": 768,
"text_config": {
"_name_or_path": "",
"add_cross_attention": false,
"architectures": null,
"attention_dropout": 0.0,
"bad_words_ids": null,
"bos_token_id": 0,
"chunk_size_feed_forward": 0,
"cross_attention_hidden_size": null,
"decoder_start_token_id": null,
"diversity_penalty": 0.0,
"do_sample": false,
"dropout": 0.0,
"early_stopping": false,
"encoder_no_repeat_ngram_size": 0,
"eos_token_id": 2,
"exponential_decay_length_penalty": null,
"finetuning_task": null,
"forced_bos_token_id": null,
"forced_eos_token_id": null,
"hidden_act": "quick_gelu",
"hidden_size": 768,
"id2label": {
"0": "LABEL_0",
"1": "LABEL_1"
},
"initializer_factor": 1.0,
"initializer_range": 0.02,
"intermediate_size": 3072,
"is_decoder": false,
"is_encoder_decoder": false,
"label2id": {
"LABEL_0": 0,
"LABEL_1": 1
},
"layer_norm_eps": 1e-05,
"length_penalty": 1.0,
"max_length": 20,
"max_position_embeddings": 77,
"min_length": 0,
"model_type": "clip_text_model",
"no_repeat_ngram_size": 0,
"num_attention_heads": 12,
"num_beam_groups": 1,
"num_beams": 1,
"num_hidden_layers": 12,
"num_return_sequences": 1,
"output_attentions": false,
"output_hidden_states": false,
"output_scores": false,
"pad_token_id": 1,
"prefix": null,
"problem_type": null,
"pruned_heads": {},
"remove_invalid_values": false,
"repetition_penalty": 1.0,
"return_dict": true,
"return_dict_in_generate": false,
"sep_token_id": null,
"task_specific_params": null,
"temperature": 1.0,
"tie_encoder_decoder": false,
"tie_word_embeddings": true,
"tokenizer_class": null,
"top_k": 50,
"top_p": 1.0,
"torch_dtype": null,
"torchscript": false,
"transformers_version": "4.21.0.dev0",
"typical_p": 1.0,
"use_bfloat16": false,
"vocab_size": 49408
},
"text_config_dict": {
"hidden_size": 768,
"intermediate_size": 3072,
"num_attention_heads": 12,
"num_hidden_layers": 12
},
"torch_dtype": "float32",
"transformers_version": null,
"vision_config": {
"_name_or_path": "",
"add_cross_attention": false,
"architectures": null,
"attention_dropout": 0.0,
"bad_words_ids": null,
"bos_token_id": null,
"chunk_size_feed_forward": 0,
"cross_attention_hidden_size": null,
"decoder_start_token_id": null,
"diversity_penalty": 0.0,
"do_sample": false,
"dropout": 0.0,
"early_stopping": false,
"encoder_no_repeat_ngram_size": 0,
"eos_token_id": null,
"exponential_decay_length_penalty": null,
"finetuning_task": null,
"forced_bos_token_id": null,
"forced_eos_token_id": null,
"hidden_act": "quick_gelu",
"hidden_size": 1024,
"id2label": {
"0": "LABEL_0",
"1": "LABEL_1"
},
"image_size": 224,
"initializer_factor": 1.0,
"initializer_range": 0.02,
"intermediate_size": 4096,
"is_decoder": false,
"is_encoder_decoder": false,
"label2id": {
"LABEL_0": 0,
"LABEL_1": 1
},
"layer_norm_eps": 1e-05,
"length_penalty": 1.0,
"max_length": 20,
"min_length": 0,
"model_type": "clip_vision_model",
"no_repeat_ngram_size": 0,
"num_attention_heads": 16,
"num_beam_groups": 1,
"num_beams": 1,
"num_hidden_layers": 24,
"num_return_sequences": 1,
"output_attentions": false,
"output_hidden_states": false,
"output_scores": false,
"pad_token_id": null,
"patch_size": 14,
"prefix": null,
"problem_type": null,
"pruned_heads": {},
"remove_invalid_values": false,
"repetition_penalty": 1.0,
"return_dict": true,
"return_dict_in_generate": false,
"sep_token_id": null,
"task_specific_params": null,
"temperature": 1.0,
"tie_encoder_decoder": false,
"tie_word_embeddings": true,
"tokenizer_class": null,
"top_k": 50,
"top_p": 1.0,
"torch_dtype": null,
"torchscript": false,
"transformers_version": "4.21.0.dev0",
"typical_p": 1.0,
"use_bfloat16": false
},
"vision_config_dict": {
"hidden_size": 1024,
"intermediate_size": 4096,
"num_attention_heads": 16,
"num_hidden_layers": 24,
"patch_size": 14
}
}
-20
View File
@@ -1,20 +0,0 @@
{
"crop_size": 224,
"do_center_crop": true,
"do_convert_rgb": true,
"do_normalize": true,
"do_resize": true,
"feature_extractor_type": "CLIPFeatureExtractor",
"image_mean": [
0.48145466,
0.4578275,
0.40821073
],
"image_std": [
0.26862954,
0.26130258,
0.27577711
],
"resample": 3,
"size": 224
}
+2 -2
View File
@@ -67,11 +67,11 @@
-1
],
[
279,
246,
74
],
[
246,
279,
74
],
[
File diff suppressed because it is too large Load Diff
File diff suppressed because one or more lines are too long
@@ -1,526 +0,0 @@
{
"height": 512,
"width": 512,
"keypoints": [
[
259,
75
],
[
261,
127
],
[
222,
133
],
[
208,
196
],
[
199,
253
],
[
296,
132
],
[
316,
192
],
[
333,
246
],
[
241,
249
],
[
242,
348
],
[
255,
448
],
[
277,
248
],
[
296,
344
],
[
322,
448
],
[
251,
69
],
[
268,
67
],
[
237,
77
],
[
280,
73
],
[
239,
71
],
[
239,
77
],
[
241,
81
],
[
243,
87
],
[
243,
91
],
[
247,
95
],
[
251,
96
],
[
255,
100
],
[
261,
100
],
[
264,
98
],
[
270,
96
],
[
272,
93
],
[
276,
89
],
[
278,
83
],
[
278,
79
],
[
278,
73
],
[
278,
69
],
[
243,
66
],
[
245,
66
],
[
249,
64
],
[
251,
64
],
[
255,
64
],
[
262,
64
],
[
266,
62
],
[
268,
62
],
[
272,
62
],
[
274,
64
],
[
259,
69
],
[
259,
71
],
[
259,
73
],
[
261,
77
],
[
255,
79
],
[
259,
79
],
[
261,
79
],
[
262,
79
],
[
264,
79
],
[
247,
69
],
[
249,
69
],
[
251,
67
],
[
255,
69
],
[
251,
69
],
[
249,
69
],
[
264,
69
],
[
266,
67
],
[
268,
67
],
[
272,
67
],
[
268,
69
],
[
266,
69
],
[
251,
87
],
[
255,
83
],
[
259,
83
],
[
261,
83
],
[
262,
83
],
[
264,
83
],
[
268,
85
],
[
266,
87
],
[
264,
91
],
[
261,
91
],
[
257,
91
],
[
253,
89
],
[
251,
87
],
[
255,
85
],
[
261,
85
],
[
264,
85
],
[
268,
85
],
[
264,
87
],
[
261,
89
],
[
255,
89
],
[
251,
69
],
[
268,
67
],
[
333,
249
],
[
328,
255
],
[
324,
262
],
[
322,
270
],
[
320,
278
],
[
331,
270
],
[
329,
278
],
[
324,
278
],
[
322,
276
],
[
335,
272
],
[
329,
278
],
[
326,
278
],
[
322,
276
],
[
337,
272
],
[
331,
278
],
[
328,
278
],
[
326,
276
],
[
337,
270
],
[
333,
276
],
[
329,
276
],
[
328,
276
],
[
199,
256
],
[
203,
263
],
[
207,
267
],
[
209,
275
],
[
209,
283
],
[
199,
277
],
[
203,
285
],
[
207,
283
],
[
211,
281
],
[
197,
279
],
[
201,
285
],
[
207,
283
],
[
209,
281
],
[
196,
277
],
[
201,
283
],
[
205,
283
],
[
209,
281
],
[
197,
277
],
[
201,
281
],
[
203,
281
],
[
205,
279
]
]
}
Binary file not shown.

Before

Width:  |  Height:  |  Size: 87 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 86 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 86 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 1.8 MiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 1.8 MiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 1.8 MiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 134 KiB

@@ -1,872 +0,0 @@
{
"last_node_id": 28,
"last_link_id": 35,
"nodes": [
{
"id": 2,
"type": "LoadImage",
"pos": [
1968,
-883
],
"size": {
"0": 315,
"1": 314
},
"flags": {},
"order": 0,
"mode": 0,
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
1
],
"shape": 3,
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": null,
"shape": 3
}
],
"title": "Image1",
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"image1.png",
"image"
]
},
{
"id": 3,
"type": "LoadImage",
"pos": [
1969,
-514
],
"size": {
"0": 315,
"1": 314
},
"flags": {},
"order": 1,
"mode": 0,
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
2
],
"shape": 3,
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": null,
"shape": 3
}
],
"title": "Image2",
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"image2.png",
"image"
]
},
{
"id": 4,
"type": "LoadImage",
"pos": [
1969,
-151
],
"size": {
"0": 315,
"1": 314
},
"flags": {},
"order": 2,
"mode": 0,
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
3
],
"shape": 3,
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": null,
"shape": 3
}
],
"title": "Image3",
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"image3.png",
"image"
]
},
{
"id": 5,
"type": "LoadImage",
"pos": [
2370,
-600
],
"size": {
"0": 315,
"1": 314
},
"flags": {},
"order": 3,
"mode": 0,
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
5
],
"shape": 3,
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": null,
"shape": 3
}
],
"title": "ATR1",
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"atr1.png",
"image"
]
},
{
"id": 6,
"type": "LoadImage",
"pos": [
2367,
-236
],
"size": {
"0": 315,
"1": 314
},
"flags": {},
"order": 4,
"mode": 0,
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
6
],
"shape": 3,
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": null,
"shape": 3
}
],
"title": "ATR2",
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"atr2.png",
"image"
]
},
{
"id": 14,
"type": "ImageBatch",
"pos": [
2745,
-584
],
"size": {
"0": 210,
"1": 46
},
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "image1",
"type": "IMAGE",
"link": 5
},
{
"name": "image2",
"type": "IMAGE",
"link": 6
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
7
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "ImageBatch"
}
},
{
"id": 7,
"type": "LoadImage",
"pos": [
2372,
129
],
"size": {
"0": 315,
"1": 314
},
"flags": {},
"order": 5,
"mode": 0,
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
8
],
"shape": 3,
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": null,
"shape": 3
}
],
"title": "ATR3",
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"atr3.png",
"image"
]
},
{
"id": 11,
"type": "ImageBatch",
"pos": [
2390,
-848
],
"size": {
"0": 210,
"1": 46
},
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"name": "image1",
"type": "IMAGE",
"link": 1
},
{
"name": "image2",
"type": "IMAGE",
"link": 2
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
4
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "ImageBatch"
}
},
{
"id": 12,
"type": "ImageBatch",
"pos": [
2981,
-582
],
"size": {
"0": 210,
"1": 46
},
"flags": {},
"order": 11,
"mode": 0,
"inputs": [
{
"name": "image1",
"type": "IMAGE",
"link": 7
},
{
"name": "image2",
"type": "IMAGE",
"link": 8
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
10
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "ImageBatch"
}
},
{
"id": 8,
"type": "LoadImage",
"pos": [
2786,
-268
],
"size": {
"0": 315,
"1": 314
},
"flags": {},
"order": 6,
"mode": 0,
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
11
],
"shape": 3,
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": null,
"shape": 3
}
],
"title": "Mask",
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"mask (30).png",
"image"
]
},
{
"id": 15,
"type": "RepeatImageBatch",
"pos": [
2793,
-367
],
"size": {
"0": 315,
"1": 58
},
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"name": "image",
"type": "IMAGE",
"link": 11
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
18
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "RepeatImageBatch"
},
"widgets_values": [
3
]
},
{
"id": 13,
"type": "ImageBatch",
"pos": [
2636,
-847
],
"size": {
"0": 210,
"1": 46
},
"flags": {},
"order": 10,
"mode": 0,
"inputs": [
{
"name": "image1",
"type": "IMAGE",
"link": 4
},
{
"name": "image2",
"type": "IMAGE",
"link": 3
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
16
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "ImageBatch"
}
},
{
"id": 26,
"type": "Image To Mask",
"pos": [
3667,
-11
],
"size": {
"0": 315,
"1": 58
},
"flags": {},
"order": 15,
"mode": 0,
"inputs": [
{
"name": "image",
"type": "IMAGE",
"link": 32
}
],
"outputs": [
{
"name": "MASK",
"type": "MASK",
"links": [
33
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "Image To Mask"
},
"widgets_values": [
"intensity"
]
},
{
"id": 27,
"type": "InpaintPreprocessor",
"pos": [
3678,
99
],
"size": {
"0": 210,
"1": 46
},
"flags": {},
"order": 16,
"mode": 0,
"inputs": [
{
"name": "image",
"type": "IMAGE",
"link": 34
},
{
"name": "mask",
"type": "MASK",
"link": 33
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
35
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "InpaintPreprocessor"
}
},
{
"id": 17,
"type": "PreviewImage",
"pos": [
3679,
-294
],
"size": {
"0": 686.6637573242188,
"1": 246
},
"flags": {},
"order": 14,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 19
}
],
"properties": {
"Node name for S&R": "PreviewImage"
}
},
{
"id": 28,
"type": "PreviewImage",
"pos": [
3675,
196
],
"size": {
"0": 786.8681640625,
"1": 246
},
"flags": {},
"order": 17,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 35
}
],
"properties": {
"Node name for S&R": "PreviewImage"
}
},
{
"id": 16,
"type": "PreviewImage",
"pos": [
3661,
-574
],
"size": {
"0": 714.0955810546875,
"1": 246
},
"flags": {},
"order": 13,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 17
}
],
"properties": {
"Node name for S&R": "PreviewImage"
}
},
{
"id": 1,
"type": "tri3d-extract-parts-batch",
"pos": [
3280,
-482
],
"size": {
"0": 315,
"1": 530
},
"flags": {},
"order": 12,
"mode": 0,
"inputs": [
{
"name": "batch_images",
"type": "IMAGE",
"link": 16
},
{
"name": "batch_segs",
"type": "IMAGE",
"link": 10
},
{
"name": "batch_secondaries",
"type": "IMAGE",
"link": 18
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
17,
34
],
"shape": 3,
"slot_index": 0
},
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
19,
32
],
"shape": 3,
"slot_index": 1
}
],
"properties": {
"Node name for S&R": "tri3d-extract-parts-batch"
},
"widgets_values": [
20,
false,
false,
true,
true,
false,
false,
false,
false,
false,
false,
false,
false,
false,
false,
false,
false,
false,
false
]
}
],
"links": [
[
1,
2,
0,
11,
0,
"IMAGE"
],
[
2,
3,
0,
11,
1,
"IMAGE"
],
[
3,
4,
0,
13,
1,
"IMAGE"
],
[
4,
11,
0,
13,
0,
"IMAGE"
],
[
5,
5,
0,
14,
0,
"IMAGE"
],
[
6,
6,
0,
14,
1,
"IMAGE"
],
[
7,
14,
0,
12,
0,
"IMAGE"
],
[
8,
7,
0,
12,
1,
"IMAGE"
],
[
10,
12,
0,
1,
1,
"IMAGE"
],
[
11,
8,
0,
15,
0,
"IMAGE"
],
[
16,
13,
0,
1,
0,
"IMAGE"
],
[
17,
1,
0,
16,
0,
"IMAGE"
],
[
18,
15,
0,
1,
2,
"IMAGE"
],
[
19,
1,
1,
17,
0,
"IMAGE"
],
[
32,
1,
1,
26,
0,
"IMAGE"
],
[
33,
26,
0,
27,
1,
"MASK"
],
[
34,
1,
0,
27,
0,
"IMAGE"
],
[
35,
27,
0,
28,
0,
"IMAGE"
]
],
"groups": [],
"config": {},
"extra": {},
"version": 0.4
}
-526
View File
@@ -1,526 +0,0 @@
{
"height": 512,
"width": 512,
"keypoints": [
[
-1,
-1
],
[
264,
125
],
[
298,
125
],
[
310,
193
],
[
314,
254
],
[
230,
125
],
[
216,
192
],
[
211,
252
],
[
286,
244
],
[
292,
339
],
[
294,
448
],
[
240,
247
],
[
242,
336
],
[
246,
444
],
[
-1,
-1
],
[
-1,
-1
],
[
279,
74
],
[
246,
74
],
[
240,
66
],
[
240,
72
],
[
240,
76
],
[
242,
82
],
[
242,
86
],
[
246,
88
],
[
248,
92
],
[
252,
94
],
[
258,
94
],
[
262,
94
],
[
266,
94
],
[
270,
92
],
[
274,
90
],
[
276,
86
],
[
278,
82
],
[
280,
78
],
[
282,
74
],
[
248,
58
],
[
250,
56
],
[
252,
56
],
[
256,
56
],
[
260,
56
],
[
268,
58
],
[
272,
58
],
[
274,
60
],
[
278,
62
],
[
280,
64
],
[
264,
64
],
[
262,
66
],
[
262,
68
],
[
262,
68
],
[
258,
74
],
[
258,
74
],
[
262,
74
],
[
264,
74
],
[
266,
74
],
[
250,
62
],
[
252,
62
],
[
256,
62
],
[
258,
62
],
[
254,
64
],
[
252,
62
],
[
268,
64
],
[
272,
64
],
[
274,
64
],
[
276,
68
],
[
274,
68
],
[
270,
66
],
[
252,
80
],
[
256,
78
],
[
258,
78
],
[
260,
78
],
[
262,
78
],
[
266,
80
],
[
268,
82
],
[
266,
84
],
[
262,
84
],
[
260,
84
],
[
256,
84
],
[
254,
82
],
[
252,
80
],
[
256,
78
],
[
260,
78
],
[
264,
80
],
[
268,
82
],
[
264,
84
],
[
260,
82
],
[
256,
82
],
[
254,
62
],
[
272,
64
],
[
211,
255
],
[
221,
262
],
[
226,
268
],
[
229,
272
],
[
232,
278
],
[
217,
279
],
[
223,
286
],
[
229,
288
],
[
232,
291
],
[
214,
281
],
[
221,
288
],
[
223,
291
],
[
229,
291
],
[
211,
281
],
[
217,
286
],
[
223,
288
],
[
223,
291
],
[
211,
281
],
[
214,
286
],
[
217,
288
],
[
221,
288
],
[
314.2191598676356,
257
],
[
308.2191598676356,
262
],
[
306.2191598676356,
266
],
[
303.2191598676356,
270
],
[
301.2191598676356,
276
],
[
312.2191598676356,
279
],
[
311.2191598676356,
284
],
[
308.2191598676356,
290
],
[
303.2191598676356,
292
],
[
313.2191598676356,
279
],
[
312.2191598676356,
284
],
[
311.2191598676356,
290
],
[
307.2191598676356,
292
],
[
314.2191598676356,
279
],
[
313.2191598676356,
284
],
[
312.2191598676356,
289
],
[
311.2191598676356,
293
],
[
317.2191598676356,
279
],
[
314.2191598676356,
283
],
[
313.2191598676356,
287
],
[
312.2191598676356,
290
]
]
}
-526
View File
@@ -1,526 +0,0 @@
{
"height": 512,
"width": 512,
"keypoints": [
[
-1,
-1
],
[
252,
130
],
[
302,
130
],
[
316,
213
],
[
323,
284
],
[
202,
130
],
[
191,
211
],
[
187,
282
],
[
282,
278
],
[
280,
380
],
[
280,
484
],
[
222,
278
],
[
226,
382
],
[
222,
480
],
[
-1,
-1
],
[
-1,
-1
],
[
286,
72
],
[
227,
73
],
[
257,
59
],
[
257,
63
],
[
257,
65
],
[
253,
67
],
[
-2,
-2
],
[
-2,
-2
],
[
-2,
-2
],
[
-2,
-2
],
[
-2,
-2
],
[
279,
79
],
[
-2,
-2
],
[
243,
67
],
[
253,
71
],
[
257,
69
],
[
-2,
-2
],
[
-2,
-2
],
[
-2,
-2
],
[
275,
61
],
[
279,
59
],
[
255,
55
],
[
-2,
-2
],
[
-2,
-2
],
[
255,
59
],
[
255,
59
],
[
255,
59
],
[
255,
57
],
[
255,
57
],
[
255,
59
],
[
255,
59
],
[
255,
61
],
[
255,
63
],
[
277,
67
],
[
275,
67
],
[
277,
67
],
[
277,
67
],
[
277,
67
],
[
255,
59
],
[
255,
57
],
[
257,
59
],
[
257,
61
],
[
257,
59
],
[
255,
59
],
[
255,
59
],
[
257,
59
],
[
257,
59
],
[
-2,
-2
],
[
255,
59
],
[
-2,
-2
],
[
-2,
-2
],
[
-2,
-2
],
[
-2,
-2
],
[
279,
71
],
[
243,
67
],
[
279,
73
],
[
277,
71
],
[
279,
71
],
[
-2,
-2
],
[
-2,
-2
],
[
-2,
-2
],
[
-2,
-2
],
[
-2,
-2
],
[
-2,
-2
],
[
279,
69
],
[
279,
71
],
[
277,
71
],
[
279,
71
],
[
277,
69
],
[
-2,
-2
],
[
250,
56
],
[
263,
63
],
[
186.65255255416977,
285
],
[
196.65255255416977,
291
],
[
199.65255255416977,
298
],
[
201.65255255416977,
306
],
[
201.65255255416977,
310
],
[
194.65255255416977,
308
],
[
196.65255255416977,
312
],
[
199.65255255416977,
316
],
[
203.65255255416977,
318
],
[
190.65255255416977,
308
],
[
192.65255255416977,
314
],
[
194.65255255416977,
318
],
[
201.65255255416977,
320
],
[
186.65255255416977,
308
],
[
186.65255255416977,
314
],
[
190.65255255416977,
318
],
[
196.65255255416977,
320
],
[
182.65255255416977,
308
],
[
184.65255255416977,
312
],
[
186.65255255416977,
316
],
[
190.65255255416977,
318
],
[
323,
287
],
[
313,
298
],
[
310,
305
],
[
310,
312
],
[
310,
317
],
[
320,
315
],
[
317,
319
],
[
313,
324
],
[
308,
326
],
[
322,
315
],
[
320,
322
],
[
315,
326
],
[
308,
326
],
[
324,
317
],
[
320,
322
],
[
315,
326
],
[
310,
326
],
[
327,
317
],
[
324,
322
],
[
320,
324
],
[
315,
326
]
]
}
-526
View File
@@ -1,526 +0,0 @@
{
"height": 512,
"width": 512,
"keypoints": [
[
-1,
-1
],
[
254,
118
],
[
292,
118
],
[
306,
194
],
[
310,
262
],
[
216,
118
],
[
200,
193
],
[
195,
260
],
[
278,
251
],
[
286,
357
],
[
285.99326159010275,
462.45102588957826
],
[
228,
254
],
[
230,
353
],
[
231.99326159010275,
458.45102588957826
],
[
-1,
-1
],
[
-1,
-1
],
[
277,
62
],
[
241,
62
],
[
234,
53
],
[
234,
60
],
[
234,
64
],
[
236,
71
],
[
236,
75
],
[
241,
77
],
[
243,
82
],
[
247,
84
],
[
254,
84
],
[
258,
84
],
[
263,
84
],
[
267,
82
],
[
272,
80
],
[
274,
75
],
[
276,
71
],
[
278,
66
],
[
281,
62
],
[
243,
44
],
[
245,
42
],
[
247,
42
],
[
252,
42
],
[
256,
42
],
[
265,
44
],
[
269,
44
],
[
272,
46
],
[
276,
49
],
[
278,
51
],
[
261,
51
],
[
258,
53
],
[
258,
55
],
[
258,
55
],
[
254,
62
],
[
254,
62
],
[
258,
62
],
[
261,
62
],
[
263,
62
],
[
245,
49
],
[
247,
49
],
[
252,
49
],
[
254,
49
],
[
249,
51
],
[
247,
49
],
[
265,
51
],
[
269,
51
],
[
272,
51
],
[
274,
55
],
[
272,
55
],
[
267,
53
],
[
247,
69
],
[
252,
66
],
[
254,
66
],
[
256,
66
],
[
258,
66
],
[
263,
69
],
[
265,
71
],
[
263,
73
],
[
258,
73
],
[
256,
73
],
[
252,
73
],
[
249,
71
],
[
247,
69
],
[
252,
66
],
[
256,
66
],
[
261,
69
],
[
265,
71
],
[
261,
73
],
[
256,
71
],
[
252,
71
],
[
249,
49
],
[
269,
51
],
[
195,
262
],
[
206,
270
],
[
210,
276
],
[
213,
280
],
[
216,
285
],
[
201,
287
],
[
208,
294
],
[
213,
296
],
[
216,
299
],
[
197,
289
],
[
206,
296
],
[
208,
299
],
[
213,
299
],
[
195,
289
],
[
201,
294
],
[
208,
296
],
[
208,
299
],
[
195,
289
],
[
197,
294
],
[
201,
296
],
[
206,
296
],
[
310,
265
],
[
304,
270
],
[
302,
274
],
[
299,
279
],
[
297,
284
],
[
308,
287
],
[
307,
292
],
[
304,
299
],
[
299,
301
],
[
309,
287
],
[
308,
292
],
[
307,
299
],
[
303,
301
],
[
310,
287
],
[
309,
292
],
[
308,
298
],
[
307,
302
],
[
313,
287
],
[
310,
291
],
[
309,
296
],
[
308,
299
]
]
}
-526
View File
@@ -1,526 +0,0 @@
{
"height": 512,
"width": 512,
"keypoints": [
[
234,
85
],
[
264,
125
],
[
298,
125
],
[
310,
193
],
[
314,
254
],
[
230,
125
],
[
216,
192
],
[
211,
252
],
[
286,
244
],
[
292,
339
],
[
294,
448
],
[
240,
247
],
[
242,
336
],
[
246,
444
],
[
-1,
-1
],
[
238,
72
],
[
-1,
-1
],
[
262,
70
],
[
212,
82
],
[
212,
88
],
[
212,
92
],
[
214,
98
],
[
214,
102
],
[
218,
104
],
[
220,
108
],
[
224,
110
],
[
230,
110
],
[
234,
110
],
[
238,
110
],
[
242,
108
],
[
246,
106
],
[
248,
102
],
[
250,
98
],
[
252,
94
],
[
254,
90
],
[
220,
74
],
[
222,
72
],
[
224,
72
],
[
228,
72
],
[
232,
72
],
[
240,
74
],
[
244,
74
],
[
246,
76
],
[
250,
78
],
[
252,
80
],
[
236,
80
],
[
234,
82
],
[
232,
84
],
[
255,
85
],
[
230,
90
],
[
230,
90
],
[
234,
90
],
[
236,
90
],
[
238,
90
],
[
222,
78
],
[
224,
78
],
[
228,
78
],
[
230,
78
],
[
226,
80
],
[
224,
78
],
[
240,
80
],
[
244,
80
],
[
246,
80
],
[
248,
84
],
[
246,
84
],
[
242,
82
],
[
224,
96
],
[
228,
94
],
[
230,
94
],
[
232,
94
],
[
234,
94
],
[
238,
96
],
[
240,
98
],
[
238,
100
],
[
234,
100
],
[
232,
100
],
[
228,
100
],
[
226,
98
],
[
224,
96
],
[
228,
94
],
[
232,
94
],
[
236,
96
],
[
240,
98
],
[
236,
100
],
[
232,
98
],
[
228,
98
],
[
226,
78
],
[
244,
80
],
[
211,
255
],
[
221,
262
],
[
226,
268
],
[
229,
272
],
[
232,
278
],
[
217,
279
],
[
223,
286
],
[
229,
288
],
[
232,
291
],
[
214,
281
],
[
221,
288
],
[
223,
291
],
[
229,
291
],
[
211,
281
],
[
217,
286
],
[
223,
288
],
[
223,
291
],
[
211,
281
],
[
214,
286
],
[
217,
288
],
[
221,
288
],
[
314,
257
],
[
308,
262
],
[
306,
266
],
[
303,
270
],
[
301,
276
],
[
312,
279
],
[
311,
284
],
[
308,
290
],
[
303,
292
],
[
313,
279
],
[
312,
284
],
[
311,
290
],
[
307,
292
],
[
314,
279
],
[
313,
284
],
[
312,
289
],
[
311,
293
],
[
317,
279
],
[
314,
283
],
[
313,
287
],
[
312,
290
]
]
}
-526
View File
@@ -1,526 +0,0 @@
{
"height": 512,
"width": 512,
"keypoints": [
[
298,
85
],
[
264,
125
],
[
298,
125
],
[
310,
193
],
[
314,
254
],
[
230,
125
],
[
216,
192
],
[
211,
252
],
[
286,
244
],
[
292,
339
],
[
294,
448
],
[
240,
247
],
[
242,
336
],
[
246,
444
],
[
293.3658724790831,
75
],
[
-1,
-1
],
[
270.3658724790831,
73
],
[
-1,
-1
],
[
276.3658724790831,
82
],
[
276.3658724790831,
88
],
[
276.3658724790831,
92
],
[
278.3658724790831,
98
],
[
278.3658724790831,
102
],
[
282.3658724790831,
104
],
[
284.3658724790831,
108
],
[
288.3658724790831,
110
],
[
294.3658724790831,
110
],
[
298.3658724790831,
110
],
[
302.3658724790831,
110
],
[
306.3658724790831,
108
],
[
310.3658724790831,
106
],
[
312.3658724790831,
102
],
[
314.3658724790831,
98
],
[
316.3658724790831,
94
],
[
318.3658724790831,
90
],
[
284.3658724790831,
74
],
[
286.3658724790831,
72
],
[
288.3658724790831,
72
],
[
292.3658724790831,
72
],
[
296.3658724790831,
72
],
[
304.3658724790831,
74
],
[
308.3658724790831,
74
],
[
310.3658724790831,
76
],
[
314.3658724790831,
78
],
[
316.3658724790831,
80
],
[
300.3658724790831,
80
],
[
298.3658724790831,
82
],
[
296.3658724790831,
84
],
[
319.3658724790831,
85
],
[
294.3658724790831,
90
],
[
294.3658724790831,
90
],
[
298.3658724790831,
90
],
[
300.3658724790831,
90
],
[
302.3658724790831,
90
],
[
286.3658724790831,
78
],
[
288.3658724790831,
78
],
[
292.3658724790831,
78
],
[
294.3658724790831,
78
],
[
290.3658724790831,
80
],
[
288.3658724790831,
78
],
[
304.3658724790831,
80
],
[
308.3658724790831,
80
],
[
310.3658724790831,
80
],
[
312.3658724790831,
84
],
[
310.3658724790831,
84
],
[
306.3658724790831,
82
],
[
288.3658724790831,
96
],
[
292.3658724790831,
94
],
[
294.3658724790831,
94
],
[
296.3658724790831,
94
],
[
298.3658724790831,
94
],
[
302.3658724790831,
96
],
[
304.3658724790831,
98
],
[
302.3658724790831,
100
],
[
298.3658724790831,
100
],
[
296.3658724790831,
100
],
[
292.3658724790831,
100
],
[
290.3658724790831,
98
],
[
288.3658724790831,
96
],
[
292.3658724790831,
94
],
[
296.3658724790831,
94
],
[
300.3658724790831,
96
],
[
304.3658724790831,
98
],
[
300.3658724790831,
100
],
[
296.3658724790831,
98
],
[
292.3658724790831,
98
],
[
290.3658724790831,
78
],
[
308.3658724790831,
80
],
[
211,
255
],
[
221,
262
],
[
226,
268
],
[
229,
272
],
[
232,
278
],
[
217,
279
],
[
223,
286
],
[
229,
288
],
[
232,
291
],
[
214,
281
],
[
221,
288
],
[
223,
291
],
[
229,
291
],
[
211,
281
],
[
217,
286
],
[
223,
288
],
[
223,
291
],
[
211,
281
],
[
214,
286
],
[
217,
288
],
[
221,
288
],
[
314,
257
],
[
308,
262
],
[
306,
266
],
[
303,
270
],
[
301,
276
],
[
312,
279
],
[
311,
284
],
[
308,
290
],
[
303,
292
],
[
313,
279
],
[
312,
284
],
[
311,
290
],
[
307,
292
],
[
314,
279
],
[
313,
284
],
[
312,
289
],
[
311,
293
],
[
317,
279
],
[
314,
283
],
[
313,
287
],
[
312,
290
]
]
}
-316
View File
@@ -1,316 +0,0 @@
#!/usr/bin/python3
import torch
import cv2
import numpy as np
#!/usr/bin/python3
def from_torch_image(image):
image = image.cpu().numpy() * 255.0
image = np.clip(image, 0, 255).astype(np.uint8)
return image
def to_torch_image(image):
image = image.astype(dtype=np.float32)
image /= 255.0
image = torch.from_numpy(image)
return image
def scaled_paste_2(
image_background,
image_foreground,
mask_foreground,
scale_factor,
height_factor=1.2,
):
print('DEBUG scaled_paste 0 ', image_background.shape,
image_foreground.shape, mask_foreground.shape, scale_factor,
height_factor)
height = image_foreground.shape[0] * height_factor
print('DEBUG scaled_paste 1 ', height)
max_0 = max(image_background.shape[0], height)
max_1 = max(image_background.shape[1], image_foreground.shape[1])
print('DEBUG scaled_paste 2 ', max_0, max_1)
ratio_0 = max_0 / image_background.shape[0]
ratio_1 = max_1 / image_background.shape[1]
ratio_max = max(ratio_0, ratio_1) * scale_factor
print('DEBUG scaled_paste 2 ', ratio_0, ratio_1, ratio_max)
size_0 = int(image_background.shape[0] * scale_factor)
size_1 = int(image_background.shape[1] * scale_factor)
print('DEBUG scaled_paste 3 ', size_0, size_1)
image_background = cv2.resize(image_background, (size_1, size_0),
cv2.INTER_CUBIC)
print('DEBUG scaled_paste 4 ', image_background.shape)
bg_h = image_background.shape[0]
fg_h = image_foreground.shape[0]
end_0 = int(bg_h - (bg_h * (height_factor-1)))
# end_0 = int(image_background.shape[0])
begin_0 = max(0, int(end_0 - fg_h))
# end_0 = int(begin_0 + image_foreground.shape[0])
fg_start_height = fg_h - (end_0 - begin_0)
print('DEBUG scaled_paste 5 ', begin_0, end_0, fg_start_height)
end_1 = image_background.shape[1]
begin_1 = end_1 - image_foreground.shape[1]
begin_1 = int(begin_1 / 2)
end_1 = int(begin_1 + image_foreground.shape[1])
print('DEBUG scaled_paste 6 ', begin_1, end_1)
image_reference = image_background[begin_0:end_0, begin_1:end_1, :]
image_foreground = image_foreground[fg_start_height:,:,:]
mask_foreground = mask_foreground[fg_start_height:,:]
print('DEBUG scaled_paste 7 ', image_reference.shape, image_foreground.shape, mask_foreground.shape)
for i in range(3):
image_reference[:, :,
i] = (mask_foreground * image_foreground[:, :, i]) + (
(1 - mask_foreground) * image_reference[:, :, i])
return image_background
def scaled_paste(
image_background,
image_foreground,
mask_foreground,
scale_factor,
height_factor=1.2,
):
print('DEBUG scaled_paste 0 ', image_background.shape,
image_foreground.shape, mask_foreground.shape, scale_factor,
height_factor)
height = image_foreground.shape[0] * height_factor
print('DEBUG scaled_paste 1 ', height)
max_0 = max(image_background.shape[0], height)
max_1 = max(image_background.shape[1], image_foreground.shape[1])
print('DEBUG scaled_paste 2 ', max_0, max_1)
ratio_0 = max_0 / image_background.shape[0]
ratio_1 = max_1 / image_background.shape[1]
ratio_max = max(ratio_0, ratio_1) * scale_factor
print('DEBUG scaled_paste 2 ', ratio_0, ratio_1, ratio_max)
size_0 = int(image_background.shape[0] * ratio_max) + 1
size_1 = int(image_background.shape[1] * ratio_max) + 1
print('DEBUG scaled_paste 3 ', size_0, size_1)
image_background = cv2.resize(image_background, (size_1, size_0),
cv2.INTER_CUBIC)
print('DEBUG scaled_paste 4 ', image_background.shape)
end_0 = int(image_background.shape[0])
begin_0 = int(end_0 - height)
end_0 = int(begin_0 + image_foreground.shape[0])
print('DEBUG scaled_paste 5 ', begin_0, end_0)
end_1 = image_background.shape[1]
begin_1 = end_1 - image_foreground.shape[1]
begin_1 = int(begin_1 / 2)
end_1 = int(begin_1 + image_foreground.shape[1])
print('DEBUG scaled_paste 6 ', begin_1, end_1)
image_reference = image_background[begin_0:end_0, begin_1:end_1, :]
for i in range(3):
image_reference[:, :,
i] = (mask_foreground * image_foreground[:, :, i]) + (
(1 - mask_foreground) * image_reference[:, :, i])
return image_background
#!/usr/bin/python3
class main_scaled_paste_2():
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image_background": ("IMAGE", ),
"image_foreground": ("IMAGE", ),
"mask_foreground": ("MASK", ),
"scale_factor": ("FLOAT", {
"default": 1.2,
"min": 1,
"max": 10,
"step": 0.05
}),
"height_factor": ("FLOAT", {
"default": 1.01,
"min": 1,
"max": 8,
"step": 0.05
}),
},
}
FUNCTION = "run"
RETURN_TYPES = ("IMAGE", )
CATEGORY = "TRI3D"
def run(
self,
image_background,
image_foreground,
mask_foreground,
scale_factor,
height_factor,
):
print('DEBUG 0 ', image_background.shape, image_foreground.shape,
mask_foreground.shape)
image_background = from_torch_image(image_background)
image_foreground = from_torch_image(image_foreground)
mask_foreground = mask_foreground.cpu().numpy()
image_output = scaled_paste_2(
image_background[0],
image_foreground[0],
mask_foreground[0],
scale_factor,
height_factor,
)
print('DEBUG 1 ', image_output.shape)
image_output = to_torch_image(image=image_output)
print('DEBUG 2 ', image_output.shape)
image_output = image_output.unsqueeze(0)
print('DEBUG 3 ', image_output.shape)
return (image_output, )
class main_scaled_paste():
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image_background": ("IMAGE", ),
"image_foreground": ("IMAGE", ),
"mask_foreground": ("MASK", ),
"scale_factor": ("FLOAT", {
"default": 1.2,
"min": 1,
"max": 10,
"step": 0.05
}),
"height_factor": ("FLOAT", {
"default": 1.01,
"min": 1,
"max": 8,
"step": 0.05
}),
},
}
FUNCTION = "run"
RETURN_TYPES = ("IMAGE", )
CATEGORY = "TRI3D"
def run(
self,
image_background,
image_foreground,
mask_foreground,
scale_factor,
height_factor,
):
print('DEBUG 0 ', image_background.shape, image_foreground.shape,
mask_foreground.shape)
image_background = from_torch_image(image_background)
image_foreground = from_torch_image(image_foreground)
mask_foreground = mask_foreground.cpu().numpy()
image_output = scaled_paste(
image_background[0],
image_foreground[0],
mask_foreground[0],
scale_factor,
height_factor,
)
print('DEBUG 1 ', image_output.shape)
image_output = to_torch_image(image=image_output)
print('DEBUG 2 ', image_output.shape)
image_output = image_output.unsqueeze(0)
print('DEBUG 3 ', image_output.shape)
return (image_output, )
#!/usr/bin/python3
# mask = cv2.imread('/home/asd/DATASETS/BG_SWAP_HACK_TEST/FOREGROUND_MASK.png',
# cv2.IMREAD_GRAYSCALE)
# mask = mask.astype(dtype=np.float32) / 255.0
# image_background = scaled_paste(
# image_background=cv2.imread(
# '/home/asd/DATASETS/BG_SWAP_HACK_TEST/BACKGROUND_DEPTH.png',
# cv2.IMREAD_COLOR),
# image_foreground=cv2.imread(
# '/home/asd/DATASETS/BG_SWAP_HACK_TEST/FOREGROUND_DEPTH.png',
# cv2.IMREAD_COLOR),
# mask_foreground=mask,
# scale_factor=2,
# height_factor=1.05,
# )
# cv2.imwrite('tmp.png', image_background)
NODE_CLASS_MAPPINGS = {
'main_scaled_paste': main_scaled_paste,
}
NODE_DISPLAY_NAME_MAPPINGS = {
'main_scaled_paste': 'main_scaled_paste',
}
-610
View File
@@ -1,610 +0,0 @@
#!/usr/bin/python3
import cv2
import math
import matplotlib.pyplot as plt
import numpy as np
import torch
#!/usr/bin/python3
def from_torch_image(image):
image = image.cpu().numpy() * 255.0
image = np.clip(image, 0, 255).astype(np.uint8)
return image
def to_torch_image(image):
image = image.astype(dtype=np.float32)
image /= 255.0
image = torch.from_numpy(image)
return image
def do_custom_threshhold(image, value):
image = image.astype(dtype=np.float64)
image = 255 * (image - value) / (255 - value)
image = np.clip(image, 0, 255)
image = image.astype(dtype=np.uint8)
return image
def scaled_paste(
image_background,
image_foreground,
mask_foreground,
scale_factor=1.2,
height_factor=1.05,
):
height = image_foreground.shape[0] * height_factor
max_0 = max(image_background.shape[0], height)
max_1 = max(image_background.shape[1], image_foreground.shape[1])
ratio_0 = max_0 / image_background.shape[0]
ratio_1 = max_1 / image_background.shape[1]
ratio_max = max(ratio_0, ratio_1) * scale_factor
size_0 = int(image_background.shape[0] * ratio_max) + 1
size_1 = int(image_background.shape[1] * ratio_max) + 1
image_background = cv2.resize(image_background, (size_1, size_0),
cv2.INTER_CUBIC)
end_0 = int(image_background.shape[0])
begin_0 = int(end_0 - height)
end_0 = int(begin_0 + image_foreground.shape[0])
end_1 = image_background.shape[1]
begin_1 = end_1 - image_foreground.shape[1]
begin_1 = int(begin_1 / 2)
end_1 = int(begin_1 + image_foreground.shape[1])
image_reference = image_background[begin_0:end_0, begin_1:end_1, :]
for i in range(3):
image_reference[:, :,
i] = (mask_foreground * image_foreground[:, :, i]) + (
(1 - mask_foreground) * image_reference[:, :, i])
return image_background
def do_bg_swap(
bkg_image,
subject_image,
mask_image,
threshhold_hist,
scale_factor=1.2,
height_factor=1.05,
):
blank_subject_image = np.zeros(subject_image.shape, dtype=np.uint8)
blank_subject_image += 255
blank_background_image = np.zeros(bkg_image.shape, dtype=np.uint8)
blank_subject_mask = np.zeros(
(subject_image.shape[0], subject_image.shape[1]), dtype=np.float64)
blank_subject_mask += 1
mask_image_3channel = subject_image.copy()
for i in range(3):
mask_image_3channel[:, :, i] = mask_image
mask_image = mask_image.astype(dtype=np.float64)
mask_image /= 255.0
result_image = scaled_paste(
image_background=bkg_image,
image_foreground=subject_image,
mask_foreground=mask_image,
scale_factor=scale_factor,
height_factor=height_factor,
)
luminosity_image = scaled_paste(
image_background=blank_background_image + 255,
image_foreground=subject_image,
mask_foreground=blank_subject_mask,
scale_factor=scale_factor,
height_factor=height_factor,
)
final_mask = scaled_paste(
image_background=blank_background_image,
image_foreground=mask_image_3channel,
mask_foreground=mask_image,
scale_factor=scale_factor,
height_factor=height_factor,
)
result_image_lab = cv2.cvtColor(src=result_image, code=cv2.COLOR_RGB2LAB)
luminosity_image_lab = cv2.cvtColor(src=luminosity_image,
code=cv2.COLOR_RGB2LAB)[:, :, 0]
luminosity_image_lab_flip = 255 - luminosity_image_lab
luminosity_image_lab_flip = do_custom_threshhold(
image=luminosity_image_lab_flip, value=threshhold_hist)
luminosity_image_lab_flip_full = luminosity_image_lab_flip.copy()
luminosity_image_lab_flip *= 1 - (final_mask[:, :, 0]
> 127.5).astype(dtype=np.uint8)
for i in range(3):
result_image[:, :,
i] = (result_image[:, :, i] *
(1 - (luminosity_image_lab_flip / 255.0))).astype(
dtype=np.uint8)
return (result_image, luminosity_image_lab_flip_full)
def find_threshold(image_input, threshold=0.0001):
image_input_L = cv2.cvtColor(image_input, cv2.COLOR_RGB2LAB)[:, :,
0].flatten()
image_input_L = 255 - image_input_L
hist = np.histogram(image_input_L, range(0, 256, 1))
values = hist[0]
values = values.astype(dtype=np.float64)
values /= len(image_input_L)
for i in range(0, values.shape[0], 1):
lhd = 0
rhd = 0
if i > 0:
lhd = values[i] - values[i - 1]
if i < values.shape[0] - 1:
rhd = values[i + 1] - values[i]
print(lhd, rhd)
if max(lhd, rhd) > threshold:
return i
def get_mu_sigma(array_input, mask_input):
array_input = array_input.astype(dtype=np.float32).flatten()
mask_input = mask_input.astype(dtype=np.float32).flatten()
sum = np.sum(mask_input)
mean = np.sum(array_input * mask_input) / sum
array_input -= mean
array_input *= mask_input
sigma = math.sqrt(np.sum(np.square(array_input)) / sum)
return mean, sigma
def renormalize_array_main(array_input, mask_input, mu, sigma):
array_input_original = array_input.copy()
in_mu, in_sigma = get_mu_sigma(array_input, mask_input)
array_input = (((array_input - in_mu) / in_sigma) * sigma) + mu
array_input_original = (array_input_original *
(1 - mask_input)) + (array_input * mask_input)
return array_input_original
#!/usr/bin/python3
class simple_bg_swap:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"bkg_image": ("IMAGE", ),
"subject_image": ("IMAGE", ),
"subject_mask": ("MASK", ),
"threshhold_hist": (
"INT",
{
"default": 150,
"min": 0, #Minimum value
"max": 255, #Maximum value
"step": 1, #Slider's step
"display":
"number" # Cosmetic only: display as "number" or "slider"
}),
"scale_factor": (
"FLOAT",
{
"default": 1.2,
"min": 0.0,
"max": 10.0,
"step": 0.01,
"round":
0.001, #The value represeting the precision to round to, will be set to the step value by default. Can be set to False to disable rounding.
"display": "number"
}),
"height_factor": (
"FLOAT",
{
"default": 1.05,
"min": 1.0,
"max": 8.0,
"step": 0.01,
"round":
0.001, #The value represeting the precision to round to, will be set to the step value by default. Can be set to False to disable rounding.
"display": "number"
}),
},
}
RETURN_TYPES = (
"IMAGE",
"MASK",
)
RETURN_NAMES = (
"output bg swapped image",
"shadow layer",
)
FUNCTION = "test"
#OUTPUT_NODE = False
CATEGORY = "TRI3D"
def test(
self,
bkg_image,
subject_image,
subject_mask,
threshhold_hist,
scale_factor,
height_factor,
):
mask_image = subject_mask
bkg_image = from_torch_image(image=bkg_image)
subject_image = from_torch_image(image=subject_image)
mask_image = from_torch_image(image=mask_image)
batch_size = bkg_image.shape[0]
ret = []
ret_lum = []
if (subject_image.shape[0] == batch_size) and (mask_image.shape[0]
== batch_size):
for i in range(batch_size):
result, luminosity = do_bg_swap(
bkg_image[i],
subject_image[i],
mask_image[i],
threshhold_hist,
scale_factor,
height_factor,
)
result = to_torch_image(result)
result = result.unsqueeze(0)
ret.append(result)
luminosity = to_torch_image(luminosity)
luminosity = luminosity.unsqueeze(0)
ret_lum.append(luminosity)
else:
print(
'input format is not correct, got different batch sizes for each input image'
)
ret = torch.cat(ret, dim=0)
ret_lum = torch.cat(ret_lum, dim=0)
return (
ret,
ret_lum,
)
class get_threshold_for_bg_swap:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"subject_image": ("IMAGE", ),
"gradient_threshold": (
"FLOAT",
{
"default": 0.0001,
"min": 0.0,
"max": 1.0,
"step": 0.00001,
"round":
0.000001, #The value represeting the precision to round to, will be set to the step value by default. Can be set to False to disable rounding.
"display": "number"
}),
},
}
RETURN_TYPES = ("INT", )
RETURN_NAMES = ("output histogram threshold", )
FUNCTION = "test"
#OUTPUT_NODE = False
CATEGORY = "TRI3D"
def test(
self,
subject_image,
gradient_threshold,
):
subject_image = from_torch_image(image=subject_image)
return (find_threshold(subject_image[0],
threshold=gradient_threshold), )
class RGB_2_LAB:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"input_RGB_image": ("IMAGE", ),
},
}
RETURN_TYPES = ("MASK", "MASK", "MASK")
RETURN_NAMES = ("L", "A", "B")
FUNCTION = "test"
#OUTPUT_NODE = False
CATEGORY = "TRI3D"
def test(self, input_RGB_image):
print('input_RGB_image.shape', input_RGB_image.shape)
input_RGB_image = from_torch_image(image=input_RGB_image)
ret_L = []
ret_A = []
ret_B = []
for i in range(input_RGB_image.shape[0]):
tmp = cv2.cvtColor(input_RGB_image[i], cv2.COLOR_RGB2LAB)
ret_L.append(to_torch_image(image=tmp[:, :, 0]).unsqueeze(0))
ret_A.append(to_torch_image(image=tmp[:, :, 1]).unsqueeze(0))
ret_B.append(to_torch_image(image=tmp[:, :, 2]).unsqueeze(0))
ret_L = torch.cat(ret_L, dim=0)
ret_A = torch.cat(ret_A, dim=0)
ret_B = torch.cat(ret_B, dim=0)
print(
'LAB output',
ret_L.shape,
ret_A.shape,
ret_B.shape,
)
return (ret_L, ret_A, ret_B)
class LAB_2_RGB:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"input_L": ("MASK", ),
"input_A": ("MASK", ),
"input_B": ("MASK", ),
},
}
RETURN_TYPES = ("IMAGE", )
RETURN_NAMES = ("Output RGB image", )
FUNCTION = "test"
#OUTPUT_NODE = False
CATEGORY = "TRI3D"
def test(self, input_L, input_A, input_B):
batch_size = input_L.shape[0]
print(input_L.shape, input_A.shape, input_B.shape)
ret = []
if (input_A.shape[0] == batch_size) and (input_B.shape[0]
== batch_size):
for i in range(batch_size):
input_L_NP = from_torch_image(image=input_L[i])
input_A_NP = from_torch_image(image=input_A[i])
input_B_NP = from_torch_image(image=input_B[i])
Y_MAX = input_L_NP.shape[0]
X_MAX = input_L_NP.shape[1]
if (input_A_NP.shape[0]
== Y_MAX) and (input_B_NP.shape[0] == Y_MAX) and (
(input_A_NP.shape[1] == X_MAX) and
(input_B_NP.shape[1] == X_MAX)):
image = np.zeros((Y_MAX, X_MAX, 3), dtype=np.uint8)
image[:, :, 0] = input_L_NP
image[:, :, 1] = input_A_NP
image[:, :, 2] = input_B_NP
image = cv2.cvtColor(image, cv2.COLOR_LAB2RGB)
image = to_torch_image(image).unsqueeze(0)
print('image.shape')
print(image.shape)
ret.append(image)
else:
print('Resolution of different layers donot match')
else:
print('batch size of different layers donot match')
ret = torch.cat(ret, dim=0)
print('ret.shape', ret.shape)
return (ret, )
class get_mean_and_standard_deviation:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"input_array": ("MASK", ),
"input_mask": ("MASK", ),
},
}
RETURN_TYPES = (
"FLOAT",
"FLOAT",
)
RETURN_NAMES = (
"Mean",
"Standard deviation",
)
FUNCTION = "test"
CATEGORY = "TRI3D"
def test(self, input_array, input_mask):
input_array = input_array.cpu().numpy()
input_mask = input_mask.cpu().numpy()
mean, sigma = get_mu_sigma(array_input=input_array[0],
mask_input=input_mask[0])
return (
mean,
sigma,
)
class renormalize_array:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"input_array": ("MASK", ),
"input_mask": ("MASK", ),
"input_mean": (
"FLOAT",
{
"default": 1.0,
"min": 0.0,
"max": 2.0,
"step": 0.001,
"round":
0.000001, #The value represeting the precision to round to, will be set to the step value by default. Can be set to False to disable rounding.
"display": "number"
}),
"input_standard_deviation": (
"FLOAT",
{
"default": 1.0,
"min": 0.0,
"max": 2.0,
"step": 0.001,
"round":
0.000001, #The value represeting the precision to round to, will be set to the step value by default. Can be set to False to disable rounding.
"display": "number"
}),
},
}
RETURN_TYPES = ("MASK", )
RETURN_NAMES = ("Output array as mask", )
FUNCTION = "test"
#OUTPUT_NODE = False
CATEGORY = "TRI3D"
def test(
self,
input_array,
input_mask,
input_mean,
input_standard_deviation,
):
batch_size = input_array.shape[0]
ret = []
if input_mask.shape[0] == batch_size:
for i in range(batch_size):
tmp = renormalize_array_main(
array_input=input_array[i].cpu().numpy(),
mask_input=input_mask[i].cpu().numpy(),
mu=input_mean,
sigma=input_standard_deviation)
tmp = torch.from_numpy(tmp)
tmp = tmp.unsqueeze(0)
ret.append(tmp)
else:
print('batch size of different layers donot match')
ret = torch.cat(ret, dim=0)
return (ret, )
-820
View File
@@ -1,820 +0,0 @@
import numpy as np
import torch
import json
import cv2
# {0, "Nose"},
# // {1, "Neck"},
# // {2, "RShoulder"},
# // {3, "RElbow"},
# // {4, "RWrist"},
# // {5, "LShoulder"},
# // {6, "LElbow"},
# // {7, "LWrist"},
# // {8, "MidHip"},
# // {9, "RHip"},
# // {10, "RKnee"},
# // {11, "RAnkle"},
# // {12, "LHip"},
# // {13, "LKnee"},
# // {14, "LAnkle"},
# // {15, "REye"},
# // {16, "LEye"},
# // {17, "REar"},
# // {18, "LEar"},
# // {19, "LBigToe"},
# // {20, "LSmallToe"},
# // {21, "LHeel"},
# // {22, "RBigToe"},
# // {23, "RSmallToe"},
# // {24, "RHeel"},
# // {25, "Background"}
class TRI3D_SmartBox:
def from_torch_image(self, image):
image = image.cpu().numpy() * 255.0
image = np.clip(image, 0, 255).astype(np.uint8)
return image
def to_torch_image(self, image):
image = image.astype(dtype=np.float32)
image /= 255.0
image = torch.from_numpy(image)
return image
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE", ),
"keypoints_json": ("STRING", {"multiline": True}),
},
}
FUNCTION = "run"
RETURN_TYPES = ("IMAGE", )
CATEGORY = "TRI3D"
def extract_torso_keypoints(self, keypoints):
# Indices for torso-related keypoints
torso_indices = [8, 9, 10, 11, 12, 13]
return [keypoints[i] for i in torso_indices]
def run(self, image, keypoints_json):
kp_data = json.loads(open(keypoints_json, 'r').read())
original_height, original_width = kp_data['height'], kp_data['width']
torso_keypoints = self.extract_torso_keypoints(kp_data['keypoints'])
# Convert Torch image to OpenCV format
cv_image = self.from_torch_image(image)
# Remove the batch dimension if present
if len(cv_image.shape) == 4:
cv_image = cv_image[0]
# Adjust keypoints to match the image dimensions
adjusted_keypoints = self.adjust_keypoints(torso_keypoints, cv_image.shape, original_height, original_width)
# Fill the area below the hip line
filled_image = self.fill_below_hip(cv_image, adjusted_keypoints)
# Convert back to Torch format
torch_image = self.to_torch_image(filled_image)
# Add the batch dimension back
torch_image = torch_image.unsqueeze(0)
return (torch_image,)
def adjust_keypoints(self, keypoints, image_shape, original_height, original_width):
image_height, image_width = image_shape[:2]
scale_x = image_width / original_width
scale_y = image_height / original_height
adjusted_keypoints = [
(int(x * scale_x), int(y * scale_y)) for x, y in keypoints
]
return adjusted_keypoints
def fill_below_hip(self, image, keypoints):
# Correct the indices for hip keypoints
# Assuming indices 8 and 11 are for left and right hips
# print(keypoints,"hip keypoints")
try:
valid_y_coords = [kp[1] for kp in keypoints if kp[1] >= 0]
hip_y = min(valid_y_coords) if valid_y_coords else 0
except:
hip_y = 0
if hip_y == 0:
return image
# Find the bounding box of the mask below the hip line
mask = image[:, :, 0] # Assuming single-channel mask
below_hip = mask[hip_y:, :]
contours, _ = cv2.findContours(below_hip, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
cnt = 0
for contour in contours:
x, y, w, h = cv2.boundingRect(contour)
# print(cnt, x,y,w,h, cv2.contourArea(contour), "cnt,x,y,w,h,area")
cnt+=1
# cv2.rectangle(image, (x, y + hip_y), (x + w, y + h + hip_y), (255, 255, 255), -1)
contours = [contour for contour in contours if cv2.contourArea(contour) > 20]
if len(contours) == 0:
return image
# Combine all contours into one
all_contours = np.vstack(contours)
# Calculate a single bounding rectangle for all contours
x, y, w, h = cv2.boundingRect(all_contours)
# print(x,y,w,h, "x,y,w,h")
cv2.rectangle(image, (x, y + hip_y), (x + w, y + h + hip_y), (255, 255, 255), -1)
return image
class TRI3D_Skip_HeadMask:
def from_torch_image(self, image):
image = image.cpu().numpy() * 255.0
image = np.clip(image, 0, 255).astype(np.uint8)
return image
def to_torch_image(self, image):
image = image.astype(dtype=np.float32)
image /= 255.0
image = torch.from_numpy(image)
return image
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE", ),
"head_mask": ("IMAGE", ),
},
}
FUNCTION = "run"
RETURN_TYPES = ("IMAGE", )
CATEGORY = "TRI3D"
def run(self, image, head_mask):
# Convert Torch images to OpenCV format
cv_image = self.from_torch_image(image)
cv_head_mask = self.from_torch_image(head_mask)
# Remove the batch dimension if present
if len(cv_image.shape) == 4:
cv_image = cv_image[0]
if len(cv_head_mask.shape) == 4:
cv_head_mask = cv_head_mask[0]
# Find the lowest point in the head mask
mask = cv_head_mask[:, :, 0] # Assuming single-channel mask
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
lowest_y = 0
for contour in contours:
for point in contour:
x, y = point[0]
if y > lowest_y:
lowest_y = y
# Black out everything above the lowest point
cv_image[:lowest_y, :] = 0
# Convert back to Torch format
torch_image = self.to_torch_image(cv_image)
# Add the batch dimension back
torch_image = torch_image.unsqueeze(0)
return (torch_image,)
class TRI3D_Skip_HeadMask_AddNeck:
def adjust_keypoints(self, keypoints, image_shape, original_height, original_width):
image_height, image_width = image_shape[:2]
scale_x = image_width / original_width
scale_y = image_height / original_height
adjusted_keypoints = [
(int(x * scale_x), int(y * scale_y)) for x, y in keypoints
]
return adjusted_keypoints
def from_torch_image(self, image):
image = image.cpu().numpy() * 255.0
image = np.clip(image, 0, 255).astype(np.uint8)
return image
def to_torch_image(self, image):
image = image.astype(dtype=np.float32)
image /= 255.0
image = torch.from_numpy(image)
return image
def extract_neck_keypoint(self, keypoints):
# Indices for torso-related keypoints
neck_indices = [1]
return [keypoints[i] for i in neck_indices]
def extract_ear_keypoints(self, keypoints):
# Indices for ear keypoints (17=right ear, 18=left ear)
ear_indices = [17, 18]
return [keypoints[i] for i in ear_indices]
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE", ),
"head_mask": ("IMAGE", ),
"keypoints_json": ("STRING", {"multiline": True}),
"ratio_aggression": ("FLOAT", {"default": 0.5, "min": 0, "max": 1, "step": 0.01}),
"neck_width_factor": ("FLOAT", {"default": 0.8, "min": 0.1, "max": 1.5, "step": 0.05}),
},
}
FUNCTION = "run"
RETURN_TYPES = ("IMAGE", )
CATEGORY = "TRI3D"
def run(self, image, head_mask, keypoints_json, ratio_aggression, neck_width_factor):
# Convert Torch images to OpenCV format
cv_image = self.from_torch_image(image)
cv_head_mask = self.from_torch_image(head_mask)
# Remove the batch dimension if present
if len(cv_image.shape) == 4:
cv_image = cv_image[0]
if len(cv_head_mask.shape) == 4:
cv_head_mask = cv_head_mask[0]
kp_data = json.loads(open(keypoints_json, 'r').read())
original_height, original_width = kp_data['height'], kp_data['width']
neck_keypoints = self.extract_neck_keypoint(kp_data['keypoints'])
# Make a copy of the original image
result_image = cv_image.copy()
# Adjust keypoints to match the image dimensions
adjusted_neck_keypoints = self.adjust_keypoints(neck_keypoints, cv_image.shape, original_height, original_width)
# Find the lowest point and face dimensions in the head mask
mask = cv_head_mask[:, :, 0] # Assuming single-channel mask
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
# Find the chin point (lowest point) and calculate face properties
lowest_y = 0
face_center_x = cv_image.shape[1] // 2 # Default to center of image
face_width = cv_image.shape[1] // 3 # Default face width
if contours:
# Find the lowest point (chin)
for contour in contours:
for point in contour:
x, y = point[0]
if y > lowest_y:
lowest_y = y
# Calculate face bounding box and center of gravity
x, y, w, h = cv2.boundingRect(contours[0])
face_width = w
# Calculate center of gravity of the face mask
M = cv2.moments(contours[0])
if M["m00"] != 0:
face_center_x = int(M["m10"] / M["m00"])
else:
face_center_x = x + w // 2
# Calculate weighted average point between neck and chin
neck_y = adjusted_neck_keypoints[0][1]
if neck_y <= 0:
neck_y = lowest_y
average_y = int((neck_y * ratio_aggression + lowest_y * (1 - ratio_aggression)))
print(neck_y, lowest_y, "neck_y, lowest_y")
print(average_y, "average_y")
# ZONE 1: Black out everything above the chin point
result_image[:lowest_y, :] = 0
# ZONE 2: Create a triangle for the neck area
if lowest_y < average_y: # Only process if there's a gap between chin and average_y
# Create a mask for Zone 2
zone2_mask = np.zeros_like(cv_image[:,:,0])
# Create a triangle with apex at weighted average point and base at chin level
# Apply the neck width factor to the face width
neck_width = int(face_width * neck_width_factor)
triangle_half_width = neck_width // 2
# Create polygon points for the triangle
triangle_points = np.array([
[face_center_x, average_y], # Apex at weighted average point
[face_center_x - triangle_half_width, lowest_y], # Left base point at chin level
[face_center_x + triangle_half_width, lowest_y] # Right base point at chin level
], dtype=np.int32)
# Fill the triangle in the mask
cv2.fillPoly(zone2_mask, [triangle_points], 255)
# Apply the mask only to the region between chin and weighted average
for y in range(lowest_y, average_y):
for x in range(cv_image.shape[1]):
if zone2_mask[y, x] > 0:
result_image[y, x] = 0
# ZONE 3: Area below weighted average point is left as is
# No action needed for this zone
# Convert back to Torch format
torch_image = self.to_torch_image(result_image)
# Add the batch dimension back
torch_image = torch_image.unsqueeze(0)
return (torch_image,)
class TRI3D_Image_extend:
def from_torch_image(self, image):
image = image.cpu().numpy() * 255.0
image = np.clip(image, 0, 255).astype(np.uint8)
return image
def to_torch_image(self, image):
image = image.astype(dtype=np.float32)
image /= 255.0
image = torch.from_numpy(image)
return image
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"face_mask": ("IMAGE", ),
"image": ("IMAGE", ),
"ratio": ("FLOAT", {"default": 1.5, "min": 1.2, "max": 2, "step": 0.01}),
},
}
FUNCTION = "run"
RETURN_TYPES = ("IMAGE", "IMAGE", )
RETURN_NAMES = ("image", "mask_image", )
CATEGORY = "TRI3D"
def run(self, face_mask, image, ratio):
cv_face_mask = self.from_torch_image(face_mask)
cv_image = self.from_torch_image(image)
# Remove the batch dimension if present
if len(cv_image.shape) == 4:
cv_image = cv_image[0]
if len(cv_face_mask.shape) == 4:
cv_face_mask = cv_face_mask[0]
mask = cv_face_mask[:, :, 0] # Assuming single-channel mask
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
lowest_y = 0
highest_y = cv_image.shape[0]
for contour in contours:
for point in contour:
x, y = point[0]
if y > lowest_y:
lowest_y = y
if y < highest_y:
highest_y = y
y_below_face = cv_image.shape[0] - lowest_y
y_face = lowest_y-highest_y
# Only extend if the space below face is less than 1.5 times face height
target_below_face = int(y_face * ratio)
# print("y_face", y_face)
# print("lowest_y", lowest_y)
# print("highest_y", highest_y)
# print("target_below_face", target_below_face)
# print("y_below_face", y_below_face)
original_height = cv_image.shape[0]
original_width = cv_image.shape[1]
if y_below_face < target_below_face:
y_extend = target_below_face - y_below_face
# Calculate how much to extend horizontally to maintain aspect ratio
new_height = original_height + y_extend
new_width = int(original_width * (new_height / original_height))
x_extend = new_width - original_width
x_extend_left = x_extend // 2
x_extend_right = x_extend - x_extend_left
# Extend the image in all necessary directions
cv_image = cv2.copyMakeBorder(
cv_image,
0, y_extend, # top, bottom
x_extend_left, x_extend_right, # left, right
cv2.BORDER_CONSTANT,
value=[0, 0, 0]
)
# Create extension mask
extension_mask = np.zeros_like(cv_image)
# Make extended portions white
extension_mask[original_height:, :] = 255 # bottom extension
extension_mask[:, :x_extend_left] = 255 # left extension
extension_mask[:, -x_extend_right:] = 255 # right extension
else:
extension_mask = np.zeros_like(cv_image)
# Convert both images back to torch format
torch_image = self.to_torch_image(cv_image)
torch_mask = self.to_torch_image(extension_mask)
# Add batch dimension to both
torch_image = torch_image.unsqueeze(0)
torch_mask = torch_mask.unsqueeze(0)
return (torch_image, torch_mask)
class TRI3D_Smart_Depth:
def from_torch_image(self, image):
image = image.cpu().numpy() * 255.0
image = np.clip(image, 0, 255).astype(np.uint8)
return image
def to_torch_image(self, image):
image = image.astype(dtype=np.float32)
image /= 255.0
image = torch.from_numpy(image)
return image
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE", ),
"keypoints_json": ("STRING", {"multiline": True}),
},
}
FUNCTION = "run"
RETURN_TYPES = ("IMAGE", )
CATEGORY = "TRI3D"
def extract_torso_keypoints(self, keypoints):
# Indices for torso-related keypoints
torso_indices = [8, 9, 10, 11, 12, 13]
return [keypoints[i] for i in torso_indices]
def run(self, image, keypoints_json):
kp_data = json.loads(open(keypoints_json, 'r').read())
original_height, original_width = kp_data['height'], kp_data['width']
torso_keypoints = self.extract_torso_keypoints(kp_data['keypoints'])
# Convert Torch image to OpenCV format
cv_image = self.from_torch_image(image)
# Remove the batch dimension if present
if len(cv_image.shape) == 4:
cv_image = cv_image[0]
# Adjust keypoints to match the image dimensions
adjusted_keypoints = self.adjust_keypoints(torso_keypoints, cv_image.shape, original_height, original_width)
# Fill the area below the hip line
filled_image = self.fill_below_hip(cv_image, adjusted_keypoints)
# Convert back to Torch format
torch_image = self.to_torch_image(filled_image)
# Add the batch dimension back
torch_image = torch_image.unsqueeze(0)
return (torch_image,)
def adjust_keypoints(self, keypoints, image_shape, original_height, original_width):
image_height, image_width = image_shape[:2]
scale_x = image_width / original_width
scale_y = image_height / original_height
adjusted_keypoints = [
(int(x * scale_x), int(y * scale_y)) for x, y in keypoints
]
return adjusted_keypoints
def fill_below_hip(self, image, keypoints):
# Correct the indices for hip keypoints
# Assuming indices 8 and 11 are for left and right hips
# print(keypoints,"hip keypoints")
try:
valid_y_coords = [kp[1] for kp in keypoints if kp[1] >= 0]
hip_y = min(valid_y_coords) if valid_y_coords else 0
except:
hip_y = 0
if hip_y == 0:
return image
# Find the bounding box of the mask below the hip line
mask = image[:, :, 0] # Assuming single-channel mask
below_hip = mask[hip_y:, :]
contours, _ = cv2.findContours(below_hip, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
cnt = 0
for contour in contours:
x, y, w, h = cv2.boundingRect(contour)
# print(cnt, x,y,w,h, cv2.contourArea(contour), "cnt,x,y,w,h,area")
cnt+=1
# cv2.rectangle(image, (x, y + hip_y), (x + w, y + h + hip_y), (255, 255, 255), -1)
contours = [contour for contour in contours if cv2.contourArea(contour) > 0]
if len(contours) == 0:
return image
# Combine all contours into one
all_contours = np.vstack(contours)
# Calculate a single bounding rectangle for all contours
x, y, w, h = cv2.boundingRect(all_contours)
# print(x,y,w,h, "x,y,w,h")
cv2.rectangle(image, (x, y + hip_y), (x + w, y + h + hip_y), (0, 0, 0), -1)
return image
class TRI3D_NarrowfyImage:
def from_torch_image(self, image):
image = image.cpu().numpy() * 255.0
image = np.clip(image, 0, 255).astype(np.uint8)
return image
def to_torch_image(self, image):
image = image.astype(dtype=np.float32)
image /= 255.0
image = torch.from_numpy(image)
return image
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE", ),
"mask": ("IMAGE", ),
"aspect_ratio": ("FLOAT", {"default": 0.33, "min": 0.25, "max": 1, "step": 0.01}),
"border_margin": ("INT", {"default": 15, "min": 10, "max": 100, "step": 1}),
},
}
FUNCTION = "run"
RETURN_TYPES = ("IMAGE", "IMAGE", "INT", "INT",)
RETURN_NAMES = ("cropped_image", "cropped_mask", "cropped_width", "cropped_height",)
CATEGORY = "TRI3D"
def run(self, image, mask, aspect_ratio, border_margin):
# Convert to CV format and remove batch dimension
cv_image = self.from_torch_image(image)[0]
cv_mask = self.from_torch_image(mask)[0]
# Find bounding box of the mask
mask_channel = cv_mask[:, :, 0]
contours, _ = cv2.findContours(mask_channel, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
if not contours:
return image, mask, aspect_ratio
# Filter contours by area
significant_contours = [cnt for cnt in contours if cv2.contourArea(cnt) > 100]
if not significant_contours:
return image, mask, aspect_ratio
# Get combined bounding box for all significant contours
x_min = float('inf')
y_min = float('inf')
x_max = 0
y_max = 0
for contour in significant_contours:
x, y, w, h = cv2.boundingRect(contour)
x_min = min(x_min, x)
y_min = min(y_min, y)
x_max = max(x_max, x + w)
y_max = max(y_max, y + h)
# Calculate final width and height with margin
margin = border_margin
x = max(0, x_min - margin) # Ensure we don't go below 0
y = max(0, y_min - margin)
w = min(cv_image.shape[1] - x, (x_max - x_min) + 2 * margin) # Ensure we don't exceed image width
h = min(cv_image.shape[0] - y, (y_max - y_min) + 2 * margin) # Ensure we don't exceed image height
# Crop both image and mask to bounding box
cropped_image = cv_image[y:y+h, x:x+w]
cropped_mask = cv_mask[y:y+h, x:x+w]
# Calculate required height for aspect ratio 1/3
min_height = w * 1/aspect_ratio
if h < min_height:
height_extend = min_height - h
# Extend image with black pixels
extended_image = cv2.copyMakeBorder(
cropped_image,
0, int(height_extend), # top, bottom
0, 0, # left, right
cv2.BORDER_CONSTANT,
value=[0, 0, 0]
)
# Create mask with white pixels only in extended region
extended_mask = cv2.copyMakeBorder(
np.zeros_like(cropped_mask), # Start with black base
0, int(height_extend), # top, bottom
0, 0, # left, right
cv2.BORDER_CONSTANT,
value=[255, 255, 255] # White extension
)
cropped_image = extended_image
cropped_mask = extended_mask
# Convert back to torch format and add batch dimension
torch_image = self.to_torch_image(cropped_image).unsqueeze(0)
torch_mask = self.to_torch_image(cropped_mask).unsqueeze(0)
return (torch_image, torch_mask,w,h)
class TRI3D_CropAndExtend:
def from_torch_image(self, image):
image = image.cpu().numpy() * 255.0
image = np.clip(image, 0, 255).astype(np.uint8)
return image
def to_torch_image(self, image):
image = image.astype(dtype=np.float32)
image /= 255.0
image = torch.from_numpy(image)
return image
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"garment_image": ("IMAGE",),
"garment_mask": ("IMAGE",),
"human_image": ("IMAGE",),
"human_mask": ("IMAGE",),
"margin": ("INT", {"default": 10, "min": 0, "max": 50}),
},
}
FUNCTION = "run"
RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "IMAGE", "INT", "INT",)
RETURN_NAMES = ("cropped_garment", "cropped_garment_mask", "cropped_human", "cropped_human_mask", "cropped_width", "cropped_height",)
def run(self, garment_image, garment_mask, human_image, human_mask, margin):
# Convert to CV format and remove batch dimension
cv_garment = self.from_torch_image(garment_image)[0]
cv_garment_mask = self.from_torch_image(garment_mask)[0]
cv_human = self.from_torch_image(human_image)[0]
cv_human_mask = self.from_torch_image(human_mask)[0]
# Process garment
mask_channel = cv_garment_mask[:, :, 0]
contours, _ = cv2.findContours(mask_channel, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
if not contours:
return garment_image, garment_mask, human_image, human_mask, cv_garment.shape[1], cv_garment.shape[0]
# Get bounding box with margin
x, y, w, h = cv2.boundingRect(contours[0])
x = max(0, x - margin)
y = max(0, y - margin)
w = min(cv_garment.shape[1] - x, w + 2 * margin)
h = min(cv_garment.shape[0] - y, h + 2 * margin)
# Store the cropped dimensions before extension
cropped_width = w
cropped_height = h
# Crop garment and its mask
cropped_garment = cv_garment[y:y+h, x:x+w]
cropped_garment_mask = cv_garment_mask[y:y+h, x:x+w]
# Calculate required height for aspect ratio 1/3
min_height = w * 3
if h < min_height:
height_extend = min_height - h
# Extend garment image and mask
extended_garment = cv2.copyMakeBorder(
cropped_garment,
0, int(height_extend),
0, 0,
cv2.BORDER_CONSTANT,
value=[0, 0, 0]
)
extended_garment_mask = cv2.copyMakeBorder(
cropped_garment_mask,
0, int(height_extend),
0, 0,
cv2.BORDER_CONSTANT,
value=[255, 255, 255]
)
cropped_garment = extended_garment
cropped_garment_mask = extended_garment_mask
# Process human image similarly
mask_channel = cv_human_mask[:, :, 0]
contours, _ = cv2.findContours(mask_channel, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
if contours:
x, y, w, h = cv2.boundingRect(contours[0])
x = max(0, x - margin)
y = max(0, y - margin)
w = min(cv_human.shape[1] - x, w + 2 * margin)
h = min(cv_human.shape[0] - y, h + 2 * margin)
cropped_human = cv_human[y:y+h, x:x+w]
cropped_human_mask = cv_human_mask[y:y+h, x:x+w]
min_height = w * 3
if h < min_height:
height_extend = min_height - h
extended_human = cv2.copyMakeBorder(
cropped_human,
0, int(height_extend),
0, 0,
cv2.BORDER_CONSTANT,
value=[0, 0, 0]
)
extended_human_mask = cv2.copyMakeBorder(
cropped_human_mask,
0, int(height_extend),
0, 0,
cv2.BORDER_CONSTANT,
value=[255, 255, 255]
)
cropped_human = extended_human
cropped_human_mask = extended_human_mask
# Convert back to torch format and add batch dimension
torch_garment = self.to_torch_image(cropped_garment).unsqueeze(0)
torch_garment_mask = self.to_torch_image(cropped_garment_mask).unsqueeze(0)
torch_human = self.to_torch_image(cropped_human).unsqueeze(0)
torch_human_mask = self.to_torch_image(cropped_human_mask).unsqueeze(0)
return (torch_garment, torch_garment_mask, torch_human, torch_human_mask, cropped_width, cropped_height)
-409
View File
@@ -1,409 +0,0 @@
import torch, cv2, json
import numpy as np
def from_torch_image(image):
image = image.cpu().numpy() * 255.0
image = np.clip(image, 0, 255).astype(np.uint8)
return image
def to_torch_image(image):
image = image.astype(dtype=np.float32)
image /= 255.0
image = torch.from_numpy(image)
return image
class TRI3D_clean_mask():
"""For the given mask and threshold area, remove all patches in the mask with area smaller than threshold"""
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"masks": ("MASK", ),
"threshold":("FLOAT",{"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01})
}
}
FUNCTION = "run"
RETURN_TYPES = ("MASK", "BOOL")
RETURN_NAMES = ("mask", "cleaned")
CATEGORY = "TRI3D"
def run(self, masks, threshold):
batch_results = []
for mask in masks:
mask = from_torch_image(mask)
mask = np.where(mask < 127, 0, 255).astype(np.uint8)
h,w = mask.shape[:2]
total_area = h*w
# num_labels, labels = cv2.connectedComponents(mask)
region_mask = np.zeros_like(mask)
# for label in range(1, num_labels):
# area_percent = (np.sum(labels == label)/ total_area) * 100
# if area_percent < threshold:
# continue
# region_mask[labels == label] = 255
less_than_threshold = True
area_percent = (np.sum(mask == 255)/ total_area) * 100
if area_percent > threshold:
region_mask[mask == 255] = 255
less_than_threshold = False
region_mask = to_torch_image(region_mask)
batch_results.append(region_mask.squeeze(0))
batch_results = torch.stack(batch_results)
return (batch_results, less_than_threshold)
class TRI3D_extract_pose_part():
"""
For the given pose, extract region around body parts, region can be defined by % of image size
"""
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE", ),
"pose_json": ("STRING",{"default" : "dwpose/keypoints/input.json"}),
"width_pad": ("FLOAT",{"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01}),
"height_pad": ("FLOAT",{"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01}),
"shoulders":("BOOLEAN", {
"default": False
})
}
}
FUNCTION = "run"
RETURN_TYPES = ("IMAGE", "STRING")
RETURN_NAMES = ("image", "coords")
CATEGORY = "TRI3D"
def get_frame_coords(self,point1, point2):
x1, y1 = point1
x2, y2 = point2
xmin, xmax, ymin, ymax = min(x1, x2), max(x1, x2), min(y1, y2), max(y1, y2)
for i in [xmin, xmax, ymin, ymax]:
if i < 0:
return None
return [xmin, xmax, ymin, ymax]
def run(self, image, pose_json, width_pad, height_pad, shoulders):
"""
image : input image
width_pad: % of image width you want to apply on both size of pose body part
height_pad: % of image width you want to apply on both size of pose body part
rest of them are body parts
"""
image = from_torch_image(image[0])
batch_result = []
input_pose = json.load(open(pose_json))
keypoints = input_pose['keypoints']
og_h, og_w = image.shape[:2]
ph, pw = [input_pose['height'], input_pose['width']]
for i,point in enumerate(keypoints):
x,y = point
y = int((y/ph)*og_h)
x = int((x/pw)*og_w)
keypoints[i] = [x, y]
width_offset = int(og_w * (width_pad) / 100)
height_offset = int(og_h * (height_pad) / 100)
xmin, xmax, ymin, ymax = [0, og_w, 0, og_h]
part_to_coords = {
"shoulders":self.get_frame_coords(keypoints[2], keypoints[5])
}
if shoulders:
print(part_to_coords["shoulders"])
if part_to_coords["shoulders"] != None:
new_xmin, new_xmax, new_ymin, new_ymax = part_to_coords["shoulders"]
xmin, xmax, ymin, ymax = new_xmin, new_xmax, new_ymin, new_ymax
xmin = max(0, xmin - width_offset)
xmax = min(og_w, xmax + width_offset)
ymin = max(0, ymin - height_offset)
ymax = min(og_h, ymax + height_offset)
image = image[ymin:ymax, xmin:xmax, :].astype(np.uint8)
image = to_torch_image(image)
batch_result.append(image)
batch_result = torch.stack(batch_result)
print("final_coords", xmin, xmax, ymin, ymax)
coords = ",".join([str(xmin), str(xmax), str(ymin), str(ymax)])
return batch_result, coords
class TRI3D_position_pose_part():
"""
put back extracted parts on OG image
"""
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"og_image": ("IMAGE", ),
"extracted_image": ("IMAGE", ),
"coords": ("STRING",{"default" : "xmin, xmax, ymin, ymax"}),
}
}
FUNCTION = "run"
RETURN_TYPES = ("IMAGE", )
RETURN_NAMES = ("image", )
CATEGORY = "TRI3D"
def run(self, og_image, extracted_image, coords):
batch_result = []
og_image = from_torch_image(og_image[0])
extracted_image = from_torch_image(extracted_image[0])
xmin, xmax, ymin, ymax = [int(i) for i in coords.split(",")]
og_image[ymin:ymax, xmin:xmax, :] = extracted_image
og_image = to_torch_image(og_image).unsqueeze(0)
batch_result.append(og_image)
batch_result = torch.stack(batch_result)
return batch_result
class TRI3D_fill_mask():
"""
fill mask with the neighbouring pixels
"""
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE", ),
"mask": ("MASK", ),
"negative_mask": ("MASK", ),
"offset":("FLOAT",{"default": 1, "min": 0.0, "max": 100.0, "step": 0.01})
}
}
FUNCTION = "run"
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("image",)
CATEGORY = "TRI3D"
def run(self, image, mask, negative_mask, offset):
image = from_torch_image(image[0])
mask = mask[0].cpu().numpy()
mask = np.expand_dims(mask, -1)
mh, mw, _ = mask.shape
inverse_mask = np.ones_like(mask) - mask
negative_mask = negative_mask[0].cpu().numpy()
indices = np.where(mask > 0)
offset = offset / 100
source = image.copy()
for y,x in zip(indices[0],indices[1]):
x_off = min(mw-1, int(x + offset * mw))
if negative_mask[y][x_off] == 0: #check if pixles on right are outside body
source[y][x] = image[y][x_off]
else:
x_off = max(0, int(x - offset * mw)) #check if pixles on left are outside body
if negative_mask[y][x_off] == 0:
source[y][x] = image[y][x_off]
else:
y_off = max(0, int(y - offset * mh))
if negative_mask[y_off][x] == 0: #check if pixles on top are outside body
source[y][x] = image[y_off][x]
else:
y_off = min(mh-1, int(y + offset * mh))
if negative_mask[y_off][x] == 0: #check if pixles on bottom are outside body
source[y][x] = image[y_off][x]
image = mask * source + inverse_mask * image
image = to_torch_image(image).unsqueeze(0)
return (image,)
class TRI3D_is_only_trouser:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"pose_json_file": ("STRING", {
"default": "dwpose/keypoints"
})
}
}
RETURN_TYPES = ("BOOLEAN", )
FUNCTION = "main"
CATEGORY = "TRI3D"
def main(self, pose_json_file):
pose = json.load(open(pose_json_file))
height = pose['height']
width = pose['width']
keypoints = pose['keypoints']
points = [0,14,15,16,17,2,1,5]
point_to_part = {0:'nose',14:"left eye",15:"right eye",16:"left ear",17:"right ear",2:"left shoulder",1:"neck",5:"right shoulder"}
all_negative = True #if all face and shoulder points are negative means it is a bottom shot
for point in points:
x,y = keypoints[point]
if x > 0 and y > 0:
all_negative = False
print(f"{point_to_part[point]} exist")
return (all_negative,)
class TRI3D_extract_facer_mask:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"background": ("BOOLEAN", {
"default": False
}),
'hair':("BOOLEAN", {
"default": False
}),
'lower_lip':("BOOLEAN", {
"default": False
}),
'inner_mouth':("BOOLEAN", {
"default": False
}),
'upper_lip':("BOOLEAN", {
"default": False
}),
'nose':("BOOLEAN", {
"default": False
}),
'left_eyebrow':("BOOLEAN", {
"default": False
}),
'right_eyebrow':("BOOLEAN", {
"default": False
}),
'left_eye':("BOOLEAN", {
"default": False
}),
'right_eye':("BOOLEAN", {
"default": False
}),
'face':("BOOLEAN", {
"default": False
})
}
}
RETURN_TYPES = ("MASK", )
FUNCTION = "main"
CATEGORY = "TRI3D"
def main(self, image, background, hair, lower_lip, inner_mouth, upper_lip, nose, left_eyebrow, right_eyebrow, left_eye, right_eye, face):
image = from_torch_image(image[0])
h,w,_ = image.shape
mask = np.zeros_like(image)
label_to_rgb = {'background':[0,0,0], 'face':[0,138,255], 'right_eye':[180, 255, 0], 'left_eye':[42, 255, 0], 'right_eyebrow':[0, 255, 96],
'left_eyebrow':[0,255,234], 'nose':[255, 192, 0], 'upper_lip':[255, 54, 0], 'inner_mouth':[255, 0, 84], 'lower_lip':[255, 0, 222],
'hair':[150,0,255]}
if background:
temp = np.all(image == label_to_rgb['background'], axis=-1)
idcs = np.where(temp==True)
mask[idcs] = 255
if face:
temp = np.all(image == label_to_rgb['face'], axis=-1)
idcs = np.where(temp==True)
mask[idcs] = 255
if right_eye:
temp = np.all(image == label_to_rgb['right_eye'], axis=-1)
idcs = np.where(temp==True)
mask[idcs] = 255
if left_eye:
temp = np.all(image == label_to_rgb['left_eye'], axis=-1)
idcs = np.where(temp==True)
mask[idcs] = 255
if right_eyebrow:
temp = np.all(image == label_to_rgb['right_eyebrow'], axis=-1)
idcs = np.where(temp==True)
mask[idcs] = 255
if left_eyebrow:
temp = np.all(image == label_to_rgb['left_eyebrow'], axis=-1)
idcs = np.where(temp==True)
mask[idcs] = 255
if nose:
temp = np.all(image == label_to_rgb['nose'], axis=-1)
idcs = np.where(temp==True)
mask[idcs] = 255
if upper_lip:
temp = np.all(image == label_to_rgb['upper_lip'], axis=-1)
idcs = np.where(temp==True)
mask[idcs] = 255
if inner_mouth:
temp = np.all(image == label_to_rgb['inner_mouth'], axis=-1)
idcs = np.where(temp==True)
mask[idcs] = 255
if lower_lip:
temp = np.all(image == label_to_rgb['lower_lip'], axis=-1)
idcs = np.where(temp==True)
mask[idcs] = 255
if hair:
temp = np.all(image == label_to_rgb['hair'], axis=-1)
idcs = np.where(temp==True)
mask[idcs] = 255
mask = to_torch_image(mask[:,:,0]).unsqueeze(0)
return (mask,)