Initial commit

29982f2 22 days ago

7.87 kB

// CED-Tiny (Consistent Ensemble Distillation ViT) AudioSet audio tagger. Self-contained Unity Sentis inference.
using System;
using System.Collections.Generic;
using System.IO;
using System.Threading.Tasks;
using Unity.InferenceEngine;
using UnityEngine;

namespace SentisModels
{
// Mono waveform -> 527 AudioSet class probabilities (multi-label: independent sigmoids, they do not sum to 1) and a
// 192-d clip embedding. The log-mel frontend (STFT as a fixed-basis Conv1d) is baked into the model, so the input is
// the raw 16000 Hz waveform in [-1, 1]; one inference covers a 3 s window. Other sample rates are resampled here.
public sealed class CedTiny : IDisposable
{
    public const int SampleRate = 16000;
    public const int WindowSamples = 48000;
    public const int ClassCount = 527;
    public const int EmbeddingSize = 192;
    const string DefaultModelFile = "ced_tiny_fp32.sentis";
    const string DefaultLabelFile = "audioset_labels.txt";

    public struct Prediction
    {
        public int Index;       // AudioSet class index (ontology order, class_labels_indices.csv)
        public string Label;
        public float Score;     // sigmoid probability
    }

    readonly BackendType m_BackendType;

    Worker m_Worker;
    Tensor<float> m_InputTensor;
    string[] m_Labels;
    readonly float[] m_Window = new float[WindowSamples];
    readonly float[] m_Stream = new float[WindowSamples];   // sliding buffer for Push / ClassifyLatest

    // Embedding of the most recent Classify call (penultimate features; use for your own classes / similarity).
    public float[] Embedding { get; } = new float[EmbeddingSize];
    public IReadOnlyList<string> Labels => m_Labels;

    public CedTiny(BackendType backendType = BackendType.GPUCompute)
    {
        m_BackendType = backendType == BackendType.GPUCompute && !SystemInfo.supportsComputeShaders
            ? BackendType.CPU
            : backendType;
    }

    public bool IsLoaded => m_Worker != null;

    // modelRoot: absolute path to the folder holding ced_tiny_fp32.sentis and audioset_labels.txt.
    public void Load(string modelRoot, string modelFile = DefaultModelFile)
    {
        Load(ModelLoader.Load(Path.Combine(modelRoot, modelFile)), File.ReadAllLines(Path.Combine(modelRoot, DefaultLabelFile)));
    }

    // For platforms where the files must be read as bytes / text (e.g. Android StreamingAssets).
    public void Load(byte[] modelBytes, string labelsText)
    {
        using var stream = new MemoryStream(modelBytes);
        Load(ModelLoader.Load(stream), labelsText.Split(new[] { "\r\n", "\n" }, StringSplitOptions.RemoveEmptyEntries));
    }

    void Load(Model model, string[] labels)
    {
        if (labels.Length != ClassCount)
            throw new ArgumentException($"expected {ClassCount} labels, got {labels.Length}");
        DisposeResources();
        m_Labels = labels;
        m_Worker = new Worker(model, m_BackendType);
        m_InputTensor = new Tensor<float>(new TensorShape(1, WindowSamples));
        Array.Clear(m_Stream, 0, m_Stream.Length);
    }

    // One silent pass so the first real call doesn't pay shader compilation / allocation.
    public void Warmup()
    {
        Classify(new float[WindowSamples]);
    }

    // samples: mono PCM at sampleRate. The LAST 3 s are classified (zero-padded in front when shorter).
    public float[] Classify(float[] samples, int sampleRate = SampleRate)
    {
        Schedule(samples, sampleRate);
        return Read();
    }

    // Async readback variant; prefer this for per-frame use on GPU backends (no main-thread stall).
    public async Task<float[]> ClassifyAsync(float[] samples, int sampleRate = SampleRate)
    {
        Schedule(samples, sampleRate);
        using var scores = await (m_Worker.PeekOutput("scores") as Tensor<float>).ReadbackAndCloneAsync();
        using var embedding = await (m_Worker.PeekOutput("embedding") as Tensor<float>).ReadbackAndCloneAsync();
        embedding.DownloadToArray().CopyTo(Embedding, 0);
        return scores.DownloadToArray();
    }

