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,)