From 75ba9bfb252fd2969f5ac2016d18b321ffbbc3eb Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Fri, 23 May 2025 05:14:55 +0000 Subject: [PATCH] percentage and cutbymask changes --- __init__.py | 3 ++ cut_by_mask_aspect_ratio.py | 76 +++++++++++++++++++++++-------------- mask_area_percentage.py | 68 +++++++++++++++++++++++++++++++++ 3 files changed, 119 insertions(+), 28 deletions(-) create mode 100644 mask_area_percentage.py diff --git a/__init__.py b/__init__.py index 627d71c..bc63f1f 100644 --- a/__init__.py +++ b/__init__.py @@ -21,6 +21,7 @@ from utility_nodes import TRI3D_extract_facer_mask from .AEMatter import (load_AEMatter_Model, run_AEMatter_inference) from .light_layer import main_light_layer from .remove_small_mask_islands import TRI3D_RemoveSmallMaskIslands +from .mask_area_percentage import TRI3D_MaskAreaPercentage from .image_stack import ( @@ -3771,6 +3772,7 @@ NODE_CLASS_MAPPINGS = { "tri3d_Remove_Small_Mask_Islands": TRI3D_RemoveSmallMaskIslands, "tri3d_CutByMaskAspectRatio": TRI3D_CutByMaskAspectRatio, "tri3d_StringContains": TRI3D_StringContains, + "tri3d_MaskAreaPercentage": TRI3D_MaskAreaPercentage, } @@ -3846,4 +3848,5 @@ NODE_DISPLAY_NAME_MAPPINGS = { "tri3d_Remove_Small_Mask_Islands": "Remove Small Mask Islands" + " v" + VERSION, "tri3d_CutByMaskAspectRatio": "Cut by mask aspect ratio" + " v" + VERSION, "tri3d_StringContains": "String contains" + " v" + VERSION, + "tri3d_MaskAreaPercentage": "Mask Area Percentage" + " v" + VERSION, } diff --git a/cut_by_mask_aspect_ratio.py b/cut_by_mask_aspect_ratio.py index b90d450..3502761 100644 --- a/cut_by_mask_aspect_ratio.py +++ b/cut_by_mask_aspect_ratio.py @@ -97,38 +97,55 @@ class TRI3D_CutByMaskAspectRatio: # Calculate current aspect ratio current_aspect_ratio = width / height - # Crop the image to the original bounding box - cropped_image = cv_image[y_min:y_max, x_min:x_max] - # Adjust width to match the target aspect ratio while keeping height constant if current_aspect_ratio < target_aspect_ratio: - # Current width is too narrow, add white padding - # Calculate the intermediate dimensions for resizing - intermediate_height = height - intermediate_width = width # Keep original width + # Current width is too narrow - need to extend it + # Calculate the required width for the target aspect ratio + required_width = int(height * target_aspect_ratio) + width_difference = required_width - width - # Calculate the final canvas size with padding - canvas_height = intermediate_height - canvas_width = int(canvas_height * target_aspect_ratio) + # Calculate how much to extend on each side + left_extend = width_difference // 2 + right_extend = width_difference - left_extend - # Calculate padding on each side - padding_x = (canvas_width - intermediate_width) // 2 + # Calculate new potential boundaries + new_x_min = x_min - left_extend + new_x_max = x_max + right_extend - # Create canvas with padding color - num_channels = cropped_image.shape[2] if len(cropped_image.shape) == 3 else 1 - if num_channels == 1: - canvas = np.full((canvas_height, canvas_width), padding_color, dtype=np.uint8) + # Check if the new boundaries are within the original image + left_padding_needed = abs(min(0, new_x_min)) + right_padding_needed = max(0, new_x_max - cv_image.shape[1]) + + # Adjust boundaries to be within the original image + new_x_min = max(0, new_x_min) + new_x_max = min(cv_image.shape[1], new_x_max) + + # Get the portion of the original image within valid boundaries + extended_image = cv_image[y_min:y_max, new_x_min:new_x_max] + + # If we need padding (i.e., extension goes beyond image boundaries) + if left_padding_needed > 0 or right_padding_needed > 0: + # Create canvas with padding color + num_channels = extended_image.shape[2] if len(extended_image.shape) == 3 else 1 + if num_channels == 1: + canvas = np.full((height, required_width), padding_color, dtype=np.uint8) + else: + canvas = np.full((height, required_width, num_channels), padding_color, dtype=np.uint8) + + # Calculate the position to place the extended image + place_x = left_padding_needed + + # Place the extended image on the canvas + if num_channels == 1: + canvas[:, place_x:place_x+extended_image.shape[1]] = extended_image + else: + canvas[:, place_x:place_x+extended_image.shape[1], :] = extended_image + + # Use the canvas as our cropped image + cropped_image = canvas else: - canvas = np.full((canvas_height, canvas_width, num_channels), padding_color, dtype=np.uint8) - - # Place the cropped image on the canvas - if num_channels == 1: - canvas[:, padding_x:padding_x+intermediate_width] = cropped_image - else: - canvas[:, padding_x:padding_x+intermediate_width, :] = cropped_image - - # Set the final image to be the padded canvas - cropped_image = canvas + # No padding needed, use the extended image + cropped_image = extended_image elif current_aspect_ratio > target_aspect_ratio: # Current width is too wide, crop it @@ -139,8 +156,11 @@ class TRI3D_CutByMaskAspectRatio: left_crop = width_difference // 2 right_crop = width_difference - left_crop - # Apply the crop to the cropped image - cropped_image = cropped_image[:, left_crop:width-right_crop] + # Apply the crop + cropped_image = cv_image[y_min:y_max, x_min+left_crop:x_max-right_crop] + else: + # Aspect ratio is already correct + cropped_image = cv_image[y_min:y_max, x_min:x_max] # Resize the cropped/padded image to the target dimensions using Lanczos interpolation resized_image = cv2.resize(cropped_image, (target_width, target_height), interpolation=cv2.INTER_LANCZOS4) diff --git a/mask_area_percentage.py b/mask_area_percentage.py new file mode 100644 index 0000000..3dddda3 --- /dev/null +++ b/mask_area_percentage.py @@ -0,0 +1,68 @@ +import os +import cv2 +import numpy as np +import torch + +class TRI3D_MaskAreaPercentage: + """ + ComfyUI node that calculates the percentage of white pixels in an image + relative to the total image area. + """ + + def from_torch_image(self, image): + """Convert a torch tensor image to numpy array for OpenCV processing""" + image = image.cpu().numpy() * 255.0 + image = np.clip(image, 0, 255).astype(np.uint8) + return image + + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "threshold": ("INT", {"default": 127, "min": 0, "max": 255, "step": 1}), + }, + } + + FUNCTION = "run" + RETURN_TYPES = ("FLOAT", "INT", "INT",) + RETURN_NAMES = ("percentage", "white_pixels", "total_pixels",) + CATEGORY = "TRI3D" + + def run(self, image, threshold=127): + # Convert Torch image to OpenCV format + cv_image = self.from_torch_image(image) + + # Remove batch dimension if present + if len(cv_image.shape) == 4: + cv_image = cv_image[0] + + # Convert to grayscale if it's a color image + if len(cv_image.shape) == 3 and cv_image.shape[2] > 1: + gray_image = cv2.cvtColor(cv_image, cv2.COLOR_RGB2GRAY) + else: + gray_image = cv_image[:, :, 0] + + # Calculate total number of pixels + total_pixels = gray_image.shape[0] * gray_image.shape[1] + + # Count white pixels (pixels with values above threshold) + _, binary_image = cv2.threshold(gray_image, threshold, 255, cv2.THRESH_BINARY) + white_pixels = cv2.countNonZero(binary_image) + + # Calculate percentage of white pixels + percentage = (white_pixels / total_pixels) * 100.0 + + return (percentage, white_pixels, total_pixels,) + +# # Node registration for ComfyUI +# NODE_CLASS_MAPPINGS = { +# "TRI3D_MaskAreaPercentage": TRI3D_MaskAreaPercentage +# } + +# NODE_DISPLAY_NAME_MAPPINGS = { +# "TRI3D_MaskAreaPercentage": "TRI3D Mask Area Percentage" +# } \ No newline at end of file