From 0955881e1663328d02fbcbdf09837003cdff7452 Mon Sep 17 00:00:00 2001 From: Pheat-AI Date: Mon, 9 Sep 2024 18:17:43 -0700 Subject: [PATCH] Update calculate_center_bounding_box.py changed this node to canny image cropper node --- calculate_center_bounding_box.py | 83 ++++++++++++-------------------- 1 file changed, 32 insertions(+), 51 deletions(-) diff --git a/calculate_center_bounding_box.py b/calculate_center_bounding_box.py index 2465488..90d0f64 100644 --- a/calculate_center_bounding_box.py +++ b/calculate_center_bounding_box.py @@ -1,72 +1,53 @@ -import comfy.utils -import numpy as np import torch -class MaskBoundingBox: +class CannyImageCropper: + def __init__(self): + pass + @classmethod - def INPUT_TYPES(s): + def INPUT_TYPES(cls): return { "required": { - "mask": ("MASK",), - "padding": ("INT", { "default": 0, "min": 0, "max": 4096, "step": 1, }), - "blur": ("INT", { "default": 0, "min": 0, "max": 256, "step": 1, }), + "canny_image": ("IMAGE",), # Canny image input + "x": ("INT",), # x-coordinate of the bounding box + "y": ("INT",), # y-coordinate of the bounding box + "width": ("INT",), # Width of the bounding box + "height": ("INT",), # Height of the bounding box + "padding": ("INT", { "default": 5, "min": 0, "max": 100, "step": 1 }), # Padding to add to the bounding box }, - "optional": { - "image_optional": ("IMAGE",), - } } - RETURN_TYPES = ("MASK", "IMAGE", "INT", "INT", "INT", "INT", "INT", "INT") - RETURN_NAMES = ("MASK", "IMAGE", "x", "y", "width", "height", "center_x", "center_y") - FUNCTION = "execute" - CATEGORY = "REMADE/mask" + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("cropped_canny_image",) + FUNCTION = "crop_canny_image" - def execute(self, mask, padding, blur, image_optional=None): - if mask.dim() == 2: - mask = mask.unsqueeze(0) + CATEGORY = "REMADE/Image" - if image_optional is None: - image_optional = mask.unsqueeze(3).repeat(1, 1, 1, 3) + def crop_canny_image(self, canny_image, x, y, width, height, padding): + # Apply padding to the bounding box to expand the cropping area + x -= padding + y -= padding + width += 2 * padding + height += 2 * padding - # Resize the image if it's not the same size as the mask - if image_optional.shape[1:] != mask.shape[1:]: - image_optional = comfy.utils.common_upscale(image_optional.permute([0, 3, 1, 2]), mask.shape[2], mask.shape[1], upscale_method='bicubic', crop='center').permute([0, 2, 3, 1]) + # Ensure that the crop area doesn't exceed the image boundaries + x = max(0, x) + y = max(0, y) + x2 = min(canny_image.shape[2], x + width) + y2 = min(canny_image.shape[1], y + height) - # Match batch size - if image_optional.shape[0] < mask.shape[0]: - image_optional = torch.cat((image_optional, image_optional[-1].unsqueeze(0).repeat(mask.shape[0] - image_optional.shape[0], 1, 1, 1)), dim=0) - elif image_optional.shape[0] > mask.shape[0]: - image_optional = image_optional[:mask.shape[0]] + # Crop the canny image based on the expanded bounding box coordinates + cropped_canny_image = canny_image[:, y:y2, x:x2, :] - # Blur the mask - if blur > 0: - if blur % 2 == 0: - blur += 1 - mask = T.functional.gaussian_blur(mask.unsqueeze(1), blur).squeeze(1) - - # Calculate bounding box coordinates - _, y, x = torch.where(mask) - x1 = max(0, x.min().item() - padding) - x2 = min(mask.shape[2], x.max().item() + 1 + padding) - y1 = max(0, y.min().item() - padding) - y2 = min(mask.shape[1], y.max().item() + 1 + padding) - - # Crop the mask and image to the bounding box - mask = mask[:, y1:y2, x1:x2] - image_optional = image_optional[:, y1:y2, x1:x2, :] - - # Calculate the center coordinates - center_x = x1 + (x2 - x1) // 2 - center_y = y1 + (y2 - y1) // 2 - - return (mask, image_optional, x1, y1, x2 - x1, y2 - y1, center_x, center_y) + # Return the cropped canny image + return (cropped_canny_image,) # A dictionary that contains all nodes you want to export with their names NODE_CLASS_MAPPINGS = { - "MaskBoundingBox": MaskBoundingBox + "CannyImageCropper": CannyImageCropper } # A dictionary that contains the friendly/human-readable titles for the nodes NODE_DISPLAY_NAME_MAPPINGS = { - "MaskBoundingBox": "Mask Bounding Box with Center" + "CannyImageCropper": "Canny Image Cropper" }