Started cloth segmentation debugging

.gitignore

@@ -11,5 +11,3 @@ cloth-segmentation/model/cloth_segm.pth
 
 dwpose/keypoints/
 huggingface/
-
-safetychecker/model.safetensors

__init__.py

@@ -12,25 +12,17 @@ import folder_paths
 from PIL import Image, ImageOps
 sys.path.append(tri3d_custom_nodes_path)
 from scaled_paste import main_scaled_paste
-from scaled_paste import main_scaled_paste_2
 from simple_bg_swap import (simple_bg_swap, get_threshold_for_bg_swap, RGB_2_LAB, LAB_2_RGB, get_mean_and_standard_deviation, renormalize_array)
 from distribution_reshape import (simple_rescale_histogram, get_histogram_limits)
-from utility_nodes import TRI3D_clean_mask, TRI3D_extract_pose_part, TRI3D_position_pose_part, TRI3D_fill_mask, TRI3D_is_only_trouser, TRI3D_extract_facer_mask
-from utility_nodes import TRI3D_extract_facer_mask
+
 
 from .AEMatter import (load_AEMatter_Model, run_AEMatter_inference)
+from .MVANet_inference import (load_MVANet_Model, run_MVANet_inference)
+from .MVANet_inference import (load_MVANet_Model, run_MVANet_inference)
+from .segment_face import main_face_segment
 from .light_layer import main_light_layer
 
 
-from .image_stack import (
-    H_Stack_Images,
-    SaveImage_absolute,
-    SaveText_absolute,
-    Wait_And_Read_File,
-)
-
-
-
 def from_torch_image(image):
     image = image.squeeze().cpu().numpy() * 255.0
     image = np.clip(image, 0, 255).astype(np.uint8)
@@ -166,7 +158,7 @@ def run_transparent_background(path_dir_input,
 
     command = [
         'transparent-background', '--source', path_dir_input,
-        '--dest', path_dir_output, '--type', 'rgba',
+        '--dest', path_dir_output, '--jit', '--type', 'rgba',
         '--mode', 'base', '--ckpt', path_file_model
     ]
 
@@ -238,7 +230,7 @@ class TRI3DLEVINDABHICLOTHSEGBATCH:
             },
         }
 
-    RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE")
+    RETURN_TYPES = ("IMAGE", )
     FUNCTION = "main"
     CATEGORY = "TRI3D"
 
@@ -292,33 +284,21 @@ class TRI3DLEVINDABHICLOTHSEGBATCH:
 
         # Collect and return the results
 
-        mask0_batch = []
-        mask1_batch = []
-        mask2_batch = []
+        
+
+        batch_results = []
         for i in range(images.shape[0]):
-            mask0_path = os.path.join(LSEG_OUTPUT_PATH, f"{i}__mask0.png")
-            mask1_path = os.path.join(LSEG_OUTPUT_PATH, f"{i}__mask1.png")
-            mask2_path = os.path.join(LSEG_OUTPUT_PATH, f"{i}__mask2.png")
-            mask0_img = cv2.imread(mask0_path, cv2.IMREAD_UNCHANGED)
-            mask1_img = cv2.imread(mask1_path, cv2.IMREAD_UNCHANGED)
-            mask2_img = cv2.imread(mask2_path, cv2.IMREAD_UNCHANGED)
-            # Ensure single channel, convert to 3 channel if needed for consistency
-            if mask0_img is not None and len(mask0_img.shape) == 2:
-                mask0_img = cv2.cvtColor(mask0_img, cv2.COLOR_GRAY2RGB)
-            if mask1_img is not None and len(mask1_img.shape) == 2:
-                mask1_img = cv2.cvtColor(mask1_img, cv2.COLOR_GRAY2RGB)
-            if mask2_img is not None and len(mask2_img.shape) == 2:
-                mask2_img = cv2.cvtColor(mask2_img, cv2.COLOR_GRAY2RGB)
-            mask0_tensor = cv2_img_to_tensor(mask0_img).squeeze(0)
-            mask1_tensor = cv2_img_to_tensor(mask1_img).squeeze(0)
-            mask2_tensor = cv2_img_to_tensor(mask2_img).squeeze(0)
-            mask0_batch.append(mask0_tensor)
-            mask1_batch.append(mask1_tensor)
-            mask2_batch.append(mask2_tensor)
-        mask0_batch = torch.stack(mask0_batch)
-        mask1_batch = torch.stack(mask1_batch)
-        mask2_batch = torch.stack(mask2_batch)
-        return (mask0_batch, mask1_batch, mask2_batch)
+            cv2_segm = cv2.imread(LSEG_OUTPUT_PATH + f'image{i}.png', cv2.IMREAD_UNCHANGED)  # Read PNG with alpha channel
+            cv2_segm = cv2.cvtColor(cv2_segm, cv2.COLOR_BGRA2RGBA)  # Convert from BGRA to RGBA
+            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, )
 
 
 
@@ -1953,8 +1933,6 @@ class TRI3DDWPose_Preprocessor:
         out_image = torch.stack(out_image_list, dim=0)
         del model
 
-        # print(save_file_path, "save_file_path")
-
         return (out_image, save_file_path)
 
 
@@ -2680,12 +2658,9 @@ class TRI3D_reLUM:
                 mu_2, sigma_2 = get_mu_sigma(array_input=image_2[:, :, i],
                                              mask_input=mask_2)
 
-                sigma_calculated = sigma_1 * factor_sigma[i]
-                if factor_sigma[i] < 0:
-                    sigma_calculated = sigma_2
                 image_2[:, :, i] = (
                     ((image_2[:, :, i] - mu_2) / sigma_2) *
-                    sigma_calculated) + (mu_1 * factor_mean[i])
+                    (sigma_1 * factor_sigma[i])) + (mu_1 * factor_mean[i])
 
             image_2 = np.clip(image_2, 0, 255)
             image_2 = image_2.astype(dtype=np.uint8)
@@ -2860,14 +2835,9 @@ class TRI3D_recolor_LAB:
                 mu_2, sigma_2 = get_mu_sigma(array_input=image_2[:, :, i],
                                              mask_input=mask_2)
 
-                sigma_calculated = sigma_1 * factor_sigma[i]
-                if factor_sigma[i] < 0:
-                    sigma_calculated = sigma_2
-
                 image_2[:, :, i] = (
                     ((image_2[:, :, i] - mu_2) / sigma_2) *
-                    sigma_calculated) + (mu_1 * factor_mean[i])
-
+                    (sigma_1 * factor_sigma[i])) + (mu_1 * factor_mean[i])
 
             image_2 = np.clip(image_2, 0, 255)
             image_2 = image_2.astype(dtype=np.uint8)
@@ -3698,6 +3668,7 @@ class TRI3D_BGREMOVE_MEGA():
 
             batch_results = torch.stack(batch_results)
             batch_results_masks = torch.stack(batch_results_masks)
+
             return (batch_results,batch_results_masks)
                 
 
@@ -3706,8 +3677,6 @@ class TRI3D_BGREMOVE_MEGA():
 
 
 from photoroom import TRI3D_photoroom_bgremove_api
-from smart_box import TRI3D_SmartBox, TRI3D_Skip_HeadMask, TRI3D_Skip_HeadMask_AddNeck, TRI3D_Image_extend, TRI3D_Smart_Depth, TRI3D_NarrowfyImage
-from nsfw import TRI3DNSFWFilter
 
 # A dictionary that contains all nodes you want to export with their names
 # NOTE: names should be globally unique
@@ -3743,7 +3712,6 @@ NODE_CLASS_MAPPINGS = {
     "tri3d-composite-image-splitter": TRI3DCompositeImageSplitter,
     'tri3d-main_transparent_background': main_transparent_background,
     'tri3d-scaled-paste': main_scaled_paste,
-    'tri3d-scaled-paste_unsafe': main_scaled_paste_2,
     'tri3d-luminosity-match': TRI3D_reLUM,
     'tri3d-simple_bg_swap': simple_bg_swap,
     'tri3d-get_threshold_for_bg_swap': get_threshold_for_bg_swap,
@@ -3754,31 +3722,17 @@ NODE_CLASS_MAPPINGS = {
     "tri3d-simple_rescale_histogram": simple_rescale_histogram,
     "tri3d-get_histogram_limits": get_histogram_limits,
     "tri3d-clear-memory": clear_memory,
+    "tri3d-load_MVANet_Model": load_MVANet_Model,
+    "tri3d-run_MVANet_inference": run_MVANet_inference,
     'tri3d-load_AEMatter_Model': load_AEMatter_Model,
     'tri3d-run_AEMatter_inference': run_AEMatter_inference,
     "tri3d-bgremove-mega" :TRI3D_BGREMOVE_MEGA,
+    'tri3d-facer_face_segment' : main_face_segment,
     'tri3d-flexible_color_extract' : main_light_layer,
-    'tri3d-clean_mask': TRI3D_clean_mask,
-    "tri3d-extract_pose_part": TRI3D_extract_pose_part,
-    "tri3d_position_pose_part":TRI3D_position_pose_part,
-    "tri3d_fill_mask": TRI3D_fill_mask,
-    "tri3d_is_only_trouser": TRI3D_is_only_trouser,
-    "tri3d_extract_facer_mask":TRI3D_extract_facer_mask,
-    "tri3d_H_Stack_Images": H_Stack_Images,
-    "tri3d_SaveImage_absolute":SaveImage_absolute,
-    "tri3d_SaveText_absolute":SaveText_absolute,
-    "tri3d_Wait_And_Read_File":Wait_And_Read_File,
-    "tri3d_SmartBox": TRI3D_SmartBox,
-    "tri3d_Skip_HeadMask": TRI3D_Skip_HeadMask,
-    "tri3d_Skip_HeadMask_AddNeck": TRI3D_Skip_HeadMask_AddNeck,
-    "tri3d_Image_extend": TRI3D_Image_extend,
-    "tri3d_Smart_Depth": TRI3D_Smart_Depth,
-    "tri3d_NSFWFilter": TRI3DNSFWFilter,
-    "tri3d_NarrowfyImage": TRI3D_NarrowfyImage,
 }
 
 
-VERSION = "4.9.0"
+VERSION = "4.2.1"
 # A dictionary that contains the friendly/humanly readable titles for the nodes
 NODE_DISPLAY_NAME_MAPPINGS = {
     "tri3d-photoroom-bgremove-api": "Photoroom BG Remove" + " v" + VERSION,
@@ -3814,7 +3768,6 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "tri3d-composite-image-splitter": "Composite Image Splitter" + " v" + VERSION,
     'tri3d-main_transparent_background': 'Transparent Background' + " v" + VERSION,
     'tri3d-scaled-paste': 'Scaled paste' + " v" + VERSION,
-    'tri3d-scaled-paste_unsafe': 'Scaled paste unsafe' + " v" + VERSION,
     'tri3d-luminosity-match': 'Luminosity match' + " v" + VERSION,
     'tri3d-simple_bg_swap':  'Simple bg swap' + " v" + VERSION,
     'tri3d-get_threshold_for_bg_swap':  'Get threshold for bg swap' + " v" + VERSION,
@@ -3825,25 +3778,11 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "tri3d-simple_rescale_histogram": 'Rescale the layer to have given max and min values' + " v" + VERSION,
     "tri3d-get_histogram_limits": 'Calculate max and min values for rescaling histogram' + " v" + VERSION,
     "tri3d-clear_memory": 'Clear Memory' + " v" + VERSION,
+    "tri3d-load_MVANet_Model": 'Load MVANet Model' + ' v' + VERSION,
+    "tri3d-run_MVANet_inference": 'Run MVANet inference' + ' v' + VERSION,
     'tri3d-load_AEMatter_Model': 'Load AEMatter Model' + ' v' + VERSION,
     'tri3d-run_AEMatter_inference': 'Run AEMatter inference' + ' v' + VERSION,
     "tri3d-bgremove-mega": "BG Remove Mega" + " v" + VERSION,
+    'tri3d-facer_face_segment': "Segment face using facer" + " v" + VERSION,
     'tri3d-flexible_color_extract': "Flexible color extract" + " v" + VERSION,
-    'tri3d-clean_mask': "Clear small patches" + " v" + VERSION,
-    "tri3d-extract_pose_part": "Extract pose part" + " v" + VERSION,
-    "tri3d_position_pose_part": "Position pose part" + " v" + VERSION,
-    "tri3d_fill_mask": "Fill mask" + " v" + VERSION,
-    "tri3d_is_only_trouser": "Is only trouser" + " v" + VERSION,
-    "tri3d_extract_facer_mask": "Extract facer mask" + " v" + VERSION,
-    "tri3d_H_Stack_Images": "Stack images for cat vton with flux" + " v" + VERSION,
-    "tri3d_SaveImage_absolute": "Save image to an absolute path and provide text optional to control execution order" + " v" + VERSION,
-    "tri3d_SaveText_absolute": "Save text to an absolute path and provide text optional to control execution order " + " v" + VERSION,
-    "tri3d_Wait_And_Read_File": "Wait and read text file, optional control from text " + " v" + VERSION,
-    "tri3d_SmartBox": "Smart Box" + " v" + VERSION,
-    "tri3d_Skip_HeadMask": "Skip Head Mask" + " v" + VERSION,
-    "tri3d_Skip_HeadMask_AddNeck": "Skip Head Mask and add neck" + " v" + VERSION,
-    "tri3d_NSFWFilter": "TRI3D NSFW Filter" + " v" + VERSION,
-    "tri3d_Image_extend": "Image extend" + " v" + VERSION,
-    "tri3d_Smart_Depth": "Smart Depth" + " v" + VERSION,
-    "tri3d_NarrowfyImage": "Narrowfy Image" + " v" + VERSION,
 }

cloth-segmentation/app.py

@@ -1,47 +1,35 @@
 import PIL
+import cv2
 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)    
