From 4ef3d762b7ef5bf850fbd6caddf9ea1f29d26bbb Mon Sep 17 00:00:00 2001
From: Anson Kao <ansonkao@ansons-mbp.home>
Date: Tue, 28 Nov 2023 18:56:43 -0500
Subject: [PATCH] Transform works for both image and mask now!

---
 .DS_Store                 | Bin 6148 -> 0 bytes
 .gitignore                |   1 +
 comfyui_geometry_nodes.py |  80 +++++++++++++++++++++++++++-----------
 3 files changed, 59 insertions(+), 22 deletions(-)
 delete mode 100644 .DS_Store
 create mode 100644 .gitignore

diff --git a/.DS_Store b/.DS_Store
deleted file mode 100644
index 4460dffdea9c6698b010f5b2d21adaff32b6fe36..0000000000000000000000000000000000000000
GIT binary patch
literal 0
HcmV?d00001

literal 6148
zcmeHKOHKko5Pdz82w@@abYU`y`$@dPfUqD@7jgjPCm{lZ0!DUo1ux(&JdKCJSKS>V
z4Z1YORFUr2U0ttw>P<3D12DB^{{pB0D6$DgC3bsE)}`%O<UO)zBu8_c^}6G%%ef>}
z`}+#WyDK2U5F;#t_gCfG+uE8fq>MIyBTOaMjuXtV3Rv%qO}?;B(L;s_?y<s>%r^HL
zSkHi2p6G2eV}?Gxb3VyVMspryJZ2V&o{}$5A{jC~ql2;KPchtnCUX{4a)>K5aft?Q
z&_W}Kehv03;g*qA)NxHrSW7=(tOG`cbq(-n>y^>OKdscaz)PrfGzs`eRH8%Hj1>)7
z-Gs9_9`;ukeFSVy)l>mhKo$5=0iM}n#j-=~RRL8%71$^s--nn@F!h)_w4V-k_6k5O
z&~1%l``4csi-VYY%pKB0Gmc7hRO7E0#@`%_qjNp=@wr1shp{^|kKft&8;bE44-?yP
znAD;6s(>o6t-yhsR^<F|e1HGnCh1HSPzC;#0;W)J)oU!tpRGH^$ypn--LZ*DT<&l~
k;lLinjFqGKgv}cLu?&c*$J`-3H2ouBWzbF)_*Dfy0b}2WMgRZ+

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..e43b0f9
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1 @@
+.DS_Store
diff --git a/comfyui_geometry_nodes.py b/comfyui_geometry_nodes.py
index 8a445b8..6310c8a 100644
--- a/comfyui_geometry_nodes.py
+++ b/comfyui_geometry_nodes.py
@@ -14,8 +14,8 @@ def register_node(identifier: str, display_name: str):
 
     return decorator
 
-@register_node("TransformImageToMatchMask", "Transform Image to Match Mask")
-class TransformImageToMatchMask:
+@register_node("TransformTemplateOntoFaceMask", "Transform Template onto Face Mask")
+class TransformTemplateOntoFaceMask:
     """
     Takes a mask as input, and calculates the centroid.
     Useful to find the center of a shape in the mask.
@@ -28,14 +28,15 @@ class TransformImageToMatchMask:
     def INPUT_TYPES(s):
         return {
             "required": {
-                "mask": ("MASK",),
-                "image": ("IMAGE",),
+                "face_mask": ("MASK",),
+                "template_image": ("IMAGE",),
+                "template_mask": ("MASK",),
             },
         }
 
-    RETURN_TYPES = ("IMAGE",)
+    RETURN_TYPES = ("IMAGE", "MASK")
     CATEGORY = "Geometry"
-    FUNCTION = "transform_image"
+    FUNCTION = "transform_template"
 
     def convert_image_from_tensor_to_numpy(self, image_tensor):
         # Convert from PyTorch tensor to NumPy array
@@ -49,6 +50,15 @@ class TransformImageToMatchMask:
 
         return image_np
 
+    def convert_mask_from_tensor_to_numpy(self, mask_tensor):
+        # Convert from PyTorch tensor to NumPy array
+        mask_np = mask_tensor.squeeze().numpy()
+
+        # Convert from normalized floats back to uint8
+        mask_np = (mask_np * 255).astype(np.uint8)
+
+        return mask_np
+
     def convert_image_from_numpy_to_tensor(self, image_np):
         # Convert from BGR to RGB
         image_rgb = cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB)
@@ -64,10 +74,22 @@ class TransformImageToMatchMask:
 
         return image_tensor
 
-    def calculate_transformation(self, mask, image):
+    def convert_mask_from_numpy_to_tensor(self, mask_np):
+        # Normalize the pixel values to [0.0, 1.0]
+        mask_normalized = mask_np.astype(np.float32) / 255.0
+
+        # Convert to a PyTorch tensor
+        mask_tensor = torch.from_numpy(mask_normalized)
+
+        # Add a batch dimension with [None,] or .unsqueeze(0)
+        mask_tensor = mask_tensor[None,]
+
+        return mask_tensor
+
+    def calculate_transformation(self, face_mask, template_mask):
         # Convert the tensor to a numpy array and remove the batch and color dimensions
         # The resulting array will have shape [height, width]
-        np_image = mask.squeeze().numpy()
+        np_image = self.convert_mask_from_tensor_to_numpy(face_mask)
 
         # Assuming the object is white and the background is black
         # Create a binary image (you might need to adjust the thresholding logic based on your image)
@@ -108,8 +130,8 @@ class TransformImageToMatchMask:
         ellipse_width = int(minor_axis_length / 2)
 
         # Starting and ending dimensions
-        height_out, width_out = mask.shape[1:3]
-        height_in, width_in = image.shape[1:3]
+        height_out, width_out = face_mask.shape[1:3]
+        height_in, width_in = template_mask.shape[1:3]
 
         # Calculate the scale relative to the original template size
         scale_x = ellipse_width / 72
@@ -172,22 +194,31 @@ class TransformImageToMatchMask:
             matrix,
         )
 
-    def transform_image(self, mask, image):
-        centroid_x, centroid_y, ellipse_length, ellipse_width, rotation, matrix = self.calculate_transformation(mask, image)
+    def transform_template(self, face_mask, template_image, template_mask):
+        centroid_x, centroid_y, ellipse_length, ellipse_width, rotation, matrix = self.calculate_transformation(face_mask, template_mask)
 
         # Create a blank canvas the same size as the mask
-        height_in, width_in = image.shape[1:3]
-        height_out, width_out = mask.shape[1:3]
-        canvas = np.zeros((height_out, width_out, 3), np.uint8)
+        height_in, width_in = template_image.shape[1:3]
+        height_out, width_out = face_mask.shape[1:3]
+        image_canvas = np.zeros((height_out, width_out, 3), np.uint8)
+        mask_canvas = np.zeros((height_out, width_out), np.uint8)
 
         # Drop the image into the center of this canvas
         x_offset = (width_out - width_in) // 2
         y_offset = (height_out - height_in) // 2
-        image_np = self.convert_image_from_tensor_to_numpy(image)
-        canvas[y_offset:y_offset + height_in, x_offset:x_offset + width_in] = image_np
+        template_image_np = self.convert_image_from_tensor_to_numpy(template_image)
+        template_mask_np = self.convert_mask_from_tensor_to_numpy(template_mask)
+
+        # Invert the mask before transforming so unmasked space around the canvas is maintained automatically
+        template_mask_np = 255.0 - template_mask_np
+
+        # Drop the input template image and mask into the center of the output canvases
+        image_canvas[y_offset:y_offset + height_in, x_offset:x_offset + width_in] = template_image_np
+        mask_canvas[y_offset:y_offset + height_in, x_offset:x_offset + width_in] = template_mask_np
 
         # Apply the affine transformation
-        transformed_image = cv2.warpAffine(canvas, matrix[:2], (width_out, height_out))
+        transformed_template_image = cv2.warpAffine(image_canvas, matrix[:2], (width_out, height_out))
+        transformed_template_mask = cv2.warpAffine(mask_canvas, matrix[:2], (width_out, height_out))
 
         # Draw the ellipse to preview the calculation
         debug_calculation = False
@@ -195,9 +226,14 @@ class TransformImageToMatchMask:
             centroid = (centroid_x, centroid_y)
             axes = (ellipse_length, ellipse_width)
             angle = rotation * (180 / np.pi)
-            cv2.ellipse(transformed_image, centroid, axes, angle, 0, 350, (0, 255, 0), 2)        
-            cv2.circle(transformed_image, (centroid_x, centroid_y), radius=5, color=(255, 0, 0), thickness=-5)
+            cv2.ellipse(transformed_template_image, centroid, axes, angle, 0, 350, (0, 255, 0), 2)        
+            cv2.circle(transformed_template_image, (centroid_x, centroid_y), radius=5, color=(255, 0, 0), thickness=-5)
+
+        # Invert the mask back to normal
+        transformed_template_mask = 255.0 - transformed_template_mask
 
         # Convert the image back to a tensor
-        final_image = self.convert_image_from_numpy_to_tensor(transformed_image)
-        return (final_image,)
+        final_image = self.convert_image_from_numpy_to_tensor(transformed_template_image)
+        final_mask = self.convert_mask_from_numpy_to_tensor(transformed_template_mask)
+
+        return (final_image, final_mask)