chameleon/Assets/Scripts/Editor/TextureNormalMapGenerator.cs

using System;
using System.Diagnostics;
using System.IO;
using System.Text;
using System.Threading;
using UnityEditor;
using UnityEngine;
using Debug = UnityEngine.Debug;

namespace Chameleon.Editor
{
    /// <summary>
    /// Right-click a texture asset -> Process -> Generate Normal + Height Map (DeepBump).
    /// Runs the DeepBump ML CLI (Tools/DeepBump, own venv) twice: color_to_normals estimates a
    /// normal map from the color/albedo image, normals_to_height then derives a height/displacement
    /// map from that normal map (DeepBump has no direct color->height path). Both land next to the
    /// original as _normal.png (Texture Type = Normal Map) and _height.png (linear, sRGB off).
    /// DeepBump outputs OpenGL convention, which matches Unity.
    /// </summary>
    public static class TextureNormalMapGenerator
    {
        const string MenuPath = "Assets/Process/Generate Normal + Height Map (DeepBump)";
        const string DeepBumpDir = "Tools/DeepBump";
        // Inputs above this get Lanczos-downscaled first: DeepBump cost is quadratic in resolution
        // (256px tiles) and fewer tiles also means fewer visible tile seams. Normal-map detail is
        // low-frequency anyway; strength is tuned via _BumpScale.
        const int MaxInputSize = 2048;

        static string ProjectRoot => Path.GetDirectoryName(Application.dataPath);

        [MenuItem(MenuPath, true)]
        static bool Validate()
            => Selection.GetFiltered<Texture2D>(SelectionMode.Assets).Length > 0;

        [MenuItem(MenuPath, false, 2001)]
        static void Process()
        {
            string deepBump = Path.Combine(ProjectRoot, DeepBumpDir);
            string python = Path.Combine(deepBump, "venv/Scripts/python.exe");
            if (!File.Exists(python))
            {
                Debug.LogError($"[DeepBump] venv python not found at '{python}'. " +
                               "Set up Tools/DeepBump (clone + venv with numpy/onnxruntime/imageio) first.");
                return;
            }

            var mainThread = SynchronizationContext.Current;
            foreach (var texture in Selection.GetFiltered<Texture2D>(SelectionMode.Assets))
                ProcessTexture(texture, deepBump, python, mainThread);
        }

        static void ProcessTexture(Texture2D texture, string deepBump, string python, SynchronizationContext mainThread)
        {
            string assetPath = AssetDatabase.GetAssetPath(texture);
            if (string.IsNullOrEmpty(assetPath) || !File.Exists(assetPath))
            {
                Debug.LogError($"[DeepBump] No source file for '{texture.name}'.");
                return;
            }

            string directory = Path.GetDirectoryName(assetPath)?.Replace('\\', '/');
            string baseName = Path.GetFileNameWithoutExtension(assetPath);
            string normalAssetPath = AssetDatabase.GenerateUniqueAssetPath($"{directory}/{baseName}_normal.png");
            string heightAssetPath = AssetDatabase.GenerateUniqueAssetPath($"{directory}/{baseName}_height.png");

            string inputFull = Path.GetFullPath(assetPath);
            string normalFull = Path.GetFullPath(normalAssetPath);
            string heightFull = Path.GetFullPath(heightAssetPath);
            string tempInput = DownscaleIfNeeded(assetPath, inputFull, python, deepBump);
            string itemName = Path.GetFileName(assetPath);

            RunCli(python, deepBump, tempInput ?? inputFull, normalFull, "color_to_normals",
                "DeepBump Normal Map", itemName, mainThread, (normalOk, normalLog) =>
            {
                DeleteTemp(tempInput);
                if (!normalOk)
                {
                    if (File.Exists(normalFull)) File.Delete(normalFull);
                    Debug.LogError($"[DeepBump] Normal map generation failed for '{normalAssetPath}':\n{normalLog}");
                    return;
                }
                ImportNormalMap(normalAssetPath);

                // Height is derived from the just-generated normal map (full-res PNG on disk, before
                // Unity's import size cap can shrink it).
                RunCli(python, deepBump, normalFull, heightFull, "normals_to_height",
                    "DeepBump Height Map", itemName, mainThread, (heightOk, heightLog) =>
                {
                    if (!heightOk)
                    {
                        if (File.Exists(heightFull)) File.Delete(heightFull);
                        Debug.LogError($"[DeepBump] Height map generation failed for '{heightAssetPath}' " +
                                       $"(the normal map was created fine):\n{heightLog}");
                        return;
                    }
                    ImportHeightMap(heightAssetPath);
                });
            });
        }

