Merge pull request #60 from TRI3D-LC/saveflattened

flattened

.env_sample

@@ -1 +1,3 @@
-CLIPDROP_API_KEY=
+CLIPDROP_API_KEY=
+COMFY_PYTHON_PATH=/home/ubuntu/.conda/envs/comfy/bin/python
+PHOTOROOM_API_KEY=3603b83dfa1846bc3c7270ead7876

.gitignore

@@ -5,5 +5,11 @@ venv
 .DS_Store
 checkpoints/
 checkpoint/
+.env
+.pth
+cloth-segmentation/model/cloth_segm.pth
 
-dwpose/keypoints/
+dwpose/keypoints/
+huggingface/
+
+safetychecker/model.safetensors

atr_node/simple_extractor.py

@@ -88,9 +88,29 @@ 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':
@@ -102,7 +122,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()

cloth-segmentation/app.py

@@ -0,0 +1,29 @@
+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):
+    palette = get_palette(4)
+    cloth_seg = generate_mask(img, net=net,device=device)
+    return cloth_seg
+
+INPUT_PATH = "./input/"
+OUTPUT_PATH = "./output/"
+
+import os 
+for cur_image in os.listdir(INPUT_PATH):
+    img = PIL.Image.open(INPUT_PATH + cur_image)
+    cloth_seg = run(img)
+    cloth_seg.save(OUTPUT_PATH + cur_image, format="PNG")

cloth-segmentation/model/ex.txt

@@ -0,0 +1 @@
+/*upload model */

cloth-segmentation/network.py

@@ -0,0 +1,560 @@
+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

cloth-segmentation/options.py

@@ -0,0 +1,12 @@
+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()

cloth-segmentation/process.py

@@ -0,0 +1,235 @@
+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'):
+    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
+
+    # 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)
+
+    # Create an RGBA image for the output
+    extracted_garment = Image.new("RGBA", img_size)
+    original_img = img.resize(img_size)  # Resize the processed image back to original size
+    extracted_garment.paste(original_img, mask=binary_mask_img)
+
+    # Save the garment image with transparency
+    garment_path = os.path.join(output_dir, 'extracted_garment.png')
+    extracted_garment.save(garment_path, format="PNG")
+
+    return 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)

cut_by_mask_aspect_ratio.py

