Add size_as parameter to ImageOutput

chaiNNer-org · Jun 30, 2024 · 66f584d · 66f584d
1 parent afe35c3
commit 66f584d
Show file tree

Hide file tree

Showing 17 changed files with 54 additions and 97 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)
 

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_dimension/resize/resize.py b/backend/src/packages/chaiNNer_standard/image_dimension/resize/resize.py
@@ -25,6 +25,11 @@ class ImageResizeMode(Enum):
     ABSOLUTE = 1
 
 
+class InputDimension(Enum):
+    WIDTH = 0
+    HEIGHT = 1
+
+
 @resize_group.register(
     schema_id="chainner:image:resize",
     name="Resize",
@@ -49,8 +54,25 @@ class ImageResizeMode(Enum):
             ).with_id(2),
         ),
         if_enum_group(1, ImageResizeMode.ABSOLUTE)(
-            NumberInput("Width", min=1, default=1, unit="px").with_id(3),
-            NumberInput("Height", min=1, default=1, unit="px").with_id(4),
+            BoolInput("Maintain Aspect Ratio", default=False).with_id(7),
+            if_group(Condition.bool(7, True))(
+                EnumInput(
+                    InputDimension,
+                    label="Dimension",
+                    default=InputDimension.WIDTH,
+                    label_style="inline",
+                ).with_id(8),
+            ),
+            if_group(
+                Condition.bool(7, False) | Condition.enum(8, InputDimension.WIDTH)
+            )(
+                NumberInput("Width", min=1, default=1, unit="px").with_id(3),
+            ),
+            if_group(
+                Condition.bool(7, False) | Condition.enum(8, InputDimension.HEIGHT)
+            )(
+                NumberInput("Height", min=1, default=1, unit="px").with_id(4),
+            ),
         ),
         ResizeFilterInput().with_id(5),
         if_group(Condition.type(0, "Image { channels: 4 } "))(
@@ -109,6 +131,8 @@ def resize_node(
     img: np.ndarray,
     mode: ImageResizeMode,
     scale: float,
+    maintain_aspect_ratio: bool,
+    dimension: InputDimension,
     width: int,
     height: int,
     filter: ResizeFilter,

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: