Improved threshold nodes (#2091)

backend/src/nodes/properties/inputs/image_dropdown_inputs.py

@@ -1,5 +1,3 @@
-import cv2
-
 import navi
 
 from ...impl.color.convert_data import (
@@ -97,72 +95,6 @@ def BorderType::getOutputChannels(type: BorderType, channels: uint, color: Color
     )
 
 
-def ThresholdInput() -> DropDownInput:
-    """Threshold type option dropdown"""
-    return DropDownInput(
-        input_type="ThresholdType",
-        label="Threshold Type",
-        options=[
-            {
-                "option": "Binary",
-                "value": cv2.THRESH_BINARY,
-            },
-            {
-                "option": "Binary (Inverted)",
-                "value": cv2.THRESH_BINARY_INV,
-            },
-            {
-                "option": "Truncated",
-                "value": cv2.THRESH_TRUNC,
-            },
-            {
-                "option": "To Zero",
-                "value": cv2.THRESH_TOZERO,
-            },
-            {
-                "option": "To Zero (Inverted)",
-                "value": cv2.THRESH_TOZERO_INV,
-            },
-        ],
-    )
-
-
-def AdaptiveThresholdInput() -> DropDownInput:
-    """Adaptive Threshold type option dropdown"""
-    return DropDownInput(
-        input_type="AdaptiveThresholdType",
-        label="Threshold Type",
-        options=[
-            {
-                "option": "Binary",
-                "value": cv2.THRESH_BINARY,
-            },
-            {
-                "option": "Binary (Inverted)",
-                "value": cv2.THRESH_BINARY_INV,
-            },
-        ],
-    )
-
-
-def AdaptiveMethodInput() -> DropDownInput:
-    """Adaptive method border option dropdown"""
-    return DropDownInput(
-        input_type="AdaptiveMethod",
-        label="Adaptive Method",
-        options=[
-            {
-                "option": "Mean - C",
-                "value": cv2.ADAPTIVE_THRESH_MEAN_C,
-            },
-            {
-                "option": "Gaussian - C",
-                "value": cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
-            },
-        ],
-    )
-
-
 def NormalChannelInvertInput() -> DropDownInput:
     return DropDownInput(
         input_type="NormalChannelInvert",

backend/src/packages/chaiNNer_standard/image_adjustment/adjustments/threshold.py

@@ -1,15 +1,35 @@
 from __future__ import annotations
 
+from enum import Enum
+from typing import Dict
+
 import cv2
 import numpy as np
-from sanic.log import logger
 
-from nodes.properties.inputs import ImageInput, SliderInput, ThresholdInput
+from nodes.groups import if_enum_group
+from nodes.properties.inputs import EnumInput, ImageInput, SliderInput
 from nodes.properties.outputs import ImageOutput
 
 from .. import adjustments_group
 
 
+class ThresholdType(Enum):
+    BINARY = cv2.THRESH_BINARY
+    BINARY_INV = cv2.THRESH_BINARY_INV
+    TRUNC = cv2.THRESH_TRUNC
+    TO_ZERO = cv2.THRESH_TOZERO
+    TO_ZERO_INV = cv2.THRESH_TOZERO_INV
+
+
+_THRESHOLD_TYPE_LABELS: Dict[ThresholdType, str] = {
+    ThresholdType.BINARY: "Binary",
+    ThresholdType.BINARY_INV: "Binary (Inverted)",
+    ThresholdType.TRUNC: "Truncated",
+    ThresholdType.TO_ZERO: "To Zero",
+    ThresholdType.TO_ZERO_INV: "To Zero (Inverted)",
+}
+
+
 @adjustments_group.register(
     schema_id="chainner:image:threshold",
     name="Threshold",
@@ -24,29 +44,33 @@
             precision=1,
             controls_step=1,
         ),
-        SliderInput(
-            "Maximum Value",
-            maximum=100,
-            default=100,
-            precision=1,
-            controls_step=1,
+        EnumInput(
+            ThresholdType,
+            "Threshold Type",
+            default=ThresholdType.BINARY,
+            option_labels=_THRESHOLD_TYPE_LABELS,
+        ).with_id(3),
+        if_enum_group(3, (ThresholdType.BINARY, ThresholdType.BINARY_INV))(
+            SliderInput(
+                "Maximum Value",
+                maximum=100,
+                default=100,
+                precision=1,
+                controls_step=1,
+            ).with_id(2),
         ),
-        ThresholdInput(),
     ],
     outputs=[ImageOutput(image_type="Input0")],
 )
 def threshold_node(
-    img: np.ndarray, thresh: float, maxval: float, thresh_type: int
+    img: np.ndarray,
+    threshold: float,
+    thresh_type: ThresholdType,
+    max_value: float,
 ) -> np.ndarray:
-    """Takes an image and applies a threshold to it"""
-
-    logger.debug(f"thresh {thresh}, maxval {maxval}, type {thresh_type}")
-
-    real_thresh = thresh / 100
-    real_maxval = maxval / 100
-
-    logger.debug(f"real_thresh {real_thresh}, real_maxval {real_maxval}")
+    threshold /= 100
+    max_value /= 100
 
-    _, result = cv2.threshold(img, real_thresh, real_maxval, thresh_type)
+    _, result = cv2.threshold(img, threshold, max_value, thresh_type.value)
 
     return result

backend/src/packages/chaiNNer_standard/image_adjustment/adjustments/threshold_adaptive.py

@@ -1,65 +1,101 @@
 from __future__ import annotations
 
+from enum import Enum
+from typing import Dict
+
 import cv2
 import numpy as np
 
 from nodes.impl.image_utils import to_uint8
-from nodes.properties.inputs import (
-    AdaptiveMethodInput,
-    AdaptiveThresholdInput,
-    ImageInput,
-    NumberInput,
-    SliderInput,
-)
+from nodes.properties.inputs import EnumInput, ImageInput, NumberInput, SliderInput
 from nodes.properties.outputs import ImageOutput
 
 from .. import adjustments_group
 
 
+class AdaptiveThresholdType(Enum):
+    BINARY = cv2.THRESH_BINARY
+    BINARY_INV = cv2.THRESH_BINARY_INV
+
+
+_THRESHOLD_TYPE_LABELS: Dict[AdaptiveThresholdType, str] = {
+    AdaptiveThresholdType.BINARY: "Binary",
+    AdaptiveThresholdType.BINARY_INV: "Binary (Inverted)",
+}
+
+
+class AdaptiveMethod(Enum):
+    MEAN = cv2.ADAPTIVE_THRESH_MEAN_C
+    GAUSSIAN = cv2.ADAPTIVE_THRESH_GAUSSIAN_C
+
+
+_ADAPTIVE_METHOD_LABELS: Dict[AdaptiveMethod, str] = {
+    AdaptiveMethod.MEAN: "Mean",
+    AdaptiveMethod.GAUSSIAN: "Gaussian",
+}
+
+
 @adjustments_group.register(
     schema_id="chainner:image:threshold_adaptive",
     name="Threshold (Adaptive)",
     description="Similar to regular threshold, but determines the threshold for a pixel based on a small region around it.",
     icon="MdAutoGraph",
     inputs=[
         ImageInput(channels=1),
+        EnumInput(
+            AdaptiveThresholdType,
+            "Threshold Type",
+            default=AdaptiveThresholdType.BINARY,
+            option_labels=_THRESHOLD_TYPE_LABELS,
+        ).with_id(3),
         SliderInput(
             "Maximum Value",
             maximum=100,
             default=100,
             precision=1,
             controls_step=1,
+        ).with_id(1),
+        EnumInput(
+            AdaptiveMethod,
+            "Adaptive Method",
+            default=AdaptiveMethod.MEAN,
+            option_labels=_ADAPTIVE_METHOD_LABELS,
+        ).with_id(2),
+        NumberInput("Block Radius", default=1, minimum=1).with_id(4),
+        NumberInput(
+            "Constant Subtraction",
+            minimum=-100,
+            maximum=100,
+            default=0,
+            precision=1,
+        )
+        .with_id(5)
+        .with_docs(
+            "A constant value that is subtracted from the automatically determined adaptive threshold.",
+            "Assuming that **Threshold Type** is *Binary*, then higher values will result in more white pixels and lower values will result in more black pixels.",
         ),
-        AdaptiveMethodInput(),
-        AdaptiveThresholdInput(),
-        NumberInput("Block Radius", default=1, minimum=1),
-        NumberInput("Mean Subtraction"),
     ],
     outputs=[ImageOutput(image_type="Input0")],
     limited_to_8bpc=True,
 )
 def threshold_adaptive_node(
     img: np.ndarray,
-    maxval: float,
-    adaptive_method: int,
-    thresh_type: int,
+    threshold_type: AdaptiveThresholdType,
+    max_value: float,
+    adaptive_method: AdaptiveMethod,
     block_radius: int,
-    c: int,
+    c: float,
 ) -> np.ndarray:
-    """Takes an image and applies an adaptive threshold to it"""
-
     # Adaptive threshold requires uint8 input
     img = to_uint8(img, normalized=True)
 
-    real_maxval = maxval / 100 * 255
+    max_value = max_value / 100 * 255
 
-    result = cv2.adaptiveThreshold(
+    return cv2.adaptiveThreshold(
         img,
-        real_maxval,
-        adaptive_method,
-        thresh_type,
+        max_value,
+        adaptive_method.value,
+        threshold_type.value,
         block_radius * 2 + 1,
-        c,
+        round(c / 100 * 255),
     )
-
-    return result

src/common/types/chainner-scope.ts

@@ -92,15 +92,11 @@ let OnnxGenericModel = OnnxModel {
 struct IteratorAuto;
 
 // various inputs
-struct AdaptiveMethod;
-struct AdaptiveThresholdType;
 struct ColorSpace { channels: 1 | 3 | 4, supportsAlpha: bool }
 struct DdsFormat;
 struct DdsMipMaps;
-struct ImageExtension;
 struct NormalChannelInvert;
 struct RotateInterpolationMode;
-struct ThresholdType;
 struct TileSize;
 struct VideoPreset;
 struct VideoType;