Adds HSV Threshold Mask
actually update examples as well
This commit is contained in:
@@ -166,7 +166,7 @@ class Blur:
|
||||
|
||||
return (blurred,)
|
||||
|
||||
class CannyEdgeDetection:
|
||||
class CannyEdgeMask:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@@ -717,6 +717,60 @@ class Glow:
|
||||
def add_glow(self, img, blurred_img, intensity):
|
||||
return img + blurred_img * intensity
|
||||
|
||||
class HSVThresholdMask:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
"low_threshold": ("FLOAT", {
|
||||
"default": 0.2,
|
||||
"min": 0,
|
||||
"max": 1,
|
||||
"step": 0.1
|
||||
}),
|
||||
"high_threshold": ("FLOAT", {
|
||||
"default": 0.7,
|
||||
"min": 0,
|
||||
"max": 1,
|
||||
"step": 0.1
|
||||
}),
|
||||
"hsv_channel": (["hue", "saturation", "value"],),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "hsv_threshold"
|
||||
|
||||
CATEGORY = "postprocessing/Masks"
|
||||
|
||||
def hsv_threshold(self, image: torch.Tensor, low_threshold: float, high_threshold: float, hsv_channel: str):
|
||||
batch_size, height, width, _ = image.shape
|
||||
result = torch.zeros(batch_size, height, width)
|
||||
|
||||
if hsv_channel == "hue":
|
||||
channel = 0
|
||||
low_threshold, high_threshold = int(low_threshold * 180), int(high_threshold * 180)
|
||||
elif hsv_channel == "saturation":
|
||||
channel = 1
|
||||
low_threshold, high_threshold = int(low_threshold * 255), int(high_threshold * 255)
|
||||
elif hsv_channel == "value":
|
||||
channel = 2
|
||||
low_threshold, high_threshold = int(low_threshold * 255), int(high_threshold * 255)
|
||||
|
||||
for b in range(batch_size):
|
||||
tensor_image = (image[b].numpy().copy() * 255).astype(np.uint8)
|
||||
hsv_image = cv2.cvtColor(tensor_image, cv2.COLOR_RGB2HSV)
|
||||
|
||||
mask = cv2.inRange(hsv_image[:, :, channel], low_threshold, high_threshold)
|
||||
tensor = torch.from_numpy(mask).float() / 255.
|
||||
result[b] = tensor
|
||||
|
||||
return (result,)
|
||||
|
||||
class KuwaharaBlur:
|
||||
def __init__(self):
|
||||
pass
|
||||
@@ -1265,7 +1319,7 @@ NODE_CLASS_MAPPINGS = {
|
||||
"ArithmeticBlend": ArithmeticBlend,
|
||||
"Blend": Blend,
|
||||
"Blur": Blur,
|
||||
"CannyEdgeDetection": CannyEdgeDetection,
|
||||
"CannyEdgeMask": CannyEdgeMask,
|
||||
"ChromaticAberration": ChromaticAberration,
|
||||
"ColorCorrect": ColorCorrect,
|
||||
"ColorTint": ColorTint,
|
||||
@@ -1273,6 +1327,7 @@ NODE_CLASS_MAPPINGS = {
|
||||
"DodgeAndBurn": DodgeAndBurn,
|
||||
"FilmGrain": FilmGrain,
|
||||
"Glow": Glow,
|
||||
"HSVThresholdMask": HSVThresholdMask,
|
||||
"KuwaharaBlur": KuwaharaBlur,
|
||||
"Parabolize": Parabolize,
|
||||
"PencilSketch": PencilSketch,
|
||||
|
||||
Reference in New Issue
Block a user