        /// <summary>
        /// Runs one DeepBump CLI module as a background process with its own progress item;
        /// onDone(success, stderr) is invoked on the main thread.
        /// </summary>
        static void RunCli(string python, string deepBump, string inputFull, string outputFull, string module,
            string progressTitle, string itemName, SynchronizationContext mainThread, Action<bool, string> onDone)
        {
            int progressId = Progress.Start(progressTitle, itemName);

            var startInfo = new ProcessStartInfo
            {
                FileName = python,
                Arguments = $"cli.py \"{inputFull}\" \"{outputFull}\" {module} --verbose",
                WorkingDirectory = deepBump, // cli.py imports its sibling modules relative to cwd
                UseShellExecute = false,
                CreateNoWindow = true,
                RedirectStandardOutput = true,
                RedirectStandardError = true,
            };
            // cli.py prints a Unicode arrow; the default cp1252 console encoding would crash it.
            startInfo.EnvironmentVariables["PYTHONIOENCODING"] = "utf-8";

            var log = new StringBuilder();

            try
            {
                var process = new Process { StartInfo = startInfo, EnableRaisingEvents = true };
                process.OutputDataReceived += (_, e) =>
                {
                    if (e.Data == null) return;
                    // Verbose progress lines are "current/total" per tile.
                    int slash = e.Data.IndexOf('/');
                    if (slash > 0
                        && float.TryParse(e.Data.Substring(0, slash), out float current)
                        && float.TryParse(e.Data.Substring(slash + 1), out float total) && total > 0)
                        Progress.Report(progressId, current / total);
                };
                process.ErrorDataReceived += (_, e) => { if (e.Data != null) log.AppendLine(e.Data); };
                process.Exited += (_, _) =>
                {
                    bool success = process.ExitCode == 0;
                    process.Dispose();
                    mainThread.Post(_ =>
                    {
                        Progress.Remove(progressId);
                        onDone(success, log.ToString());
                    }, null);
                };
                process.Start();
                process.BeginOutputReadLine();
                process.BeginErrorReadLine();
            }
            catch (Exception e)
            {
                Progress.Remove(progressId);
                onDone(false, $"Failed to start python: {e.Message}");
            }
        }

        /// <summary>
        /// Returns a temp-file path with the downscaled image if the source exceeds
        /// <see cref="MaxInputSize"/>, otherwise null (use the original). Runs synchronously
        /// (sub-second even for large photos).
        /// </summary>
        static string DownscaleIfNeeded(string assetPath, string inputFull, string python, string deepBump)
        {
            int width = 0, height = 0;
            if (AssetImporter.GetAtPath(assetPath) is TextureImporter importer)
                importer.GetSourceTextureWidthAndHeight(out width, out height);
            if (Mathf.Max(width, height) <= MaxInputSize)
                return null;

            string temp = Path.Combine(Path.GetTempPath(), $"deepbump_{Guid.NewGuid():N}.png");
            string code = "from PIL import Image; im = Image.open(r'" + inputFull + "').convert('RGB'); " +
                          $"im.thumbnail(({MaxInputSize}, {MaxInputSize}), Image.LANCZOS); " +
                          "im.save(r'" + temp + "')";
            var startInfo = new ProcessStartInfo
            {
                FileName = python,
                Arguments = $"-c \"{code}\"",
                WorkingDirectory = deepBump,
                UseShellExecute = false,
                CreateNoWindow = true,
                RedirectStandardError = true,
            };
            try
            {
                // Fully qualified: bare "Process" resolves to this class's Process() menu method.
                using var process = System.Diagnostics.Process.Start(startInfo);
                string error = process.StandardError.ReadToEnd();
                process.WaitForExit();
                if (process.ExitCode == 0 && File.Exists(temp))
                    return temp;
                Debug.LogWarning($"[DeepBump] Pre-resize failed, falling back to full resolution:\n{error}");
            }
            catch (Exception e)
            {
                Debug.LogWarning($"[DeepBump] Pre-resize failed, falling back to full resolution: {e.Message}");
            }
            DeleteTemp(temp);
            return null;
        }

        static void DeleteTemp(string tempInput)
        {
            if (tempInput != null && File.Exists(tempInput))
                File.Delete(tempInput);
        }

        static void ImportNormalMap(string assetPath)
        {
            AssetDatabase.ImportAsset(assetPath);
            if (AssetImporter.GetAtPath(assetPath) is TextureImporter importer
                && importer.textureType != TextureImporterType.NormalMap)
            {
                importer.textureType = TextureImporterType.NormalMap;
                importer.SaveAndReimport();
            }
            Ping(assetPath);
        }

        static void ImportHeightMap(string assetPath)
        {
            AssetDatabase.ImportAsset(assetPath);
            // Grayscale displacement data, not color: linear sampling.
            if (AssetImporter.GetAtPath(assetPath) is TextureImporter importer && importer.sRGBTexture)
            {
                importer.sRGBTexture = false;
                importer.SaveAndReimport();
            }
            Ping(assetPath);
        }

        static void Ping(string assetPath)
        {
            var imported = AssetDatabase.LoadAssetAtPath<Texture2D>(assetPath);
            EditorGUIUtility.PingObject(imported);
            Debug.Log($"[DeepBump] Created '{assetPath}'.", imported);
        }
    }
}