Impostors: baker now records albedo, normal, gloss, and depth

4 files changed, 196 insertions, 86 deletions

diff --git a/3ner.cginc b/3ner.cginc
index c104c68..db10781 100644
--- a/3ner.cginc
+++ b/3ner.cginc
@@ -265,8 +265,12 @@ float4 frag(v2f i, uint facing : SV_IsFrontFace) : SV_Target {
   i.normal *= facing ? 1 : -1;
 
   Pbr pbr = getPbr(i);
-#if defined(_UNLIT)
+#if defined(_DEBUG_VIEW_UNLIT)
   return pbr.albedo;
+#elif defined(_DEBUG_VIEW_WORLD_SPACE_NORMALS)
+  return float4(i.normal, 1);
+#elif defined(_DEBUG_VIEW_METALLIC_GLOSS)
+  return float4(pbr.metallic, pbr.smoothness, 0, 1);
 #endif
 
   LightData light_data;
diff --git a/3ner.shader b/3ner.shader
index 70a5174..4818c08 100644
--- a/3ner.shader
+++ b/3ner.shader
@@ -395,9 +395,11 @@ Shader "yum_food/3ner"
 
       [HideInInspector] m_start_Rendering_Options("Rendering Options", Float) = 0
 
-      //ifex _Unlit==0
-      [ThryToggle(_UNLIT)] _Unlit("Unlit", Float) = 0
-      //endex
+      [HideInInspector] m_start_Debug_Views("Debug Views", Float) = 0
+        [ThryToggle(_DEBUG_VIEW_UNLIT)] _Debug_View_Unlit("Unlit", Float) = 0
+        [ThryToggle(_DEBUG_VIEW_WORLD_SPACE_NORMALS)] _Debug_View_World_Space_Normals("World space normals", Float) = 0
+        [ThryToggle(_DEBUG_VIEW_METALLIC_GLOSS)] _Debug_View_Metallic_Gloss("Metallic gloss", Float) = 0
+      [HideInInspector] m_end_Debug_Views("Debug Views", Float) = 0
 
       [DoNotAnimate][HideInInspector] Instancing ("Instancing", Float) = 0
       //ifex Instancing==0
diff --git a/Scripts/Impostors.cs b/Scripts/Impostors.cs
index ded7e6d..e8513eb 100644
--- a/Scripts/Impostors.cs
+++ b/Scripts/Impostors.cs
@@ -125,74 +125,182 @@ public class Impostors : MonoBehaviour
     cameras = null;
   }
 
