Output normalization (#1717)

* Image output normalization

* Apply to nodes

* Added 2D-ize back

* Didn't mean to remove that

Author: RunDevelopment

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

@@ -4,7 +4,7 @@
 
 import numpy as np
 
-from ...impl.image_utils import get_h_w_c, normalize
+from ...impl.image_utils import get_h_w_c
 from ...utils.format import format_image_with_channels
 from .. import expression
 from .base_input import BaseInput
@@ -47,10 +47,12 @@ def enforce(self, value):
                 f"The input {self.label} only supports {expected} but was given {actual}."
             )
 
+        assert value.dtype == np.float32, "Expected the input image to be normalized."
+
         if c == 1 and value.ndim == 3:
             value = value[:, :, 0]
 
-        return normalize(value)
+        return value
 
 
 class VideoInput(BaseInput):

backend/src/nodes/properties/outputs/base_output.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Literal, Union
+from typing import Literal
 
 from base_types import OutputId
 
@@ -20,7 +20,7 @@ def __init__(
         self.output_type: expression.ExpressionJson = output_type
         self.label: str = label
         self.id: OutputId = OutputId(-1)
-        self.never_reason: Union[str, None] = None
+        self.never_reason: str | None = None
         self.kind: OutputKind = kind
         self.has_handle: bool = has_handle
 
@@ -34,7 +34,7 @@ def toDict(self):
             "hasHandle": self.has_handle,
         }
 
-    def with_id(self, output_id: Union[OutputId, int]):
+    def with_id(self, output_id: OutputId | int):
         self.id = OutputId(output_id)
         return self
 
@@ -54,5 +54,6 @@ def get_broadcast_data(self, _value):
     def get_broadcast_type(self, _value) -> expression.ExpressionJson | None:
         return None
 
-    def validate(self, value) -> None:
+    def enforce(self, value: object) -> object:
         assert value is not None
+        return value

backend/src/nodes/properties/outputs/file_outputs.py

@@ -13,8 +13,9 @@ def __init__(self, file_type: expression.ExpressionJson, label: str):
     def get_broadcast_data(self, value: str):
         return value
 
-    def validate(self, value) -> None:
+    def enforce(self, value) -> str:
         assert isinstance(value, str)
+        return value
 
 
 class DirectoryOutput(BaseOutput):
@@ -32,5 +33,6 @@ def __init__(self, label: str = "Directory", of_input: int | None = None):
     def get_broadcast_type(self, value: str):
         return expression.named("Directory", {"path": expression.literal(value)})
 
-    def validate(self, value) -> None:
+    def enforce(self, value) -> str:
         assert isinstance(value, str)
+        return value

backend/src/nodes/properties/outputs/generic_outputs.py