@@ -0,0 +1,183 @@
+import os
+import cv2
+import numpy as np
+import torch
+
+class TRI3D_CutByMaskAspectRatio:
+    """
+    ComfyUI node that crops an image based on a mask's bounding box,
+    adjusts the aspect ratio, and resizes to specified dimensions.
+    """
+    
+    def from_torch_image(self, image):
+        """Convert a torch tensor image to numpy array for OpenCV processing"""
+        image = image.cpu().numpy() * 255.0
+        image = np.clip(image, 0, 255).astype(np.uint8)
+        return image
+
+    def to_torch_image(self, image):
+        """Convert numpy array back to torch tensor format"""
+        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",),
+                "margin": ("INT", {"default": 10, "min": 0, "max": 100, "step": 1}),
+                "target_width": ("INT", {"default": 512, "min": 64, "max": 4096, "step": 8}),
+                "target_height": ("INT", {"default": 512, "min": 64, "max": 4096, "step": 8}),
+                "padding_color": ("INT", {"default": 255, "min": 0, "max": 255, "step": 1}),
+            },
+        }
+
+    FUNCTION = "run"
+    RETURN_TYPES = ("IMAGE",)
+    CATEGORY = "TRI3D"
+
+    def run(self, image, mask, margin, target_width, target_height, padding_color=255):
+        # Convert Torch images to OpenCV format
+        cv_image = self.from_torch_image(image)
+        cv_mask = self.from_torch_image(mask)
+        
+        # Remove batch dimension if present
+        if len(cv_image.shape) == 4:
+            cv_image = cv_image[0]
+        if len(cv_mask.shape) == 4:
+            cv_mask = cv_mask[0]
+        
+        # Convert mask to grayscale if it's not already
+        if len(cv_mask.shape) == 3 and cv_mask.shape[2] > 1:
+            mask_gray = cv2.cvtColor(cv_mask, cv2.COLOR_RGB2GRAY)
+        else:
+            mask_gray = cv_mask[:, :, 0]
+        
+        # Create binary mask
+        _, binary_mask = cv2.threshold(mask_gray, 127, 255, cv2.THRESH_BINARY)
+        
+        # Find contours in the binary mask
+        contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        
+        if not contours:
+            # If no contours found, return the original image
+            print("No contours found in mask. Returning original image.")
+            return (image,)
+        
+        # Find bounding box around all contours
+        x_min, y_min = float('inf'), float('inf')
+        x_max, y_max = 0, 0
+        
+        for contour in 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)
+        
+        # Add margin to bounding box
+        x_min = max(0, x_min - margin)
+        y_min = max(0, y_min - margin)
+        x_max = min(cv_image.shape[1], x_max + margin)
+        y_max = min(cv_image.shape[0], y_max + margin)
+        
+        # Current dimensions of the bounding box
+        height = y_max - y_min
+        width = x_max - x_min
+        
+        # Calculate the target aspect ratio (width/height)
+        target_aspect_ratio = target_width / target_height
+        
+        # Calculate current aspect ratio
+        current_aspect_ratio = width / height
+        
+        # Adjust width to match the target aspect ratio while keeping height constant
+        if current_aspect_ratio < target_aspect_ratio:
+            # Current width is too narrow - need to extend it
+            # Calculate the required width for the target aspect ratio
+            required_width = int(height * target_aspect_ratio)
+            width_difference = required_width - width
+            
+            # Calculate how much to extend on each side
+            left_extend = width_difference // 2
+            right_extend = width_difference - left_extend
+            
+            # Calculate new potential boundaries
+            new_x_min = x_min - left_extend
+            new_x_max = x_max + right_extend
+            
+            # Check if the new boundaries are within the original image
+            left_padding_needed = abs(min(0, new_x_min))
+            right_padding_needed = max(0, new_x_max - cv_image.shape[1])
+            
+            # Adjust boundaries to be within the original image
+            new_x_min = max(0, new_x_min)
+            new_x_max = min(cv_image.shape[1], new_x_max)
+            
+            # Get the portion of the original image within valid boundaries
+            extended_image = cv_image[y_min:y_max, new_x_min:new_x_max]
+            
+            # If we need padding (i.e., extension goes beyond image boundaries)
+            if left_padding_needed > 0 or right_padding_needed > 0:
+                # Create canvas with padding color
+                num_channels = extended_image.shape[2] if len(extended_image.shape) == 3 else 1
+                if num_channels == 1:
+                    canvas = np.full((height, required_width), padding_color, dtype=np.uint8)
+                else:
+                    canvas = np.full((height, required_width, num_channels), padding_color, dtype=np.uint8)
+                
+                # Calculate the position to place the extended image
+                place_x = left_padding_needed
+                
+                # Place the extended image on the canvas
+                if num_channels == 1:
+                    canvas[:, place_x:place_x+extended_image.shape[1]] = extended_image
+                else:
+                    canvas[:, place_x:place_x+extended_image.shape[1], :] = extended_image
+                
+                # Use the canvas as our cropped image
+                cropped_image = canvas
+            else:
+                # No padding needed, use the extended image
+                cropped_image = extended_image
+            
+        elif current_aspect_ratio > target_aspect_ratio:
+            # Current width is too wide, crop it
+            new_width = int(height * target_aspect_ratio)
+            width_difference = width - new_width
+            
+            # Crop equally from both sides if possible
+            left_crop = width_difference // 2
+            right_crop = width_difference - left_crop
+            
+            # Apply the crop
+            cropped_image = cv_image[y_min:y_max, x_min+left_crop:x_max-right_crop]
+        else:
+            # Aspect ratio is already correct
+            cropped_image = cv_image[y_min:y_max, x_min:x_max]
+        
+        # Resize the cropped/padded image to the target dimensions using Lanczos interpolation
+        resized_image = cv2.resize(cropped_image, (target_width, target_height), interpolation=cv2.INTER_LANCZOS4)
+        
+        # Convert back to torch format
+        torch_image = self.to_torch_image(resized_image)
+        
+        # Add batch dimension back
+        torch_image = torch_image.unsqueeze(0)
+        
+        return (torch_image,)
+
+# Node registration for ComfyUI
+NODE_CLASS_MAPPINGS = {
+    "TRI3D_CutByMaskAspectRatio": TRI3D_CutByMaskAspectRatio
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "TRI3D_CutByMaskAspectRatio": "TRI3D Cut By Mask Aspect Ratio"
+}

