Add Noise Node

backend/src/nodes/nodes/image_filter/add_noise.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+import numpy as np
+
+from . import category as ImageFilterCategory
+from ...node_base import NodeBase
+from ...node_factory import NodeFactory
+from ...properties.inputs import (
+    ImageInput,
+    SliderInput,
+    NoiseTypeDropdown,
+    NoiseColorDropdown,
+)
+from ...properties.outputs import ImageOutput
+from ...utils.noise_utils import (
+    gaussian_noise,
+    uniform_noise,
+    salt_and_pepper_noise,
+    poisson_noise,
+    speckle_noise,
+    NoiseType,
+)
+
+
+@NodeFactory.register("chainner:image:add_noise")
+class AddNoiseNode(NodeBase):
+    def __init__(self):
+        super().__init__()
+        self.description = "Add various kinds of noise to an image."
+        self.inputs = [
+            ImageInput(channels=[1, 3, 4]),
+            NoiseTypeDropdown(),
+            NoiseColorDropdown(),
+            SliderInput("Amount", minimum=0, maximum=100, default=50),
+        ]
+        self.outputs = [
+            ImageOutput(
+                image_type="""
+                Image {
+                    width: Input0.width,
+                    height: Input0.height,
+                    channels: max(
+                        Input0.channels,
+                        match Input2 { NoiseColor::Rgb => 3, NoiseColor::Gray => 1 }
+                    )
+                }"""
+            )
+        ]
+        self.category = ImageFilterCategory
+        self.name = "Add Noise"
+        self.icon = "CgEditNoise"
+        self.sub = "Noise"
+
+    def run(
+        self,
+        img: np.ndarray,
+        noise_type: str,
+        noise_color: NoiseType,
+        amount: int,
+    ) -> np.ndarray:
+        if noise_type == "gaussian":
+            return gaussian_noise(img, amount / 100, noise_color)
+        elif noise_type == "uniform":
+            return uniform_noise(img, amount / 100, noise_color)
+        elif noise_type == "salt_and_pepper":
+            return salt_and_pepper_noise(img, amount / 100, noise_color)
+        elif noise_type == "poisson":
+            return poisson_noise(img, amount / 100, noise_color)
+        elif noise_type == "speckle":
+            return speckle_noise(img, amount / 100, noise_color)
+        else:
+            raise ValueError(f"Unknown noise type: {noise_type}")

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

@@ -445,3 +445,51 @@ def NormalMappingAlphaInput() -> DropDownInput:
             },
         ],
     )
+
+
+def NoiseTypeDropdown() -> DropDownInput:
+    return DropDownInput(
+        input_type="NoiseType",
+        label="Noise Type",
+        options=[
+            {
+                "option": "Gaussian",
+                "value": "gaussian",
+            },
+            {
+                "option": "Uniform",
+                "value": "uniform",
+            },
+            {
+                "option": "Salt & Pepper",
+                "value": "salt_and_pepper",
+            },
+            {
+                "option": "Speckle",
+                "value": "speckle",
+            },
+            {
+                "option": "Poisson",
+                "value": "poisson",
+            },
+        ],
+    )
+
+
+def NoiseColorDropdown() -> DropDownInput:
+    return DropDownInput(
+        input_type="NoiseColor",
+        label="Noise Color",
+        options=[
+            {
+                "option": "Color",
+                "value": "rgb",
+                "type": "NoiseColor::Rgb",
+            },
+            {
+                "option": "Monochrome",
+                "value": "gray",
+                "type": "NoiseColor::Gray",
+            },
+        ],
+    )

backend/src/nodes/utils/noise_utils.py

