adds solarize

2023-04-02 19:33:16 -04:00
parent 8da09a5812
commit d0e69db667
3 changed files with 112 additions and 0 deletions
@@ -0,0 +1,33 @@
+import torch
+
+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,)
+
+NODE_CLASS_MAPPINGS = {
+    "Solarize": Solarize,
+}