diff --git a/__init__.py b/__init__.py index e8cc23b..e382a6e 100644 --- a/__init__.py +++ b/__init__.py @@ -1941,6 +1941,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) @@ -3692,6 +3694,7 @@ class TRI3D_BGREMOVE_MEGA(): from photoroom import TRI3D_photoroom_bgremove_api +from smart_box import TRI3D_SmartBox # A dictionary that contains all nodes you want to export with their names # NOTE: names should be globally unique @@ -3752,10 +3755,11 @@ NODE_CLASS_MAPPINGS = { "tri3d_SaveImage_absolute":SaveImage_absolute, "tri3d_SaveText_absolute":SaveText_absolute, "tri3d_Wait_And_Read_File":Wait_And_Read_File, + "tri3d_SmartBox": TRI3D_SmartBox, } -VERSION = "4.7.3" +VERSION = "4.8.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, @@ -3816,4 +3820,5 @@ NODE_DISPLAY_NAME_MAPPINGS = { "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, } diff --git a/dwpose/comfy_utils.py b/dwpose/comfy_utils.py index ae614e3..c75be41 100644 --- a/dwpose/comfy_utils.py +++ b/dwpose/comfy_utils.py @@ -274,4 +274,10 @@ def switch_to_backpose(input_keypoints, input_width): x,y = input_keypoints[i] input_keypoints[i] = [input_width - x, y] - return input_keypoints \ No newline at end of file + return input_keypoints + +def extract_torso_keypoints(keypoints): + # Indices for torso-related keypoints + torso_indices = [8, 9, 10, 11, 12, 13] + return [keypoints[i] for i in torso_indices if keypoints[i] != [-1, -1]] + diff --git a/smart_box.py b/smart_box.py new file mode 100644 index 0000000..bb04da4 --- /dev/null +++ b/smart_box.py @@ -0,0 +1,91 @@ +import numpy as np +import torch +import json +from dwpose.comfy_utils import extract_torso_keypoints +import cv2 + +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 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 = 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 + hip_y = (keypoints[0][1] + keypoints[1][1]) // 2 + + # 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) + + for contour in contours: + x, y, w, h = cv2.boundingRect(contour) + cv2.rectangle(image, (x, y + hip_y), (x + w, y + h + hip_y), (255, 255, 255), -1) + + return image + +