147 lines
3.6 KiB
C#
147 lines
3.6 KiB
C#
// Copyright (c) 2021 homuler
|
|
//
|
|
// Use of this source code is governed by an MIT-style
|
|
// license that can be found in the LICENSE file or at
|
|
// https://opensource.org/licenses/MIT.
|
|
|
|
using System;
|
|
using System.Runtime.InteropServices;
|
|
using UnityEngine;
|
|
using UnityEngine.Rendering;
|
|
using UnityEngine.UI;
|
|
|
|
namespace Mediapipe.Unity
|
|
{
|
|
#pragma warning disable IDE0065
|
|
using Color = UnityEngine.Color;
|
|
#pragma warning restore IDE0065
|
|
|
|
public class MaskAnnotation : HierarchicalAnnotation
|
|
{
|
|
[SerializeField] private RawImage _screen;
|
|
[SerializeField] private Shader _maskShader;
|
|
[SerializeField] private Texture2D _maskTexture;
|
|
[SerializeField] private Color _color = Color.blue;
|
|
[SerializeField, Range(0, 1)] private float _threshold = 0.9f;
|
|
|
|
private Material _prevMaterial;
|
|
private Material _material;
|
|
private GraphicsBuffer _maskBuffer;
|
|
|
|
private void OnEnable()
|
|
{
|
|
ApplyMaskTexture(_maskTexture, _color);
|
|
ApplyThreshold(_threshold);
|
|
}
|
|
|
|
private void OnDisable()
|
|
{
|
|
if (_prevMaterial != null)
|
|
{
|
|
ApplyMaterial(_prevMaterial);
|
|
}
|
|
}
|
|
|
|
#if UNITY_EDITOR
|
|
private void OnValidate()
|
|
{
|
|
if (!UnityEditor.PrefabUtility.IsPartOfAnyPrefab(this))
|
|
{
|
|
ApplyMaskTexture(_maskTexture, _color);
|
|
ApplyThreshold(_threshold);
|
|
}
|
|
}
|
|
#endif
|
|
|
|
private void OnDestroy()
|
|
{
|
|
if (_maskBuffer != null)
|
|
{
|
|
_maskBuffer.Release();
|
|
}
|
|
}
|
|
|
|
public void Init(int width, int height)
|
|
{
|
|
_material = new Material(_maskShader)
|
|
{
|
|
renderQueue = (int)RenderQueue.Transparent
|
|
};
|
|
|
|
_material.SetTexture("_MainTex", _screen.texture);
|
|
ApplyMaskTexture(_maskTexture, _color);
|
|
_material.SetInt("_Width", width);
|
|
_material.SetInt("_Height", height);
|
|
ApplyThreshold(_threshold);
|
|
InitMaskBuffer(width, height);
|
|
}
|
|
|
|
public void Draw(float[] mask, int width, int height)
|
|
{
|
|
if (mask == null)
|
|
{
|
|
ApplyMaterial(_prevMaterial);
|
|
return;
|
|
}
|
|
|
|
if (mask.Length != width * height)
|
|
{
|
|
throw new ArgumentException("mask size must equal width * height");
|
|
}
|
|
|
|
ApplyMaterial(_material);
|
|
_maskBuffer.SetData(mask);
|
|
}
|
|
|
|
private Texture2D CreateMonoColorTexture(Color color)
|
|
{
|
|
var texture = new Texture2D(1, 1, TextureFormat.RGBA32, false);
|
|
var textureColor = new Color32((byte)(255 * color.r), (byte)(255 * color.g), (byte)(255 * color.b), (byte)(255 * color.a));
|
|
texture.SetPixels32(new Color32[] { textureColor });
|
|
texture.Apply();
|
|
|
|
return texture;
|
|
}
|
|
|
|
private void InitMaskBuffer(int width, int height)
|
|
{
|
|
if (_maskBuffer != null)
|
|
{
|
|
_maskBuffer.Release();
|
|
}
|
|
var stride = Marshal.SizeOf(typeof(float));
|
|
_maskBuffer = new GraphicsBuffer(GraphicsBuffer.Target.Structured, width * height, stride);
|
|
_material.SetBuffer("_MaskBuffer", _maskBuffer);
|
|
}
|
|
|
|
private void ApplyMaterial(Material material)
|
|
{
|
|
if (_prevMaterial == null)
|
|
{
|
|
// backup
|
|
_prevMaterial = _screen.material;
|
|
}
|
|
if (_screen.material != material)
|
|
{
|
|
_screen.material = material;
|
|
}
|
|
}
|
|
|
|
private void ApplyMaskTexture(Texture maskTexture, Color maskColor)
|
|
{
|
|
if (_material != null)
|
|
{
|
|
_material.SetTexture("_MaskTex", maskTexture == null ? CreateMonoColorTexture(maskColor) : maskTexture);
|
|
}
|
|
}
|
|
|
|
private void ApplyThreshold(float threshold)
|
|
{
|
|
if (_material != null)
|
|
{
|
|
_material.SetFloat("_Threshold", threshold);
|
|
}
|
|
}
|
|
}
|
|
}
|