commit 56e932e99709f6babb7b580b54169402fdc749db Author: var1ableX Date: Fri Jul 26 20:59:47 2024 -0400 initial commit diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..56f4019 --- /dev/null +++ b/__init__.py @@ -0,0 +1,11 @@ +from .maskToCrop import GetMaskDimensions + +# Node class mappings +NODE_CLASS_MAPPINGS = { + "GetMaskDimensions": GetMaskDimensions, +} + +# Node display name mappings +NODE_DISPLAY_NAME_MAPPINGS = { + "GetMaskDimensions": "Get Mask Dimensions", +} \ No newline at end of file diff --git a/maskToCrop.py b/maskToCrop.py new file mode 100644 index 0000000..665495b --- /dev/null +++ b/maskToCrop.py @@ -0,0 +1,53 @@ +import torch + +class GetMaskDimensions: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "mask": ("MASK",), + } + } + + RETURN_TYPES = ("INT", "INT", "INT", "INT",) + RETURN_NAMES = ("width", "height", "x", "y",) + FUNCTION = "execute" + CATEGORY = "image" + + def execute(self, mask): + def get_mask_dimensions(mask): + # Assuming mask_tensor is of shape (1, H, W) + mask_array = mask[0] + + # Find all coordinates where the mask is non-zero + coordinates = torch.nonzero(mask_array, as_tuple=False) + + if coordinates.numel() == 0: + # If there are no non-zero coordinates, return a zero-sized rectangle + return 0, 0, 0, 0 + + # Get the minimum and maximum x and y coordinates + min_coords = torch.min(coordinates, dim=0)[0] + max_coords = torch.max(coordinates, dim=0)[0] + + # Extract bounding box coordinates + top_left_x = min_coords[1].item() + top_left_y = min_coords[0].item() + bottom_right_x = max_coords[1].item() + bottom_right_y = max_coords[0].item() + + # Calculate width and height + width = bottom_right_x - top_left_x + 1 + height = bottom_right_y - top_left_y + 1 + + return width, height, top_left_x, top_left_y + + # Process the mask tensor to find cropping dimensions + width, height, top_left_x, top_left_y = get_mask_dimensions(mask) + + # Return the results as node outputs + return width, height, top_left_x, top_left_y +