@@ -16,8 +16,9 @@ def __init__(
     def get_broadcast_type(self, value: int | float):
         return expression.literal(value)
 
-    def validate(self, value) -> None:
+    def enforce(self, value) -> int | float:
         assert isinstance(value, (int, float))
+        return value
 
 
 class TextOutput(BaseOutput):
@@ -31,8 +32,9 @@ def __init__(
     def get_broadcast_type(self, value: str):
         return expression.literal(value)
 
-    def validate(self, value) -> None:
+    def enforce(self, value) -> str:
         assert isinstance(value, str)
+        return value
 
 
 def FileNameOutput(label: str = "Name", of_input: int | None = None):
@@ -49,5 +51,6 @@ class SeedOutput(BaseOutput):
     def __init__(self, label: str = "Seed"):
         super().__init__(output_type="Seed", label=label, kind="generic")
 
-    def validate(self, value) -> None:
+    def enforce(self, value) -> Seed:
         assert isinstance(value, Seed)
+        return value

backend/src/nodes/properties/outputs/numpy_outputs.py

@@ -4,7 +4,7 @@
 import cv2
 import numpy as np
 
-from ...impl.image_utils import to_uint8
+from ...impl.image_utils import normalize, to_uint8
 from ...impl.pil_utils import InterpolationMethod, resize
 from ...utils.format import format_image_with_channels
 from ...utils.utils import get_h_w_c
@@ -24,8 +24,9 @@ def __init__(
     ):
         super().__init__(output_type, label, kind=kind, has_handle=has_handle)
 
-    def validate(self, value) -> None:
+    def enforce(self, value) -> np.ndarray:
         assert isinstance(value, np.ndarray)
+        return value
 
 
 def AudioOutput():
@@ -41,6 +42,7 @@ def __init__(
         kind: OutputKind = "image",
         has_handle: bool = True,
         channels: Optional[int] = None,
+        assume_normalized: bool = False,
     ):
         super().__init__(
             expression.intersect(image_type, expression.Image(channels=channels)),
@@ -50,6 +52,7 @@ def __init__(
         )
 
         self.channels: Optional[int] = channels
+        self.assume_normalized: bool = assume_normalized
 
     def get_broadcast_data(self, value: np.ndarray):
         h, w, c = get_h_w_c(value)
@@ -63,9 +66,8 @@ def get_broadcast_type(self, value: np.ndarray):
         h, w, c = get_h_w_c(value)
         return expression.Image(width=w, height=h, channels=c)
 
-    def validate(self, value) -> None:
+    def enforce(self, value) -> np.ndarray:
         assert isinstance(value, np.ndarray)
-        assert value.dtype == np.float32
 
         _, _, c = get_h_w_c(value)
 
@@ -78,6 +80,21 @@ def validate(self, value) -> None:
                 f" Please report this bug."
             )
 
+        # flatting 3D single-channel images to 2D
+        if c == 1 and value.ndim == 3:
+            value = value[:, :, 0]
+
+        if self.assume_normalized:
+            assert value.dtype == np.float32, (
+                f"The output {self.label} did not return a normalized image."
+                f" This is a bug in the implementation of the node."
+                f" Please report this bug."
+                f"\n\nTo the author of this node: Either use `normalize` or remove `assume_normalized=True` from this output."
+            )
+            return value
+
+        return normalize(value)
+
 
 def preview_encode(
     img: np.ndarray,

backend/src/packages/chaiNNer_standard/image/batch_processing/video_frame_iterator.py

@@ -10,7 +10,7 @@
 import numpy as np
 from sanic.log import logger
 
-from nodes.impl.image_utils import normalize, to_uint8
+from nodes.impl.image_utils import to_uint8
 from nodes.properties.inputs import (
     BoolInput,
     DirectoryInput,
@@ -73,7 +73,7 @@ class Writer:
 def VideoFrameIteratorFrameLoaderNode(
     img: np.ndarray, idx: int, video_dir: str, video_name: str
 ) -> Tuple[np.ndarray, int, str, str]:
-    return normalize(img), idx, video_dir, video_name
+    return img, idx, video_dir, video_name
 
 
 @batch_processing_group.register(

backend/src/packages/chaiNNer_standard/image/create_images/create_color_gray.py

@@ -31,7 +31,8 @@
                 width="Input0",
                 height="Input1",
                 channels="1",
-            )
+            ),
+            assume_normalized=True,
         )
     ],
 )

backend/src/packages/chaiNNer_standard/image/create_images/create_color_rgb.py

@@ -45,7 +45,8 @@
                 width="Input0",
                 height="Input1",
                 channels="3",
-            )
+            ),
+            assume_normalized=True,
         )
     ],
 )

backend/src/packages/chaiNNer_standard/image/create_images/create_color_rgba.py

@@ -52,7 +52,8 @@
                 width="Input0",
                 height="Input1",
                 channels="4",
-            )
+            ),
+            assume_normalized=True,
         )
     ],
 )

backend/src/packages/chaiNNer_standard/image/create_images/create_noise.py

@@ -200,4 +200,4 @@ def create_noise_node(
                 kwargs["seed"] = (kwargs["seed"] + 1) % (2**32)
         img /= total_brightness
 
-    return np.clip(img, 0, 1)
+    return img