adds solarize

This commit is contained in:
EllangoK
2023-04-02 19:33:16 -04:00
parent 8da09a5812
commit d0e69db667
3 changed files with 112 additions and 0 deletions
+78
View File
@@ -46,6 +46,54 @@ class ArithmeticBlend:
def difference(self, img1, img2):
return torch.abs(img1 - img2)
class Blend:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image1": ("IMAGE",),
"image2": ("IMAGE",),
"blend_factor": ("FLOAT", {
"default": 0.5,
"min": 0.0,
"max": 1.0,
"step": 0.01
}),
"blend_mode": (["normal", "multiply", "screen", "overlay", "soft_light"],),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "blend_images"
CATEGORY = "postprocessing"
def blend_images(self, image1: torch.Tensor, image2: torch.Tensor, blend_factor: float, blend_mode: str):
blended_image = self.blend_mode(image1, image2, blend_mode)
blended_image = image1 * (1 - blend_factor) + blended_image * blend_factor
blended_image = torch.clamp(blended_image, 0, 1)
return (blended_image,)
def blend_mode(self, img1, img2, mode):
if mode == "normal":
return img2
elif mode == "multiply":
return img1 * img2
elif mode == "screen":
return 1 - (1 - img1) * (1 - img2)
elif mode == "overlay":
return torch.where(img1 <= 0.5, 2 * img1 * img2, 1 - 2 * (1 - img1) * (1 - img2))
elif mode == "soft_light":
return torch.where(img2 <= 0.5, img1 - (1 - 2 * img2) * img1 * (1 - img1), img1 + (2 * img2 - 1) * (self.g(img1) - img1))
else:
raise ValueError(f"Unsupported blend mode: {mode}")
def g(self, x):
return torch.where(x <= 0.25, ((16 * x - 12) * x + 4) * x, torch.sqrt(x))
class Blur:
def __init__(self):
pass
@@ -787,6 +835,34 @@ class Sharpen:
return (result,)
class Solarize:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"threshold": ("FLOAT", {
"default": 0.5,
"min": 0.0,
"max": 1.0,
"step": 0.01
}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "solarize_image"
CATEGORY = "postprocessing"
def solarize_image(self, image: torch.Tensor, threshold: float):
solarized_image = torch.where(image > threshold, 1 - image, image)
solarized_image = torch.clamp(solarized_image, 0, 1)
return (solarized_image,)
def sort_span(span, sort_by, reverse_sorting):
if sort_by == 'H':
key = lambda x: x[1][0]
@@ -880,6 +956,7 @@ def pixel_sort(img, mask, horizontal_sort=False, span_limit=None, sort_by='H', r
NODE_CLASS_MAPPINGS = {
"ArithmeticBlend": ArithmeticBlend,
"Blend": Blend,
"Blur": Blur,
"CannyEdgeDetection": CannyEdgeDetection,
"ColorCorrect": ColorCorrect,
@@ -892,4 +969,5 @@ NODE_CLASS_MAPPINGS = {
"PixelSort": PixelSort,
"Pixelize": Pixelize,
"Sharpen": Sharpen,
"Solarize": Solarize,
}