From d10c1195e60de7f1471b6dfa5f211b3677a28279 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Mon, 27 Jan 2025 10:57:51 +0000 Subject: [PATCH 1/4] TRI3D_Image_extend node to extend image for a close up image input --- __init__.py | 6 ++-- smart_box.py | 84 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 88 insertions(+), 2 deletions(-) diff --git a/__init__.py b/__init__.py index 6704fdd..b42829b 100644 --- a/__init__.py +++ b/__init__.py @@ -3694,7 +3694,7 @@ class TRI3D_BGREMOVE_MEGA(): from photoroom import TRI3D_photoroom_bgremove_api -from smart_box import TRI3D_SmartBox, TRI3D_Skip_HeadMask, TRI3D_Skip_HeadMask_AddNeck +from smart_box import TRI3D_SmartBox, TRI3D_Skip_HeadMask, TRI3D_Skip_HeadMask_AddNeck, TRI3D_Image_extend from nsfw import TRI3DNSFWFilter # A dictionary that contains all nodes you want to export with their names @@ -3759,11 +3759,12 @@ NODE_CLASS_MAPPINGS = { "tri3d_SmartBox": TRI3D_SmartBox, "tri3d_Skip_HeadMask": TRI3D_Skip_HeadMask, "tri3d_Skip_HeadMask_AddNeck": TRI3D_Skip_HeadMask_AddNeck, + "tri3d_Image_extend": TRI3D_Image_extend, "tri3d_NSFWFilter": TRI3DNSFWFilter, } -VERSION = "4.8.5" +VERSION = "4.8.6" # 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, @@ -3828,4 +3829,5 @@ NODE_DISPLAY_NAME_MAPPINGS = { "tri3d_Skip_HeadMask": "Skip Head Mask" + " v" + VERSION, "tri3d_Skip_HeadMask_AddNeck": "Skip Head Mask and add neck" + " v" + VERSION, "tri3d_NSFWFilter": "TRI3D NSFW Filter" + " v" + VERSION, + "tri3d_Image_extend": "Image extend" + " v" + VERSION, } diff --git a/smart_box.py b/smart_box.py index 08baeaa..e800394 100644 --- a/smart_box.py +++ b/smart_box.py @@ -303,3 +303,87 @@ class TRI3D_Skip_HeadMask_AddNeck: 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", ), + }, + } + + FUNCTION = "run" + RETURN_TYPES = ("IMAGE", "IMAGE", ) + RETURN_NAMES = ("image", "mask_image", ) + CATEGORY = "TRI3D" + + def run(self, face_mask, image): + 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 * 1.5) + 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] + if y_below_face < target_below_face: + y_extend = target_below_face - y_below_face + cv_image = cv2.copyMakeBorder(cv_image, 0, y_extend, 0, 0, cv2.BORDER_CONSTANT, value=[0, 0, 0]) + + # Create extension mask + extension_mask = np.zeros_like(cv_image) + extension_mask[original_height:, :] = 255 # Make extended portion white + + 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) From b430d01c5b85045545e1171e83ced343afe2ca83 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Mon, 27 Jan 2025 11:48:53 +0000 Subject: [PATCH 2/4] taking ratio as an input --- smart_box.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/smart_box.py b/smart_box.py index e800394..5dbd06e 100644 --- a/smart_box.py +++ b/smart_box.py @@ -326,6 +326,7 @@ class TRI3D_Image_extend: "required": { "face_mask": ("IMAGE", ), "image": ("IMAGE", ), + "ratio": ("FLOAT", {"default": 1.5, "min": 1.2, "max": 2, "step": 0.01}), }, } @@ -359,7 +360,7 @@ class TRI3D_Image_extend: 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 * 1.5) + target_below_face = int(y_face * ratio) print("y_face", y_face) print("lowest_y", lowest_y) print("highest_y", highest_y) From 15ae5d9fef72704d22da3ac15f020575c0ff81b6 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Mon, 27 Jan 2025 12:02:54 +0000 Subject: [PATCH 3/4] rectified bug of not importing ratio and commented print statements --- smart_box.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/smart_box.py b/smart_box.py index 5dbd06e..9b0e0b4 100644 --- a/smart_box.py +++ b/smart_box.py @@ -335,7 +335,7 @@ class TRI3D_Image_extend: RETURN_NAMES = ("image", "mask_image", ) CATEGORY = "TRI3D" - def run(self, face_mask, image): + def run(self, face_mask, image, ratio): cv_face_mask = self.from_torch_image(face_mask) cv_image = self.from_torch_image(image) @@ -361,11 +361,11 @@ class TRI3D_Image_extend: # 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) + # 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] if y_below_face < target_below_face: From bd3dbad41c40690705fb0358fa7ffa21ad5a5d06 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Mon, 27 Jan 2025 12:07:43 +0000 Subject: [PATCH 4/4] modified code to extend image to preserve original aspect ratio --- smart_box.py | 24 ++++++++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) diff --git a/smart_box.py b/smart_box.py index 9b0e0b4..cb7f732 100644 --- a/smart_box.py +++ b/smart_box.py @@ -368,13 +368,33 @@ class TRI3D_Image_extend: # 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 - cv_image = cv2.copyMakeBorder(cv_image, 0, y_extend, 0, 0, cv2.BORDER_CONSTANT, value=[0, 0, 0]) + + # 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) - extension_mask[original_height:, :] = 255 # Make extended portion white + # 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)