Sky-Kim's picture
Initial commit
2e7837a
Raw
History Blame Contribute Delete
3.14 kB
using System.Collections.Generic;
using System.Threading.Tasks;
using Unity.InferenceEngine;
using UnityEngine;
using SentisModels;
namespace OnDeviceAgent.Inference
{
public abstract class YoloDetector : MonoBehaviour
{
[SerializeField, Tooltip("Optional StreamingAssets-relative model override; blank uses the subclass default.")]
string m_ModelPath = string.Empty;
[SerializeField] BackendType m_BackendType = BackendType.GPUCompute;
[SerializeField, Range(0.05f, 0.95f)] float m_ConfidenceThreshold = 0.25f;
[SerializeField, Tooltip("Log the raw model output shapes on the first inference (diagnostic).")]
bool m_LogOutputShapeOnce = true;
[SerializeField, Tooltip("Run one dummy inference when the model loads so the GPUCompute shaders " +
"compile up front; otherwise the first real detection stalls while every kernel is compiled.")]
bool m_WarmupOnEnable = true;
[SerializeField, Tooltip("Swap R<->B before inference. YOLOX is typically trained on BGR while Sentis " +
"ToTensor produces RGB. If detections are wrong (false positives for unrelated classes), toggle this.")]
bool m_SwapRedBlue = true;
protected abstract string DefaultModelPath { get; }
YoloxDetector m_Core;
public struct Detection
{
public int ClassId;
public string ClassName;
public float Confidence;
public Rect BoxXyxy;
}
public bool IsReady => m_Core != null && m_Core.IsReady;
void OnEnable()
{
LoadModel();
}
void OnDisable()
{
m_Core?.Dispose();
m_Core = null;
}
void LoadModel()
{
m_Core?.Dispose();
m_Core = new YoloxDetector(
m_BackendType,
m_ConfidenceThreshold,
m_SwapRedBlue,
m_WarmupOnEnable,
m_LogOutputShapeOnce);
try
{
var root = ModelRootProvisioner.EnsureModelRoot("Model/YOLO");
m_Core.Load(root, ModelFileRelativeToRoot());
}
catch (System.Exception exception)
{
Debug.LogException(exception);
}
}
string GetModelPath()
{
return string.IsNullOrWhiteSpace(m_ModelPath) ? DefaultModelPath : m_ModelPath.Trim();
}
// Strips the "Model/YOLO/" package-root prefix from the configured path to get the file name the package
// Load expects relative to its modelRoot.
string ModelFileRelativeToRoot()
{
const string prefix = "Model/YOLO/";
var path = GetModelPath().Replace('\\', '/');
return path.StartsWith(prefix, System.StringComparison.Ordinal) ? path.Substring(prefix.Length) : path;
}
public async Task<List<Detection>> DetectAsync(Texture frame)
{
if (m_Core == null)
return new List<Detection>();
var raw = await m_Core.DetectAsync(frame);
var result = new List<Detection>(raw.Count);
foreach (var d in raw)
result.Add(new Detection { ClassId = d.ClassId, ClassName = d.ClassName, Confidence = d.Confidence, BoxXyxy = d.BoxXyxy });
return result;
}
}
}