From 4906ba4a32fd9f243602e4a17d1cbe73a3da4b5e Mon Sep 17 00:00:00 2001
From: Michael Schmidt <mitchi5000.ms@googlemail.com>
Date: Sun, 3 Dec 2023 18:27:40 +0100
Subject: [PATCH] Add strengthen normals node (#2372)

---
 backend/src/nodes/impl/normals/addition.py    | 34 +++++++++++++++
 .../normal_map/strengthen_normals.py          | 42 +++++++++++++++++++
 2 files changed, 76 insertions(+)
 create mode 100644 backend/src/packages/chaiNNer_standard/material_textures/normal_map/strengthen_normals.py

diff --git a/backend/src/nodes/impl/normals/addition.py b/backend/src/nodes/impl/normals/addition.py
index a7833cf93..f90688705 100644
--- a/backend/src/nodes/impl/normals/addition.py
+++ b/backend/src/nodes/impl/normals/addition.py
@@ -97,3 +97,37 @@ def add_normals(
         return __angles(xyz1, xyz2, f1, f2)
     else:
         raise AssertionError(f"Invalid normal addition method {method}")
+
+
+def strengthen_normals(method: AdditionMethod, n: np.ndarray, f: float) -> XYZ:
+    """
+    Same as `add_normals`, but with `n2` being the same as `n1` and `f2=0`.
+    """
+    # Convert BGR to XY
+    x = n[:, :, 2] * 2 - 1
+    y = n[:, :, 1] * 2 - 1
+
+    x, y, z = normalize_normals(x, y)
+
+    if method is AdditionMethod.PARTIAL_DERIVATIVES:
+        # Slopes aren't defined for z=0, so set to near-zero decimal
+        z = np.maximum(z, 0.001, out=z)
+
+        n_f = f / z
+
+        x = x * n_f
+        y = y * n_f
+
+        l_r = 1 / np.sqrt(np.square(x) + np.square(y) + 1)
+        x *= l_r
+        y *= l_r
+        z = l_r
+
+        return x, y, z
+    elif method is AdditionMethod.ANGLES:
+        return normalize_normals(
+            np.sin(__clamp_angles(np.arcsin(x) * f)),
+            np.sin(__clamp_angles(np.arcsin(y) * f)),
+        )
+    else:
+        raise AssertionError(f"Invalid normal addition method {method}")
diff --git a/backend/src/packages/chaiNNer_standard/material_textures/normal_map/strengthen_normals.py b/backend/src/packages/chaiNNer_standard/material_textures/normal_map/strengthen_normals.py
new file mode 100644
index 000000000..ce919f418
--- /dev/null
+++ b/backend/src/packages/chaiNNer_standard/material_textures/normal_map/strengthen_normals.py
@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+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
+from nodes.properties.outputs import ImageOutput
+
+from .. import normal_map_group
+
+
+@normal_map_group.register(
+    schema_id="chainner:image:strengthen_normals",
+    name="Strengthen Normals",
+    description=[
+        "Strengths and weakens the normals in the given normal map. Only the R and G channels of the input image will be used. The output normal map is guaranteed to be normalized.",
+        "Conceptually, this node is equivalent to `chainner:image:add_normals` with the strength of the second normal map set to 0.",
+    ],
+    icon="MdExpand",
+    inputs=[
+        ImageInput("Normal Map", channels=[3, 4]),
+        SliderInput("Strength", maximum=400, default=100),
+        EnumInput(
+            AdditionMethod,
+            label="Method",
+            default=AdditionMethod.PARTIAL_DERIVATIVES,
+        ),
+    ],
+    outputs=[
+        ImageOutput(
+            "Normal Map",
+            image_type=navi.Image(size_as="Input0"),
+            channels=3,
+        ),
+    ],
+)
+def strengthen_normals_node(
+    n: np.ndarray, strength: int, method: AdditionMethod
+) -> np.ndarray:
+    return xyz_to_bgr(strengthen_normals(method, n, strength / 100))