distribution_reshape.py

@@ -0,0 +1,189 @@
+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"
+}

dwpose/comfy_utils.py

@@ -274,4 +274,6 @@ 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
+
+

dwpose_conversion.py

@@ -0,0 +1,111 @@
+import os
+import json
+import torch
+import numpy as np
+import folder_paths
+
+print("Loading TRI3D_SavePoseKeypointsJSON module")
+
+class SaveFlattenedPoseKpsAsJsonFile:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "pose_kps": ("POSE_KEYPOINT",),
+                "file_path": ("STRING", {"default": "dwpose/keypoints/input.json"})
+            }
+        }
+    RETURN_TYPES = (
+        "STRING",
+    )
+    FUNCTION = "save_flattened_pose_kps"
+    OUTPUT_NODE = True
+    CATEGORY = "ControlNet Preprocessors/Pose Keypoint Postprocess"
+
+    def __init__(self):
+        self.output_dir = folder_paths.get_output_directory()
+        self.type = "output"
+        self.prefix_append = ""
+
+    def _flatten_openpose_dict(self, pose_dict: dict) -> dict:
+        """
+        Converts a single OpenPose dictionary into flattened format.
+        """
+        # Get canvas dimensions from the input dictionary
+        H = pose_dict.get('canvas_height', 512)
+        W = pose_dict.get('canvas_width', 512)
+        
+        flat_keypoints = []
+
+        # Check if any person was detected
+        if not pose_dict.get('people'):
+            # If no people, return a list of 130 invalid keypoints
+            flat_keypoints.extend([[-1, -1]] * 130)
+            return {"height": H, "width": W, "keypoints": flat_keypoints}
+
+        person = pose_dict['people'][0]  # Process the first person found
+
+        # Helper function to process each body part
+        def process_part(keypoints_data, expected_length):
+            processed_kps = []
+            if keypoints_data:
+                # Iterate in steps of 3 (x, y, confidence)
+                for i in range(0, len(keypoints_data), 3):
+                    x, y, conf = keypoints_data[i], keypoints_data[i+1], keypoints_data[i+2]
+                    # Use confidence score to check for validity. If 0, it's a missing point.
+                    if conf > 0:
+                        processed_kps.append([x, y])
+                    else:
+                        processed_kps.append([-1, -1])
+            
+            # Ensure the list has the exact expected length
+            while len(processed_kps) < expected_length:
+                processed_kps.append([-1, -1])
+            
+            return processed_kps
+
+        # Process parts in order: body -> face -> left hand -> right hand
+        body_kps = process_part(person.get('pose_keypoints_2d'), 18)
+        face_kps = process_part(person.get('face_keypoints_2d'), 70)
+        left_hand_kps = process_part(person.get('hand_left_keypoints_2d'), 21)
+        right_hand_kps = process_part(person.get('hand_right_keypoints_2d'), 21)
+
+        # Combine all parts into the final flat list
+        flat_keypoints.extend(body_kps)
+        flat_keypoints.extend(face_kps)
+        flat_keypoints.extend(left_hand_kps)
+        flat_keypoints.extend(right_hand_kps)
+        
+        return {"height": H, "width": W, "keypoints": flat_keypoints}
+
+    def save_flattened_pose_kps(self, pose_kps, file_path):
+        # filename_prefix += self.prefix_append
+        
+        # # Get the save path using the first pose keypoint's dimensions
+        # full_output_folder, filename, counter, subfolder, filename_prefix = \
+        #     folder_paths.get_save_image_path(filename_prefix, self.output_dir, 
+        #                                   pose_kps[0]["canvas_width"], 
+        #                                   pose_kps[0]["canvas_height"])
+        
+        # Process each pose keypoint in the batch
+        flattened_poses = []
+        for pose_dict in pose_kps:
+            flattened_data = self._flatten_openpose_dict(pose_dict)
+            flattened_poses.append(flattened_data)
+            
+        # # Save the flattened data
+        # file = f"{filename}_{counter:05}.json"
+        # save_path = os.path.join(full_output_folder, file)
+
+        cur_file_dir = os.path.dirname(os.path.realpath(__file__))
+        save_path = os.path.join(cur_file_dir,
+                                          file_path)
+        
+        with open(save_path, 'w') as f:
+            if len(flattened_poses) == 1:
+                json.dump(flattened_poses[0], f, indent=4)  # Save single pose directly
+            else:
+                json.dump(flattened_poses, f, indent=4)  # Save batch as array
+                
+        print(f"Saved flattened pose keypoints to: {save_path}")
+        return (save_path,)