    // Streaming: append new microphone / clip samples (mono, any rate) to the sliding window.
    public void Push(float[] samples, int sampleRate = SampleRate)
    {
        var x = sampleRate == SampleRate ? samples : Resample(samples, sampleRate, SampleRate);
        if (x.Length >= WindowSamples)
        {
            Array.Copy(x, x.Length - WindowSamples, m_Stream, 0, WindowSamples);
            return;
        }
        Array.Copy(m_Stream, x.Length, m_Stream, 0, WindowSamples - x.Length);
        Array.Copy(x, 0, m_Stream, WindowSamples - x.Length, x.Length);
    }

    // Classifies the last 3 s pushed with Push (call e.g. every 0.25-0.5 s for a sliding window).
    public float[] ClassifyLatest() => Classify(m_Stream);

    public Task<float[]> ClassifyLatestAsync() => ClassifyAsync(m_Stream);

    // Highest-scoring classes, best first.
    public List<Prediction> Top(float[] scores, int k = 5, float minScore = 0f)
    {
        var order = new int[scores.Length];
        for (var i = 0; i < order.Length; i++)
            order[i] = i;
        Array.Sort(order, (a, b) => scores[b].CompareTo(scores[a]));
        var result = new List<Prediction>();
        for (var i = 0; i < order.Length && result.Count < k; i++)
        {
            if (scores[order[i]] < minScore)
                break;
            result.Add(new Prediction { Index = order[i], Label = m_Labels[order[i]], Score = scores[order[i]] });
        }
        return result;
    }

    // Average the channels of interleaved PCM (AudioClip.GetData / Microphone data).
    public static float[] ToMono(float[] interleaved, int channels)
    {
        if (channels <= 1)
            return interleaved;
        var mono = new float[interleaved.Length / channels];
        for (var i = 0; i < mono.Length; i++)
        {
            var sum = 0f;
            for (var c = 0; c < channels; c++)
                sum += interleaved[i * channels + c];
            mono[i] = sum / channels;
        }
        return mono;
    }

    // Linear-interpolation resampler (adequate for tagging; use a band-limited resampler for best fidelity).
    public static float[] Resample(float[] src, int srcRate, int dstRate)
    {
        if (srcRate == dstRate || src.Length == 0)
            return src;
        var n = (int)((long)src.Length * dstRate / srcRate);
        var dst = new float[n];
        var step = (double)srcRate / dstRate;
        for (var i = 0; i < n; i++)
        {
            var p = i * step;
            var j = (int)p;
            var t = (float)(p - j);
            dst[i] = j + 1 < src.Length ? src[j] + (src[j + 1] - src[j]) * t : src[src.Length - 1];
        }
        return dst;
    }

    void Schedule(float[] samples, int sampleRate)
    {
        if (m_Worker == null)
            throw new InvalidOperationException("CedTiny.Load must be called before Classify.");
        var x = sampleRate == SampleRate ? samples : Resample(samples, sampleRate, SampleRate);
        var n = Math.Min(x.Length, WindowSamples);
        Array.Clear(m_Window, 0, WindowSamples - n);
        Array.Copy(x, x.Length - n, m_Window, WindowSamples - n, n);
        m_InputTensor.Upload(m_Window);
        m_Worker.Schedule(m_InputTensor);
    }

    float[] Read()
    {
        using var scores = (m_Worker.PeekOutput("scores") as Tensor<float>).ReadbackAndClone();
        using var embedding = (m_Worker.PeekOutput("embedding") as Tensor<float>).ReadbackAndClone();
        embedding.DownloadToArray().CopyTo(Embedding, 0);
        return scores.DownloadToArray();
    }

    void DisposeResources()
    {
        m_Worker?.Dispose();
        m_Worker = null;
        m_InputTensor?.Dispose();
        m_InputTensor = null;
    }