+  struct BakePass
+  {
+    public string name;
+    public string keyword;
+    public Texture2D atlas;
+    public bool isDepth;
+
+    public BakePass(string name, string keyword, Texture2D atlas, bool isDepth = false)
+    {
+      this.name = name;
+      this.keyword = keyword;
+      this.atlas = atlas;
+      this.isDepth = isDepth;
+    }
+  }
+
+  void RenderAtlasPass(Texture2D atlas, RenderTexture colorRT, RenderTexture depthOnlyRT, Material depthBlitMat = null)
+  {
+    bool isDepth = depthBlitMat != null;
+    RenderTexture linearDepthRT = isDepth ? RenderTexture.GetTemporary(cameraResolution, cameraResolution, 0, RenderTextureFormat.RFloat) : null;
+
+    int idx = 0;
+    for (int y = 0; y < gridResolution; y++)
+    {
+      for (int x = 0; x < gridResolution; x++)
+      {
+        Camera cam = cameras[idx++];
+        cam.SetTargetBuffers(colorRT.colorBuffer, depthOnlyRT.depthBuffer);
+        cam.Render();
+
+        if (isDepth)
+        {
+          depthBlitMat.SetTexture("_DepthTex", depthOnlyRT);
+          Graphics.Blit(null, linearDepthRT, depthBlitMat);
+          RenderTexture.active = linearDepthRT;
+        }
+        else
+        {
+          RenderTexture.active = colorRT;
+        }
+
+        Texture2D temp = new Texture2D(cameraResolution, cameraResolution, isDepth ? TextureFormat.RFloat : TextureFormat.RGBA32, false);
+        temp.ReadPixels(new Rect(0, 0, cameraResolution, cameraResolution), 0, 0);
+        temp.Apply();
+        atlas.SetPixels(x * cameraResolution, y * cameraResolution, cameraResolution, cameraResolution, temp.GetPixels());
+        DestroyImmediate(temp);
+      }
+    }
+    atlas.Apply();
+
+    if (linearDepthRT != null) RenderTexture.ReleaseTemporary(linearDepthRT);
+  }
+
+  void SetMaterialDebugKeyword(string keyword, bool enabled)
+  {
+    if (originalMesh == null || string.IsNullOrEmpty(keyword)) return;
+    foreach (Renderer r in originalMesh.GetComponentsInChildren<Renderer>(true))
+    {
+      foreach (Material mat in r.sharedMaterials)
+      {
+        if (mat != null)
+        {
+          if (enabled)
+            mat.EnableKeyword(keyword);
+          else
+            mat.DisableKeyword(keyword);
+        }
+      }
+    }
+  }
+
+  void ExecutePassesSequentially(BakePass[] passes, RenderTexture colorRT, RenderTexture depthOnlyRT, Material depthBlitMat, int passIndex = 0)
+  {
+    if (passIndex >= passes.Length)
+    {
+      // All passes complete
+      RenderTexture.active = null;
+      RenderTexture.ReleaseTemporary(colorRT);
+      RenderTexture.ReleaseTemporary(depthOnlyRT);
+      DestroyImmediate(depthBlitMat);
+
+      SaveAndConfigureTextures(passes[0].atlas, passes[1].atlas, passes[2].atlas, passes[3].atlas);
+      return;
+    }
+
+    BakePass pass = passes[passIndex];
+    Debug.Log($"Baking {pass.name} pass...");
+
+    SetMaterialDebugKeyword(pass.keyword, true);
+
+    UnityEditor.EditorApplication.delayCall += () => {
+      RenderAtlasPass(pass.atlas, colorRT, depthOnlyRT, pass.isDepth ? depthBlitMat : null);
+      SetMaterialDebugKeyword(pass.keyword, false);
+      ExecutePassesSequentially(passes, colorRT, depthOnlyRT, depthBlitMat, passIndex + 1);
+    };
+  }
+
   public void BakeTexture()
   {
     SetRenderersEnabled(true);
     if (cameras == null || cameras.Length != gridResolution * gridResolution || cameras[0] == null) CreateCameras();
 
-    // Load depth blit shader
     Shader depthBlitShader = Shader.Find("Hidden/yum_food/DepthBlit");
     if (depthBlitShader == null) { Debug.LogError("DepthBlit shader not found"); return; }
     Material depthBlitMat = new Material(depthBlitShader);
 
-    // Render atlas
     int size = cameraResolution * gridResolution;
-    Texture2D colorAtlas = new Texture2D(size, size, TextureFormat.RGBA32, false);
-    Texture2D depthAtlas = new Texture2D(size, size, TextureFormat.RFloat, false);
+    BakePass[] passes = new BakePass[]
+    {
+      new BakePass("albedo", "_DEBUG_VIEW_UNLIT", new Texture2D(size, size, TextureFormat.RGBA32, false)),
+      new BakePass("normals", "_DEBUG_VIEW_WORLD_SPACE_NORMALS", new Texture2D(size, size, TextureFormat.RGBA32, false)),
+      new BakePass("metallic/gloss", "_DEBUG_VIEW_METALLIC_GLOSS", new Texture2D(size, size, TextureFormat.RGBA32, false)),
+      new BakePass("depth", "", new Texture2D(size, size, TextureFormat.RFloat, false), true)
+    };
 
-    // Create RT with depth buffer that can be sampled as texture
     RenderTextureDescriptor desc = new RenderTextureDescriptor(cameraResolution, cameraResolution, RenderTextureFormat.ARGB32, 24);
     desc.sRGB = true;
     RenderTexture colorRT = RenderTexture.GetTemporary(desc);
 
-    // Separate depth texture
     RenderTextureDescriptor depthDesc = new RenderTextureDescriptor(cameraResolution, cameraResolution, RenderTextureFormat.Depth, 24);
     RenderTexture depthOnlyRT = RenderTexture.GetTemporary(depthDesc);
 
-    // Output for linearized depth
-    RenderTexture linearDepthRT = RenderTexture.GetTemporary(cameraResolution, cameraResolution, 0, RenderTextureFormat.RFloat);
+    // Ensure all debug keywords start disabled
+    foreach (var pass in passes) SetMaterialDebugKeyword(pass.keyword, false);
 
-    int idx = 0;
-    for (int y = 0; y < gridResolution; y++)
+    ExecutePassesSequentially(passes, colorRT, depthOnlyRT, depthBlitMat);
+  }
+
+  struct TextureExportSettings
+  {
+    public string suffix;
+    public bool isEXR;
+    public bool mipmaps;
+    public bool sRGB;
+    public FilterMode filter;
+    public bool alphaTransparency;
+    public bool uncompressed;
+
+    public TextureExportSettings(string suffix, bool isEXR = false, bool mipmaps = true, bool sRGB = true,
+                                  FilterMode filter = FilterMode.Trilinear, bool alphaTransparency = false, bool uncompressed = false)
     {
-      for (int x = 0; x < gridResolution; x++)
-      {
-        Camera cam = cameras[idx++];
+      this.suffix = suffix;
+      this.isEXR = isEXR;
+      this.mipmaps = mipmaps;
+      this.sRGB = sRGB;
+      this.filter = filter;
+      this.alphaTransparency = alphaTransparency;
+      this.uncompressed = uncompressed;
+    }
+  }
 
-        // Render to color + depth buffers simultaneously
-        cam.SetTargetBuffers(colorRT.colorBuffer, depthOnlyRT.depthBuffer);
-        cam.Render();
+  void SaveAndConfigureTexture(Texture2D atlas, TextureExportSettings settings, string baseName, out string path)
+  {
+    path = Path.Combine(OutputFolder, $"{baseName}_{settings.suffix}.{(settings.isEXR ? "exr" : "png")}");
+    byte[] data = settings.isEXR ? atlas.EncodeToEXR(Texture2D.EXRFlags.OutputAsFloat) : atlas.EncodeToPNG();
+    File.WriteAllBytes(Path.Combine(Application.dataPath, "..", path), data);
+    DestroyImmediate(atlas);
+  }
 
-        // Read color
-        RenderTexture.active = colorRT;
-        Texture2D colorTemp = new Texture2D(cameraResolution, cameraResolution);
-        colorTemp.ReadPixels(new Rect(0, 0, cameraResolution, cameraResolution), 0, 0);
-        colorTemp.Apply();
-        colorAtlas.SetPixels(x * cameraResolution, y * cameraResolution, cameraResolution, cameraResolution, colorTemp.GetPixels());
-        DestroyImmediate(colorTemp);
-
-        // Blit depth buffer through linearization shader
-        depthBlitMat.SetTexture("_DepthTex", depthOnlyRT);
-        Graphics.Blit(null, linearDepthRT, depthBlitMat);
-
-        // Read linearized depth
-        RenderTexture.active = linearDepthRT;
-        Texture2D depthTemp = new Texture2D(cameraResolution, cameraResolution, TextureFormat.RFloat, false);
-        depthTemp.ReadPixels(new Rect(0, 0, cameraResolution, cameraResolution), 0, 0);
-        depthTemp.Apply();
-        depthAtlas.SetPixels(x * cameraResolution, y * cameraResolution, cameraResolution, cameraResolution, depthTemp.GetPixels());
-        DestroyImmediate(depthTemp);
-      }
+  void ConfigureTextureImporter(string path, TextureExportSettings settings)
+  {
+    TextureImporter importer = AssetImporter.GetAtPath(path) as TextureImporter;
+    if (importer != null)
+    {
+      importer.mipmapEnabled = settings.mipmaps;
+      importer.sRGBTexture = settings.sRGB;
+      importer.wrapMode = TextureWrapMode.Clamp;
+      importer.filterMode = settings.filter;
+      if (settings.alphaTransparency) importer.alphaIsTransparency = true;
+      if (settings.uncompressed) importer.textureCompression = TextureImporterCompression.Uncompressed;
+      importer.SaveAndReimport();
     }
-    colorAtlas.Apply();
-    depthAtlas.Apply();
-    RenderTexture.active = null;
-    RenderTexture.ReleaseTemporary(colorRT);
-    RenderTexture.ReleaseTemporary(depthOnlyRT);
-    RenderTexture.ReleaseTemporary(linearDepthRT);
-    DestroyImmediate(depthBlitMat);
-
-    // Save
+  }
+
+  void SaveAndConfigureTextures(Texture2D albedoAtlas, Texture2D normalAtlas, Texture2D metallicGlossAtlas, Texture2D depthAtlas)
+  {
     if (!AssetDatabase.IsValidFolder(OutputFolder))
     {
       if (!AssetDatabase.IsValidFolder("Assets/yum_food/3ner"))
@@ -200,44 +308,31 @@ public class Impostors : MonoBehaviour
       AssetDatabase.CreateFolder("Assets/yum_food/3ner", "Impostor_Generated");
     }
 
-    string name = gameObject.name.Replace(" ", "_");
-    string colorPath = Path.Combine(OutputFolder, $"{name}_atlas.png");
-    string depthPath = Path.Combine(OutputFolder, $"{name}_depth.exr");
+    string baseName = gameObject.name.Replace(" ", "_");
 
-    File.WriteAllBytes(Path.Combine(Application.dataPath, "..", colorPath), colorAtlas.EncodeToPNG());
-    File.WriteAllBytes(Path.Combine(Application.dataPath, "..", depthPath), depthAtlas.EncodeToEXR(Texture2D.EXRFlags.OutputAsFloat));
-    DestroyImmediate(colorAtlas);
-    DestroyImmediate(depthAtlas);
+    var exportSettings = new (Texture2D atlas, TextureExportSettings settings, string materialProp)[]
+    {
+      (albedoAtlas, new TextureExportSettings("albedo", mipmaps: true, sRGB: true, alphaTransparency: true), "_ImpostorAtlas"),
+      (normalAtlas, new TextureExportSettings("normal", mipmaps: true, sRGB: false), "_ImpostorNormalAtlas"),
+      (metallicGlossAtlas, new TextureExportSettings("metallic_gloss", mipmaps: true, sRGB: false), "_ImpostorMetallicGlossAtlas"),
+      (depthAtlas, new TextureExportSettings("depth", isEXR: true, mipmaps: false, sRGB: false, filter: FilterMode.Bilinear, uncompressed: true), "_ImpostorDepthAtlas")
+    };
 
-    AssetDatabase.Refresh();
+    string[] paths = new string[exportSettings.Length];
+    for (int i = 0; i < exportSettings.Length; i++)
+      SaveAndConfigureTexture(exportSettings[i].atlas, exportSettings[i].settings, baseName, out paths[i]);
 
-    // Configure color atlas importer
-    TextureImporter colorImporter = AssetImporter.GetAtPath(colorPath) as TextureImporter;
-    if (colorImporter != null)
-    {
-      colorImporter.mipmapEnabled = true;
-      colorImporter.alphaIsTransparency = true;
-      colorImporter.wrapMode = TextureWrapMode.Clamp;
-      colorImporter.filterMode = FilterMode.Trilinear;
-      colorImporter.SaveAndReimport();
-    }
+    AssetDatabase.Refresh();
 
-    // Configure depth atlas importer
-    TextureImporter depthImporter = AssetImporter.GetAtPath(depthPath) as TextureImporter;
-    if (depthImporter != null)
-    {
-      depthImporter.mipmapEnabled = false;
-      depthImporter.sRGBTexture = false;  // Linear data
-      depthImporter.wrapMode = TextureWrapMode.Clamp;
-      depthImporter.filterMode = FilterMode.Bilinear;
-      depthImporter.textureCompression = TextureImporterCompression.Uncompressed;
-      depthImporter.SaveAndReimport();
-    }
+    for (int i = 0; i < paths.Length; i++)
+      ConfigureTextureImporter(paths[i], exportSettings[i].settings);
 
     // Create impostor
-    Texture2D colorTex = AssetDatabase.LoadAssetAtPath<Texture2D>(colorPath);
-    Texture2D depthTex = AssetDatabase.LoadAssetAtPath<Texture2D>(depthPath);
-    if (colorTex != null)
+    Texture2D[] textures = new Texture2D[paths.Length];
+    for (int i = 0; i < paths.Length; i++)
+      textures[i] = AssetDatabase.LoadAssetAtPath<Texture2D>(paths[i]);
+
+    if (textures[0] != null)
     {
       DestroyExistingImpostor();
 
@@ -245,11 +340,11 @@ public class Impostors : MonoBehaviour
       if (shader == null) { Debug.LogError("Shader not found"); return; }
 
       impostorMaterial = new Material(shader);
-      impostorMaterial.SetTexture("_ImpostorAtlas", colorTex);
-      impostorMaterial.SetTexture("_ImpostorDepthAtlas", depthTex);
+      for (int i = 0; i < textures.Length; i++)
+        impostorMaterial.SetTexture(exportSettings[i].materialProp, textures[i]);
       impostorMaterial.SetInt("_GridResolution", gridResolution);
       impostorMaterial.SetFloat("_SphereRadius", sphere_radius_);
-      AssetDatabase.CreateAsset(impostorMaterial, Path.Combine(OutputFolder, $"{name}_mat.mat"));
+      AssetDatabase.CreateAsset(impostorMaterial, Path.Combine(OutputFolder, $"{baseName}_mat.mat"));
 
       impostorObject = GameObject.CreatePrimitive(PrimitiveType.Quad);
       impostorObject.name = "Impostor";
@@ -259,6 +354,7 @@ public class Impostors : MonoBehaviour
       impostorObject.GetComponent<MeshRenderer>().sharedMaterial = impostorMaterial;
 
       SetRenderersEnabled(false);
+      Debug.Log("Impostor baking complete!");
     }
   }
 
diff --git a/features.cginc b/features.cginc
index 9d1c17b..d9bc95e 100644
--- a/features.cginc
+++ b/features.cginc
@@ -37,8 +37,16 @@
 #pragma shader_feature_local _VERTEX_DEFORMATION_TESSELLATION
 //endex
 
-//ifex _Unlit==0
-#pragma shader_feature_local _UNLIT
+//ifex _Debug_View_Unlit==0
+#pragma shader_feature_local _DEBUG_VIEW_UNLIT
+//endex
+
+//ifex _Debug_View_World_Space_Normals==0
+#pragma shader_feature_local _DEBUG_VIEW_WORLD_SPACE_NORMALS
+//endex
+
+//ifex _Debug_View_Metallic_Gloss==0
+#pragma shader_feature_local _DEBUG_VIEW_METALLIC_GLOSS
 //endex
 
 //ifex _Instance_Distance_Culling_Enabled==0