Merge pull request #53 from TRI3D-LC/image_entend

Image entend

__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,
 }

smart_box.py

@@ -303,3 +303,108 @@ 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", ),
+                "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)