+    net = load_seg_model(checkpoint_path, device=device)
     return net
 
+
 net = initialize_and_load_models()
 
-def run(img, image_id, output_dir):
+
+def run(img):
     palette = get_palette(4)
-    mask0, mask1, mask2, cloth_seg = generate_mask(img, net=net, device=device, image_id=image_id, output_dir=output_dir)
-    return mask0, mask1, mask2, cloth_seg
+    cloth_seg = generate_mask(img, net=net, device=device)
+    return cloth_seg
+
 
 INPUT_PATH = "./input/"
 OUTPUT_PATH = "./output/"
 
-import os 
-mask0_paths = []
-mask1_paths = []
-mask2_paths = []
-cloth_paths = []
-
-for idx, cur_image in enumerate(os.listdir(INPUT_PATH)):
+import os
+for cur_image in os.listdir(INPUT_PATH):
     img = PIL.Image.open(INPUT_PATH + cur_image)
-    mask0, mask1, mask2, cloth_seg = run(img, image_id=idx, output_dir=OUTPUT_PATH)
-    # Save masks and cloth_seg with unique names (already saved in generate_mask)
-    mask0_path = os.path.join(OUTPUT_PATH, f"{idx}__mask0.png")
-    mask1_path = os.path.join(OUTPUT_PATH, f"{idx}__mask1.png")
-    mask2_path = os.path.join(OUTPUT_PATH, f"{idx}__mask2.png")
-    cloth_path = os.path.join(OUTPUT_PATH, f"{idx}__extracted_garment.png")
-    mask0_paths.append(mask0_path)
-    mask1_paths.append(mask1_path)
-    mask2_paths.append(mask2_path)
-    cloth_paths.append(cloth_path)
+    cloth_seg = run(img)
 
-print("Mask0 batch:", mask0_paths)
-print("Mask1 batch:", mask1_paths)
-print("Mask2 batch:", mask2_paths)
-print("Garment batch:", cloth_paths)
+    cv2.imwrite(OUTPUT_PATH + cur_image,
+                cv2.cvtColor(src=cloth_seg, code=cv2.COLOR_RGB2BGR))
+    # cloth_seg.save(OUTPUT_PATH + cur_image, format="PNG")

cloth-segmentation/process.py

@@ -14,6 +14,19 @@ import torchvision.transforms as transforms
 from collections import OrderedDict
 from options import opt
 
+import einops
+
+def do_recolor(vis_seg_probs, n_classes):
+    val = int(255 / n_classes)
+    not_visible = (vis_seg_probs == 0).astype(dtype=np.uint8)
+    not_visible = 1 - not_visible
+    not_visible *= 255
+    vis_seg_probs *= val
+    ret = np.array((vis_seg_probs, not_visible, not_visible), np.uint8)
+    ret = einops.rearrange(ret, 'c h w -> h w c')
+    ret = cv2.cvtColor(ret, cv2.COLOR_HSV2RGB_FULL)
+    return ret
+
 
 def load_checkpoint(model, checkpoint_path):
     if not os.path.exists(checkpoint_path):
@@ -101,69 +114,59 @@ def apply_transform(img):
 from PIL import Image
 
 
-def generate_mask(input_image, net, device='cpu', image_id=None, output_dir=None):
+def generate_mask(input_image, net, device='cpu'):
     img = input_image
     img_size = img.size
-    img = img.resize((768, 768), Image.BICUBIC)
+    # img = img.resize((768, 768), Image.BICUBIC)
     image_tensor = apply_transform(img)
     image_tensor = torch.unsqueeze(image_tensor, 0)
 
-    # Allow output_dir override for batch processing
-    if output_dir is None:
-        output_dir = os.path.join(opt.output, 'extracted_garment')
+    output_dir = os.path.join(opt.output, 'extracted_garment')
     os.makedirs(output_dir, exist_ok=True)
 
+    print('#### DEBUG START ####')
     with torch.no_grad():
         output_tensor = net(image_tensor.to(device))
+        print(output_tensor[0].shape)
         output_tensor = F.log_softmax(output_tensor[0], dim=1)
         output_tensor = torch.max(output_tensor, dim=1, keepdim=True)[1]
         output_tensor = torch.squeeze(output_tensor, dim=0)
         output_arr = output_tensor.cpu().numpy()
 
-    # Create and save individual masks for classes 1, 2, 3
-    classes_of_interest = [1, 2, 3]
-    mask_imgs = []
-    for idx, cls in enumerate(classes_of_interest):
-        mask = np.zeros_like(output_arr, dtype=np.uint8)
-        mask[output_arr == cls] = 255
-        if mask.ndim > 2:
-            mask = mask.squeeze()
-        if mask.ndim != 2:
-            raise ValueError(f"mask{idx} must be a 2-dimensional array")
-        mask_img = Image.fromarray(mask, mode='L').resize(img_size, Image.BICUBIC)
-        # Save with unique name if image_id is provided
-        if image_id is not None:
-            mask_path = os.path.join(output_dir, f'{image_id}__mask{idx}.png')
-        else:
-            mask_path = os.path.join(output_dir, f'mask{idx}.png')
-        mask_img.save(mask_path, format="PNG")
-        print(f"Saved mask{idx} at: {mask_path}")
-        mask_imgs.append(mask_img)
+    print(output_arr.shape)
+    image_tmp = do_recolor(vis_seg_probs = output_arr.squeeze(0), n_classes = 4)
+    print(image_tmp.shape)
+    print('#### DEBUG STOP ####')
 
-    # Create a binary mask where selected classes are 1, others are 0
-    binary_mask = np.zeros_like(output_arr, dtype=np.uint8)
-    for cls in classes_of_interest:
-        binary_mask[output_arr == cls] = 255
-    if binary_mask.ndim > 2:
-        binary_mask = binary_mask.squeeze()
-    if binary_mask.ndim != 2:
-        raise ValueError("binary_mask must be a 2-dimensional array")
-    binary_mask_img = Image.fromarray(binary_mask, mode='L').resize(img_size, Image.BICUBIC)
+    garment_path = os.path.join(output_dir, 'extracted_garment.png')
+    cv2.imwrite(garment_path, cv2.cvtColor(src = image_tmp, code = cv2.COLOR_RGB2BGR))
+    return image_tmp
 
-    # Create an RGBA image for the output
-    extracted_garment = Image.new("RGBA", img_size)
-    original_img = img.resize(img_size)
-    extracted_garment.paste(original_img, mask=binary_mask_img)
+    # # 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
 
-    # Save the garment image with transparency
-    if image_id is not None:
-        garment_path = os.path.join(output_dir, f'{image_id}__extracted_garment.png')
-    else:
-        garment_path = os.path.join(output_dir, 'extracted_garment.png')
-    extracted_garment.save(garment_path, format="PNG")
-    print(f"Saved extracted garment at: {garment_path}")
+    # # 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
 
-    return (*mask_imgs, extracted_garment)
 
 # def generate_mask(input_image, net, device='cpu'):
 #     img = input_image
@@ -254,4 +257,4 @@ if __name__ == '__main__':
     parser.add_argument('--checkpoint_path', type=str, default='model/cloth_segm.pth', help='Path to the checkpoint file')
     args = parser.parse_args()
 
-    main(args)
+    main(args)

dwpose/comfy_utils.py

@@ -274,6 +274,4 @@ def switch_to_backpose(input_keypoints, input_width):
         x,y = input_keypoints[i]
         input_keypoints[i] = [input_width - x, y]
 
-    return input_keypoints
-
-
+    return input_keypoints

image_stack.py

