From 00da1910634fbf314d407608efb281ae6f7f1ba2 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Thu, 18 Sep 2025 21:55:59 +0300
Subject: [PATCH] Add BlockifyMask

---
 __init__.py         |  1 +
 nodes/mask_nodes.py | 67 +++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 66 insertions(+), 2 deletions(-)

diff --git a/__init__.py b/__init__.py
index 7ba74db..817ddf6 100644
--- a/__init__.py
+++ b/__init__.py
@@ -25,6 +25,7 @@ NODE_CONFIG = {
     "DrawMaskOnImage": {"class": DrawMaskOnImage, "name": "Draw Mask On Image"},
     "DownloadAndLoadCLIPSeg": {"class": DownloadAndLoadCLIPSeg, "name": "(Down)load CLIPSeg"},
     "BatchCLIPSeg": {"class": BatchCLIPSeg, "name": "Batch CLIPSeg"},
+    "BlockifyMask": {"class": BlockifyMask, "name": "Create Block Mask"},
     "ColorToMask": {"class": ColorToMask, "name": "Color To Mask"},
     "CreateGradientMask": {"class": CreateGradientMask, "name": "Create Gradient Mask"},
     "CreateTextMask": {"class": CreateTextMask, "name": "Create Text Mask"},
diff --git a/nodes/mask_nodes.py b/nodes/mask_nodes.py
index 9b5172d..69bbfa4 100644
--- a/nodes/mask_nodes.py
+++ b/nodes/mask_nodes.py
@@ -1516,7 +1516,7 @@ class DrawMaskOnImage:
     RETURN_TYPES = ("IMAGE", )
     RETURN_NAMES = ("images",)
     FUNCTION = "apply"
-    CATEGORY = "KJNodes/image"
+    CATEGORY = "KJNodes/masking"
     DESCRIPTION = "Applies the provided masks to the input images."
 
     def apply(self, image, mask, color):
@@ -1559,4 +1559,67 @@ class DrawMaskOnImage:
 
         out_rgb = torch.stack(output_images, dim=0)
         
-        return (out_rgb, )
\ No newline at end of file
+        return (out_rgb, )
+
+
+class BlockifyMask:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+                    "masks": ("MASK",),
+                    "block_size": ("INT", {"default": 32, "min": 8, "max": 512, "step": 1, "tooltip": "Size of blocks in pixels (smaller = smaller blocks)"}),
+                }
+                }
+
+    RETURN_TYPES = ("MASK", )
+    RETURN_NAMES = ("mask",)
+    FUNCTION = "process"
+    CATEGORY = "KJNodes/masking"
+    DESCRIPTION = "Creates a block mask by dividing the bounding box of each mask into blocks of the specified size and filling in blocks that contain any part of the original mask."
+
+    def process(self, masks, block_size):
+        batch_size = masks.shape[0]
+        result_masks = []
+        
+        for i in range(batch_size):
+            mask = masks[i]
+            
+            # Find bounding box using tensor operations
+            nonzero_coords = torch.nonzero(mask, as_tuple=True)
+            if len(nonzero_coords[0]) == 0:  # Empty mask
+                result_masks.append(mask)
+                continue
+                
+            y_coords, x_coords = nonzero_coords
+            y_min, y_max = y_coords.min(), y_coords.max()
+            x_min, x_max = x_coords.min(), x_coords.max()
+            
+            bbox_width = x_max - x_min + 1
+            bbox_height = y_max - y_min + 1
+            
+            # Calculate number of blocks that fit
+            w_divisions = max(1, bbox_width // block_size)
+            h_divisions = max(1, bbox_height // block_size)
+            
+            # Calculate actual block sizes (might be slightly larger than block_size)
+            w_slice = bbox_width // w_divisions
+            h_slice = bbox_height // h_divisions
+            
+            # Create output mask (copy of input)
+            output_mask = mask.clone()
+            
+            # Process grid cells
+            for w_start in range(x_min, x_max + 1, w_slice):
+                w_end = min(w_start + w_slice, x_max + 1)
+                for h_start in range(y_min, y_max + 1, h_slice):
+                    h_end = min(h_start + h_slice, y_max + 1)
+                    
+                    # Check if this cell contains any mask content
+                    cell_region = mask[h_start:h_end, w_start:w_end]
+                    if cell_region.sum() > 0:
+                        # Fill the entire cell
+                        output_mask[h_start:h_end, w_start:w_end] = 1.0
+            
+            result_masks.append(output_mask)
+        
+        return torch.stack(result_masks, dim=0),
\ No newline at end of file