percentage and cutbymask changes
This commit is contained in:
Ubuntu
@@ -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,
|
||||
}
|
||||
|
||||
+48
-28
@@ -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)
|
||||
|
||||
@@ -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"
|
||||
# }
|
||||
Reference in New Issue
Block a user