From d6bd841d7fa11e2a05e8cfd0a816db3a25221b13 Mon Sep 17 00:00:00 2001 From: var1ableX Date: Fri, 13 Sep 2024 21:23:12 -0400 Subject: [PATCH] v2.0 --- __init__.py | 11 +++- maskToCrop.py | 53 --------------- unmasked.py | 120 ++++++++++++++++++++++++++++++++++ web/js/getRandomDimensions.js | 68 +++++++++++++++++++ 4 files changed, 197 insertions(+), 55 deletions(-) delete mode 100644 maskToCrop.py create mode 100644 unmasked.py create mode 100644 web/js/getRandomDimensions.js diff --git a/__init__.py b/__init__.py index 56f4019..958e39c 100644 --- a/__init__.py +++ b/__init__.py @@ -1,11 +1,18 @@ -from .maskToCrop import GetMaskDimensions +from .unmasked import GetMaskDimensions, isMaskEmpty, GetRandomDimensions # Node class mappings NODE_CLASS_MAPPINGS = { "GetMaskDimensions": GetMaskDimensions, + "isMaskEmpty": isMaskEmpty, + "GetRandomDimensions": GetRandomDimensions } # Node display name mappings NODE_DISPLAY_NAME_MAPPINGS = { "GetMaskDimensions": "Get Mask Dimensions", -} \ No newline at end of file + "isMaskEmpty": "Is Mask Empty", + "GetRandomDimensions": "Get Random Dimensions" +} + +WEB_DIRECTORY = "./web" +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] \ No newline at end of file diff --git a/maskToCrop.py b/maskToCrop.py deleted file mode 100644 index 665495b..0000000 --- a/maskToCrop.py +++ /dev/null @@ -1,53 +0,0 @@ -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 - diff --git a/unmasked.py b/unmasked.py new file mode 100644 index 0000000..23a9b42 --- /dev/null +++ b/unmasked.py @@ -0,0 +1,120 @@ +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 + +class isMaskEmpty: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "mask": ("MASK",), + } + } + + RETURN_TYPES = ("BOOLEAN",) + RETURN_NAMES = ("boolean",) + FUNCTION = "execute" + CATEGORY = "image" + + def execute(self, mask): + def is_mask_empty(mask): + return not torch.any(mask == 1).item() + + isEmpty = is_mask_empty(mask) + return (isEmpty,) # Return the boolean value as a tuple + +import random + +class GetRandomDimensions: + def __init__(self): + self.is_changed_enabled = True + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "min_width": ("INT", {"default": 768, "step": 16, "display": "number"}), + "min_height": ("INT", {"default": 768, "step": 16, "display": "number"}), + "max_width": ("INT", {"default": 1280, "step": 16, "display": "number"}), + "max_height": ("INT", {"default": 1280, "step": 16, "display": "number"}), + "randomize": ("BOOLEAN", {"default": False}), + } + } + + RETURN_TYPES = ("INT", "INT") + RETURN_NAMES = ("width", "height",) + FUNCTION = "execute" + CATEGORY = "image" + OUTPUT_NODE = True + + def execute(self, min_width, min_height, max_width, max_height, randomize ): + def get_random_dimensions(min_val, max_val): + min_div_16 = min_val // 16 + max_div_16 = max_val // 16 + if min_div_16 > max_div_16: + raise ValueError("No numbers in the specified range are divisible by 16.") + return random.randint(min_div_16, max_div_16) * 16 + + if randomize: + width = get_random_dimensions(min_width, max_width) + height = get_random_dimensions(min_height, max_height) + else: + width, height = max_width, max_height + + # Setting the resolution string + text = f"{width}x{height}" + #return width, height + return {"ui": {"text": text}, + "result": (width, height, text)} diff --git a/web/js/getRandomDimensions.js b/web/js/getRandomDimensions.js new file mode 100644 index 0000000..d51c885 --- /dev/null +++ b/web/js/getRandomDimensions.js @@ -0,0 +1,68 @@ +import { app } from "../../../scripts/app.js"; +import { ComfyWidgets } from "../../../scripts/widgets.js"; + +// Add new colors to LGraphCanvas.node_colors +function addNewColors() { + // Completely replacing the node_colors object + LGraphCanvas.node_colors = { + red: { color: "#C62828", bgcolor: "#E53935", groupcolor: "#B71C1C" }, // Red: 800, 600, 900 + pink: { color: "#AD1457", bgcolor: "#D81B60", groupcolor: "#880E4F" }, // Pink: 800, 600, 900 + purple: { color: "#6A1B9A", bgcolor: "#8E24AA", groupcolor: "#4A148C" }, // Purple: 800, 600, 900 + indigo: { color: "#283593", bgcolor: "#3949AB", groupcolor: "#1A237E" }, // Indigo: 800, 600, 900 + blue: { color: "#1565C0", bgcolor: "#1E88E5", groupcolor: "#0D47A1" }, // Blue: 800, 600, 900 + cyan: { color: "#00838F", bgcolor: "#00ACC1", groupcolor: "#006064" }, // Cyan: 800, 600, 900 + teal: { color: "#00695C", bgcolor: "#00897B", groupcolor: "#004D40" }, // Teal: 800, 600, 900 + green: { color: "#2E7D32", bgcolor: "#43A047", groupcolor: "#1B5E20" }, // Green: 800, 600, 900 + "light green": { color: "#558B2F", bgcolor: "#7CB342", groupcolor: "#33691E" }, // Light Green: 800, 600, 900 + "deep orange": { color: "#D84315", bgcolor: "#F4511E", groupcolor: "#BF360C" }, // Deep Orange: 800, 600, 900 + grey: { color: "#424242", bgcolor: "#757575", groupcolor: "#212121" }, // Grey: 800, 600, 900 + + // Adding the 6 extra colors + "blue grey": { color: "#455A64", bgcolor: "#607D8B", groupcolor: "#263238" }, // Blue Grey + lime: { color: "#9E9D24", bgcolor: "#C0CA33", groupcolor: "#827717" }, // Lime + "light blue": { color: "#0277BD", bgcolor: "#039BE5", groupcolor: "#01579B" }, // Light Blue + "steel blue": { color: "#4682B4", bgcolor: "#5A9BD3", groupcolor: "#2B5A8A" }, // Steel Blue + "turquoise": { color: "#00796B", bgcolor: "#009688", groupcolor: "#004D40" }, // Turquoise + navy: { color: "#003366", bgcolor: "#004080", groupcolor: "#001f3f" }, // Navy Blue + }; +} + + +// Function to enhance nodes with text display +function addTextDisplay(nodeType) { + const onNodeCreated = nodeType.prototype.onNodeCreated; + nodeType.prototype.onNodeCreated = function () { + const r = onNodeCreated?.apply(this, arguments); + const w = ComfyWidgets["STRING"](this, "display", ["STRING", { multiline: true, placeholder: " " }], app).widget; + w.inputEl.readOnly = true; + w.inputEl.style.opacity = 0.7; + w.inputEl.style.cursor = "auto"; + return r; + }; + + const onExecuted = nodeType.prototype.onExecuted; + nodeType.prototype.onExecuted = function (message) { + onExecuted?.apply(this, arguments); + + for (const widget of this.widgets) { + if (widget.type === "customtext" && widget.name === "display" && widget.inputEl.readOnly === true) { + widget.value = message.text.join(''); + } + } + + this.onResize?.(this.size); + }; +} + +app.registerExtension({ + name: "rka.GetRandomDimensions", + setup() { + // Call the function to add new colors + addNewColors(); + }, + async beforeRegisterNodeDef(nodeType, nodeData, app) { + if (["GetRandomDimensions"].includes(nodeData.name)) { + addTextDisplay(nodeType); + } + }, +});