@@ -0,0 +1,135 @@
+from typing import Literal, Callable, List
+import numpy as np
+
+from .utils import get_h_w_c
+from .image_utils import as_target_channels
+
+
+def __add_noises(
+    image: np.ndarray,
+    noise_gen: Callable[[int, int], List[np.ndarray]],
+    combine: Callable[[np.ndarray, List[np.ndarray]], np.ndarray],
+) -> np.ndarray:
+    img = image
+    h, w, c = get_h_w_c(img)
+    assert c != 2, "Noise cannot be added to 2-channel images."
+
+    if c > 3:
+        img = img[:, :, :3]
+
+    noises = noise_gen(h, w)
+    assert len(noises) > 0
+
+    max_channels = min(c, 3)
+    for n in noises:
+        noise_channels = get_h_w_c(n)[2]
+        assert noise_channels in (1, 3), "Noise must be a grayscale or RGB image."
+        max_channels = max(max_channels, noise_channels)
+
+    noises = [as_target_channels(n, max_channels) for n in noises]
+    img = as_target_channels(img, max_channels)
+
+    result = combine(img, noises)
+
+    if c > 3:
+        result = np.concatenate([result, image[:, :, 3:]], axis=2)
+
+    return np.clip(result, 0, 1)
+
+
+def __add_noise(
+    image: np.ndarray,
+    noise_gen: Callable[[int, int], np.ndarray],
+    combine: Callable[[np.ndarray, np.ndarray], np.ndarray] = lambda i, n: i + n,
+) -> np.ndarray:
+    return __add_noises(
+        image,
+        lambda h, w: [noise_gen(h, w)],
+        lambda i, n: combine(i, n[0]),
+    )
+
+
+NoiseType = Literal["gray", "rgb"]
+
+
+# Applies gaussian noise to an image
+def gaussian_noise(
+    image: np.ndarray,
+    amount: float,
+    noise_type: NoiseType,
+    seed: int = 0,
+) -> np.ndarray:
+    rng = np.random.default_rng(seed)
+    noise_c = 3 if noise_type == "rgb" else 1
+    return __add_noise(
+        image,
+        lambda h, w: rng.normal(0, amount, (h, w, noise_c)),
+    )
+
+
+# Applies uniform noise to an image
+def uniform_noise(
+    image: np.ndarray,
+    amount: float,
+    noise_type: NoiseType,
+    seed: int = 0,
+) -> np.ndarray:
+    rng = np.random.default_rng(seed)
+    noise_c = 3 if noise_type == "rgb" else 1
+    return __add_noise(
+        image,
+        lambda h, w: rng.uniform(-amount, amount, (h, w, noise_c)),
+    )
+
+
+# Applies salt and pepper noise to an image
+def salt_and_pepper_noise(
+    image: np.ndarray,
+    amount: float,
+    noise_type: NoiseType,
+    seed: int = 0,
+) -> np.ndarray:
+    def gen_noise(h: int, w: int):
+        rng = np.random.default_rng(seed)
+        noise_c = 3 if noise_type == "rgb" else 1
+        amt = amount / 2
+        pepper = rng.choice([0, 1], (h, w, noise_c), p=[amt, 1 - amt])
+        salt = rng.choice([0, 1], (h, w, noise_c), p=[1 - amt, amt])
+        return [pepper, salt]
+
+    def combine(i: np.ndarray, n: List[np.ndarray]):
+        pepper, salt = n
+        return np.where(salt == 1, 1, np.where(pepper == 0, 0, i))
+
+    return __add_noises(image, gen_noise, combine)
+
+
+# Applies poisson noise to an image
+def poisson_noise(
+    image: np.ndarray,
+    amount: float,
+    noise_type: NoiseType,
+    seed: int = 0,
+) -> np.ndarray:
+    rng = np.random.default_rng(seed)
+    noise_c = 3 if noise_type == "rgb" else 1
+    return __add_noise(
+        image,
+        lambda h, w: rng.poisson(amount, (h, w, noise_c)),
+    )
+
+
+# Applies speckle noise to an image
+def speckle_noise(
+    image: np.ndarray,
+    amount: float,
+    noise_type: NoiseType,
+    seed: int = 0,
+) -> np.ndarray:
+    rng = np.random.default_rng(seed)
+    noise_c = 3 if noise_type == "rgb" else 1
+    return __add_noise(
+        image,
+        lambda h, w: rng.normal(0, amount, (h, w, noise_c)),
+        lambda i, n: i + i * n,
+    )

src/common/types/chainner-scope.ts

@@ -83,6 +83,7 @@ struct TileMode;
 struct TransferColorspace;
 struct VideoType;
 struct VideoPreset;
+struct NoiseType;
 
 enum BorderType { ReflectMirror, Wrap, Replicate, Black, Transparent }
 enum FillColor { Auto, Black, Transparent }
@@ -93,6 +94,7 @@ enum PaddingAlignment { Start, End, Center }
 enum ResizeCondition { Both, Upscale, Downscale }
 enum RotateSizeChange { Crop, Expand }
 enum SideSelection { Width, Height, Shorter, Longer }
+enum NoiseColor { Rgb, Gray }
 
 def FillColor::getOutputChannels(fill: FillColor, channels: uint) {
     match fill {