@@ -1,203 +0,0 @@
-#!/usr/bin/python3
-
-from PIL import Image, ImageOps, ImageSequence, ImageFile
-from PIL.PngImagePlugin import PngInfo
-import cv2
-import hashlib
-import json
-import logging
-import math
-import numpy as np
-import os
-import random
-import safetensors.torch
-import sys
-import time
-import torch
-import traceback
-
-
-def load_image(path):
-
-    return torch.from_numpy(cv2.imread(
-        path, cv2.IMREAD_COLOR)).to(dtype=torch.float32) / 255.0
-
-
-def do_stack(img1, img2):
-
-    dim = max(max(img1.shape[0], img2.shape[0]), img1.shape[1] + img2.shape[1])
-
-    out = torch.zeros((dim, dim, 3), dtype=img1.dtype, device=img1.device) + 1
-
-    diff1 = (out.shape[0] - img1.shape[0]) // 2
-    diff2 = (out.shape[0] - img2.shape[0]) // 2
-
-    part0 = 0
-    part1 = img1.shape[1]
-    part2 = img2.shape[1] + img1.shape[1]
-
-    out[diff1:diff1 + img1.shape[0], part0:part1, :] = img1
-    out[diff2:diff2 + img2.shape[0], part1:part2, :] = img2
-
-    return out
-
-
-def save_image(image, outpath):
-
-    cv2.imwrite(outpath,
-                (image * 255).to(dtype=torch.uint8).detach().cpu().numpy())
-
-
-class H_Stack_Images:
-
-    def __init__(self):
-        pass
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "image_L": ("IMAGE", ),
-                "image_R": ("IMAGE", ),
-            },
-        }
-
-    RETURN_TYPES = ("IMAGE", )
-    FUNCTION = "test"
-    CATEGORY = "TRI3D"
-
-    def test(self, image_L, image_R):
-
-        return (do_stack(img1=image_L[0], img2=image_R[0]).unsqueeze(0), )
-
-
-class SaveImage_absolute:
-
-    def __init__(self):
-        pass
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "images": ("IMAGE", {
-                    "tooltip": "The images to save."
-                }),
-                "absolute_filename": ("STRING", {
-                    "default":
-                    "image.png",
-                    "tooltip":
-                    "The absolute path to the file to save."
-                })
-            },
-        }
-
-    RETURN_TYPES = ("STRING", )
-    RETURN_NAMES = ("text to control order", )
-    FUNCTION = "save_images"
-
-    OUTPUT_NODE = True
-
-    CATEGORY = "image"
-    DESCRIPTION = "Saves the input images to an absolute path."
-
-    def save_images(self, images, absolute_filename):
-        i = 255.0 * images[0].cpu().numpy()
-        img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
-        img.save(absolute_filename)
-
-        return (absolute_filename, )
-
-
-class SaveText_absolute:
-
-    def __init__(self):
-        pass
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "text": ("STRING", {
-                    "multiline": True,
-                    "dynamicPrompts": True,
-                    "tooltip": "Text to be saved to the file."
-                }),
-                "absolute_filename": ("STRING", {
-                    "default":
-                    "image.txt",
-                    "tooltip":
-                    "The absolute path to the file to save."
-                })
-            },
-            "optional": {
-                "text_opt": ("STRING", {
-                    "multiline":
-                    True,
-                    "dynamicPrompts":
-                    True,
-                    "tooltip":
-                    "Text to provide order when necessary (to create work files after txt files)."
-                }),
-            }
-        }
-
-    RETURN_TYPES = ("STRING", )
-    RETURN_NAMES = ("same text as input", )
-    FUNCTION = "save_text"
-
-    OUTPUT_NODE = True
-
-    CATEGORY = "text"
-    DESCRIPTION = "Saves the input text to an absolute path."
-
-    def save_text(self, text, absolute_filename, text_opt=''):
-        open(absolute_filename, "w").write(text)
-        return (text, )
-
-
-class Wait_And_Read_File:
-
-    def __init__(self):
-        pass
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "absolute_filename": ("STRING", {
-                    "default":
-                    "image.txt",
-                    "tooltip":
-                    "The absolute path to the file to read."
-                })
-            },
-            "optional": {
-                "text": ("STRING", {
-                    "multiline":
-                    True,
-                    "dynamicPrompts":
-                    True,
-                    "tooltip":
-                    "Text to provide order when necessary (to wait on done file)."
-                }),
-            }
-        }
-
-    RETURN_TYPES = ("STRING", )
-    RETURN_NAMES = ("text from file", )
-    FUNCTION = "read_text"
-
-    OUTPUT_NODE = True
-
-    CATEGORY = "text"
-    DESCRIPTION = "Saves the input text to an absolute path."
-
-    def read_text(self, absolute_filename, text=''):
-        while not os.path.exists(absolute_filename):
-            time.sleep(0.1)
-
-        res = open(absolute_filename, "r").read()
-        os.unlink(absolute_filename)
-
-        return (res, )

nsfw.py

@@ -1,166 +0,0 @@
-from __future__ import annotations
-from weakref import ref as WeakRef
-from pathlib import Path
-from tqdm import tqdm
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from torch import Tensor
-from transformers import CLIPImageProcessor, CLIPConfig, CLIPVisionModel, PreTrainedModel
-from kornia.filters import box_blur
-
-
-
-def cosine_similarity(image_embeds: Tensor, text_embeds: Tensor):
-    if image_embeds.dim() == 2 and text_embeds.dim() == 2:
-        image_embeds = image_embeds.unsqueeze(1)
-    return F.cosine_similarity(image_embeds, text_embeds, dim=-1)
-
-
-class CLIPSafetyChecker(PreTrainedModel):
-    # https://huggingface.co/CompVis/stable-diffusion-safety-checker
-    # Adapted from:
-    # https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/safety_checker.py
-
-    config_class = CLIPConfig
-    _no_split_modules = ["CLIPEncoderLayer"]
-
-    def __init__(self, config: CLIPConfig):
-        super().__init__(config)
-        projdim = config.projection_dim
-
-        self.vision_model = CLIPVisionModel(config.vision_config)
-        self.visual_projection = nn.Linear(config.vision_config.hidden_size, projdim, bias=False)
-
-        self.concept_embeds = nn.Parameter(torch.ones(17, projdim), requires_grad=False)
-        self.special_care_embeds = nn.Parameter(torch.ones(3, projdim), requires_grad=False)
-        self.concept_embeds_weights = nn.Parameter(torch.ones(17), requires_grad=False)
-        self.special_care_embeds_weights = nn.Parameter(torch.ones(3), requires_grad=False)
-
-    def forward(self, clip_input, images: Tensor, sensitivity: float, alternate_image: Tensor):
-        with torch.no_grad():
-            image_batch = self.vision_model(clip_input)[1]
-            image_embeds = self.visual_projection(image_batch)
-            sensitivity = -0.1 + 0.14 * sensitivity
-
-            special_cos_dist = cosine_similarity(image_embeds, self.special_care_embeds)
-            special_scores_threshold = self.special_care_embeds_weights.unsqueeze(0)
-            special_scores = special_cos_dist - special_scores_threshold + sensitivity
-
-            if torch.any(special_scores > 0):
-                sensitivity = sensitivity + 0.01
-
-            cos_dist = cosine_similarity(image_embeds, self.concept_embeds)
-            concept_threshold = self.concept_embeds_weights.unsqueeze(0)
-            concept_scores = cos_dist - concept_threshold + sensitivity
-
-            is_nsfw = [torch.any(concept_scores[i] > 0) for i in range(concept_scores.shape[0])]
-            is_nsfw = [x.item() for x in is_nsfw]
-            return self.filter_images(images, alternate_image, is_nsfw)
-
-    def filter_images(self, images: Tensor, alternate_image: Tensor, is_nsfw: list[bool]):
-        if not any(is_nsfw):
-            return images
-
-        images = images.clone()
-        for idx, nsfw in enumerate(is_nsfw):
-            if nsfw:
-                # Resize alternate image to match original image dimensions
-                resized_alternate = F.interpolate(
-                    alternate_image[idx:idx+1],  # Add batch dimension
-                    size=(images[idx].shape[1], images[idx].shape[2]),  # Target height, width
-                    mode='bilinear',
-                    align_corners=False
-                )
-                images[idx] = resized_alternate.squeeze(0)  # Remove batch dimension
-        return images
-
-
-class CachedModels:
-    _instance: WeakRef | None = None
-
-    def __init__(self):
-        model_dir = Path(__file__).parent / "safetychecker"
-        model_file = model_dir / "model.safetensors"
-        if not model_file.exists():
-            self.download(
-                "https://huggingface.co/CompVis/stable-diffusion-safety-checker/resolve/refs%2Fpr%2F41/model.safetensors",
-                target=model_file,
-            )
-        self.feature_extractor = CLIPImageProcessor.from_pretrained(model_dir)
-        self.safety_checker = CLIPSafetyChecker.from_pretrained(model_dir)
-
-    @classmethod
-    def load(cls):
-        models = cls._instance and cls._instance()
-        if models is None:
-            models = cls()
-            cls._instance = WeakRef(models)
-        return models
-
-    def download(self, url: str, target: Path):
-        import requests
-
-        try:
-            target_temp = target.with_suffix(".download")
-            with requests.get(url, stream=True) as response:
-                text = "NSFWFilter model download"
-                total = int(response.headers.get("content-length", 0))
-                pbar = tqdm(None, total=total, unit="b", unit_scale=True, desc=text)
-                with open(target_temp, "wb") as f:
-                    for chunk in response.iter_content(chunk_size=8192):
-                        f.write(chunk)
-                        pbar.update(len(chunk))
-                pbar.close()
-            target_temp.rename(target)
-        except Exception as e:
-            raise RuntimeError(
-                f"NSFWFilter: Failed to download safety-checker model from {url} to target location {target}: {e}"
-            ) from e
-
-
-
-def to_bchw(image: torch.Tensor):
-    if image.ndim == 3:
-        image = image.unsqueeze(0)
-    return image.movedim(-1, 1)
-
-
-def to_bhwc(image: torch.Tensor):
-    return image.movedim(1, -1)
-
-
-def mask_batch(mask: torch.Tensor):
-    if mask.ndim == 2:
-        mask = mask.unsqueeze(0)
-    return mask
-
-
-class TRI3DNSFWFilter:
-    models: CachedModels
-
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {
-            "required": {
-                "image": ("IMAGE",),
-                "alternate_image": ("IMAGE",),
-                "sensitivity": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.10}),
-            },
-        }
-
-    RETURN_TYPES = ("IMAGE",)
-    FUNCTION = "check"
-    CATEGORY = "TRI3D NSFW"
-
-    def __init__(self):
-        self.models = CachedModels.load()
-
-    def check(self, image, alternate_image,sensitivity):
-        image = to_bchw(image)
-        alternate_image = to_bchw(alternate_image)
-        input = self.models.feature_extractor(image, do_rescale=False, return_tensors="pt")
-        filtered = self.models.safety_checker(
-            images=image, clip_input=input.pixel_values, sensitivity=sensitivity, alternate_image=alternate_image
-        )
-        return (to_bhwc(filtered),)

safetychecker/config.json