image_stack.py

@@ -0,0 +1,203 @@
+#!/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, )

light_layer.py

@@ -0,0 +1,202 @@
+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',
+}

mask_area_percentage.py

@@ -0,0 +1,68 @@
+import os
+import cv2
+import numpy as np
+import torch
+
+class TRI3D_MaskAreaPercentage:
+    """
+    ComfyUI node that calculates the percentage of white pixels in an image
+    relative to the total image area.
+    """
+    
+    def from_torch_image(self, image):
+        """Convert a torch tensor image to numpy array for OpenCV processing"""
+        image = image.cpu().numpy() * 255.0
+        image = np.clip(image, 0, 255).astype(np.uint8)
+        return image
+
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "threshold": ("INT", {"default": 127, "min": 0, "max": 255, "step": 1}),
+            },
+        }
+
+    FUNCTION = "run"
+    RETURN_TYPES = ("FLOAT", "INT", "INT",)
+    RETURN_NAMES = ("percentage", "white_pixels", "total_pixels",)
+    CATEGORY = "TRI3D"
+
+    def run(self, image, threshold=127):
+        # Convert Torch image to OpenCV format
+        cv_image = self.from_torch_image(image)
+        
+        # Remove batch dimension if present
+        if len(cv_image.shape) == 4:
+            cv_image = cv_image[0]
+        
+        # Convert to grayscale if it's a color image
+        if len(cv_image.shape) == 3 and cv_image.shape[2] > 1:
+            gray_image = cv2.cvtColor(cv_image, cv2.COLOR_RGB2GRAY)
+        else:
+            gray_image = cv_image[:, :, 0]
+        
+        # Calculate total number of pixels
+        total_pixels = gray_image.shape[0] * gray_image.shape[1]
+        
+        # Count white pixels (pixels with values above threshold)
+        _, binary_image = cv2.threshold(gray_image, threshold, 255, cv2.THRESH_BINARY)
+        white_pixels = cv2.countNonZero(binary_image)
+        
+        # Calculate percentage of white pixels
+        percentage = (white_pixels / total_pixels) * 100.0
+        
+        return (percentage, white_pixels, total_pixels,)
+
+# # Node registration for ComfyUI
+# NODE_CLASS_MAPPINGS = {
+#     "TRI3D_MaskAreaPercentage": TRI3D_MaskAreaPercentage
+# }
+
+# NODE_DISPLAY_NAME_MAPPINGS = {
+#     "TRI3D_MaskAreaPercentage": "TRI3D Mask Area Percentage"
+# } 

nsfw.py

@@ -0,0 +1,166 @@
+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),)

photoroom.py

@@ -0,0 +1,131 @@
+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,)
+            
+
+
+        
+      
+

remove_small_mask_islands.py

