Merge branch 'main' into sequence_types

chaiNNer-org · Jun 30, 2024 · dff9227 · dff9227
2 parents a2dcaef + 2503261
commit dff9227
Show file tree

Hide file tree

Showing 21 changed files with 143 additions and 105 deletions.
diff --git a/backend/src/nodes/properties/outputs/numpy_outputs.py b/backend/src/nodes/properties/outputs/numpy_outputs.py
@@ -46,11 +46,13 @@ class ImageOutput(NumPyOutput):
     def __init__(
         self,
         label: str = "Image",
+        *,
         image_type: navi.ExpressionJson = "Image",
         kind: OutputKind = "generic",
         has_handle: bool = True,
         channels: int | None = None,
         shape_as: int | InputId | None = None,
+        size_as: int | InputId | None = None,
         assume_normalized: bool = False,
     ):
         # narrow down type
@@ -60,6 +62,10 @@ def __init__(
             )
         if shape_as is not None:
             image_type = navi.intersect_with_error(image_type, f"Input{shape_as}")
+        if size_as is not None:
+            image_type = navi.intersect_with_error(
+                image_type, navi.Image(size_as=f"Input{size_as}")
+            )
 
         super().__init__(image_type, label, kind=kind, has_handle=has_handle)
 
@@ -159,7 +165,7 @@ def __init__(
     ):
         super().__init__(
             label,
-            image_type,
+            image_type=image_type,
             kind=kind,
             has_handle=has_handle,
             assume_normalized=assume_normalized,

diff --git a/backend/src/packages/chaiNNer_onnx/onnx/processing/remove_background.py b/backend/src/packages/chaiNNer_onnx/onnx/processing/remove_background.py
@@ -2,7 +2,6 @@
 
 import numpy as np
 
-import navi
 from api import NodeContext
 from nodes.groups import Condition, if_group
 from nodes.impl.onnx.model import OnnxRemBg
@@ -51,7 +50,7 @@
             """,
             channels=4,
         ),
-        ImageOutput("Mask", image_type=navi.Image(size_as="Input0"), channels=1),
+        ImageOutput("Mask", size_as=0, channels=1),
     ],
     node_context=True,
 )

diff --git a/backend/src/packages/chaiNNer_standard/image_adjustment/adjustments/opacity.py b/backend/src/packages/chaiNNer_standard/image_adjustment/adjustments/opacity.py
@@ -2,7 +2,6 @@
 
 import numpy as np
 
-import navi
 from api import KeyInfo
 from nodes.impl.pil_utils import convert_to_bgra
 from nodes.properties.inputs import ImageInput, SliderInput
@@ -22,11 +21,7 @@
         SliderInput("Opacity", max=100, default=100, precision=1, step=1, unit="%"),
     ],
     outputs=[
-        ImageOutput(
-            image_type=navi.Image(size_as="Input0"),
-            channels=4,
-            assume_normalized=True,
-        )
+        ImageOutput(size_as=0, channels=4, assume_normalized=True),
     ],
     key_info=KeyInfo.number(1),
 )

diff --git a/backend/src/packages/chaiNNer_standard/image_channel/all/merge_channels.py b/backend/src/packages/chaiNNer_standard/image_channel/all/merge_channels.py
@@ -27,8 +27,8 @@
     ],
     outputs=[
         ImageOutput(
+            size_as=0,
             image_type=navi.Image(
-                size_as="Input0",
                 channels="""
                     match (
                         Input0.channels
@@ -41,7 +41,7 @@
                         int(4..) => 4
                     }
                     """,
-            )
+            ),
         )
     ],
     deprecated=True,

diff --git a/backend/src/packages/chaiNNer_standard/image_channel/all/separate_rgba.py b/backend/src/packages/chaiNNer_standard/image_channel/all/separate_rgba.py
@@ -2,7 +2,6 @@
 
 import numpy as np
 
-import navi
 from nodes.properties.inputs import ImageInput
 from nodes.properties.outputs import ImageOutput
 from nodes.utils.utils import get_h_w_c
@@ -20,35 +19,17 @@
     icon="MdCallSplit",
     inputs=[ImageInput()],
     outputs=[
-        ImageOutput(
-            "R Channel",
-            image_type=navi.Image(size_as="Input0"),
-            channels=1,
-            assume_normalized=True,
-        )
+        ImageOutput("R Channel", size_as=0, channels=1, assume_normalized=True)
         .with_docs("The red channel.")
         .with_id(2),
-        ImageOutput(
-            "G Channel",
-            image_type=navi.Image(size_as="Input0"),
-            channels=1,
-            assume_normalized=True,
-        )
+        ImageOutput("G Channel", size_as=0, channels=1, assume_normalized=True)
         .with_docs("The green channel.")
         .with_id(1),
-        ImageOutput(
-            "B Channel",
-            image_type=navi.Image(size_as="Input0"),
-            channels=1,
-            assume_normalized=True,
-        )
+        ImageOutput("B Channel", size_as=0, channels=1, assume_normalized=True)
         .with_docs("The blue channel.")
         .with_id(0),
         ImageOutput(
-            "A Channel",
-            image_type=navi.Image(size_as="Input0"),
-            channels=1,
-            assume_normalized=True,
+            "A Channel", size_as=0, channels=1, assume_normalized=True
         ).with_docs("The alpha (transparency mask) channel."),
     ],
 )

diff --git a/backend/src/packages/chaiNNer_standard/image_channel/misc/chroma_key.py b/backend/src/packages/chaiNNer_standard/image_channel/misc/chroma_key.py
@@ -6,7 +6,6 @@
 import numpy as np
 import pymatting
 
-import navi
 from nodes.groups import if_enum_group, linked_inputs_group
 from nodes.impl.color.color import Color
 from nodes.properties.inputs import (
@@ -62,10 +61,7 @@ class KeyMethod(Enum):
         ),
     ],
     outputs=[
-        ImageOutput(
-            image_type=navi.Image(size_as="Input0"),
-            channels=4,
-        ),
+        ImageOutput(size_as=0, channels=4),
     ],
 )
 def chroma_key_node(

diff --git a/backend/src/packages/chaiNNer_standard/image_channel/misc/fill_alpha.py b/backend/src/packages/chaiNNer_standard/image_channel/misc/fill_alpha.py
@@ -9,7 +9,6 @@
     fill_alpha_nearest_color,
 )
 
-import navi
 from nodes.properties.inputs import EnumInput, ImageInput
 from nodes.properties.outputs import ImageOutput
 
@@ -32,16 +31,8 @@ class AlphaFillMethod(Enum):
         EnumInput(AlphaFillMethod, label="Fill Method"),
     ],
     outputs=[
-        ImageOutput(
-            "RGB",
-            image_type=navi.Image(size_as="Input0"),
-            channels=3,
-        ),
-        ImageOutput(
-            "Alpha",
-            image_type=navi.Image(size_as="Input0"),
-            channels=1,
-        ),
+        ImageOutput("RGB", size_as=0, channels=3),
+        ImageOutput("Alpha", size_as=0, channels=1),
     ],
 )
 def fill_alpha_node(

diff --git a/backend/src/packages/chaiNNer_standard/image_channel/transparency/split_transparency.py b/backend/src/packages/chaiNNer_standard/image_channel/transparency/split_transparency.py
@@ -2,7 +2,6 @@
 
 import numpy as np
 
-import navi
 from nodes.impl.image_utils import as_target_channels
 from nodes.properties.inputs import ImageInput
 from nodes.properties.outputs import ImageOutput
@@ -18,18 +17,8 @@
     icon="MdCallSplit",
     inputs=[ImageInput(channels=[1, 3, 4])],
     outputs=[
-        ImageOutput(
-            "RGB",
-            image_type=navi.Image(size_as="Input0"),
-            channels=3,
-            assume_normalized=True,
-        ),
-        ImageOutput(
-            "Alpha",
-            image_type=navi.Image(size_as="Input0"),
-            channels=1,
-            assume_normalized=True,
-        ),
+        ImageOutput("RGB", size_as=0, channels=3, assume_normalized=True),
+        ImageOutput("Alpha", size_as=0, channels=1, assume_normalized=True),
     ],
 )
 def split_transparency_node(img: np.ndarray) -> tuple[np.ndarray, np.ndarray]:

diff --git a/backend/src/packages/chaiNNer_standard/image_filter/miscellaneous/canny_edge_detection.py b/backend/src/packages/chaiNNer_standard/image_filter/miscellaneous/canny_edge_detection.py
@@ -3,7 +3,6 @@
 import cv2
 import numpy as np
 
-import navi
 from nodes.impl.image_utils import to_uint8
 from nodes.properties.inputs import ImageInput, NumberInput
 from nodes.properties.outputs import ImageOutput
@@ -23,7 +22,7 @@
         NumberInput("Lower Threshold", min=0, default=100),
         NumberInput("Upper Threshold", min=0, default=300),
     ],
-    outputs=[ImageOutput(image_type=navi.Image(size_as="Input0"), channels=1)],
+    outputs=[ImageOutput(size_as=0, channels=1)],
 )
 def canny_edge_detection_node(
     img: np.ndarray,

diff --git a/backend/src/packages/chaiNNer_standard/image_utility/miscellaneous/change_color_model.py b/backend/src/packages/chaiNNer_standard/image_utility/miscellaneous/change_color_model.py
@@ -45,8 +45,8 @@
     ],
     outputs=[
         ImageOutput(
+            size_as=0,
             image_type=navi.Image(
-                size_as="Input0",
                 channels="""
                     if Input2.supportsAlpha and Input3 {
                         4

diff --git a/backend/src/packages/chaiNNer_standard/material_textures/conversion/metal_to_specular.py b/backend/src/packages/chaiNNer_standard/material_textures/conversion/metal_to_specular.py
@@ -2,7 +2,6 @@
 
 import numpy as np
 
-import navi
 from nodes.impl.resize import ResizeFilter, resize
 from nodes.properties.inputs import ImageInput
 from nodes.properties.outputs import ImageOutput
@@ -63,11 +62,7 @@ def metal_to_spec(
     ],
     outputs=[
         ImageOutput("Diffuse", shape_as=0),
-        ImageOutput(
-            "Specular",
-            image_type=navi.Image(size_as="Input1"),
-            channels=3,
-        ),
+        ImageOutput("Specular", size_as=1, channels=3),
         ImageOutput(
             "Gloss",
             image_type="""

diff --git a/backend/src/packages/chaiNNer_standard/material_textures/conversion/specular_to_metal.py b/backend/src/packages/chaiNNer_standard/material_textures/conversion/specular_to_metal.py
@@ -2,7 +2,6 @@
 
 import numpy as np
 
-import navi
 from nodes.impl.resize import ResizeFilter, resize
 from nodes.properties.inputs import ImageInput, SliderInput
 from nodes.properties.outputs import ImageOutput
@@ -89,11 +88,7 @@ def spec_to_metal(
     ],
     outputs=[
         ImageOutput("Albedo", shape_as=0),
-        ImageOutput(
-            "Metal",
-            image_type=navi.Image(size_as="Input1"),
-            channels=1,
-        ),
+        ImageOutput("Metal", size_as=1, channels=1),
         ImageOutput(
             "Roughness",
             image_type="""

diff --git a/backend/src/packages/chaiNNer_standard/material_textures/normal_map/balance_normals.py b/backend/src/packages/chaiNNer_standard/material_textures/normal_map/balance_normals.py
@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+import numpy as np
+
+import navi
+from nodes.impl.normals.util import gr_to_xyz, normalize_normals, xyz_to_bgr
+from nodes.properties.inputs import ImageInput
+from nodes.properties.outputs import ImageOutput
+
+from .. import normal_map_group
+
+
+@normal_map_group.register(
+    schema_id="chainner:image:balance_normals",
+    name="Balance Normals",
+    description=[
+        "This ensures that the average of all normals is pointing straight up. The input normal map is normalized before this operation is applied. The output normal map is guaranteed to be normalized.",
+    ],
+    icon="MdExpand",
+    inputs=[
+        ImageInput("Normal Map", channels=[3, 4]),
+    ],
+    outputs=[
+        ImageOutput("Normal Map", image_type=navi.Image(size_as="Input0"), channels=3),
+    ],
+)
+def balance_normals_node(n: np.ndarray) -> np.ndarray:
+    x, y, _ = gr_to_xyz(n)
+
+    x -= np.mean(x)
+    y -= np.mean(y)
+
+    return xyz_to_bgr(normalize_normals(x, y))
diff --git a/backend/src/packages/chaiNNer_standard/material_textures/normal_map/normal_map_generator.py b/backend/src/packages/chaiNNer_standard/material_textures/normal_map/normal_map_generator.py
@@ -240,8 +240,8 @@ def normalize(x: np.ndarray, y: np.ndarray):
     outputs=[
         ImageOutput(
             "Normal Map",
+            size_as=0,
             image_type=navi.Image(
-                size_as="Input0",
                 channels="match Input7 { AlphaOutput::None => 3, _ => 4 }",
             ),
         ),

diff --git a/backend/src/packages/chaiNNer_standard/material_textures/normal_map/normalize_normals.py b/backend/src/packages/chaiNNer_standard/material_textures/normal_map/normalize_normals.py
@@ -4,7 +4,6 @@
 
 import numpy as np
 
-import navi
 from nodes.impl.normals.util import gr_to_xyz, xyz_to_bgr
 from nodes.properties.inputs import EnumInput, ImageInput
 from nodes.properties.outputs import ImageOutput
@@ -42,11 +41,7 @@ class BChannel(Enum):
         ),
     ],
     outputs=[
-        ImageOutput(
-            "Normal Map",
-            image_type=navi.Image(size_as="Input0"),
-            channels=3,
-        ),
+        ImageOutput("Normal Map", size_as=0, channels=3),
     ],
 )
 def normalize_normals_node(img: np.ndarray, b: BChannel) -> np.ndarray:

diff --git a/backend/src/packages/chaiNNer_standard/material_textures/normal_map/scale_normals.py b/backend/src/packages/chaiNNer_standard/material_textures/normal_map/scale_normals.py
@@ -2,7 +2,6 @@
 
 import numpy as np
 
-import navi
 from nodes.impl.normals.addition import AdditionMethod, strengthen_normals
 from nodes.impl.normals.util import xyz_to_bgr
 from nodes.properties.inputs import EnumInput, ImageInput, SliderInput
@@ -29,11 +28,7 @@
         ),
     ],
     outputs=[
-        ImageOutput(
-            "Normal Map",
-            image_type=navi.Image(size_as="Input0"),
-            channels=3,
-        ),
+        ImageOutput("Normal Map", size_as=0, channels=3),
     ],
 )
 def scale_normals_node(

diff --git a/docs/nodes.md b/docs/nodes.md
@@ -213,7 +213,7 @@ Let's take a look at the inputs and outputs of the Create Border node. This node
 
 ### Documentation
 
-Inputs and outputs both allow optional documentation to be added to add more information about it to the Node Documentation modal. You can add documentation to an input or output using the `.with_docs()` method, like so: 
+Inputs and outputs both allow optional documentation to be added to add more information about it to the Node Documentation modal. You can add documentation to an input or output using the `.with_docs()` method, like so:
 
 ```python
     inputs=[
@@ -303,10 +303,7 @@ from .. import adjustments_group
         ),
     ],
     outputs=[
-        ImageOutput(
-            image_type=navi.Image(size_as="Input0"),
-            channels=4,
-        )
+        ImageOutput(size_as=0, channels=4)
     ],
 )
 def opacity_node(img: np.ndarray, opacity: float) -> np.ndarray: