class TRI3DExtractPartsMaskBatch: def __init__(self): pass @classmethod def INPUT_TYPES(cls): return { "required": { "batch_images": ("IMAGE",), "batch_segs": ("IMAGE",), "batch_secondary": ("IMAGE",), "right_leg": ("BOOLEAN", {"default": False}), "right_hand": ("BOOLEAN", {"default": True}), "head": ("BOOLEAN", {"default": False}), "hair": ("BOOLEAN", {"default": False}), "left_shoe" : ("BOOLEAN", {"default": False}), "bag" : ("BOOLEAN", {"default": False}), "background" : ("BOOLEAN", {"default": False}), "dress" : ("BOOLEAN", {"default": False}), "left_leg": ("BOOLEAN", {"default": False}), "right_shoe" : ("BOOLEAN", {"default": False}), "left_hand": ("BOOLEAN", {"default": True}), "upper_garment" : ("BOOLEAN", {"default": False}), "lower_garment" : ("BOOLEAN", {"default": False}), "belt" : ("BOOLEAN", {"default": False}), "skirt" : ("BOOLEAN", {"default": False}), "hat" : ("BOOLEAN", {"default": False}), "sunglasses" : ("BOOLEAN", {"default": False}), "scarf" : ("BOOLEAN", {"default": False}), }, } RETURN_TYPES = ("IMAGE","IMAGE","IMAGE") FUNCTION = "main" CATEGORY = "TRI3D" def main(self, batch_images, batch_segs,batch_secondary, right_leg,right_hand, head, hair, left_shoe,bag,background,dress,left_leg,right_shoe,left_hand, upper_garment,lower_garment,belt,skirt,hat,sunglasses,scarf): import cv2 import numpy as np import torch from pprint import pprint def get_segment_counts(segm): # Load the segmentation image # Reshape the image array to be 2D reshaped = segm.reshape(-1, segm.shape[-1]) # Find unique vectors and their counts unique_vectors, counts = np.unique(reshaped, axis=0, return_counts=True) segment_counts = list(zip(unique_vectors, counts)) pprint(segment_counts) return segment_counts def generate_mask(seg_img, color_code_list): seg_mask = np.zeros_like(seg_img[:,:,0], dtype=np.uint8) for color in color_code_list: lowerb = np.array(color, dtype=np.uint8) upperb = np.array(color, dtype=np.uint8) temp_mask = cv2.inRange(seg_img, lowerb, upperb) seg_mask = cv2.bitwise_or(seg_mask, temp_mask) # mask_3channel = cv2.merge([seg_mask, seg_mask, seg_mask]) return seg_mask def tensor_to_cv2_img(tensor, remove_alpha=False): i = 255. * tensor.squeeze(0).cpu().numpy() # This will give us (H, W, C) img = np.clip(i, 0, 255).astype(np.uint8) return img def cv2_img_to_tensor(img): img = img.astype(np.float32) / 255.0 img = torch.from_numpy(img)[None,] return img masks = [] extracted_images = [] extracted_secondaries = [] for i in range(batch_images.shape[0]): seg = batch_segs[i] cv2_seg = tensor_to_cv2_img(seg) color_code_list = [] #################ATR MAPPING################# if right_leg: color_code_list.append([192,0,128]) if right_hand: color_code_list.append([192,128,128]) if head: color_code_list.append([192,128,0]) if hair: color_code_list.append([0,128,0]) if left_shoe: color_code_list.append([192,0,0]) if bag: color_code_list.append([0,64,0]) if background: color_code_list.append([0,0,0]) if dress: color_code_list.append([128,128,128]) if left_leg: color_code_list.append([64,0,128]) if right_shoe: color_code_list.append([64,128,0]) if left_hand: color_code_list.append([64,128,128]) if upper_garment: color_code_list.append([0,0,128]) if lower_garment: color_code_list.append([0,128,128]) if belt: color_code_list.append([64,0,0]) if skirt: color_code_list.append([128,0,128]) if hat: color_code_list.append([128,0,0]) if sunglasses: color_code_list.append([128,128,0]) if scarf: color_code_list.append([128,64,0]) # get_segment_counts(cv2_seg) mask = generate_mask(cv2_seg, color_code_list) mask_3channel = cv2.merge([mask,mask,mask]) tensor_mask = cv2_img_to_tensor(mask_3channel) cv2_image = tensor_to_cv2_img(batch_images[i]) cv2_secondary = tensor_to_cv2_img(batch_secondary[i]) extracted_image = cv2.bitwise_and(cv2_image, cv2_image, mask=mask) # Use the single_channel_mask here extracted_secondary = cv2.bitwise_and(cv2_secondary, cv2_secondary, mask=mask) # Use the single_channel_mask here tensor_extracted_image = cv2_img_to_tensor(extracted_image) tensor_extracted_secondary = cv2_img_to_tensor(extracted_secondary) extracted_images.append(tensor_extracted_image.squeeze(0)) extracted_secondaries.append(tensor_extracted_secondary.squeeze(0)) print(tensor_mask.shape,"tensor_mask.shape") masks.append(tensor_mask.squeeze(0)) # Convert the masks to tensors batch_masks = torch.stack(masks) batch_imgs = torch.stack(extracted_images) batch_secondaries = torch.stack(extracted_secondaries) print(batch_masks.shape,"batch_masks.shape") return (batch_masks,batch_imgs, batch_secondaries) class TRI3DExtractPartsBatch: def __init__(self): pass @classmethod def INPUT_TYPES(s): return { "required": { "batch_images": ("IMAGE",), "batch_segs" : ("IMAGE",), "batch_secondaries" : ("IMAGE",), "margin" : ("INT", {"default": 15, "min": 0 }), "right_leg": ("BOOLEAN", {"default": False}), "right_hand": ("BOOLEAN", {"default": True}), "head": ("BOOLEAN", {"default": False}), "hair": ("BOOLEAN", {"default": False}), "left_shoe" : ("BOOLEAN", {"default": False}), "bag" : ("BOOLEAN", {"default": False}), "background" : ("BOOLEAN", {"default": False}), "dress" : ("BOOLEAN", {"default": False}), "left_leg": ("BOOLEAN", {"default": False}), "right_shoe" : ("BOOLEAN", {"default": False}), "left_hand": ("BOOLEAN", {"default": True}), "upper_garment" : ("BOOLEAN", {"default": False}), "lower_garment" : ("BOOLEAN", {"default": False}), "belt" : ("BOOLEAN", {"default": False}), "skirt" : ("BOOLEAN", {"default": False}), "hat" : ("BOOLEAN", {"default": False}), "sunglasses" : ("BOOLEAN", {"default": False}), "scarf" : ("BOOLEAN", {"default": False}), }, } RETURN_TYPES = ("IMAGE","IMAGE",) FUNCTION = "main" CATEGORY = "TRI3D" def main(self, batch_images, batch_segs, batch_secondaries, margin, right_leg,right_hand, head, hair, left_shoe,bag,background,dress,left_leg,right_shoe,left_hand, upper_garment,lower_garment,belt,skirt,hat,sunglasses,scarf): import cv2 import numpy as np import torch from pprint import pprint def get_segment_counts(segm): # Load the segmentation image # Reshape the image array to be 2D reshaped = segm.reshape(-1, segm.shape[-1]) # Find unique vectors and their counts unique_vectors, counts = np.unique(reshaped, axis=0, return_counts=True) segment_counts = list(zip(unique_vectors, counts)) pprint(segment_counts) return segment_counts def bounded_image(seg_img, color_code_list, input_img): import cv2 import numpy as np # Create a mask for hands seg_img = cv2.resize(seg_img,(input_img.shape[1],input_img.shape[0]),interpolation=cv2.INTER_NEAREST) hand_mask = np.zeros_like(seg_img[:,:,0]) for color in color_code_list: lowerb = np.array(color, dtype=np.uint8) upperb = np.array(color, dtype=np.uint8) temp_mask = cv2.inRange(seg_img, lowerb, upperb) hand_mask = cv2.bitwise_or(hand_mask, temp_mask) # Find contours to get the bounding box of the hands contours, _ = cv2.findContours(hand_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) # If no contours were found, just return None if not contours: return None # Combine all contours to find encompassing bounding box all_points = np.concatenate(contours, axis=0) x, y, w, h = cv2.boundingRect(all_points) print(x,y,w,h,"x,y,w,h") x = max(x - margin, 0) y = max(y - margin, 0) w = min(w + 2*margin, input_img.shape[1] - x) # Ensure width does not exceed image boundary h = min(h + 2*margin, input_img.shape[0] - y) # Ensure height does not exceed image boundary print(x,y,w,h,"x,y,w,h") print(input_img.shape,"input_img.shape") # Extract the region from the original image that contains both hands hand_region = input_img[y:y+h, x:x+w] return hand_region def tensor_to_cv2_img(tensor, remove_alpha=False): i = 255. * tensor.squeeze(0).cpu().numpy() # This will give us (H, W, C) img = np.clip(i, 0, 255).astype(np.uint8) return img def cv2_img_to_tensor(img): img = img.astype(np.float32) / 255.0 img = torch.from_numpy(img)[None,] return img batch_results = [] images = [] secondaries = [] for i in range(batch_images.shape[0]): image = batch_images[i] seg = batch_segs[i] cv2_image = tensor_to_cv2_img(image) cv2_secondary = tensor_to_cv2_img(batch_secondaries[i]) cv2_seg = tensor_to_cv2_img(seg) color_code_list = [] #################ATR MAPPING################# if right_leg: color_code_list.append([192,0,128]) if right_hand: color_code_list.append([192,128,128]) if head: color_code_list.append([192,128,0]) if hair: color_code_list.append([0,128,0]) if left_shoe: color_code_list.append([192,0,0]) if bag: color_code_list.append([0,64,0]) if background: color_code_list.append([0,0,0]) if dress: color_code_list.append([128,128,128]) if left_leg: color_code_list.append([64,0,128]) if right_shoe: color_code_list.append([64,128,0]) if left_hand: color_code_list.append([64,128,128]) if upper_garment: color_code_list.append([0,0,128]) if lower_garment: color_code_list.append([0,128,128]) if belt: color_code_list.append([64,0,0]) if skirt: color_code_list.append([128,0,128]) if hat: color_code_list.append([128,0,0]) if sunglasses: color_code_list.append([128,128,0]) if scarf: color_code_list.append([128,64,0]) bimage = bounded_image(cv2_seg, color_code_list, cv2_image) bsecondary = bounded_image(cv2_seg, color_code_list, cv2_secondary) # Handle case when bimage is None to avoid error during conversion to tensor if bimage is not None: images.append(bimage) else: black_img = np.zeros_like(cv2_image) images.append(black_img) if bsecondary is not None: secondaries.append(bsecondary) else: black_img = np.zeros_like(cv2_secondary) secondaries.append(black_img) # Get max height and width max_height = max(img.shape[0] for img in images) max_width = max(img.shape[1] for img in images) batch_results = [] batch_secondaries = [] for img in images: # Resize the image to max height and width resized_img = cv2.resize(img, (max_width, max_height), interpolation=cv2.INTER_AREA) tensor_img = cv2_img_to_tensor(resized_img) batch_results.append(tensor_img.squeeze(0)) for sec in secondaries: # Resize the image to max height and width resized_sec = cv2.resize(sec, (max_width, max_height), interpolation=cv2.INTER_AREA) tensor_sec = cv2_img_to_tensor(resized_sec) batch_secondaries.append(tensor_sec.squeeze(0)) batch_results = torch.stack(batch_results) batch_secondaries = torch.stack(batch_secondaries) print(batch_results.shape,"batch_results.shape") return (batch_results,batch_secondaries) class TRI3DPositionPartsBatch: def __init__(self): pass @classmethod def INPUT_TYPES(s): return { "required": { "batch_images": ("IMAGE",), "batch_segs" : ("IMAGE",), "batch_handimgs" : ("IMAGE",), "margin" : ("INT", {"default": 15, "min": 0 }), "right_leg": ("BOOLEAN", {"default": False}), "right_hand": ("BOOLEAN", {"default": True}), "head": ("BOOLEAN", {"default": False}), "hair": ("BOOLEAN", {"default": False}), "left_shoe" : ("BOOLEAN", {"default": False}), "bag" : ("BOOLEAN", {"default": False}), "background" : ("BOOLEAN", {"default": False}), "dress" : ("BOOLEAN", {"default": False}), "left_leg": ("BOOLEAN", {"default": False}), "right_shoe" : ("BOOLEAN", {"default": False}), "left_hand": ("BOOLEAN", {"default": True}), "upper_garment" : ("BOOLEAN", {"default": False}), "lower_garment" : ("BOOLEAN", {"default": False}), "belt" : ("BOOLEAN", {"default": False}), "skirt" : ("BOOLEAN", {"default": False}), "hat" : ("BOOLEAN", {"default": False}), "sunglasses" : ("BOOLEAN", {"default": False}), "scarf" : ("BOOLEAN", {"default": False}), }, } RETURN_TYPES = ("IMAGE",) FUNCTION = "main" CATEGORY = "TRI3D" def main(self, batch_images, batch_segs, batch_handimgs, margin,right_leg,right_hand, head, hair, left_shoe,bag,background,dress,left_leg,right_shoe,left_hand, upper_garment,lower_garment,belt,skirt,hat,sunglasses,scarf): import cv2 import numpy as np import torch from pprint import pprint def bounded_image_points(seg_img, color_code_list, input_img): import cv2 import numpy as np # Create a mask for hands seg_img = cv2.resize(seg_img,(input_img.shape[1],input_img.shape[0]),interpolation=cv2.INTER_NEAREST) hand_mask = np.zeros_like(seg_img[:,:,0]) for color in color_code_list: lowerb = np.array(color, dtype=np.uint8) upperb = np.array(color, dtype=np.uint8) temp_mask = cv2.inRange(seg_img, lowerb, upperb) hand_mask = cv2.bitwise_or(hand_mask, temp_mask) # Find contours to get the bounding box of the hands contours, _ = cv2.findContours(hand_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) # If no contours were found, just return None if not contours: return None # Combine all contours to find encompassing bounding box all_points = np.concatenate(contours, axis=0) x, y, w, h = cv2.boundingRect(all_points) x = max(x - margin, 0) y = max(y - margin, 0) w = min(w + 2*margin, input_img.shape[1] - x) # Ensure width does not exceed image boundary h = min(h + 2*margin, input_img.shape[0] - y) # Ensure height does not exceed image boundary return (x,y,w,h) def tensor_to_cv2_img(tensor, remove_alpha=False): i = 255. * tensor.squeeze(0).cpu().numpy() # This will give us (H, W, C) img = np.clip(i, 0, 255).astype(np.uint8) return img def cv2_img_to_tensor(img): img = img.astype(np.float32) / 255.0 img = torch.from_numpy(img)[None,] return img batch_results = [] for i in range(batch_images.shape[0]): image = batch_images[i] seg = batch_segs[i] handimg = batch_handimgs[i] cv2_image = tensor_to_cv2_img(image) cv2_seg = tensor_to_cv2_img(seg) color_code_list = [] #################ATR MAPPING################# if right_leg: color_code_list.append([192,0,128]) if right_hand: color_code_list.append([192,128,128]) if head: color_code_list.append([192,128,0]) if hair: color_code_list.append([0,128,0]) if left_shoe: color_code_list.append([192,0,0]) if bag: color_code_list.append([0,64,0]) if background: color_code_list.append([0,0,0]) if dress: color_code_list.append([128,128,128]) if left_leg: color_code_list.append([64,0,128]) if right_shoe: color_code_list.append([64,128,0]) if left_hand: color_code_list.append([64,128,128]) if upper_garment: color_code_list.append([0,0,128]) if lower_garment: color_code_list.append([0,128,128]) if belt: color_code_list.append([64,0,0]) if skirt: color_code_list.append([128,0,128]) if hat: color_code_list.append([128,0,0]) if sunglasses: color_code_list.append([128,128,0]) if scarf: color_code_list.append([128,64,0]) positions = bounded_image_points(cv2_seg, color_code_list, cv2_image) cv2_handimg = tensor_to_cv2_img(handimg) cv2_handimg = cv2.resize(cv2_handimg, (positions[2], positions[3]), interpolation=cv2.INTER_NEAREST) cv2_image[positions[1]:positions[1]+positions[3], positions[0]:positions[0]+positions[2]] = cv2_handimg b_tensor_img = cv2_img_to_tensor(cv2_image) batch_results.append(b_tensor_img.squeeze(0)) batch_results = torch.stack(batch_results) return (batch_results,) class TRI3DATRParseBatch: def __init__(self): pass @classmethod def INPUT_TYPES(s): return { "required": { "images": ("IMAGE",), }, } RETURN_TYPES = ("IMAGE",) FUNCTION = "main" CATEGORY = "TRI3D" def main(self, images): import cv2 import numpy as np import torch import os import shutil def tensor_to_cv2_img(tensor, remove_alpha=False): i = 255. * tensor.cpu().numpy() # This will give us (H, W, C) img = np.clip(i, 0, 255).astype(np.uint8) return img def cv2_img_to_tensor(img): img = img.astype(np.float32) / 255.0 img = torch.from_numpy(img)[None,] return img ATR_PATH = 'custom_nodes/tri3d-comfyui-nodes/atr_node/' ATR_INPUT_PATH = ATR_PATH + 'input/' ATR_OUTPUT_PATH = ATR_PATH + 'output/' # Create the input directory if it does not exist shutil.rmtree(ATR_INPUT_PATH, ignore_errors=True) os.makedirs(ATR_INPUT_PATH, exist_ok=True) shutil.rmtree(ATR_OUTPUT_PATH, ignore_errors=True) os.makedirs(ATR_OUTPUT_PATH, exist_ok=True) for i in range(images.shape[0]): image = images[i] cv2_image = tensor_to_cv2_img(image) cv2_image = cv2.cvtColor(cv2_image, cv2.COLOR_BGR2RGB) cv2.imwrite(ATR_INPUT_PATH + f"image{i}.png", cv2_image) # Run the ATR model cwd = os.getcwd() os.chdir(ATR_PATH) os.system("python simple_extractor.py --dataset atr --model-restore 'checkpoints/atr.pth' --input-dir input --output-dir output") os.chdir(cwd) # Collect and return the results batch_results = [] for i in range(images.shape[0]): cv2_segm = cv2.imread(ATR_OUTPUT_PATH + f'image{i}.png') cv2_segm = cv2.cvtColor(cv2_segm, cv2.COLOR_BGR2RGB) b_tensor_img = cv2_img_to_tensor(cv2_segm) batch_results.append(b_tensor_img.squeeze(0)) batch_results = torch.stack(batch_results) return (batch_results,) class TRI3DExtractHand: def __init__(self): pass @classmethod def INPUT_TYPES(s): return { "required": { "image": ("IMAGE",), "seg" : ("IMAGE",), "margin" : ("INT", {"default": 15, "min": 0 }), "left_hand" : ("BOOLEAN", {"default": True}), "right_hand" : ("BOOLEAN", {"default": True}), "head" : ("BOOLEAN", {"default": False}), "hair" : ("BOOLEAN", {"default": False}), "left_leg" : ("BOOLEAN", {"default": False}), "right_leg" : ("BOOLEAN", {"default": False}), }, } RETURN_TYPES = ("IMAGE",) FUNCTION = "main" CATEGORY = "TRI3D" def main(self, image,seg,margin,left_hand,right_hand,head,hair,left_leg,right_leg): import cv2 import numpy as np import torch from pprint import pprint def get_segment_counts(segm): # Load the segmentation image # Reshape the image array to be 2D reshaped = segm.reshape(-1, segm.shape[-1]) # Find unique vectors and their counts unique_vectors, counts = np.unique(reshaped, axis=0, return_counts=True) segment_counts = list(zip(unique_vectors, counts)) pprint(segment_counts) return segment_counts def bounded_image(seg_img, color_code_list, input_img): import cv2 import numpy as np # Create a mask for hands seg_img = cv2.resize(seg_img,(input_img.shape[1],input_img.shape[0]),interpolation=cv2.INTER_NEAREST) hand_mask = np.zeros_like(seg_img[:,:,0]) for color in color_code_list: lowerb = np.array(color, dtype=np.uint8) upperb = np.array(color, dtype=np.uint8) temp_mask = cv2.inRange(seg_img, lowerb, upperb) hand_mask = cv2.bitwise_or(hand_mask, temp_mask) # Find contours to get the bounding box of the hands contours, _ = cv2.findContours(hand_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) # If no contours were found, just return None if not contours: return None # Combine all contours to find encompassing bounding box all_points = np.concatenate(contours, axis=0) x, y, w, h = cv2.boundingRect(all_points) print(x,y,w,h,"x,y,w,h") x = max(x - margin, 0) y = max(y - margin, 0) w = min(w + 2*margin, input_img.shape[1] - x) # Ensure width does not exceed image boundary h = min(h + 2*margin, input_img.shape[0] - y) # Ensure height does not exceed image boundary print(x,y,w,h,"x,y,w,h") print(input_img.shape,"input_img.shape") # Extract the region from the original image that contains both hands hand_region = input_img[y:y+h, x:x+w] return hand_region def tensor_to_cv2_img(tensor, remove_alpha=False): i = 255. * tensor.squeeze(0).cpu().numpy() # This will give us (H, W, C) img = np.clip(i, 0, 255).astype(np.uint8) return img def cv2_img_to_tensor(img): img = img.astype(np.float32) / 255.0 img = torch.from_numpy(img)[None,] return img cv2_image = tensor_to_cv2_img(image) cv2_seg = tensor_to_cv2_img(seg) get_segment_counts(cv2_seg) color_code_list = [] if left_hand: color_code_list.append([64,128,128]) if right_hand: color_code_list.append([192,128,128]) if head: color_code_list.append([192,128,0]) if hair: color_code_list.append([0,128,0]) if left_leg: color_code_list.append([192,0,0]) if right_leg: color_code_list.append([64,128,0]) # color_code_list = [[64,128,128], [192,128,128]] bimage = bounded_image(cv2_seg,color_code_list,cv2_image) b_tensor_img = cv2_img_to_tensor(bimage) return (b_tensor_img,) class TRI3DATRParse: def __init__(self): pass @classmethod def INPUT_TYPES(s): return { "required": { "image": ("IMAGE",), }, } RETURN_TYPES = ("IMAGE",) FUNCTION = "main" CATEGORY = "TRI3D" def main(self, image): import cv2 import numpy as np import torch import os import shutil from pprint import pprint def tensor_to_cv2_img(tensor, remove_alpha=False): i = 255. * tensor.squeeze(0).cpu().numpy() # This will give us (H, W, C) img = np.clip(i, 0, 255).astype(np.uint8) return img def cv2_img_to_tensor(img): img = img.astype(np.float32) / 255.0 img = torch.from_numpy(img)[None,] return img cv2_image = tensor_to_cv2_img(image) cv2_image = cv2.cvtColor(cv2_image, cv2.COLOR_BGR2RGB) ATR_PATH = 'custom_nodes/tri3d-comfyui-nodes/atr_node/' ATR_INPUT_PATH = ATR_PATH + 'input/' ATR_OUTPUT_PATH = ATR_PATH + 'output/' # Create the input directory if it does not exist shutil.rmtree(ATR_INPUT_PATH, ignore_errors=True) os.makedirs(ATR_INPUT_PATH, exist_ok=True) shutil.rmtree(ATR_OUTPUT_PATH, ignore_errors=True) os.makedirs(ATR_OUTPUT_PATH, exist_ok=True) cv2.imwrite(ATR_INPUT_PATH + "image.png",cv2_image) # Run the ATR model cwd = os.getcwd() os.chdir(ATR_PATH) os.system("python simple_extractor.py --dataset atr --model-restore 'checkpoints/atr.pth' --input-dir input --output-dir output") # Load the segmentation image os.chdir(cwd) cv2_segm = cv2.imread(ATR_OUTPUT_PATH + 'image.png') cv2_segm = cv2.cvtColor(cv2_segm, cv2.COLOR_BGR2RGB) b_tensor_img = cv2_img_to_tensor(cv2_segm) return (b_tensor_img,) class TRI3DPositiontHands: def __init__(self): pass @classmethod def INPUT_TYPES(s): return { "required": { "image": ("IMAGE",), "seg" : ("IMAGE",), "handimg" : ("IMAGE",), "margin" : ("INT", {"default": 15, "min": 0 }), "left_hand" : ("BOOLEAN", {"default": True}), "right_hand" : ("BOOLEAN", {"default": True}), "head" : ("BOOLEAN", {"default": False}), "hair" : ("BOOLEAN", {"default": False}), "left_leg" : ("BOOLEAN", {"default": False}), "right_leg" : ("BOOLEAN", {"default": False}), }, } RETURN_TYPES = ("IMAGE",) FUNCTION = "main" CATEGORY = "TRI3D" def main(self, image,seg,handimg,margin,left_hand,right_hand,head,hair,left_leg,right_leg): import cv2 import numpy as np import torch from pprint import pprint def bounded_image_points(seg_img, color_code_list, input_img): import cv2 import numpy as np # Create a mask for hands seg_img = cv2.resize(seg_img,(input_img.shape[1],input_img.shape[0]),interpolation=cv2.INTER_NEAREST) hand_mask = np.zeros_like(seg_img[:,:,0]) for color in color_code_list: lowerb = np.array(color, dtype=np.uint8) upperb = np.array(color, dtype=np.uint8) temp_mask = cv2.inRange(seg_img, lowerb, upperb) hand_mask = cv2.bitwise_or(hand_mask, temp_mask) # Find contours to get the bounding box of the hands contours, _ = cv2.findContours(hand_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) # If no contours were found, just return None if not contours: return None # Combine all contours to find encompassing bounding box all_points = np.concatenate(contours, axis=0) x, y, w, h = cv2.boundingRect(all_points) x = max(x - margin, 0) y = max(y - margin, 0) w = min(w + 2*margin, input_img.shape[1] - x) # Ensure width does not exceed image boundary h = min(h + 2*margin, input_img.shape[0] - y) # Ensure height does not exceed image boundary return (x,y,w,h) def tensor_to_cv2_img(tensor, remove_alpha=False): i = 255. * tensor.squeeze(0).cpu().numpy() # This will give us (H, W, C) img = np.clip(i, 0, 255).astype(np.uint8) return img def cv2_img_to_tensor(img): img = img.astype(np.float32) / 255.0 img = torch.from_numpy(img)[None,] return img cv2_image = tensor_to_cv2_img(image) cv2_seg = tensor_to_cv2_img(seg) # cv2_seg = cv2.resize(cv2_seg,(cv2_image.shape[1],cv2_image.shape[0]),interpolation=cv2.INTER_NEAREST) # 128 128 64 / 128 128 192 # color_code_list = [[128,128,64], [128,128,192]] # color_code_list = [[64,128,128], [192,128,128]] color_code_list = [] if left_hand: color_code_list.append([64,128,128]) if right_hand: color_code_list.append([192,128,128]) if head: color_code_list.append([192,128,0]) if hair: color_code_list.append([0,128,0]) if left_leg: color_code_list.append([192,0,0]) if right_leg: color_code_list.append([64,128,0]) positions = bounded_image_points(cv2_seg,color_code_list,cv2_image) print(positions,"positions") try: cv2_handimg = tensor_to_cv2_img(handimg) #Resize cv2_handimg to positions print("before resizing ",cv2_handimg.shape,"handimg.shape") cv2_handimg = cv2.resize(cv2_handimg,(positions[2],positions[3]),interpolation=cv2.INTER_NEAREST) print(positions,"positions") print(cv2_image.shape,"cv2img.shape") print(cv2_handimg.shape,"handimg.shape") #position cv2_handimg in cv2_image cv2_image[positions[1]:positions[1]+positions[3],positions[0]:positions[0]+positions[2]] = cv2_handimg except Exception as e: print(e) pass b_tensor_img = cv2_img_to_tensor(cv2_image) return (b_tensor_img,) class TRI3DFuzzification: def __init__(self): pass @classmethod def INPUT_TYPES(s): return { "required": { "input": ("IMAGE",), "inputseg" : ("IMAGE",), "controlnetoutput": ("IMAGE",), "controlnetoutputseg" : ("IMAGE",), }, } RETURN_TYPES = ("IMAGE",) FUNCTION = "main" CATEGORY = "TRI3D" def main(self, input,inputseg,controlnetoutput,controlnetoutputseg): import cv2 import numpy as np import torch from pprint import pprint from scipy.spatial import distance def combined_image(img1, img2, mask): if img1.shape[2] == 4: # Convert it from four channels to three channels img1 = cv2.cvtColor(img1, cv2.COLOR_BGRA2BGR) mask_inv = cv2.bitwise_not(mask) # Normalize the masks to the range [0, 1] mask = cv2.normalize(mask, None, alpha=0, beta=1, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_32F) mask_inv = cv2.normalize(mask_inv, None, alpha=0, beta=1, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_32F) # Convert images and masks to the same data type img1 = img1.astype(np.float32) img2 = img2.astype(np.float32) mask = mask.astype(np.float32) mask_inv = mask_inv.astype(np.float32) img1 = cv2.resize(img1, (mask.shape[1], mask.shape[0]),interpolation=cv2.INTER_NEAREST) img2 = cv2.resize(img2, (mask.shape[1], mask.shape[0]),interpolation=cv2.INTER_NEAREST) # Check if img1 (and hence img2) have more than one channel (e.g., RGB images) if len(img1.shape) > 2: # Convert mask and mask_inv to the same number of channels as img1 mask = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR) mask_inv = cv2.cvtColor(mask_inv, cv2.COLOR_GRAY2BGR) # Use the masks to get the weighted regions of each image img1_masked = cv2.multiply(img1, mask_inv) img2_masked = cv2.multiply(img2, mask) # Combine the two images combined = cv2.add(img1_masked, img2_masked).astype(np.uint8) return combined def fuzzify(img): if len(img.shape) == 3: img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) # Threshold the image _, thresholded = cv2.threshold(img, 1, 255, cv2.THRESH_BINARY) # Find contours in the thresholded image contours, _ = cv2.findContours(thresholded, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) # # Approximate contours to reduce number of points # epsilon_factor = 0.02 # can be adjusted, higher values mean more simplification # simplified_contours = [cv2.approxPolyDP(cnt, epsilon_factor * cv2.arcLength(cnt, True), True) for cnt in contours] # Define brush size brush_radius = int(15*img.shape[0]/1024.0) mask = np.zeros_like(img) for contour in contours: # print("contour shape - ",contour.shape) contour_points = contour.squeeze(1) # Determine the bounding box around the contour and expand it by the brush radius x_min, y_min = np.min(contour_points, axis=0) - brush_radius x_max, y_max = np.max(contour_points, axis=0) + brush_radius # Clip the coordinates to the image boundaries x_min, y_min = max(0, x_min), max(0, y_min) x_max, y_max = min(img.shape[1]-1, x_max), min(img.shape[0]-1, y_max) # # Create a grid of coordinates within this bounding box # ys, xs = np.ogrid[y_min:y_max+1, x_min:x_max+1] # grid_coords = np.column_stack((xs.ravel(), ys.ravel())) ys, xs = np.mgrid[y_min:y_max+1, x_min:x_max+1] grid_coords = np.column_stack((xs.flatten(), ys.flatten())) # Compute distances for pixels inside the bounding box distances = distance.cdist(grid_coords, contour_points, 'euclidean') min_distances = distances.min(axis=1) brush_effect = np.clip((1 - min_distances / brush_radius) * 255, 0, 255) mask[grid_coords[:, 1], grid_coords[:, 0]] = np.maximum(mask[grid_coords[:, 1], grid_coords[:, 0]], brush_effect) # Save the result final_mask = np.maximum(img, mask.astype(img.dtype)) return final_mask def get_segment_counts(segm): # Load the segmentation image # Reshape the image array to be 2D reshaped = segm.reshape(-1, segm.shape[-1]) # Find unique vectors and their counts unique_vectors, counts = np.unique(reshaped, axis=0, return_counts=True) segment_counts = list(zip(unique_vectors, counts)) pprint(segment_counts) # array([0, 0, 0], dtype=uint8), 421768), #background # (array([ 0, 0, 128], dtype=uint8), 291418), #upper garment # (array([ 0, 128, 0], dtype=uint8), 9393), #hair # (array([ 0, 128, 128], dtype=uint8), 50277), #lower garment # (array([ 64, 128, 128], dtype=uint8), 14548), #left hand # (array([192, 128, 0], dtype=uint8), 33325), #face # (array([192, 128, 128], dtype=uint8), 14855)] #right hand # (array([ 0, 0, 64], dtype=uint8), 20918), # (array([ 0, 0, 192], dtype=uint8), 20264), #left shoe # (array([ 0, 128, 0], dtype=uint8), 64359), # (array([ 0, 128, 64], dtype=uint8), 21031), #right shoe # (array([ 0, 128, 192], dtype=uint8), 76005), # (array([128, 0, 0], dtype=uint8), 102761), # (array([128, 0, 64], dtype=uint8), 89881), # (array([128, 0, 192], dtype=uint8), 93931), # (array([128, 128, 0], dtype=uint8), 39445), # (array([128, 128, 64], dtype=uint8), 44930), # (array([128, 128, 192], dtype=uint8), 59772)] # color_code_list = [[64,128,128], [192,128,128]] #left and right hands # color_code_list = [[192,128,0]] #face # color_code_list = [[0,128,0]] #hair # color_code_list = [[0,128,128]] #lower garment # color_code_list = [[0,0,128]] #upper garment # color_code_list = [[0,0,0]] #background return segment_counts def blend_images(cv2_input, cv2_inputseg, cv2_controlnetoutput, cv2_controlnetoutputseg, color_code_dict): # Helper function to create masks def get_mask_from_colors(image, color_list): mask = np.zeros_like(image[:,:,0]) for color in color_list: lowerb = np.array(color, dtype=np.uint8) upperb = np.array(color, dtype=np.uint8) temp_mask = cv2.inRange(image, lowerb, upperb) mask = cv2.bitwise_or(mask, temp_mask) return cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR) input_facehair = get_mask_from_colors(cv2_inputseg, color_code_dict['face_hair']) input_background = get_mask_from_colors(cv2_inputseg, color_code_dict['background']) input_rest = cv2.bitwise_not(cv2.add(input_facehair, input_background)) controlnet_facehair = get_mask_from_colors(cv2_controlnetoutputseg, color_code_dict['face_hair']) # Initial Image: Set it to input_rest blended_image = np.copy(cv2_input) blended_image = cv2.bitwise_and(blended_image, input_rest) # blended_image = combined_image(cv2_input, np.zeros_like(cv2_input), fuzzify(input_rest)) # Stage 2: Overlay face and hair pixels from controlnet_facehair onto blended_image face_hair_region = cv2.bitwise_and(cv2_controlnetoutput, controlnet_facehair) inverse_face_hair_mask = cv2.bitwise_not(controlnet_facehair) blended_without_facehair = cv2.bitwise_and(blended_image, inverse_face_hair_mask) blended_image = cv2.add(blended_without_facehair, face_hair_region) # blended_facehair = combined_image(blended_image, cv2_controlnetoutput, fuzzify(controlnet_facehair)) # Stage 3: Overlay the remaining pixels with white remaining_mask = cv2.bitwise_not(cv2.add(input_rest, controlnet_facehair)) inverse_remaining_mask = cv2.bitwise_not(remaining_mask) white_fill = np.ones_like(cv2_input) * 255 # Create an image filled with white white_region = cv2.bitwise_and(white_fill, remaining_mask) blended_without_remaining = cv2.bitwise_and(blended_image, inverse_remaining_mask) blended_image = cv2.add(blended_without_remaining, white_region) # remaining_mask = cv2.bitwise_not(cv2.add(input_rest, controlnet_facehair)) # blended_image = combined_image(blended_facehair, cv2_input, fuzzify(remaining_mask)) return blended_image def tensor_to_cv2_img(tensor, remove_alpha=False): i = 255. * tensor.squeeze(0).cpu().numpy() # This will give us (H, W, C) img = np.clip(i, 0, 255).astype(np.uint8) return img def cv2_img_to_tensor(img): img = img.astype(np.float32) / 255.0 img = torch.from_numpy(img)[None,] return img cv2_input = tensor_to_cv2_img(input) cv2_inputseg = tensor_to_cv2_img(inputseg) cv2_controlnetoutput = tensor_to_cv2_img(controlnetoutput) cv2_controlnetoutputseg = tensor_to_cv2_img(controlnetoutputseg) cv2_inputseg = cv2.resize(cv2_inputseg,(cv2_input.shape[1],cv2_input.shape[0]),interpolation=cv2.INTER_NEAREST) cv2_controlnetoutput = cv2.resize(cv2_controlnetoutput,(cv2_input.shape[1],cv2_input.shape[0]),interpolation=cv2.INTER_NEAREST) cv2_controlnetoutputseg = cv2.resize(cv2_controlnetoutputseg,(cv2_input.shape[1],cv2_input.shape[0]),interpolation=cv2.INTER_NEAREST) # # cv2_seg = cv2.resize(cv2_seg,(cv2_image.shape[1],cv2_image.shape[0]),interpolation=cv2.INTER_NEAREST) # color_code_list = [[192,128,0],[0,128,0]] #face and hair color_code_dict = { 'face_hair' : [[192,128,0],[0,128,0]], 'background' : [[0,0,0]], } bimage = blend_images(cv2_input,cv2_inputseg,cv2_controlnetoutput,cv2_controlnetoutputseg,color_code_dict) # bimage = bounded_image(cv2_inputseg,color_code_list,cv2_input) output_img = cv2_img_to_tensor(bimage) return (output_img,) # A dictionary that contains all nodes you want to export with their names # NOTE: names should be globally unique NODE_CLASS_MAPPINGS = { "tri3d-extract-hand": TRI3DExtractHand, "tri3d-position-hands": TRI3DPositiontHands, "tri3d-atr-parse": TRI3DATRParse, "tri3d-fuzzification": TRI3DFuzzification, "tri3d-atr-parse-batch": TRI3DATRParseBatch, "tri3d-position-parts-batch": TRI3DPositionPartsBatch, 'tri3d-extract-parts-batch': TRI3DExtractPartsBatch, 'tri3d-extract-parts-mask-batch': TRI3DExtractPartsMaskBatch, } # A dictionary that contains the friendly/humanly readable titles for the nodes NODE_DISPLAY_NAME_MAPPINGS = { "tri3d-extract-hand": "Extract Hand", "tri3d-fuzzification" : "Fuzzification", "tri3d-position-hands" : "Position Hands", "tri3d-atr-parse" : "ATR Parse", "tri3d-atr-parse-batch" : "ATR Parse Batch", "tri3d-position-parts-batch" : "Position Parts Batch", 'tri3d-extract-parts-batch': 'Extract Parts Batch', 'tri3d-extract-parts-mask-batch': 'Extract Parts Mask Batch', }