@@ -0,0 +1,117 @@
+import os
+import cv2
+import numpy as np
+import torch
+
+class TRI3D_RemoveSmallMaskIslands:
+    """
+    ComfyUI node that removes small islands of white pixels from a mask image
+    based on a specified area threshold.
+    """
+    
+    def from_torch_image(self, image):
+        """Convert a torch tensor image to numpy array for OpenCV processing"""
+        image = image.cpu().numpy() * 255.0
+        image = np.clip(image, 0, 255).astype(np.uint8)
+        return image
+
+    def to_torch_image(self, image):
+        """Convert numpy array back to torch tensor format"""
+        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", ),
+                "min_island_area": ("INT", {"default": 100, "min": 1, "max": 10000, "step": 10}),
+                "invert": ("BOOLEAN", {"default": False}),
+            },
+        }
+
+    FUNCTION = "run"
+    RETURN_TYPES = ("IMAGE", )
+    CATEGORY = "TRI3D"
+
+    def run(self, image, min_island_area, invert):
+        # Convert Torch image to OpenCV format
+        cv_image = self.from_torch_image(image)
+        
+        # Remove batch dimension if present
+        if len(cv_image.shape) == 4:
+            cv_image = cv_image[0]
+        
+        # Make a copy to work with
+        result_image = cv_image.copy()
+        
+        # Process each channel (if grayscale, it will just be one iteration)
+        height, width = cv_image.shape[:2]
+        
+        # If the image has 3 channels (RGB), convert to grayscale for contour detection
+        if len(cv_image.shape) == 3 and cv_image.shape[2] == 3:
+            # Convert to grayscale for processing
+            gray = cv2.cvtColor(cv_image, cv2.COLOR_RGB2GRAY)
+        else:
+            # Use the first channel if it's already grayscale or has alpha
+            gray = cv_image[:, :, 0]
+        
+        # Invert if needed (to work with black islands instead of white)
+        if invert:
+            gray = 255 - gray
+            
+        # Create binary image
+        _, binary = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)
+        
+        # Find contours in the binary image
+        contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        
+        # Create a blank mask for the cleaned image
+        clean_mask = np.zeros((height, width), dtype=np.uint8)
+        
+        # Draw only contours with area greater than the threshold
+        for contour in contours:
+            area = cv2.contourArea(contour)
+            if area >= min_island_area:
+                cv2.drawContours(clean_mask, [contour], 0, 255, -1)
+        
+        # Invert back if needed
+        if invert:
+            clean_mask = 255 - clean_mask
+        
+        # Apply the clean mask to each channel of the original image
+        if len(cv_image.shape) == 3 and cv_image.shape[2] == 3:
+            # RGB image
+            for i in range(3):
+                result_image[:, :, i] = cv2.bitwise_and(cv_image[:, :, i], clean_mask)
+        elif len(cv_image.shape) == 3 and cv_image.shape[2] == 4:
+            # RGBA image
+            for i in range(4):
+                result_image[:, :, i] = cv2.bitwise_and(cv_image[:, :, i], clean_mask)
+        else:
+            # Single channel image
+            result_image = cv2.bitwise_and(cv_image, clean_mask)
+            # Reshape to match expected dimensions
+            result_image = result_image.reshape(height, width, 1)
+        
+        # Convert back to torch format
+        torch_image = self.to_torch_image(result_image)
+        
+        # Add batch dimension back
+        torch_image = torch_image.unsqueeze(0)
+        
+        return (torch_image,)
+
+# # Node registration for ComfyUI
+# NODE_CLASS_MAPPINGS = {
+#     "TRI3D_RemoveSmallMaskIslands": TRI3D_RemoveSmallMaskIslands
+# }
+
+# NODE_DISPLAY_NAME_MAPPINGS = {
+#     "TRI3D_RemoveSmallMaskIslands": "TRI3D Remove Small Mask Islands"
+# }

requirements.txt

@@ -3,4 +3,9 @@ ninja
 pillow
 torch
 torchvision
-gdown
+transparent-background
+wget
+gdown
+matplotlib
+python-dotenv
+git+https://github.com/FacePerceiver/facer.git@main

safetychecker/config.json

@@ -0,0 +1,171 @@
+{
+  "_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
+  }
+}

safetychecker/preprocessor_config.json

@@ -0,0 +1,20 @@
+{
+  "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
+}

samples/3.1_resize_extract_parts/test_workflow.json

@@ -0,0 +1,872 @@
+{
+  "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
+}

samples/back_poses/backpose_kid.json

