From 7abc8bb50886d84d9ca25b62fbde54dbda26882d Mon Sep 17 00:00:00 2001
From: SuperBeastsAI <bradcmasters+superbeasts.ai@gmail.com>
Date: Sun, 7 Apr 2024 07:22:54 +1000
Subject: [PATCH] Added helper node for making mask batch with up to 12 masks.
 Includes resizing at center to ensure masks as the same size in the output

---
 superbeastsai_nodes.py | 81 ++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 79 insertions(+), 2 deletions(-)

diff --git a/superbeastsai_nodes.py b/superbeastsai_nodes.py
index 3c33af2..4c13369 100644
--- a/superbeastsai_nodes.py
+++ b/superbeastsai_nodes.py
@@ -165,10 +165,87 @@ class HDREffects:
          
         return pil2tensor(color_adjusted)
 
+class MakeResizedMaskBatch:
+    """
+    Creates a batch of masks from multiple individual masks or batches.
+    """
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "width": ("INT", {"default": 512, "min": 1, "step": 1}),
+                "height": ("INT", {"default": 768, "min": 1, "step": 1}),
+            },
+            "optional": {
+                "mask1": ("MASK",),
+                "mask2": ("MASK",),
+                "mask3": ("MASK",),
+                "mask4": ("MASK",),
+                "mask5": ("MASK",),
+                "mask6": ("MASK",),
+                "mask7": ("MASK",),
+                "mask8": ("MASK",),
+                "mask9": ("MASK",),
+                "mask10": ("MASK",),
+                "mask11": ("MASK",),
+                "mask12": ("MASK",),
+            },
+        }
+
+    RETURN_TYPES = ("MASK",)
+    FUNCTION = "append"
+    CATEGORY = "SuperBeastsAI/Masks"
+
+    def append(self, width, height, mask1=None, mask2=None, mask3=None, mask4=None, mask5=None, mask6=None,
+               mask7=None, mask8=None, mask9=None, mask10=None, mask11=None, mask12=None):
+        masks = [mask1, mask2, mask3, mask4, mask5, mask6, mask7, mask8, mask9, mask10, mask11, mask12]
+        valid_masks = [mask for mask in masks if mask is not None]
+
+        if not valid_masks:
+            raise ValueError("At least one input mask must be provided.")
+
+        cropped_masks = []
+        for mask in valid_masks:
+            if mask.ndim == 2:
+                mask = mask.unsqueeze(0).unsqueeze(0)  # Add batch and channel dimensions
+            elif mask.ndim == 3:
+                mask = mask.unsqueeze(0)  # Add batch dimension
+            elif mask.ndim != 4 or (mask.ndim == 4 and mask.shape[1] != 1):
+                raise ValueError(f"Invalid mask shape: {mask.shape}. Expected (N, 1, H, W) or (1, H, W) or (H, W).")
+
+            # Scale the mask to match the desired width while maintaining the aspect ratio
+            scale_factor = width / mask.shape[-1]
+            scaled_height = int(mask.shape[-2] * scale_factor)
+            scaled_mask = F.interpolate(mask, size=(scaled_height, width), mode='bilinear', align_corners=False)
+
+            # Perform center cropping
+            if scaled_height < height:
+                # Pad the top and bottom of the mask
+                pad_top = (height - scaled_height) // 2
+                pad_bottom = height - scaled_height - pad_top
+                cropped_mask = F.pad(scaled_mask, (0, 0, pad_top, pad_bottom), mode='constant', value=0)
+            else:
+                # Crop the center of the mask
+                crop_top = (scaled_height - height) // 2
+                crop_bottom = crop_top + height
+                cropped_mask = scaled_mask[:, :, crop_top:crop_bottom, :]
+
+            cropped_masks.append(cropped_mask)
+
+        # Concatenate the cropped masks along the batch dimension
+        result = torch.cat(cropped_masks, dim=0)
+
+        return (result,)
+
 NODE_CLASS_MAPPINGS = {
-    'HDR Effects (SuperBeasts.AI)': HDREffects
+    'HDR Effects (SuperBeasts.AI)': HDREffects,
+    'Make Resized Mask Batch (SuperBeasts.AI)': MakeResizedMaskBatch
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
-    'HDREffects': 'HDR Effects (SuperBeasts.AI)'
+    'HDREffects': 'HDR Effects (SuperBeasts.AI)',
+    'MakeResizedMaskBatch': 'Make Resized Mask Batch (SuperBeasts.AI)'
 }