com.misaki.hdrp-toon/Editor/NormalBaker/ModelOutlineJobs.cs

using System.Collections.Generic;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;
using Unity.Jobs;
using UnityEngine;

// Todo: Change this into a editor window

namespace Unity.Toonshader.Editor
{
    public class ModelOutlineJobs
    {
        public struct CollectNormalJob : IJobParallelFor
        {
            [ReadOnly] public NativeArray<Vector3> normals, vertrx;
            [NativeDisableContainerSafetyRestriction]
            public NativeArray<UnsafeParallelHashMap<Vector3, Vector3>.ParallelWriter> result;
            public CollectNormalJob(NativeArray<Vector3> normals, NativeArray<Vector3> vertrx, NativeArray<UnsafeParallelHashMap<Vector3, Vector3>.ParallelWriter> result)
            {
                this.normals = normals;
                this.vertrx = vertrx;
                this.result = result;
            }

            void IJobParallelFor.Execute(int index)
            {
                for (var i = 0; i < result.Length + 1; i++)
                {
                    if (i == result.Length)
                    {
                        Debug.LogError($"Overlapping vertices count({i})has out of bound!");
                        break;
                    }

                    if (result[i].TryAdd(vertrx[index], normals[index]))
                    {
                        break;
                    }
                }
            }
        }

        public struct BakeNormalJob : IJobParallelFor
        {
            [ReadOnly] public NativeArray<Vector3> vertrx, normals;
            [ReadOnly] public NativeArray<Vector4> tangents;
            [NativeDisableContainerSafetyRestriction]
            [ReadOnly] public NativeArray<UnsafeParallelHashMap<Vector3, Vector3>> result;
            [WriteOnly] public NativeArray<Vector2> uv2;

            public BakeNormalJob(NativeArray<Vector3> vertrx, NativeArray<Vector3> normals, NativeArray<Vector4> tangents, NativeArray<UnsafeParallelHashMap<Vector3, Vector3>> result, NativeArray<Vector2> uv2)
            {
                this.vertrx = vertrx;
                this.normals = normals;
                this.tangents = tangents;
                this.result = result;
                this.uv2 = uv2;
            }

            void IJobParallelFor.Execute(int index)
            {
                var smoothedNormals = Vector3.zero;
                for (var i = 0; i < result.Length; i++)
                {
                    if (result[i][vertrx[index]] != Vector3.zero)
                        smoothedNormals += result[i][vertrx[index]];
                    else
                        break;
                }
                smoothedNormals = smoothedNormals.normalized;

                var bitangent = (Vector3.Cross(normals[index], tangents[index]) * tangents[index].w).normalized;

                var tbn = new Matrix4x4(
                    tangents[index],
                    bitangent,
                    normals[index],
                    Vector4.zero);
                tbn = tbn.transpose;

                var bakedNormal = tbn.MultiplyVector(smoothedNormals).normalized;
                var newUV = new Vector2(bakedNormal.x, bakedNormal.y);
                uv2[index] = newUV;
            }
        }

        //[ContextMenu("Generate smooth normal to uv2")]
        //void ProcessingModel()
        //{
        //    var go = Selection.activeGameObject;
        //    if (go == null)
        //    {
        //        Debug.LogError("Select a GameObject first!");
        //        return;
        //    }

        //    Dictionary<string, Mesh> originalMesh = GetMesh(go), smoothedMesh = GetMesh(go);

        //    foreach (var item in originalMesh)
        //    {
        //        var m = item.Value;
        //        ComputeSmoothedNormalByJob(smoothedMesh[item.Key], m);
        //    }
        //}

        //void OnPreprocessModel()
        //{
        //    if (assetPath.Contains("@@@"))
        //    {
        //        ModelImporter model = assetImporter as ModelImporter;
        //        model.importNormals = ModelImporterNormals.Calculate;
        //        model.normalCalculationMode = ModelImporterNormalCalculationMode.AngleWeighted;
        //        model.normalSmoothingAngle = 180.0f;
        //        model.importAnimation = false;
        //        model.materialImportMode = ModelImporterMaterialImportMode.None;
        //    }
        //}

        //void OnPostprocessModel(GameObject g)
        //{
        //    if (!g.name.Contains("_ol") || g.name.Contains("@@@"))
        //        return;

        //    ModelImporter model = assetImporter as ModelImporter;