@@ -0,0 +1,526 @@
+{
+  "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
+    ]
+  ]
+}

samples/back_poses/backpose_kidold.json

@@ -0,0 +1,526 @@
+{
+  "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
+    ]
+  ]
+}

scaled_paste.py

@@ -0,0 +1,316 @@
+#!/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',
+}

simple_bg_swap.py

@@ -0,0 +1,610 @@
+#!/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, )

smart_box.py

@@ -0,0 +1,927 @@
+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)
+
+class TRI3D_Skip_LipMask:
+
+    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_lip_keypoints(self, keypoints):
+        # In DWPose, lips are typically keypoints in face area
+        # Assuming standard face keypoint format where lips are around indices 61-68
+        # This may need adjustment based on your specific keypoint format
+        lip_indices = range(61, 69)  # Adjust these indices based on your keypoint format
+        
+        # Filter out invalid keypoints (those with negative confidence or coordinates)
+        lip_keypoints = []
+        for idx in lip_indices:
+            if idx < len(keypoints):
+                x, y = keypoints[idx]
+                if x >= 0 and y >= 0:  # Check for valid coordinates
+                    lip_keypoints.append((x, y))
+        
+        return lip_keypoints
+
+    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 run(self, image, keypoints_json):
+        # 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]
+
+        # Make a copy of the original image
+        result_image = cv_image.copy()
+        
+        # Parse keypoints JSON
+        try:
+            kp_data = json.loads(open(keypoints_json, 'r').read())
+            original_height, original_width = kp_data['height'], kp_data['width']
+            keypoints = kp_data['keypoints']
+            
+            # Extract lip keypoints
+            lip_keypoints = self.extract_lip_keypoints(keypoints)
+            
+            # If no valid lip keypoints found, use a fallback approach
+            if not lip_keypoints:
+                # Fallback: use the nose point (index 0) as reference
+                nose_point = keypoints[0]
+                if nose_point[1] > 0:  # If y-coordinate is valid
+                    # Estimate lip position slightly below nose
+                    lip_y = int(nose_point[1] + 0.15 * cv_image.shape[0])
+                    lowest_y = lip_y
+                else:
+                    # If no valid reference point, use 1/3 of the image height
+                    lowest_y = cv_image.shape[0] // 3
+            else:
+                # Find the lowest y-coordinate among lip keypoints
+                adjusted_lip_keypoints = self.adjust_keypoints(lip_keypoints, cv_image.shape, original_height, original_width)
+                lowest_y = max([kp[1] for kp in adjusted_lip_keypoints])
+            
+            # Black out everything above the lowest lip point
+            result_image[:lowest_y, :] = 0
+            
+        except Exception as e:
+            print(f"Error processing keypoints JSON: {e}")
+            # In case of error, return the original image
+            result_image = cv_image
+        
+        # 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,)

string_check.py

@@ -0,0 +1,42 @@
+import os
+
+class TRI3D_StringContains:
+    """
+    ComfyUI node that checks if a specified string exists within another string.
+    Performs case-insensitive comparison by converting all text to lowercase.
+    """
+    
+    def __init__(self):
+        pass
+    
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "input_string": ("STRING", {"multiline": True}),
+                "search_string": ("STRING", {"default": "", "multiline": False}),
+            },
+        }
+    
+    FUNCTION = "run"
+    RETURN_TYPES = ("BOOLEAN",)
+    CATEGORY = "TRI3D"
+    
+    def run(self, input_string, search_string):
+        # Convert both strings to lowercase for case-insensitive comparison
+        input_lower = input_string.lower()
+        search_lower = search_string.lower()
+        
+        # Check if search string exists in input string
+        contains = search_lower in input_lower
+        
+        return (contains,)
+
+# # Node registration for ComfyUI
+# NODE_CLASS_MAPPINGS = {
+#     "TRI3D_StringContains": TRI3D_StringContains
+# }
+
+# NODE_DISPLAY_NAME_MAPPINGS = {
+#     "TRI3D_StringContains": "TRI3D String Contains"
+# }

utility_nodes.py

@@ -0,0 +1,409 @@
+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,)