[Unity] Replacement Shader for SceneView

00README.md

This is an example code how to use replacement shader for SceneView.

• Add this component to your Camera
• Add some GameObject which has Cube as its prefix
• Select TestReplacementShader/Enable Custom Render Mode on SceneView from Untiy Editor's main menu.
  • All Cube objects are drawn in red (replacement shader) in SceneView.
• Select TestReplacementShader/Disable Custom Render Mode on SceneView from Untiy Editor's main menu.
  • All Cube objects are drawn in normal shader.

TestReplacementShader.cs

using UnityEngine;
#if UNITY_EDITOR
using UnityEditor;
#endif
using System;
using System.Reflection;

[ExecuteInEditMode]
[ImageEffectAllowedInSceneView]
[RequireComponent(typeof(Camera))]
public class TestReplacementShader : MonoBehaviour {
    const string ReplacementShaderName = "Utj/TestReplacementShader";
    public Color backgroundColor = new Color(0f, 0f, 0f, 0f);
    public Color defaultColor = new Color(0.5f, 0.5f, 0.5f, 0f);

    Camera thisCamera_ = null;
    protected Camera ThisCamera {
        get { return thisCamera_ = thisCamera_ ?? GetComponent<Camera>(); }
    }

    static Shader replacementShader_ = null;
    protected static Shader ReplacementShader {
        get { return replacementShader_ = replacementShader_ ?? Shader.Find(ReplacementShaderName); }
    }

#if UNITY_EDITOR
    [MenuItem("TestReplacementShader/Enable Custom Render Mode on SceneView")]
    static void EnableCustomRender() {
        if(ReplacementShader != null) {
            foreach(SceneView sceneView in SceneView.sceneViews) {
                sceneView.SetSceneViewShaderReplace(ReplacementShader, null);
            }
        }
        SceneView.RepaintAll();
    }

    [MenuItem("TestReplacementShader/Disable Custom Render Mode on SceneView")]
    static void DisableCustomRender() {
        if(ReplacementShader != null) {
            foreach(SceneView sceneView in SceneView.sceneViews) {
                if(GetSceneViewReplacementShader(sceneView) == ReplacementShader) {
                    sceneView.SetSceneViewShaderReplace(null, null);
                }
            }
        }
        SceneView.RepaintAll();
    }

    static readonly FieldInfo fiSceneView_m_ReplacementShader = typeof(UnityEditor.SceneView).GetField(
        "m_ReplacementShader",
        BindingFlags.NonPublic | BindingFlags.Public | BindingFlags.Instance
    );

    static Shader GetSceneViewReplacementShader(SceneView sceneView) {
        var fi = fiSceneView_m_ReplacementShader;
        if(sceneView == null || fi == null) {
            return null;
        }
        return (Shader) fi.GetValue(sceneView);
    }
#endif

    static void TraverseAllScenesAndGameObjects(
        System.Action<UnityEngine.SceneManagement.Scene, GameObject, int> action
    ) {
        var nScene = UnityEngine.SceneManagement.SceneManager.sceneCount;
        for(int iScene = 0; iScene < nScene; ++iScene) {
            var scene = UnityEngine.SceneManagement.SceneManager.GetSceneAt(iScene);
            action(scene, null, -1);

            var rootGos = scene.GetRootGameObjects();
            foreach(var rootGo in rootGos) {
                TraverseGameObjects(scene, rootGo, action, 1);
            }
        }
    }

    static void TraverseGameObjects(
          UnityEngine.SceneManagement.Scene scene
        , GameObject go
        , System.Action<UnityEngine.SceneManagement.Scene, GameObject, int> action
        , int depth = 0
    ) {
        if(go == null) {
            return;
        }
        action(scene, go, depth);
        foreach(Transform childTransform in go.transform) {
            if(childTransform == go.transform) {
                continue;
            }
            TraverseGameObjects(scene, childTransform.gameObject, action, depth+1);
        }
    }

    void Start() {
        if(ThisCamera != null) {
            ThisCamera.clearFlags = CameraClearFlags.SolidColor;
            ThisCamera.backgroundColor = backgroundColor;
            ThisCamera.SetReplacementShader(Shader.Find(ReplacementShaderName), null);
        }

        SetMpb();
    }

    void SetMpb() {
        TraverseAllScenesAndGameObjects((scene, go, depth) => {
            if(go == null) {
                return;
            }

            Color color = defaultColor;
            if(go.name.StartsWith("Cube")) {
                color = new Color(1.0f, 0f, 0f, 0f);
            }
            Debug.LogFormat("{0} -> {1}", go.name, color);

            foreach(Renderer renderer in go.GetComponents<Renderer>()) {
                var mpb = new MaterialPropertyBlock();
                renderer.GetPropertyBlock(mpb);
                mpb.SetColor("_IdColor", color);
                renderer.SetPropertyBlock(mpb);
            }
        });
    }
}

TestReplacementShader.shader

Shader "Utj/TestReplacementShader"
{
    Properties
    {
        [PerRendererData] _IdColor  ("ID Color", Color)  = (0.4, 0.4, 0.8, 0)
    }

    SubShader
    {
        Tags { "RenderType" = "Opaque" }

        CGPROGRAM

        #pragma surface surf Lambert

        struct Input
        {
            float2 uv_MainTex;
        };

        sampler2D _MainTex;
        float3 _IdColor;

        void surf (Input IN, inout SurfaceOutput o)
        {
            half4 c = tex2D (_MainTex, IN.uv_MainTex);
            o.Albedo = half3(0, 0, 0);
            o.Alpha = c.a;
            o.Emission = _IdColor.rgb;
        }

        ENDCG
    }

    Fallback "Diffuse"
}