Merge pull request #60 from TRI3D-LC/saveflattened

flattened

.gitignore

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

__init__.py

@@ -12,15 +12,25 @@ 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 .remove_small_mask_islands import TRI3D_RemoveSmallMaskIslands
+from .mask_area_percentage import TRI3D_MaskAreaPercentage
+
+
+from .image_stack import (
+    H_Stack_Images,
+    SaveImage_absolute,
+    SaveText_absolute,
+    Wait_And_Read_File,
+)
+
 
 
 def from_torch_image(image):
@@ -158,7 +168,7 @@ def run_transparent_background(path_dir_input,
 
     command = [
         'transparent-background', '--source', path_dir_input,
-        '--dest', path_dir_output, '--jit', '--type', 'rgba',
+        '--dest', path_dir_output, '--type', 'rgba',
         '--mode', 'base', '--ckpt', path_file_model
     ]
 
@@ -1922,7 +1932,7 @@ class TRI3DDWPose_Preprocessor:
             cur_file_dir = os.path.dirname(os.path.realpath(__file__))
             save_file_path = os.path.join(cur_file_dir,
                                           filename_path)
-            json.dump(pose_dict, open(save_file_path, 'w'))
+            json.dump(pose_dict, open(save_file_path, 'w'), indent=4)
             np_result = cv2.resize(np_result, (W, H),
                                    interpolation=cv2.INTER_AREA)
             out_image_list.append(
@@ -1933,6 +1943,8 @@ 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)
 
 
@@ -2658,9 +2670,12 @@ 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_1 * factor_sigma[i])) + (mu_1 * factor_mean[i])
+                    sigma_calculated) + (mu_1 * factor_mean[i])
 
             image_2 = np.clip(image_2, 0, 255)
             image_2 = image_2.astype(dtype=np.uint8)
@@ -2835,9 +2850,14 @@ 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_1 * factor_sigma[i])) + (mu_1 * factor_mean[i])
+                    sigma_calculated) + (mu_1 * factor_mean[i])
+
 
             image_2 = np.clip(image_2, 0, 255)
             image_2 = image_2.astype(dtype=np.uint8)
@@ -3668,7 +3688,6 @@ class TRI3D_BGREMOVE_MEGA():
 
             batch_results = torch.stack(batch_results)
             batch_results_masks = torch.stack(batch_results_masks)
-
             return (batch_results,batch_results_masks)
                 
 
@@ -3677,6 +3696,11 @@ 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, TRI3D_Skip_LipMask
+from nsfw import TRI3DNSFWFilter
+from cut_by_mask_aspect_ratio import TRI3D_CutByMaskAspectRatio
+from string_check import TRI3D_StringContains
+from .dwpose_conversion import SaveFlattenedPoseKpsAsJsonFile
 
 # A dictionary that contains all nodes you want to export with their names
 # NOTE: names should be globally unique
@@ -3712,6 +3736,7 @@ 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,
@@ -3722,17 +3747,39 @@ 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_Skip_LipMask": TRI3D_Skip_LipMask,
+    "tri3d_Image_extend": TRI3D_Image_extend,
+    "tri3d_Smart_Depth": TRI3D_Smart_Depth,
+    "tri3d_NSFWFilter": TRI3DNSFWFilter,
+    "tri3d_NarrowfyImage": TRI3D_NarrowfyImage,
+    "tri3d_Skip_LipMask": TRI3D_Skip_LipMask,
+    "tri3d_Remove_Small_Mask_Islands": TRI3D_RemoveSmallMaskIslands,
+    "tri3d_CutByMaskAspectRatio": TRI3D_CutByMaskAspectRatio,
+    "tri3d_StringContains": TRI3D_StringContains,
+    "tri3d_MaskAreaPercentage": TRI3D_MaskAreaPercentage,
+    "tri3d_SaveFlattenedPoseKpsAsJsonFile": SaveFlattenedPoseKpsAsJsonFile,
+
 }
 
 
-VERSION = "4.2.1"
+VERSION = "5.1.0"
 # 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,
@@ -3768,6 +3815,7 @@ 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,
@@ -3778,11 +3826,31 @@ 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_Skip_LipMask": "Skip Lip Mask" + " 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,
+    "tri3d_Remove_Small_Mask_Islands": "Remove Small Mask Islands" + " v" + VERSION,
+    "tri3d_CutByMaskAspectRatio": "Cut by mask aspect ratio" + " v" + VERSION,
+    "tri3d_StringContains": "String contains" + " v" + VERSION,
+    "tri3d_MaskAreaPercentage": "Mask Area Percentage" + " v" + VERSION,
+    "tri3d_SaveFlattenedPoseKpsAsJsonFile": "Save Flattened Pose Keypoints as JSON File" + " v" + VERSION,
 }

cut_by_mask_aspect_ratio.py

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

dwpose/comfy_utils.py

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

dwpose_conversion.py

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

image_stack.py

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

mask_area_percentage.py

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

nsfw.py

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

remove_small_mask_islands.py

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

safetychecker/config.json

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

safetychecker/preprocessor_config.json

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

scaled_paste.py

@@ -18,6 +18,76 @@ 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,
@@ -79,6 +149,76 @@ 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

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

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

string_check.py

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

utility_nodes.py

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