From 6b10ebdc2c21f42a0d302d4abcecb7dfe83616d3 Mon Sep 17 00:00:00 2001
From: Chandler Cox <mr.chandlercox@gmail.com>
Date: Fri, 4 Oct 2024 15:29:45 -0500
Subject: [PATCH] Add options to reconstruct z channel for textures.

---
 .../TextureTool/TextureTool.DirectXTex.cpp    | 37 +++++++++++++++++++
 .../Engine/Tools/TextureTool/TextureTool.cpp  |  4 ++
 Source/Engine/Tools/TextureTool/TextureTool.h |  4 ++
 3 files changed, 45 insertions(+)

diff --git a/Source/Engine/Tools/TextureTool/TextureTool.DirectXTex.cpp b/Source/Engine/Tools/TextureTool/TextureTool.DirectXTex.cpp
index 094daefda..e62d32842 100644
--- a/Source/Engine/Tools/TextureTool/TextureTool.DirectXTex.cpp
+++ b/Source/Engine/Tools/TextureTool/TextureTool.DirectXTex.cpp
@@ -829,6 +829,43 @@ bool TextureTool::ImportTextureDirectXTex(ImageType type, const StringView& path
         SET_CURRENT_IMG(timage);
     }
 
+    // Reconstruct Z Channel
+    if (!keepAsIs & options.ReconstructZChannel)
+    {
+        auto& timage = GET_TMP_IMG();
+        bool isunorm = (DirectX::FormatDataType(currentImage->GetMetadata().format) == DirectX::FORMAT_TYPE_UNORM) != 0;
+        result = TransformImage(currentImage->GetImages(), currentImage->GetImageCount(), currentImage->GetMetadata(),
+            [&](DirectX::XMVECTOR* outPixels, const DirectX::XMVECTOR* inPixels, size_t w, size_t y)
+            {
+                static const DirectX::XMVECTORU32 s_selectz = { { { DirectX::XM_SELECT_0, DirectX::XM_SELECT_0, DirectX::XM_SELECT_1, DirectX::XM_SELECT_0 } } };
+
+                UNREFERENCED_PARAMETER(y);
+
+                for (size_t j = 0; j < w; ++j)
+                {
+                    const DirectX::XMVECTOR value = inPixels[j];
+                    DirectX::XMVECTOR z;
+                    if (isunorm)
+                    {
+                        DirectX::XMVECTOR x2 = DirectX::XMVectorMultiplyAdd(value, DirectX::g_XMTwo, DirectX::g_XMNegativeOne);
+                        x2 = DirectX::XMVectorSqrt(DirectX::XMVectorSubtract(DirectX::g_XMOne, DirectX::XMVector2Dot(x2, x2)));
+                        z = DirectX::XMVectorMultiplyAdd(x2, DirectX::g_XMOneHalf, DirectX::g_XMOneHalf);
+                    }
+                    else
+                    {
+                        z = DirectX::XMVectorSqrt(DirectX::XMVectorSubtract(DirectX::g_XMOne, DirectX::XMVector2Dot(value, value)));
+                    }
+                    outPixels[j] = XMVectorSelect(value, z, s_selectz);
+                }
+            }, timage);
+        if (FAILED(result))
+        {
+            errorMsg = String::Format(TEXT("Cannot reconstruct z channel in texture, error: {0:x}"), static_cast<uint32>(result));
+            return true;
+        }
+        SET_CURRENT_IMG(timage);
+    }
+
     // Generate mip maps chain
     if (!keepAsIs && useMipLevels && options.GenerateMipMaps)
     {
diff --git a/Source/Engine/Tools/TextureTool/TextureTool.cpp b/Source/Engine/Tools/TextureTool/TextureTool.cpp
index ff0a66177..48df4cb31 100644
--- a/Source/Engine/Tools/TextureTool/TextureTool.cpp
+++ b/Source/Engine/Tools/TextureTool/TextureTool.cpp
@@ -76,6 +76,9 @@ void TextureTool::Options::Serialize(SerializeStream& stream, const void* otherO
     stream.JKEY("InvertGreenChannel");
     stream.Bool(InvertGreenChannel);
 
+    stream.JKEY("ReconstructZChannel");
+    stream.Bool(ReconstructZChannel);
+
     stream.JKEY("Resize");
     stream.Bool(Resize);
 
@@ -134,6 +137,7 @@ void TextureTool::Options::Deserialize(DeserializeStream& stream, ISerializeModi
     GenerateMipMaps = JsonTools::GetBool(stream, "GenerateMipMaps", GenerateMipMaps);
     FlipY = JsonTools::GetBool(stream, "FlipY", FlipY);
     InvertGreenChannel = JsonTools::GetBool(stream, "InvertGreenChannel", InvertGreenChannel);
+    ReconstructZChannel = JsonTools::GetBool(stream, "ReconstructZChannel", ReconstructZChannel);
     Resize = JsonTools::GetBool(stream, "Resize", Resize);
     PreserveAlphaCoverage = JsonTools::GetBool(stream, "PreserveAlphaCoverage", PreserveAlphaCoverage);
     PreserveAlphaCoverageReference = JsonTools::GetFloat(stream, "PreserveAlphaCoverageReference", PreserveAlphaCoverageReference);
diff --git a/Source/Engine/Tools/TextureTool/TextureTool.h b/Source/Engine/Tools/TextureTool/TextureTool.h
index 85d8b2d20..179508c68 100644
--- a/Source/Engine/Tools/TextureTool/TextureTool.h
+++ b/Source/Engine/Tools/TextureTool/TextureTool.h
@@ -61,6 +61,10 @@ API_CLASS(Namespace="FlaxEngine.Tools", Static) class FLAXENGINE_API TextureTool
         API_FIELD(Attributes = "EditorOrder(71)")
         bool InvertGreenChannel = false;
 
+        // True if to invert the green channel on a normal map. Good for OpenGL to DirectX conversion.
+        API_FIELD(Attributes = "EditorOrder(72)")
+        bool ReconstructZChannel = false;
+
         // Texture size scale. Allows increasing or decreasing the imported texture resolution. Default is 1.
         API_FIELD(Attributes="EditorOrder(80), Limit(0.0001f, 1000.0f, 0.01f)")
         float Scale = 1.0f;