@@ -1,171 +0,0 @@
-{
-  "_name_or_path": "clip-vit-large-patch14/",
-  "architectures": [
-    "SafetyChecker"
-  ],
-  "initializer_factor": 1.0,
-  "logit_scale_init_value": 2.6592,
-  "model_type": "clip",
-  "projection_dim": 768,
-  "text_config": {
-    "_name_or_path": "",
-    "add_cross_attention": false,
-    "architectures": null,
-    "attention_dropout": 0.0,
-    "bad_words_ids": null,
-    "bos_token_id": 0,
-    "chunk_size_feed_forward": 0,
-    "cross_attention_hidden_size": null,
-    "decoder_start_token_id": null,
-    "diversity_penalty": 0.0,
-    "do_sample": false,
-    "dropout": 0.0,
-    "early_stopping": false,
-    "encoder_no_repeat_ngram_size": 0,
-    "eos_token_id": 2,
-    "exponential_decay_length_penalty": null,
-    "finetuning_task": null,
-    "forced_bos_token_id": null,
-    "forced_eos_token_id": null,
-    "hidden_act": "quick_gelu",
-    "hidden_size": 768,
-	"id2label": {
-      "0": "LABEL_0",
-      "1": "LABEL_1"
-    },
-    "initializer_factor": 1.0,
-    "initializer_range": 0.02,
-    "intermediate_size": 3072,
-    "is_decoder": false,
-    "is_encoder_decoder": false,
-    "label2id": {
-      "LABEL_0": 0,
-      "LABEL_1": 1
-    },
-    "layer_norm_eps": 1e-05,
-    "length_penalty": 1.0,
-    "max_length": 20,
-    "max_position_embeddings": 77,
-    "min_length": 0,
-    "model_type": "clip_text_model",
-    "no_repeat_ngram_size": 0,
-    "num_attention_heads": 12,
-    "num_beam_groups": 1,
-    "num_beams": 1,
-    "num_hidden_layers": 12,
-    "num_return_sequences": 1,
-    "output_attentions": false,
-    "output_hidden_states": false,
-    "output_scores": false,
-    "pad_token_id": 1,
-    "prefix": null,
-    "problem_type": null,
-    "pruned_heads": {},
-    "remove_invalid_values": false,
-    "repetition_penalty": 1.0,
-    "return_dict": true,
-    "return_dict_in_generate": false,
-    "sep_token_id": null,
-    "task_specific_params": null,
-    "temperature": 1.0,
-    "tie_encoder_decoder": false,
-    "tie_word_embeddings": true,
-    "tokenizer_class": null,
-    "top_k": 50,
-    "top_p": 1.0,
-    "torch_dtype": null,
-    "torchscript": false,
-    "transformers_version": "4.21.0.dev0",
-    "typical_p": 1.0,
-    "use_bfloat16": false,
-    "vocab_size": 49408
-  },
-  "text_config_dict": {
-    "hidden_size": 768,
-    "intermediate_size": 3072,
-    "num_attention_heads": 12,
-    "num_hidden_layers": 12
-  },
-  "torch_dtype": "float32",
-  "transformers_version": null,
-  "vision_config": {
-    "_name_or_path": "",
-    "add_cross_attention": false,
-    "architectures": null,
-    "attention_dropout": 0.0,
-    "bad_words_ids": null,
-    "bos_token_id": null,
-    "chunk_size_feed_forward": 0,
-    "cross_attention_hidden_size": null,
-    "decoder_start_token_id": null,
-    "diversity_penalty": 0.0,
-    "do_sample": false,
-    "dropout": 0.0,
-    "early_stopping": false,
-    "encoder_no_repeat_ngram_size": 0,
-    "eos_token_id": null,
-    "exponential_decay_length_penalty": null,
-    "finetuning_task": null,
-    "forced_bos_token_id": null,
-    "forced_eos_token_id": null,
-    "hidden_act": "quick_gelu",
-    "hidden_size": 1024,
-    "id2label": {
-      "0": "LABEL_0",
-      "1": "LABEL_1"
-    },
-    "image_size": 224,
-    "initializer_factor": 1.0,
-    "initializer_range": 0.02,
-    "intermediate_size": 4096,
-    "is_decoder": false,
-    "is_encoder_decoder": false,
-    "label2id": {
-      "LABEL_0": 0,
-      "LABEL_1": 1
-    },
-    "layer_norm_eps": 1e-05,
-    "length_penalty": 1.0,
-    "max_length": 20,
-    "min_length": 0,
-    "model_type": "clip_vision_model",
-    "no_repeat_ngram_size": 0,
-    "num_attention_heads": 16,
-    "num_beam_groups": 1,
-    "num_beams": 1,
-    "num_hidden_layers": 24,
-    "num_return_sequences": 1,
-    "output_attentions": false,
-    "output_hidden_states": false,
-    "output_scores": false,
-    "pad_token_id": null,
-    "patch_size": 14,
-    "prefix": null,
-    "problem_type": null,
-    "pruned_heads": {},
-    "remove_invalid_values": false,
-    "repetition_penalty": 1.0,
-    "return_dict": true,
-    "return_dict_in_generate": false,
-    "sep_token_id": null,
-    "task_specific_params": null,
-    "temperature": 1.0,
-    "tie_encoder_decoder": false,
-    "tie_word_embeddings": true,
-    "tokenizer_class": null,
-    "top_k": 50,
-    "top_p": 1.0,
-    "torch_dtype": null,
-    "torchscript": false,
-    "transformers_version": "4.21.0.dev0",
-    "typical_p": 1.0,
-    "use_bfloat16": false
-  },
-  "vision_config_dict": {
-    "hidden_size": 1024,
-    "intermediate_size": 4096,
-    "num_attention_heads": 16,
-    "num_hidden_layers": 24,
-    "patch_size": 14
-  }
-}

safetychecker/preprocessor_config.json

@@ -1,20 +0,0 @@
-{
-  "crop_size": 224,
-  "do_center_crop": true,
-  "do_convert_rgb": true,
-  "do_normalize": true,
-  "do_resize": true,
-  "feature_extractor_type": "CLIPFeatureExtractor",
-  "image_mean": [
-    0.48145466,
-    0.4578275,
-    0.40821073
-  ],
-  "image_std": [
-    0.26862954,
-    0.26130258,
-    0.27577711
-  ],
-  "resample": 3,
-  "size": 224
-}

scaled_paste.py

@@ -18,76 +18,6 @@ def to_torch_image(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,
@@ -149,76 +79,6 @@ def scaled_paste(
 
 
 #!/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):

segment_face.py

@@ -0,0 +1,317 @@
+#!/usr/bin/python3
+import cv2
+import numpy as np
+
+import torch
+import einops
+
+import facer
+
+
+def load_image(image_path):
+    image = cv2.imread(image_path, cv2.IMREAD_COLOR)
+    image = cv2.cvtColor(image, code=cv2.COLOR_BGR2RGB)
+    image = torch.from_numpy(image).to(dtype=torch.float32) / 255.0
+    return image
+
+
+def do_recolor(vis_seg_probs, n_classes):
+    val = int(255 / n_classes)
+    vis_seg_probs = vis_seg_probs.cpu().detach().numpy()
+    not_visible = (vis_seg_probs == 0).astype(dtype=np.uint8)
+    not_visible = 1 - not_visible
+    not_visible *= 255
+    vis_seg_probs *= val
+    ret = np.array((vis_seg_probs, not_visible, not_visible), np.uint8)
+    ret = einops.rearrange(ret, 'c h w -> h w c')
+    ret = cv2.cvtColor(ret, cv2.COLOR_HSV2BGR_FULL)
+    return ret
+
+
+def detect_face_from_tensor(image):
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    image *= 255
+    image = image.to(dtype=torch.uint8)
+    image = facer.hwc2bchw(image).to(device=device)
+    face_detector = facer.face_detector('retinaface/mobilenet', device=device)
+
+    with torch.inference_mode():
+        faces = face_detector(image)
+
+    face_parser = facer.face_parser(
+        'farl/lapa/448', device=device)  # optional "farl/celebm/448"
+
+    with torch.inference_mode():
+        faces = face_parser(image, faces)
+
+    seg_logits = faces['seg']['logits']
+    num_faces = seg_logits.shape[0]
+    print(num_faces)
+
+    if num_faces >= 1:
+
+        seg_probs = seg_logits.softmax(dim=1)  # nfaces x nclasses x h x w
+        n_classes = seg_probs.size(1)
+
+        vis_seg_probs = seg_probs.argmax(dim=1)
+        vis_seg_probs = einops.einsum(vis_seg_probs, 'b h w -> h w')
+
+        return (vis_seg_probs, n_classes, num_faces)
+
+    else:
+
+        vis_seg_probs = torch.zeros((image.shape[0], image.shape[1]),
+                                    dtype=torch.int64)
+
+        n_classes = 11
+
+        return (vis_seg_probs, n_classes, num_faces)
+
+
+def full_work_wrapper(image):
+
+    try:
+        res, n_classes, num_faces = detect_face_from_tensor(image)
+    except:
+        res = torch.zeros((image.shape[0], image.shape[1]), dtype=torch.int64)
+        n_classes = 11
+    tup = do_recolor(res, n_classes)
+    return tup
+
+
+def run_slave(input_image_path, output_image_path, tmp_file_path):
+    import os
+
+    EXEC_STRING = '''
+import os
+
+try:
+    del os.environ['AUX_ANNOTATOR_CKPTS_PATH']
+    os.unsetenv('AUX_ANNOTATOR_CKPTS_PATH')
+except:
+    print('Failed to unset AUX_ANNOTATOR_CKPTS_PATH')
+try:
+    del os.environ['AUX_ORT_PROVIDERS']
+    os.unsetenv('AUX_ORT_PROVIDERS')
+except:
+    print('Failed to unset AUX_ORT_PROVIDERS')
+try:
+    del os.environ['AUX_TEMP_DIR']
+    os.unsetenv('AUX_TEMP_DIR')
+except:
+    print('Failed to unset AUX_TEMP_DIR')
+try:
+    del os.environ['AUX_USE_SYMLINKS']
+    os.unsetenv('AUX_USE_SYMLINKS')
+except:
+    print('Failed to unset AUX_USE_SYMLINKS')
+try:
+    del os.environ['CUBLAS_WORKSPACE_CONFIG']
+    os.unsetenv('CUBLAS_WORKSPACE_CONFIG')
+except:
+    print('Failed to unset CUBLAS_WORKSPACE_CONFIG')
+try:
+    del os.environ['CUDA_MODULE_LOADING']
+    os.unsetenv('CUDA_MODULE_LOADING')
+except:
+    print('Failed to unset CUDA_MODULE_LOADING')
+try:
+    del os.environ['DWPOSE_ONNXRT_CHECKED']
+    os.unsetenv('DWPOSE_ONNXRT_CHECKED')
+except:
+    print('Failed to unset DWPOSE_ONNXRT_CHECKED')
+try:
+    del os.environ['KINETO_LOG_LEVEL']
+    os.unsetenv('KINETO_LOG_LEVEL')
+except:
+    print('Failed to unset KINETO_LOG_LEVEL')
+try:
+    del os.environ['KMP_DUPLICATE_LIB_OK']
+    os.unsetenv('KMP_DUPLICATE_LIB_OK')
+except:
+    print('Failed to unset KMP_DUPLICATE_LIB_OK')
+try:
+    del os.environ['KMP_INIT_AT_FORK']
+    os.unsetenv('KMP_INIT_AT_FORK')
+except:
+    print('Failed to unset KMP_INIT_AT_FORK')
+try:
+    del os.environ['PYTORCH_CUDA_ALLOC_CONF']
+    os.unsetenv('PYTORCH_CUDA_ALLOC_CONF')
+except:
+    print('Failed to unset PYTORCH_CUDA_ALLOC_CONF')
+try:
+    del os.environ['PYTORCH_ENABLE_MPS_FALLBACK']
+    os.unsetenv('PYTORCH_ENABLE_MPS_FALLBACK')
+except:
+    print('Failed to unset PYTORCH_ENABLE_MPS_FALLBACK')
+try:
+    del os.environ['PYTORCH_NVML_BASED_CUDA_CHECK']
+    os.unsetenv('PYTORCH_NVML_BASED_CUDA_CHECK')
+except:
+    print('Failed to unset PYTORCH_NVML_BASED_CUDA_CHECK')
+try:
+    del os.environ['TF_CPP_MIN_LOG_LEVEL']
+    os.unsetenv('TF_CPP_MIN_LOG_LEVEL')
+except:
+    print('Failed to unset TF_CPP_MIN_LOG_LEVEL')
+try:
+    del os.environ['TOKENIZERS_PARALLELISM']
+    os.unsetenv('TOKENIZERS_PARALLELISM')
+except:
+    print('Failed to unset TOKENIZERS_PARALLELISM')
+try:
+    del os.environ['TORCH_CPP_LOG_LEVEL']
+    os.unsetenv('TORCH_CPP_LOG_LEVEL')
+except:
+    print('Failed to unset TORCH_CPP_LOG_LEVEL')
+
+import torch
+import facer
+import cv2
+import einops
+import numpy as np
+import sys
+
+
+def load_image(image_path):
+    image = cv2.imread(image_path, cv2.IMREAD_COLOR)
+    image = cv2.cvtColor(image, code=cv2.COLOR_BGR2RGB)
+    image = torch.from_numpy(image).to(dtype=torch.float32) / 255.0
+    return image
+
+
+def do_recolor(vis_seg_probs, n_classes):
+    val = int(255 / n_classes)
+    vis_seg_probs = vis_seg_probs.cpu().detach().numpy()
+    not_visible = (vis_seg_probs == 0).astype(dtype=np.uint8)
+    not_visible = 1 - not_visible
+    not_visible *= 255
+    vis_seg_probs *= val
+    ret = np.array((vis_seg_probs, not_visible, not_visible), np.uint8)
+    ret = einops.rearrange(ret, 'c h w -> h w c')
+    ret = cv2.cvtColor(ret, cv2.COLOR_HSV2BGR_FULL)
+    return ret
+
+
+def detect_face_from_tensor(image):
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    image *= 255
+    image = image.to(dtype=torch.uint8)
+    image = facer.hwc2bchw(image).to(device=device)
+    face_detector = facer.face_detector('retinaface/mobilenet', device=device)
+
+    with torch.inference_mode():
+        faces = face_detector(image)
+
+    face_parser = facer.face_parser(
+        'farl/lapa/448', device=device)  # optional "farl/celebm/448"
+
+    with torch.inference_mode():
+        faces = face_parser(image, faces)
+
+    seg_logits = faces['seg']['logits']
+    seg_probs = seg_logits.softmax(dim=1)  # nfaces x nclasses x h x w
+    n_classes = seg_probs.size(1)
+
+    vis_seg_probs = seg_probs.argmax(dim=1)
+    vis_seg_probs = einops.einsum(vis_seg_probs, 'b h w -> h w')
+    return (vis_seg_probs, n_classes)
+
+
+def full_work_wrapper(image):
+    try:
+        res, n_classes = detect_face_from_tensor(image)
+        tup = do_recolor(res, n_classes)
+    except:
+        print('Warning: Failed to find a face.')
+        tup = np.zeros(image.shape, dtype=np.uint8)
+    return tup
+
+tup = full_work_wrapper(image=load_image(image_path=sys.argv[1]))
+cv2.imwrite(sys.argv[2], tup)
+'''
+
+    with open(tmp_file_path, 'w', encoding='utf-8') as f:
+        f.write(EXEC_STRING)
+
+    CMD = 'env > ~/env.txt ; python3 ' + tmp_file_path + ' ' + input_image_path + ' ' + output_image_path
+
+    print(CMD)
+    os.system(CMD)
+
+
+def run_slave_tensor(image):
+
+    import tempfile
+    import cv2
+    import os
+
+    device = image.device
+    outtype = image.dtype
+
+    path_dir = tempfile.TemporaryDirectory(
+        suffix='.dir',
+        prefix='facer.',
+        dir=None,
+        ignore_cleanup_errors=False,
+    )
+
+    path_input = path_dir.name + '/input.png'
+    path_output = path_dir.name + '/output.png'
+    path_source = path_dir.name + '/exec.py'
+
+    image = image.detach().cpu().numpy() * 255.0
+    image = image.astype(dtype=np.uint8)
+    image = cv2.cvtColor(src=image, code=cv2.COLOR_RGB2BGR)
+    cv2.imwrite(path_input, image)
+
+    run_slave(input_image_path=path_input,
+              output_image_path=path_output,
+              tmp_file_path=path_source)
+
+    os.unlink(path_input)
+    os.unlink(path_source)
+    image = cv2.imread(path_output, cv2.IMREAD_COLOR)
+    os.unlink(path_output)
+    os.rmdir(path_dir.name)
+    # image = cv2.cvtColor(src=image, code=cv2.COLOR_BGR2RGB)
+    image = image.astype(np.float32) / 255.0
+    # image = torch.from_numpy(image).to(dtype=outtype, device=device) / 255.0
+    return image
+
+
+class main_face_segment():
+
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE", ),
+                "to_run": ("BOOLEAN", )
+            },
+        }
+
+    FUNCTION = "run"
+    RETURN_TYPES = ("IMAGE", )
+    CATEGORY = "TRI3D"
+
+    def run(self, image, to_run):
+        if to_run:
+            batch_size = image.shape[0]
+            ret = []
+            for i in range(batch_size):
+                ret.append(run_slave_tensor(image[i].clone()))
+                # ret.append(full_work_wrapper(image[i].clone()))
+
+            ret = np.array(ret)
+
+            ret = torch.from_numpy(ret).to(dtype=image.dtype,
+                                           device=image.device)
+
+            return (ret, )
+        else:
+            return (torch.zeros_like(image), )