        //    string src = model.assetPath;
        //    string dst = Path.GetDirectoryName(src) + "/@@@" + Path.GetFileName(src);

        //    if (!File.Exists(Application.dataPath + "/" + dst.Substring(7)))
        //    {
        //        AssetDatabase.CopyAsset(src, dst);
        //        AssetDatabase.ImportAsset(dst);
        //    }
        //    else
        //    {
        //        var go = AssetDatabase.LoadAssetAtPath<GameObject>(dst);

        //        Dictionary<string, Mesh> originalMesh = GetMesh(g), smoothedMesh = GetMesh(go);

        //        foreach (var item in originalMesh)
        //        {
        //            var m = item.Value;
        //            ComputeSmoothedNormalByJob(smoothedMesh[item.Key], m);
        //        }

        //        AssetDatabase.DeleteAsset(dst);
        //    }
        //    AssetDatabase.Refresh();
        //}
    }

    public static class UTSNormalBakerHelper
    {
        public static Dictionary<string, Mesh> GetMesh(GameObject go)
        {
            static void AddMesh(Dictionary<string, Mesh> dictionary, string name, Mesh mesh)
            {
                if (dictionary.ContainsKey(name))
                    Debug.LogWarning($"Model:'{name}'is duplicate!");
                else
                    dictionary.Add(name, mesh);
            }

            var dic = new Dictionary<string, Mesh>();
            foreach (var item in go.GetComponentsInChildren<MeshFilter>())
                AddMesh(dic, item.name, item.sharedMesh);

            if (go.TryGetComponent<MeshFilter>(out var mf))
            {
                AddMesh(dic, mf.name.Replace("@", ""), mf.sharedMesh);
            }

            foreach (var item in go.GetComponentsInChildren<SkinnedMeshRenderer>())
            {
                AddMesh(dic, item.name, item.sharedMesh);
            }

            if (go.TryGetComponent<SkinnedMeshRenderer>(out var smr))
            {
                AddMesh(dic, smr.name.Replace("@", ""), smr.sharedMesh);
            }

            return dic;
        }

        public static void ComputeSmoothedNormalByJob(Mesh smoothedMesh, Mesh originalMesh, int maxOverlapvertices = 20)
        {
            int svc = smoothedMesh.vertexCount, ovc = originalMesh.vertexCount;

            // CollectNormalJob Data
            var normals = new NativeArray<Vector3>(smoothedMesh.normals, Allocator.Persistent);
            var vertrx = new NativeArray<Vector3>(smoothedMesh.vertices, Allocator.Persistent);
            var smoothedNormals = new NativeArray<Vector3>(svc, Allocator.Persistent);
            var result = new NativeArray<UnsafeParallelHashMap<Vector3, Vector3>>(maxOverlapvertices, Allocator.Persistent);
            var resultParallel = new NativeArray<UnsafeParallelHashMap<Vector3, Vector3>.ParallelWriter>(result.Length, Allocator.Persistent);

            // NormalBakeJob Data
            NativeArray<Vector3> normalsO = new NativeArray<Vector3>(originalMesh.normals, Allocator.Persistent),
                vertrxO = new NativeArray<Vector3>(originalMesh.vertices, Allocator.Persistent);
            var tangents = new NativeArray<Vector4>(originalMesh.tangents, Allocator.Persistent);
            var uv2 = new NativeArray<Vector2>(ovc, Allocator.Persistent);

            for (var i = 0; i < result.Length; i++)
            {
                result[i] = new UnsafeParallelHashMap<Vector3, Vector3>(svc, Allocator.Persistent);
                resultParallel[i] = result[i].AsParallelWriter();
            }

            var collectNormalJob = new ModelOutlineJobs.CollectNormalJob(normals, vertrx, resultParallel);
            var normalBakeJob = new ModelOutlineJobs.BakeNormalJob(vertrxO, normalsO, tangents, result, uv2);

            normalBakeJob.Schedule(ovc, 8, collectNormalJob.Schedule(svc, 100)).Complete();

            var _uv2 = new Vector2[ovc];
            uv2.CopyTo(_uv2);
            originalMesh.uv2 = _uv2;

            normals.Dispose();
            vertrx.Dispose();
            result.Dispose();
            smoothedNormals.Dispose();
            resultParallel.Dispose();
            normalsO.Dispose();
            vertrxO.Dispose();
            tangents.Dispose();
            uv2.Dispose();
        }
    }
}