Add size_as parameter to ImageOutput (#2965)

* Add `size_as` parameter to `ImageOutput` * Fixed type error
chaiNNer-org · Jun 30, 2024 · 2503261 · 2503261
1 parent a0fe2b9
commit 2503261
Show file tree

Hide file tree

Showing 16 changed files with 29 additions and 96 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/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: