StyleTransferPostProcessing.cs

using System.IO;
using TensorFlowLite;
using UnityEngine;

[RequireComponent(typeof(Camera))]
[ExecuteInEditMode]
public class StyleTransferPostProcessing : MonoBehaviour
{
    [SerializeField, FilePopup("*.tflite")] string predictionFileName = "style_predict_quantized_256.tflite";
    [SerializeField, FilePopup("*.tflite")] string transferFileName = "style_transfer_quantized_dynamic.tflite";
    [SerializeField] Texture2D styleImage = null;
    float[] styleBottleneck;
    [SerializeField] ComputeShader compute = null;
    StyleTransfer styleTransfer;
    [SerializeField]  Material adaptiveBlend;

    [SerializeField] private bool blendWithSource;
    void Start()
    {
        string predictionModelPath = Path.Combine(Application.streamingAssetsPath, predictionFileName);
        using (var predict = new StylePredict(predictionModelPath))
        {
            predict.Invoke(styleImage);
            styleBottleneck = predict.GetStyleBottleneck();
        }

        string transferModelPath = Path.Combine(Application.streamingAssetsPath, transferFileName);
        styleTransfer = new StyleTransfer(transferModelPath, styleBottleneck, compute);
    }
    
    void OnDestroy()
    {
        styleTransfer?.Dispose();
    }

    private void OnRenderImage(RenderTexture src, RenderTexture dest)
    {
        if (styleTransfer != null)
        {
            styleTransfer.Invoke(src);
            var renderTex = styleTransfer.GetResultTexture();
            if (blendWithSource)
            {
                Graphics.Blit(src, dest);
                Graphics.Blit(renderTex, dest, adaptiveBlend);
            }
            else
            {
                Graphics.Blit(renderTex,dest);
            }
        }
        else
        {
            Graphics.Blit(src,dest);
        }
    }
}