    public void Dispose()
    {
        DisposeResources();
    }
}
}

	// CED-Tiny (Consistent Ensemble Distillation ViT) AudioSet audio tagger. Self-contained Unity Sentis inference.
	using System;
	using System.Collections.Generic;
	using System.IO;
	using System.Threading.Tasks;
	using Unity.InferenceEngine;
	using UnityEngine;

	namespace SentisModels
	{
	// Mono waveform -> 527 AudioSet class probabilities (multi-label: independent sigmoids, they do not sum to 1) and a
	// 192-d clip embedding. The log-mel frontend (STFT as a fixed-basis Conv1d) is baked into the model, so the input is
	// the raw 16000 Hz waveform in [-1, 1]; one inference covers a 3 s window. Other sample rates are resampled here.
	public sealed class CedTiny : IDisposable
	{
	public const int SampleRate = 16000;
	public const int WindowSamples = 48000;
	public const int ClassCount = 527;
	public const int EmbeddingSize = 192;
	const string DefaultModelFile = "ced_tiny_fp32.sentis";
	const string DefaultLabelFile = "audioset_labels.txt";

	public struct Prediction
	{
	public int Index; // AudioSet class index (ontology order, class_labels_indices.csv)
	public string Label;
	public float Score; // sigmoid probability
	}

	readonly BackendType m_BackendType;

	Worker m_Worker;
	Tensor<float> m_InputTensor;
	string[] m_Labels;
	readonly float[] m_Window = new float[WindowSamples];
	readonly float[] m_Stream = new float[WindowSamples]; // sliding buffer for Push / ClassifyLatest

	// Embedding of the most recent Classify call (penultimate features; use for your own classes / similarity).
	public float[] Embedding { get; } = new float[EmbeddingSize];
	public IReadOnlyList<string> Labels => m_Labels;

	public CedTiny(BackendType backendType = BackendType.GPUCompute)
	{
	m_BackendType = backendType == BackendType.GPUCompute && !SystemInfo.supportsComputeShaders
	? BackendType.CPU
	: backendType;
	}

	public bool IsLoaded => m_Worker != null;

	// modelRoot: absolute path to the folder holding ced_tiny_fp32.sentis and audioset_labels.txt.
	public void Load(string modelRoot, string modelFile = DefaultModelFile)
	{
	Load(ModelLoader.Load(Path.Combine(modelRoot, modelFile)), File.ReadAllLines(Path.Combine(modelRoot, DefaultLabelFile)));
	}

	// For platforms where the files must be read as bytes / text (e.g. Android StreamingAssets).
	public void Load(byte[] modelBytes, string labelsText)
	{
	using var stream = new MemoryStream(modelBytes);
	Load(ModelLoader.Load(stream), labelsText.Split(new[] { "\r\n", "\n" }, StringSplitOptions.RemoveEmptyEntries));
	}

	void Load(Model model, string[] labels)
	{
	if (labels.Length != ClassCount)
	throw new ArgumentException($"expected {ClassCount} labels, got {labels.Length}");
	DisposeResources();
	m_Labels = labels;
	m_Worker = new Worker(model, m_BackendType);
	m_InputTensor = new Tensor<float>(new TensorShape(1, WindowSamples));
	Array.Clear(m_Stream, 0, m_Stream.Length);
	}

	// One silent pass so the first real call doesn't pay shader compilation / allocation.
	public void Warmup()
	{
	Classify(new float[WindowSamples]);
	}

	// samples: mono PCM at sampleRate. The LAST 3 s are classified (zero-padded in front when shorter).
	public float[] Classify(float[] samples, int sampleRate = SampleRate)
	{
	Schedule(samples, sampleRate);
	return Read();
	}

	// Async readback variant; prefer this for per-frame use on GPU backends (no main-thread stall).
	public async Task<float[]> ClassifyAsync(float[] samples, int sampleRate = SampleRate)
	{
	Schedule(samples, sampleRate);
	using var scores = await (m_Worker.PeekOutput("scores") as Tensor<float>).ReadbackAndCloneAsync();
	using var embedding = await (m_Worker.PeekOutput("embedding") as Tensor<float>).ReadbackAndCloneAsync();
	embedding.DownloadToArray().CopyTo(Embedding, 0);
	return scores.DownloadToArray();
	}

	// Streaming: append new microphone / clip samples (mono, any rate) to the sliding window.
	public void Push(float[] samples, int sampleRate = SampleRate)
	{
	var x = sampleRate == SampleRate ? samples : Resample(samples, sampleRate, SampleRate);
	if (x.Length >= WindowSamples)
	{
	Array.Copy(x, x.Length - WindowSamples, m_Stream, 0, WindowSamples);
	return;
	}
	Array.Copy(m_Stream, x.Length, m_Stream, 0, WindowSamples - x.Length);
	Array.Copy(x, 0, m_Stream, WindowSamples - x.Length, x.Length);
	}

	// Classifies the last 3 s pushed with Push (call e.g. every 0.25-0.5 s for a sliding window).
	public float[] ClassifyLatest() => Classify(m_Stream);

	public Task<float[]> ClassifyLatestAsync() => ClassifyAsync(m_Stream);

	// Highest-scoring classes, best first.
	public List<Prediction> Top(float[] scores, int k = 5, float minScore = 0f)
	{
	var order = new int[scores.Length];
	for (var i = 0; i < order.Length; i++)
	order[i] = i;
	Array.Sort(order, (a, b) => scores[b].CompareTo(scores[a]));
	var result = new List<Prediction>();
	for (var i = 0; i < order.Length && result.Count < k; i++)
	{
	if (scores[order[i]] < minScore)
	break;
	result.Add(new Prediction { Index = order[i], Label = m_Labels[order[i]], Score = scores[order[i]] });
	}
	return result;
	}

	// Average the channels of interleaved PCM (AudioClip.GetData / Microphone data).
	public static float[] ToMono(float[] interleaved, int channels)
	{
	if (channels <= 1)
	return interleaved;
	var mono = new float[interleaved.Length / channels];
	for (var i = 0; i < mono.Length; i++)
	{
	var sum = 0f;
	for (var c = 0; c < channels; c++)
	sum += interleaved[i * channels + c];
	mono[i] = sum / channels;
	}
	return mono;
	}

	// Linear-interpolation resampler (adequate for tagging; use a band-limited resampler for best fidelity).
	public static float[] Resample(float[] src, int srcRate, int dstRate)
	{
	if (srcRate == dstRate \|\| src.Length == 0)
	return src;
	var n = (int)((long)src.Length * dstRate / srcRate);
	var dst = new float[n];
	var step = (double)srcRate / dstRate;
	for (var i = 0; i < n; i++)
	{
	var p = i * step;
	var j = (int)p;
	var t = (float)(p - j);
	dst[i] = j + 1 < src.Length ? src[j] + (src[j + 1] - src[j]) * t : src[src.Length - 1];
	}
	return dst;
	}

	void Schedule(float[] samples, int sampleRate)
	{
	if (m_Worker == null)
	throw new InvalidOperationException("CedTiny.Load must be called before Classify.");
	var x = sampleRate == SampleRate ? samples : Resample(samples, sampleRate, SampleRate);
	var n = Math.Min(x.Length, WindowSamples);
	Array.Clear(m_Window, 0, WindowSamples - n);
	Array.Copy(x, x.Length - n, m_Window, WindowSamples - n, n);
	m_InputTensor.Upload(m_Window);
	m_Worker.Schedule(m_InputTensor);
	}

	float[] Read()
	{
	using var scores = (m_Worker.PeekOutput("scores") as Tensor<float>).ReadbackAndClone();
	using var embedding = (m_Worker.PeekOutput("embedding") as Tensor<float>).ReadbackAndClone();
	embedding.DownloadToArray().CopyTo(Embedding, 0);
	return scores.DownloadToArray();
	}

	void DisposeResources()
	{
	m_Worker?.Dispose();
	m_Worker = null;
	m_InputTensor?.Dispose();
	m_InputTensor = null;
	}

	public void Dispose()
	{
	DisposeResources();
	}
	}
	}