smart_box.py

@@ -1,820 +0,0 @@
-import numpy as np
-import torch
-import json
-
-import cv2
-
-#  {0,  "Nose"},
-# //     {1,  "Neck"},
-# //     {2,  "RShoulder"},
-# //     {3,  "RElbow"},
-# //     {4,  "RWrist"},
-# //     {5,  "LShoulder"},
-# //     {6,  "LElbow"},
-# //     {7,  "LWrist"},
-# //     {8,  "MidHip"},
-# //     {9,  "RHip"},
-# //     {10, "RKnee"},
-# //     {11, "RAnkle"},
-# //     {12, "LHip"},
-# //     {13, "LKnee"},
-# //     {14, "LAnkle"},
-# //     {15, "REye"},
-# //     {16, "LEye"},
-# //     {17, "REar"},
-# //     {18, "LEar"},
-# //     {19, "LBigToe"},
-# //     {20, "LSmallToe"},
-# //     {21, "LHeel"},
-# //     {22, "RBigToe"},
-# //     {23, "RSmallToe"},
-# //     {24, "RHeel"},
-# //     {25, "Background"}
-
-
-class TRI3D_SmartBox:
-    
-    
-    def from_torch_image(self, image):
-        image = image.cpu().numpy() * 255.0
-        image = np.clip(image, 0, 255).astype(np.uint8)
-        return image
-
-    def to_torch_image(self, image):
-        image = image.astype(dtype=np.float32)
-        image /= 255.0
-        image = torch.from_numpy(image)
-        return image
-
-    def __init__(self):
-        pass
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "image": ("IMAGE", ),
-                "keypoints_json": ("STRING", {"multiline": True}),
-            },
-        }
-
-    FUNCTION = "run"
-    RETURN_TYPES = ("IMAGE", )
-    CATEGORY = "TRI3D"
-
-    def extract_torso_keypoints(self, keypoints):
-        # Indices for torso-related keypoints
-        torso_indices = [8, 9, 10, 11, 12, 13]
-        return [keypoints[i] for i in torso_indices]
-
-
-    def run(self, image, keypoints_json):
-        
-        kp_data = json.loads(open(keypoints_json, 'r').read())
-        original_height, original_width = kp_data['height'], kp_data['width']
-        torso_keypoints = self.extract_torso_keypoints(kp_data['keypoints'])
-
-        # Convert Torch image to OpenCV format
-        cv_image = self.from_torch_image(image)
-        
-        # Remove the batch dimension if present
-        if len(cv_image.shape) == 4:
-            cv_image = cv_image[0]
-
-        # Adjust keypoints to match the image dimensions
-        adjusted_keypoints = self.adjust_keypoints(torso_keypoints, cv_image.shape, original_height, original_width)
-        
-
-        # Fill the area below the hip line
-        filled_image = self.fill_below_hip(cv_image, adjusted_keypoints)
-        
-
-        # Convert back to Torch format
-        torch_image = self.to_torch_image(filled_image)
-
-        # Add the batch dimension back
-        torch_image = torch_image.unsqueeze(0)
-        
-
-        return (torch_image,)
-
-    def adjust_keypoints(self, keypoints, image_shape, original_height, original_width):
-        image_height, image_width = image_shape[:2]
-        scale_x = image_width / original_width
-        scale_y = image_height / original_height
-
-        adjusted_keypoints = [
-            (int(x * scale_x), int(y * scale_y)) for x, y in keypoints
-        ]
-        return adjusted_keypoints
-
-    def fill_below_hip(self, image, keypoints):
-        # Correct the indices for hip keypoints
-        # Assuming indices 8 and 11 are for left and right hips
-        # print(keypoints,"hip keypoints")
-        try:
-            valid_y_coords = [kp[1] for kp in keypoints if kp[1] >= 0]
-            hip_y = min(valid_y_coords) if valid_y_coords else 0
-        except:
-            hip_y = 0
-        
-        if hip_y == 0:
-            return image
-
-        # Find the bounding box of the mask below the hip line
-        mask = image[:, :, 0]  # Assuming single-channel mask
-        below_hip = mask[hip_y:, :]
-        contours, _ = cv2.findContours(below_hip, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-        cnt = 0
-        
-        for contour in contours:
-            x, y, w, h = cv2.boundingRect(contour)
-            # print(cnt, x,y,w,h, cv2.contourArea(contour), "cnt,x,y,w,h,area")
-            cnt+=1
-        #     cv2.rectangle(image, (x, y + hip_y), (x + w, y + h + hip_y), (255, 255, 255), -1)
-        contours = [contour for contour in contours if cv2.contourArea(contour) > 20]
-
-        if len(contours) == 0:
-            return image
-        # Combine all contours into one
-        all_contours = np.vstack(contours)
-
-        # Calculate a single bounding rectangle for all contours
-        x, y, w, h = cv2.boundingRect(all_contours)
-        # print(x,y,w,h, "x,y,w,h")
-        cv2.rectangle(image, (x, y + hip_y), (x + w, y + h + hip_y), (255, 255, 255), -1)
-
-        return image
-
-    
-class TRI3D_Skip_HeadMask:
-    def from_torch_image(self, image):
-        image = image.cpu().numpy() * 255.0
-        image = np.clip(image, 0, 255).astype(np.uint8)
-        return image
-
-    def to_torch_image(self, image):
-        image = image.astype(dtype=np.float32)
-        image /= 255.0
-        image = torch.from_numpy(image)
-        return image
-
-    def __init__(self):
-        pass
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "image": ("IMAGE", ),
-                "head_mask": ("IMAGE", ),
-            },
-        }
-
-    FUNCTION = "run"
-    RETURN_TYPES = ("IMAGE", )
-    CATEGORY = "TRI3D"
-
-    def run(self, image, head_mask):
-        # Convert Torch images to OpenCV format
-        cv_image = self.from_torch_image(image)
-        cv_head_mask = self.from_torch_image(head_mask)
-
-        # Remove the batch dimension if present
-        if len(cv_image.shape) == 4:
-            cv_image = cv_image[0]
-        if len(cv_head_mask.shape) == 4:
-            cv_head_mask = cv_head_mask[0]
-
-        # Find the lowest point in the head mask
-        mask = cv_head_mask[:, :, 0]  # Assuming single-channel mask
-        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-        lowest_y = 0
-        for contour in contours:
-            for point in contour:
-                x, y = point[0]
-                if y > lowest_y:
-                    lowest_y = y
-
-        # Black out everything above the lowest point
-        cv_image[:lowest_y, :] = 0
-
-        # Convert back to Torch format
-        torch_image = self.to_torch_image(cv_image)
-
-        # Add the batch dimension back
-        torch_image = torch_image.unsqueeze(0)
-
-        return (torch_image,)
-
-
-    
-class TRI3D_Skip_HeadMask_AddNeck:
-
-    def adjust_keypoints(self, keypoints, image_shape, original_height, original_width):
-        image_height, image_width = image_shape[:2]
-        scale_x = image_width / original_width
-        scale_y = image_height / original_height
-
-        adjusted_keypoints = [
-            (int(x * scale_x), int(y * scale_y)) for x, y in keypoints
-        ]
-        return adjusted_keypoints
-
-    def from_torch_image(self, image):
-        image = image.cpu().numpy() * 255.0
-        image = np.clip(image, 0, 255).astype(np.uint8)
-        return image
-
-    def to_torch_image(self, image):
-        image = image.astype(dtype=np.float32)
-        image /= 255.0
-        image = torch.from_numpy(image)
-        return image
-
-
-    def extract_neck_keypoint(self, keypoints):
-        # Indices for torso-related keypoints
-        neck_indices = [1]
-        return [keypoints[i] for i in neck_indices]
-
-    def extract_ear_keypoints(self, keypoints):
-        # Indices for ear keypoints (17=right ear, 18=left ear)
-        ear_indices = [17, 18]
-        return [keypoints[i] for i in ear_indices]
-
-    def __init__(self):
-        pass
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "image": ("IMAGE", ),
-                "head_mask": ("IMAGE", ),
-                "keypoints_json": ("STRING", {"multiline": True}),
-                "ratio_aggression": ("FLOAT", {"default": 0.5, "min": 0, "max": 1, "step": 0.01}),
-                "neck_width_factor": ("FLOAT", {"default": 0.8, "min": 0.1, "max": 1.5, "step": 0.05}),
-            },
-        }
-
-    FUNCTION = "run"
-    RETURN_TYPES = ("IMAGE", )
-    CATEGORY = "TRI3D"
-
-    def run(self, image, head_mask, keypoints_json, ratio_aggression, neck_width_factor):
-        # Convert Torch images to OpenCV format
-        cv_image = self.from_torch_image(image)
-        cv_head_mask = self.from_torch_image(head_mask)
-
-        # Remove the batch dimension if present
-        if len(cv_image.shape) == 4:
-            cv_image = cv_image[0]
-        if len(cv_head_mask.shape) == 4:
-            cv_head_mask = cv_head_mask[0]
-
-        kp_data = json.loads(open(keypoints_json, 'r').read())
-        original_height, original_width = kp_data['height'], kp_data['width']
-        neck_keypoints = self.extract_neck_keypoint(kp_data['keypoints'])
-        
-        # Make a copy of the original image
-        result_image = cv_image.copy()
-        
-        # Adjust keypoints to match the image dimensions
-        adjusted_neck_keypoints = self.adjust_keypoints(neck_keypoints, cv_image.shape, original_height, original_width)
-        
-        # Find the lowest point and face dimensions in the head mask
-        mask = cv_head_mask[:, :, 0]  # Assuming single-channel mask
-        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-        
-        # Find the chin point (lowest point) and calculate face properties
-        lowest_y = 0
-        face_center_x = cv_image.shape[1] // 2  # Default to center of image
-        face_width = cv_image.shape[1] // 3     # Default face width
-        
-        if contours:
-            # Find the lowest point (chin)
-            for contour in contours:
-                for point in contour:
-                    x, y = point[0]
-                    if y > lowest_y:
-                        lowest_y = y
-            
-            # Calculate face bounding box and center of gravity
-            x, y, w, h = cv2.boundingRect(contours[0])
-            face_width = w
-            
-            # Calculate center of gravity of the face mask
-            M = cv2.moments(contours[0])
-            if M["m00"] != 0:
-                face_center_x = int(M["m10"] / M["m00"])
-            else:
-                face_center_x = x + w // 2
-        
-        # Calculate weighted average point between neck and chin
-        neck_y = adjusted_neck_keypoints[0][1]
-        if neck_y <= 0:
-            neck_y = lowest_y
-        average_y = int((neck_y * ratio_aggression + lowest_y * (1 - ratio_aggression)))
-        
-        print(neck_y, lowest_y, "neck_y, lowest_y")
-        print(average_y, "average_y")
-        
-        # ZONE 1: Black out everything above the chin point
-        result_image[:lowest_y, :] = 0
-        
-        # ZONE 2: Create a triangle for the neck area
-        if lowest_y < average_y:  # Only process if there's a gap between chin and average_y
-            # Create a mask for Zone 2
-            zone2_mask = np.zeros_like(cv_image[:,:,0])
-            
-            # Create a triangle with apex at weighted average point and base at chin level
-            # Apply the neck width factor to the face width
-            neck_width = int(face_width * neck_width_factor)
-            triangle_half_width = neck_width // 2
-            
-            # Create polygon points for the triangle
-            triangle_points = np.array([
-                [face_center_x, average_y],  # Apex at weighted average point
-                [face_center_x - triangle_half_width, lowest_y],  # Left base point at chin level
-                [face_center_x + triangle_half_width, lowest_y]   # Right base point at chin level
-            ], dtype=np.int32)
-            
-            # Fill the triangle in the mask
-            cv2.fillPoly(zone2_mask, [triangle_points], 255)
-            
-            # Apply the mask only to the region between chin and weighted average
-            for y in range(lowest_y, average_y):
-                for x in range(cv_image.shape[1]):
-                    if zone2_mask[y, x] > 0:
-                        result_image[y, x] = 0
-        
-        # ZONE 3: Area below weighted average point is left as is
-        # No action needed for this zone
-        
-        # Convert back to Torch format
-        torch_image = self.to_torch_image(result_image)
-
-        # Add the batch dimension back
-        torch_image = torch_image.unsqueeze(0)
-
-        return (torch_image,)
-
-
-class TRI3D_Image_extend:
-    def from_torch_image(self, image):
-        image = image.cpu().numpy() * 255.0
-        image = np.clip(image, 0, 255).astype(np.uint8)
-        return image
-
-    def to_torch_image(self, image):
-        image = image.astype(dtype=np.float32)
-        image /= 255.0
-        image = torch.from_numpy(image)
-        return image
-    
-    def __init__(self):
-        pass
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "face_mask": ("IMAGE", ),
-                "image": ("IMAGE", ),
-                "ratio": ("FLOAT", {"default": 1.5, "min": 1.2, "max": 2, "step": 0.01}),
-            },
-        }
-    
-    FUNCTION = "run"
-    RETURN_TYPES = ("IMAGE", "IMAGE", )
-    RETURN_NAMES = ("image", "mask_image", )
-    CATEGORY = "TRI3D"
-
-    def run(self, face_mask, image, ratio):
-        cv_face_mask = self.from_torch_image(face_mask)
-        cv_image = self.from_torch_image(image)
-
-        # Remove the batch dimension if present
-        if len(cv_image.shape) == 4:
-            cv_image = cv_image[0]
-        if len(cv_face_mask.shape) == 4:
-            cv_face_mask = cv_face_mask[0]
-        mask = cv_face_mask[:, :, 0]  # Assuming single-channel mask
-        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-        lowest_y = 0
-        highest_y = cv_image.shape[0]
-        for contour in contours:
-            for point in contour:
-                x, y = point[0]
-                if y > lowest_y:
-                    lowest_y = y
-                if y < highest_y:
-                    highest_y = y
-
-        y_below_face = cv_image.shape[0] - lowest_y
-        y_face = lowest_y-highest_y
-
-        # Only extend if the space below face is less than 1.5 times face height
-        target_below_face = int(y_face * ratio)
-        # print("y_face", y_face)
-        # print("lowest_y", lowest_y)
-        # print("highest_y", highest_y)
-        # print("target_below_face", target_below_face)
-        # print("y_below_face", y_below_face)
-
-        original_height = cv_image.shape[0]
-        original_width = cv_image.shape[1]
-        
-        if y_below_face < target_below_face:
-            y_extend = target_below_face - y_below_face
-            
-            # Calculate how much to extend horizontally to maintain aspect ratio
-            new_height = original_height + y_extend
-            new_width = int(original_width * (new_height / original_height))
-            x_extend = new_width - original_width
-            x_extend_left = x_extend // 2
-            x_extend_right = x_extend - x_extend_left
-            
-            # Extend the image in all necessary directions
-            cv_image = cv2.copyMakeBorder(
-                cv_image, 
-                0, y_extend,                    # top, bottom
-                x_extend_left, x_extend_right,  # left, right
-                cv2.BORDER_CONSTANT, 
-                value=[0, 0, 0]
-            )
-            
-            # Create extension mask
-            extension_mask = np.zeros_like(cv_image)
-            # Make extended portions white
-            extension_mask[original_height:, :] = 255  # bottom extension
-            extension_mask[:, :x_extend_left] = 255    # left extension
-            extension_mask[:, -x_extend_right:] = 255  # right extension
-
-        else:
-            extension_mask = np.zeros_like(cv_image)
-
-        # Convert both images back to torch format
-        torch_image = self.to_torch_image(cv_image)
-        torch_mask = self.to_torch_image(extension_mask)
-        
-        # Add batch dimension to both
-        torch_image = torch_image.unsqueeze(0)
-        torch_mask = torch_mask.unsqueeze(0)
-        
-        return (torch_image, torch_mask)
-
-
-class TRI3D_Smart_Depth:
-    
-    
-    def from_torch_image(self, image):
-        image = image.cpu().numpy() * 255.0
-        image = np.clip(image, 0, 255).astype(np.uint8)
-        return image
-
-    def to_torch_image(self, image):
-        image = image.astype(dtype=np.float32)
-        image /= 255.0
-        image = torch.from_numpy(image)
-        return image
-
-    def __init__(self):
-        pass
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "image": ("IMAGE", ),
-                "keypoints_json": ("STRING", {"multiline": True}),
-            },
-        }
-
-    FUNCTION = "run"
-    RETURN_TYPES = ("IMAGE", )
-    CATEGORY = "TRI3D"
-
-    def extract_torso_keypoints(self, keypoints):
-        # Indices for torso-related keypoints
-        torso_indices = [8, 9, 10, 11, 12, 13]
-        return [keypoints[i] for i in torso_indices]
-
-
-    def run(self, image, keypoints_json):
-        
-        kp_data = json.loads(open(keypoints_json, 'r').read())
-        original_height, original_width = kp_data['height'], kp_data['width']
-        torso_keypoints = self.extract_torso_keypoints(kp_data['keypoints'])
-
-        # Convert Torch image to OpenCV format
-        cv_image = self.from_torch_image(image)
-        
-        # Remove the batch dimension if present
-        if len(cv_image.shape) == 4:
-            cv_image = cv_image[0]
-
-        # Adjust keypoints to match the image dimensions
-        adjusted_keypoints = self.adjust_keypoints(torso_keypoints, cv_image.shape, original_height, original_width)
-        
-
-        # Fill the area below the hip line
-        filled_image = self.fill_below_hip(cv_image, adjusted_keypoints)
-        
-
-        # Convert back to Torch format
-        torch_image = self.to_torch_image(filled_image)
-
-        # Add the batch dimension back
-        torch_image = torch_image.unsqueeze(0)
-        
-
-        return (torch_image,)
-
-    def adjust_keypoints(self, keypoints, image_shape, original_height, original_width):
-        image_height, image_width = image_shape[:2]
-        scale_x = image_width / original_width
-        scale_y = image_height / original_height
-
-        adjusted_keypoints = [
-            (int(x * scale_x), int(y * scale_y)) for x, y in keypoints
-        ]
-        return adjusted_keypoints
-
-    def fill_below_hip(self, image, keypoints):
-        # Correct the indices for hip keypoints
-        # Assuming indices 8 and 11 are for left and right hips
-        # print(keypoints,"hip keypoints")
-        try:
-            valid_y_coords = [kp[1] for kp in keypoints if kp[1] >= 0]
-            hip_y = min(valid_y_coords) if valid_y_coords else 0
-        except:
-            hip_y = 0
-        
-        if hip_y == 0:
-            return image
-
-        # Find the bounding box of the mask below the hip line
-        mask = image[:, :, 0]  # Assuming single-channel mask
-        below_hip = mask[hip_y:, :]
-        contours, _ = cv2.findContours(below_hip, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-        cnt = 0
-        
-        for contour in contours:
-            x, y, w, h = cv2.boundingRect(contour)
-            # print(cnt, x,y,w,h, cv2.contourArea(contour), "cnt,x,y,w,h,area")
-            cnt+=1
-        #     cv2.rectangle(image, (x, y + hip_y), (x + w, y + h + hip_y), (255, 255, 255), -1)
-        contours = [contour for contour in contours if cv2.contourArea(contour) > 0]
-
-        if len(contours) == 0:
-            return image
-        # Combine all contours into one
-        all_contours = np.vstack(contours)
-
-        # Calculate a single bounding rectangle for all contours
-        x, y, w, h = cv2.boundingRect(all_contours)
-        # print(x,y,w,h, "x,y,w,h")
-        cv2.rectangle(image, (x, y + hip_y), (x + w, y + h + hip_y), (0, 0, 0), -1)
-
-        return image
-
-
-class TRI3D_NarrowfyImage:
-    def from_torch_image(self, image):
-        image = image.cpu().numpy() * 255.0
-        image = np.clip(image, 0, 255).astype(np.uint8)
-        return image
-
-    def to_torch_image(self, image):
-        image = image.astype(dtype=np.float32)
-        image /= 255.0
-        image = torch.from_numpy(image)
-        return image
-    
-    def __init__(self):
-        pass
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "image": ("IMAGE", ),
-                "mask": ("IMAGE", ),
-                "aspect_ratio": ("FLOAT", {"default": 0.33, "min": 0.25, "max": 1, "step": 0.01}),
-                "border_margin": ("INT", {"default": 15, "min": 10, "max": 100, "step": 1}),
-            },
-        }
-    
-    FUNCTION = "run"
-    RETURN_TYPES = ("IMAGE", "IMAGE", "INT", "INT",)
-    RETURN_NAMES = ("cropped_image", "cropped_mask", "cropped_width", "cropped_height",)
-    CATEGORY = "TRI3D"
-
-    def run(self, image, mask, aspect_ratio, border_margin):
-        # Convert to CV format and remove batch dimension
-        cv_image = self.from_torch_image(image)[0]
-        cv_mask = self.from_torch_image(mask)[0]
-        
-        # Find bounding box of the mask
-        mask_channel = cv_mask[:, :, 0]
-        contours, _ = cv2.findContours(mask_channel, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-        
-        if not contours:
-            return image, mask, aspect_ratio
-            
-        # Filter contours by area
-        significant_contours = [cnt for cnt in contours if cv2.contourArea(cnt) > 100]
-        
-        if not significant_contours:
-            return image, mask, aspect_ratio
-            
-        # Get combined bounding box for all significant contours
-        x_min = float('inf')
-        y_min = float('inf')
-        x_max = 0
-        y_max = 0
-        
-        for contour in significant_contours:
-            x, y, w, h = cv2.boundingRect(contour)
-            x_min = min(x_min, x)
-            y_min = min(y_min, y)
-            x_max = max(x_max, x + w)
-            y_max = max(y_max, y + h)
-        
-        # Calculate final width and height with margin
-        margin = border_margin
-        x = max(0, x_min - margin)  # Ensure we don't go below 0
-        y = max(0, y_min - margin)
-        w = min(cv_image.shape[1] - x, (x_max - x_min) + 2 * margin)  # Ensure we don't exceed image width
-        h = min(cv_image.shape[0] - y, (y_max - y_min) + 2 * margin)  # Ensure we don't exceed image height
-        
-        # Crop both image and mask to bounding box
-        cropped_image = cv_image[y:y+h, x:x+w]
-        cropped_mask = cv_mask[y:y+h, x:x+w]
-        
-        # Calculate required height for aspect ratio 1/3
-        min_height = w * 1/aspect_ratio
-        if h < min_height:
-            height_extend = min_height - h
-            
-            # Extend image with black pixels
-            extended_image = cv2.copyMakeBorder(
-                cropped_image,
-                0, int(height_extend),  # top, bottom
-                0, 0,                   # left, right
-                cv2.BORDER_CONSTANT,
-                value=[0, 0, 0]
-            )
-            
-            # Create mask with white pixels only in extended region
-            extended_mask = cv2.copyMakeBorder(
-                np.zeros_like(cropped_mask),  # Start with black base
-                0, int(height_extend),  # top, bottom
-                0, 0,                   # left, right
-                cv2.BORDER_CONSTANT,
-                value=[255, 255, 255]   # White extension
-            )
-            
-            cropped_image = extended_image
-            cropped_mask = extended_mask
-        
-        # Convert back to torch format and add batch dimension
-        torch_image = self.to_torch_image(cropped_image).unsqueeze(0)
-        torch_mask = self.to_torch_image(cropped_mask).unsqueeze(0)
-        
-        return (torch_image, torch_mask,w,h)
-
-
-
-
-class TRI3D_CropAndExtend:
-    def from_torch_image(self, image):
-        image = image.cpu().numpy() * 255.0
-        image = np.clip(image, 0, 255).astype(np.uint8)
-        return image
-
-    def to_torch_image(self, image):
-        image = image.astype(dtype=np.float32)
-        image /= 255.0
-        image = torch.from_numpy(image)
-        return image
-
-    def __init__(self):
-        pass
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "garment_image": ("IMAGE",),
-                "garment_mask": ("IMAGE",),
-                "human_image": ("IMAGE",),
-                "human_mask": ("IMAGE",),
-                "margin": ("INT", {"default": 10, "min": 0, "max": 50}),
-            },
-        }
-
-    FUNCTION = "run"
-    RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "IMAGE", "INT", "INT",)
-    RETURN_NAMES = ("cropped_garment", "cropped_garment_mask", "cropped_human", "cropped_human_mask", "cropped_width", "cropped_height",)
-    
-    def run(self, garment_image, garment_mask, human_image, human_mask, margin):
-        # Convert to CV format and remove batch dimension
-        cv_garment = self.from_torch_image(garment_image)[0]
-        cv_garment_mask = self.from_torch_image(garment_mask)[0]
-        cv_human = self.from_torch_image(human_image)[0]
-        cv_human_mask = self.from_torch_image(human_mask)[0]
-        
-        # Process garment
-        mask_channel = cv_garment_mask[:, :, 0]
-        contours, _ = cv2.findContours(mask_channel, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-        
-        if not contours:
-            return garment_image, garment_mask, human_image, human_mask, cv_garment.shape[1], cv_garment.shape[0]
-            
-        # Get bounding box with margin
-        x, y, w, h = cv2.boundingRect(contours[0])
-        x = max(0, x - margin)
-        y = max(0, y - margin)
-        w = min(cv_garment.shape[1] - x, w + 2 * margin)
-        h = min(cv_garment.shape[0] - y, h + 2 * margin)
-        
-        # Store the cropped dimensions before extension
-        cropped_width = w
-        cropped_height = h
-        
-        # Crop garment and its mask
-        cropped_garment = cv_garment[y:y+h, x:x+w]
-        cropped_garment_mask = cv_garment_mask[y:y+h, x:x+w]
-        
-        # Calculate required height for aspect ratio 1/3
-        min_height = w * 3
-        if h < min_height:
-            height_extend = min_height - h
-            
-            # Extend garment image and mask
-            extended_garment = cv2.copyMakeBorder(
-                cropped_garment,
-                0, int(height_extend),
-                0, 0,
-                cv2.BORDER_CONSTANT,
-                value=[0, 0, 0]
-            )
-            
-            extended_garment_mask = cv2.copyMakeBorder(
-                cropped_garment_mask,
-                0, int(height_extend),
-                0, 0,
-                cv2.BORDER_CONSTANT,
-                value=[255, 255, 255]
-            )
-            
-            cropped_garment = extended_garment
-            cropped_garment_mask = extended_garment_mask
-        
-        # Process human image similarly
-        mask_channel = cv_human_mask[:, :, 0]
-        contours, _ = cv2.findContours(mask_channel, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-        
-        if contours:
-            x, y, w, h = cv2.boundingRect(contours[0])
-            x = max(0, x - margin)
-            y = max(0, y - margin)
-            w = min(cv_human.shape[1] - x, w + 2 * margin)
-            h = min(cv_human.shape[0] - y, h + 2 * margin)
-            
-            cropped_human = cv_human[y:y+h, x:x+w]
-            cropped_human_mask = cv_human_mask[y:y+h, x:x+w]
-            
-            min_height = w * 3
-            if h < min_height:
-                height_extend = min_height - h
-                
-                extended_human = cv2.copyMakeBorder(
-                    cropped_human,
-                    0, int(height_extend),
-                    0, 0,
-                    cv2.BORDER_CONSTANT,
-                    value=[0, 0, 0]
-                )
-                
-                extended_human_mask = cv2.copyMakeBorder(
-                    cropped_human_mask,
-                    0, int(height_extend),
-                    0, 0,
-                    cv2.BORDER_CONSTANT,
-                    value=[255, 255, 255]
-                )
-                
-                cropped_human = extended_human
-                cropped_human_mask = extended_human_mask
-        
-        # Convert back to torch format and add batch dimension
-        torch_garment = self.to_torch_image(cropped_garment).unsqueeze(0)
-        torch_garment_mask = self.to_torch_image(cropped_garment_mask).unsqueeze(0)
-        torch_human = self.to_torch_image(cropped_human).unsqueeze(0)
-        torch_human_mask = self.to_torch_image(cropped_human_mask).unsqueeze(0)
-        
-        return (torch_garment, torch_garment_mask, torch_human, torch_human_mask, cropped_width, cropped_height)

utility_nodes.py

@@ -1,409 +0,0 @@
-import torch, cv2, json
-import numpy as np
-
-
-def from_torch_image(image):
-    image = image.cpu().numpy() * 255.0
-    image = np.clip(image, 0, 255).astype(np.uint8)
-    return image
-
-
-def to_torch_image(image):
-    image = image.astype(dtype=np.float32)
-    image /= 255.0
-    image = torch.from_numpy(image)
-    return image
-
-class TRI3D_clean_mask():
-
-    """For the given mask and threshold area, remove all patches in the mask with area smaller than threshold"""
-
-    def __init__(self):
-        pass
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "masks": ("MASK", ),
-                "threshold":("FLOAT",{"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01})
-            }
-        }
-
-    FUNCTION = "run"
-    RETURN_TYPES = ("MASK", "BOOL")
-    RETURN_NAMES = ("mask", "cleaned")
-    CATEGORY = "TRI3D"
-
-    def run(self, masks, threshold):
-        batch_results = []
-        for mask in masks:
-            mask = from_torch_image(mask)
-            mask = np.where(mask < 127, 0, 255).astype(np.uint8)
-            h,w = mask.shape[:2]
-
-            total_area = h*w
-            # num_labels, labels = cv2.connectedComponents(mask)
-            region_mask = np.zeros_like(mask)
-
-            # for label in range(1, num_labels):  
-            #     area_percent = (np.sum(labels == label)/ total_area) * 100
-            #     if area_percent < threshold:
-            #         continue
-            #     region_mask[labels == label] = 255
-            less_than_threshold = True
-            area_percent = (np.sum(mask == 255)/ total_area) * 100
-            if area_percent > threshold:
-                region_mask[mask == 255] = 255
-                less_than_threshold = False
-            region_mask = to_torch_image(region_mask)
-            batch_results.append(region_mask.squeeze(0))
-
-        batch_results = torch.stack(batch_results)
-        return (batch_results, less_than_threshold)
-        
-
-class TRI3D_extract_pose_part():
-    """
-        For the given pose, extract region around body parts, region can be defined by % of image size  
-    """
-    def __init__(self):
-        pass
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "image": ("IMAGE", ),
-                "pose_json": ("STRING",{"default" : "dwpose/keypoints/input.json"}),
-                "width_pad": ("FLOAT",{"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01}),
-                "height_pad": ("FLOAT",{"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01}),
-                "shoulders":("BOOLEAN", {
-                    "default": False
-                })
-            }
-        }
-
-
-    FUNCTION = "run"
-    RETURN_TYPES = ("IMAGE", "STRING")
-    RETURN_NAMES = ("image", "coords")
-    CATEGORY = "TRI3D"
-
-    def get_frame_coords(self,point1, point2):
-        x1, y1 = point1
-        x2, y2 = point2
-
-        xmin, xmax, ymin, ymax = min(x1, x2), max(x1, x2), min(y1, y2), max(y1, y2)
-
-        for i in [xmin, xmax, ymin, ymax]:
-            if i < 0: 
-                return None
-
-        return [xmin, xmax, ymin, ymax]
-
-    def run(self, image, pose_json, width_pad, height_pad, shoulders):
-        """
-        image : input image
-        width_pad: % of image width you want to apply on both size of pose body part
-        height_pad: % of image width you want to apply on both size of pose body part
-        rest of them are body parts
-        """
-
-        image = from_torch_image(image[0])
-        batch_result = []
-        input_pose = json.load(open(pose_json))
-        keypoints = input_pose['keypoints']
-
-        og_h, og_w = image.shape[:2]
-        ph, pw = [input_pose['height'], input_pose['width']]
-
-        for i,point in enumerate(keypoints):
-            x,y = point
-            y = int((y/ph)*og_h)
-            x = int((x/pw)*og_w)
-            keypoints[i] = [x, y]
-
-        width_offset = int(og_w * (width_pad) / 100)
-        height_offset = int(og_h * (height_pad) / 100)
-
-        xmin, xmax, ymin, ymax = [0, og_w, 0, og_h]
-
-        part_to_coords = {
-            "shoulders":self.get_frame_coords(keypoints[2], keypoints[5])
-        }
-
-        if shoulders:
-            print(part_to_coords["shoulders"])
-            if part_to_coords["shoulders"] != None:
-                new_xmin, new_xmax, new_ymin, new_ymax = part_to_coords["shoulders"]
-
-                xmin, xmax, ymin, ymax = new_xmin, new_xmax, new_ymin, new_ymax
-        
-        xmin = max(0, xmin - width_offset)
-        xmax = min(og_w, xmax + width_offset)
-        ymin = max(0, ymin - height_offset)
-        ymax = min(og_h, ymax + height_offset)
-
-        image = image[ymin:ymax, xmin:xmax, :].astype(np.uint8)
-        image = to_torch_image(image)
-        batch_result.append(image)
-        batch_result = torch.stack(batch_result)
-        print("final_coords", xmin, xmax, ymin, ymax)
-        coords = ",".join([str(xmin), str(xmax), str(ymin), str(ymax)])
-        
-        return batch_result, coords
-
-class TRI3D_position_pose_part():
-    """
-        put back extracted parts on OG image 
-    """
-    def __init__(self):
-        pass
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "og_image": ("IMAGE", ),
-                "extracted_image": ("IMAGE", ),
-                "coords": ("STRING",{"default" : "xmin, xmax, ymin, ymax"}),
-            }
-        }
-
-
-    FUNCTION = "run"
-    RETURN_TYPES = ("IMAGE", )
-    RETURN_NAMES = ("image", )
-    CATEGORY = "TRI3D"
-
-    def run(self, og_image, extracted_image, coords):
-
-        batch_result = []
-        og_image = from_torch_image(og_image[0])
-        extracted_image = from_torch_image(extracted_image[0])
-        
-        xmin, xmax, ymin, ymax = [int(i) for i in coords.split(",")]
-
-        og_image[ymin:ymax, xmin:xmax, :] = extracted_image
-
-        og_image = to_torch_image(og_image).unsqueeze(0)
-        batch_result.append(og_image)
-        batch_result = torch.stack(batch_result)
-        return batch_result
-
-
-class TRI3D_fill_mask():
-    """
-        fill mask with the neighbouring pixels
-    """
-    def __init__(self):
-        pass
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "image": ("IMAGE", ),
-                "mask": ("MASK", ),
-                "negative_mask": ("MASK", ),
-                "offset":("FLOAT",{"default": 1, "min": 0.0, "max": 100.0, "step": 0.01})
-            }
-        }
-
-    FUNCTION = "run"
-    RETURN_TYPES = ("IMAGE",)
-    RETURN_NAMES = ("image",)
-    CATEGORY = "TRI3D"
-
-    def run(self, image, mask, negative_mask, offset):
-        image = from_torch_image(image[0])
-
-        mask = mask[0].cpu().numpy()
-        mask = np.expand_dims(mask, -1)
-        mh, mw, _ = mask.shape
-
-        inverse_mask = np.ones_like(mask) - mask
-        
-        negative_mask = negative_mask[0].cpu().numpy()
-        indices = np.where(mask > 0)
-
-        offset = offset / 100
-        
-        source = image.copy()
-
-        for y,x in zip(indices[0],indices[1]):
-            x_off = min(mw-1, int(x + offset * mw))
-            if negative_mask[y][x_off] == 0:         #check if pixles on right are outside body
-                source[y][x] = image[y][x_off]
-                
-            else:
-                x_off = max(0, int(x - offset * mw))       #check if pixles on left are outside body
-                if negative_mask[y][x_off] == 0:
-                    source[y][x] = image[y][x_off]
-                else:
-                    y_off = max(0, int(y - offset * mh))
-                    if negative_mask[y_off][x] == 0:            #check if pixles on top are outside body
-                        source[y][x] = image[y_off][x]
-
-                    else:
-                        y_off = min(mh-1, int(y + offset * mh))
-                        if negative_mask[y_off][x] == 0:            #check if pixles on bottom are outside body
-                            source[y][x] = image[y_off][x]
-
-        image = mask * source + inverse_mask * image
-        image = to_torch_image(image).unsqueeze(0)
-
-        return (image,)
-
-class TRI3D_is_only_trouser:
-    def __init__(self):
-        pass
-    
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "pose_json_file": ("STRING", {
-                    "default": "dwpose/keypoints"
-                })
-            }
-        }
-
-    RETURN_TYPES = ("BOOLEAN", )
-    FUNCTION = "main"
-    CATEGORY = "TRI3D"
-
-    def main(self, pose_json_file):
-        pose = json.load(open(pose_json_file))
-        height = pose['height']
-        width = pose['width']
-        keypoints = pose['keypoints']
-
-        points = [0,14,15,16,17,2,1,5]
-        point_to_part = {0:'nose',14:"left eye",15:"right eye",16:"left ear",17:"right ear",2:"left shoulder",1:"neck",5:"right shoulder"}
-        all_negative = True          #if all face and shoulder points are negative means it is a bottom shot
-        for point in points:
-            x,y = keypoints[point]
-            if x > 0 and y > 0:
-                all_negative = False    
-                print(f"{point_to_part[point]} exist")
-        return (all_negative,)
-
-class TRI3D_extract_facer_mask:
-    def __init__(self):
-        pass
-    
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "image": ("IMAGE",),
-                "background": ("BOOLEAN", {
-                    "default": False
-                }),
-                'hair':("BOOLEAN", {
-                    "default": False
-                }),
-                'lower_lip':("BOOLEAN", {
-                    "default": False
-                }),
-                'inner_mouth':("BOOLEAN", {
-                    "default": False
-                }),
-                'upper_lip':("BOOLEAN", {
-                    "default": False
-                }),
-                'nose':("BOOLEAN", {
-                    "default": False
-                }),
-                'left_eyebrow':("BOOLEAN", {
-                    "default": False
-                }),
-                'right_eyebrow':("BOOLEAN", {
-                    "default": False
-                }),
-                'left_eye':("BOOLEAN", {
-                    "default": False
-                }),
-                'right_eye':("BOOLEAN", {
-                    "default": False
-                }),
-                'face':("BOOLEAN", {
-                    "default": False
-                })
-            }
-        }
-
-    RETURN_TYPES = ("MASK", )
-    FUNCTION = "main"
-    CATEGORY = "TRI3D"
-
-    def main(self, image, background, hair, lower_lip, inner_mouth, upper_lip, nose, left_eyebrow, right_eyebrow, left_eye, right_eye, face):
-
-        image = from_torch_image(image[0])
-        h,w,_ = image.shape
-        
-        mask = np.zeros_like(image)
-        
-        label_to_rgb = {'background':[0,0,0], 'face':[0,138,255], 'right_eye':[180, 255, 0], 'left_eye':[42, 255, 0], 'right_eyebrow':[0, 255, 96],
-                'left_eyebrow':[0,255,234], 'nose':[255, 192, 0],  'upper_lip':[255, 54, 0], 'inner_mouth':[255, 0, 84], 'lower_lip':[255, 0, 222], 
-                'hair':[150,0,255]}
-
-        if background:
-            temp = np.all(image == label_to_rgb['background'], axis=-1)
-            idcs = np.where(temp==True)
-            mask[idcs] = 255
-        
-        if face:
-            temp = np.all(image == label_to_rgb['face'], axis=-1)
-            idcs = np.where(temp==True)
-            mask[idcs] = 255
-
-        if right_eye:
-            temp = np.all(image == label_to_rgb['right_eye'], axis=-1)
-            idcs = np.where(temp==True)
-            mask[idcs] = 255
-
-        if left_eye:
-            temp = np.all(image == label_to_rgb['left_eye'], axis=-1)
-            idcs = np.where(temp==True)
-            mask[idcs] = 255
-
-        if right_eyebrow:
-            temp = np.all(image == label_to_rgb['right_eyebrow'], axis=-1)
-            idcs = np.where(temp==True)
-            mask[idcs] = 255
-
-        if left_eyebrow:
-            temp = np.all(image == label_to_rgb['left_eyebrow'], axis=-1)
-            idcs = np.where(temp==True)
-            mask[idcs] = 255
-
-        if nose:
-            temp = np.all(image == label_to_rgb['nose'], axis=-1)
-            idcs = np.where(temp==True)
-            mask[idcs] = 255
-
-        if upper_lip:
-            temp = np.all(image == label_to_rgb['upper_lip'], axis=-1)
-            idcs = np.where(temp==True)
-            mask[idcs] = 255
-
-        if inner_mouth:
-            temp = np.all(image == label_to_rgb['inner_mouth'], axis=-1)
-            idcs = np.where(temp==True)
-            mask[idcs] = 255
-
-        if lower_lip:
-            temp = np.all(image == label_to_rgb['lower_lip'], axis=-1)
-            idcs = np.where(temp==True)
-            mask[idcs] = 255
-
-        if hair:
-            temp = np.all(image == label_to_rgb['hair'], axis=-1)
-            idcs = np.where(temp==True)
-            mask[idcs] = 255
-
-        mask = to_torch_image(mask[:,:,0]).unsqueeze(0)   
-        return (mask,)