feat. finish yolo

Chen-Yulin · Apr 13, 2024 · e42be3e · e42be3e
1 parent a6ac87b
commit e42be3e
Show file tree

Hide file tree

Showing 3 changed files with 213 additions and 200 deletions.
diff --git a/Assets/Scenes/SampleScene.unity b/Assets/Scenes/SampleScene.unity
@@ -3435,7 +3435,7 @@ GameObject:
   m_Icon: {fileID: 0}
   m_NavMeshLayer: 0
   m_StaticEditorFlags: 0
-  m_IsActive: 0
+  m_IsActive: 1
 --- !u!4 &245231320
 Transform:
   m_ObjectHideFlags: 0
@@ -11091,7 +11091,7 @@ GameObject:
   - component: {fileID: 800842345}
   - component: {fileID: 800842344}
   m_Layer: 0
-  m_Name: RoomMeshToggle
+  m_Name: DetectionToggle
   m_TagString: Untagged
   m_Icon: {fileID: 0}
   m_NavMeshLayer: 0
@@ -11298,9 +11298,9 @@ MonoBehaviour:
   OnClick:
     m_PersistentCalls:
       m_Calls:
-      - m_Target: {fileID: 0}
-        m_TargetAssemblyTypeName: ToggleRoomMesh, Assembly-CSharp
-        m_MethodName: OnButtonPress
+      - m_Target: {fileID: 782628376}
+        m_TargetAssemblyTypeName: Detector, Assembly-CSharp
+        m_MethodName: ToggleDetection
         m_Mode: 1
         m_Arguments:
           m_ObjectArgument: {fileID: 0}
@@ -12980,7 +12980,7 @@ GameObject:
   m_Icon: {fileID: 0}
   m_NavMeshLayer: 0
   m_StaticEditorFlags: 0
-  m_IsActive: 1
+  m_IsActive: 0
 --- !u!4 &936113383
 Transform:
   m_ObjectHideFlags: 0

diff --git a/Assets/Script/Detection/Detector.cs b/Assets/Script/Detection/Detector.cs
@@ -1,141 +1,154 @@
-using Assets.Scripts;
-using Assets.Scripts.TextureProviders;
-using NN;
-using System;
-using System.Collections.Generic;
-using System.ComponentModel;
-using System.Runtime.CompilerServices;
-using System.Runtime.Serialization;
-using Unity.Barracuda;
-using UnityEditor;
-using UnityEngine;
-using UnityEngine.Profiling;
-using UnityEngine.UI;
-
-public class Detector : MonoBehaviour
-{
-    [Tooltip("File of YOLO model.")]
-    [SerializeField]
-    protected NNModel ModelFile;
-
-    [Tooltip("RawImage component which will be used to draw resuls.")]
-    [SerializeField]
-    protected Material material_t;
-
-    [Range(0.0f, 1f)]
-    [Tooltip("The minimum value of box confidence below which boxes won't be drawn.")]
-    [SerializeField]
-    protected float MinBoxConfidence = 0.3f;
-
-    [SerializeField]
-    protected TextureProviderType.ProviderType textureProviderType;
-
-    [SerializeReference]
-    protected TextureProvider textureProvider = null;
-
-    [Tooltip("deploy on hololens")]
-    [SerializeField]
-    protected bool holo = false;
-
-    protected NNHandler nn;
-    protected Color[] colorArray = new Color[] { Color.red, Color.green, Color.blue, Color.cyan, Color.magenta, Color.yellow };
-
-    YOLOv8 yolo;
-
-    float time = 0;
-    float cycle = 0.5f;
-
-    private void OnEnable()
-    {
-        nn = new NNHandler(ModelFile);
-        yolo = new YOLOv8Segmentation(nn);
-
-        textureProvider = GetTextureProvider(nn.model);
-        textureProvider.Start();
-    }
-
-    private void Update()
-    {
-        time += Time.deltaTime;
-        if (time > cycle)
-        {
-            time = 0;
-
-            YOLOv8OutputReader.DiscardThreshold = MinBoxConfidence;
-            Texture2D texture = GetNextTexture();
-
-            var boxes = yolo.Run(texture);
-            DrawResults(boxes, texture);
-            material_t.mainTexture = texture;
-        }
-
-    }
-
-    protected TextureProvider GetTextureProvider(Model model)
-    {
-        var firstInput = model.inputs[0];
-        int height = firstInput.shape[5];
-        int width = firstInput.shape[6];
-
-        TextureProvider provider;
-        switch (textureProviderType)
-        {
-            case TextureProviderType.ProviderType.WebCam:
-                if (holo)
-                {
-                    provider = new WebCamTextureProvider(textureProvider as WebCamTextureProvider, width, height, true);
+using Assets.Scripts;
+using Assets.Scripts.TextureProviders;
+using NN;
+using System;
+using System.Collections.Generic;
+using System.ComponentModel;
+using System.Runtime.CompilerServices;
+using System.Runtime.Serialization;
+using Unity.Barracuda;
+using UnityEditor;
+using UnityEngine;
+using UnityEngine.Profiling;
+using UnityEngine.UI;
+
+public class Detector : MonoBehaviour
+{
+    [Tooltip("File of YOLO model.")]
+    [SerializeField]
+    protected NNModel ModelFile;
+
+    [Tooltip("RawImage component which will be used to draw resuls.")]
+    [SerializeField]
+    protected Material material_t;
+
+    [Range(0.0f, 1f)]
+    [Tooltip("The minimum value of box confidence below which boxes won't be drawn.")]
+    [SerializeField]
+    protected float MinBoxConfidence = 0.3f;
+
+    [SerializeField]
+    protected TextureProviderType.ProviderType textureProviderType;
+
+    [SerializeReference]
+    protected TextureProvider textureProvider = null;
+
+    [Tooltip("deploy on hololens")]
+    [SerializeField]
+    protected bool holo = false;
+
+    protected NNHandler nn;
+    protected Color[] colorArray = new Color[] { Color.red, Color.green, Color.blue, Color.cyan, Color.magenta, Color.yellow };
+
+    YOLOv8 yolo;
+
+    float time = 0;
+    float cycle = 2f;
+
+    bool detectionOn = false;
+
+    public void ToggleDetection()
+    {
+        detectionOn = !detectionOn;
+    }
+
+    private void OnEnable()
+    {
+        nn = new NNHandler(ModelFile);
+        yolo = new YOLOv8Segmentation(nn);
+
+        textureProvider = GetTextureProvider(nn.model);
+        textureProvider.Start();
+    }
+
+    private void Update()
+    {
+        if (detectionOn)
+        {
+            time += Time.deltaTime;
+            if (time > cycle)
+            {
+                time = 0;
+
+                YOLOv8OutputReader.DiscardThreshold = MinBoxConfidence;
+                Texture2D texture = GetNextTexture();
+
+                var boxes = yolo.Run(texture);
+                DrawResults(boxes, texture);
+                material_t.mainTexture = texture;
+            }
+        }
+        else
+        {
+            time = cycle / 2f;
+        }
+    }
+
+    protected TextureProvider GetTextureProvider(Model model)
+    {
+        var firstInput = model.inputs[0];
+        int height = firstInput.shape[5];
+        int width = firstInput.shape[6];
+
+        TextureProvider provider;
+        switch (textureProviderType)
+        {
+            case TextureProviderType.ProviderType.WebCam:
+                if (holo)
+                {
+                    provider = new WebCamTextureProvider(textureProvider as WebCamTextureProvider, width, height, true);
                 }
                 else
                 {
                     provider = new WebCamTextureProvider(textureProvider as WebCamTextureProvider, width, height, false);
-                }
-                break;
-
-            case TextureProviderType.ProviderType.Video:
-                provider = new VideoTextureProvider(textureProvider as VideoTextureProvider, width, height);
-                break;
-            default:
-                throw new InvalidEnumArgumentException();
-        }
-        return provider;
-    }
-
-    protected Texture2D GetNextTexture()
-    {
-        return textureProvider.GetTexture();
-    }
-
-    void OnDisable()
-    {
-        nn.Dispose();
-        textureProvider.Stop();
-    }
-
-    protected void DrawResults(IEnumerable<ResultBox> results, Texture2D img)
-    {
-        results.ForEach(box => DrawBox(box, img));
-    }
-
-    protected virtual void DrawBox(ResultBox box, Texture2D img)
-    {
-        Color boxColor = colorArray[box.bestClassIndex % colorArray.Length];
-        int boxWidth = (int)(box.score / MinBoxConfidence);
-        TextureTools.DrawRectOutline(img, box.rect, boxColor, boxWidth, rectIsNormalized: false, revertY: true);
-    }
-
-    private void OnValidate()
-    {
-        Type t = TextureProviderType.GetProviderType(textureProviderType);
-        if (textureProvider == null || t != textureProvider.GetType())
-        {
-            if (nn == null)
-                textureProvider = RuntimeHelpers.GetUninitializedObject(t) as TextureProvider;
-            else
-            {
-                textureProvider = GetTextureProvider(nn.model);
-                textureProvider.Start();
-            }
-
-        }
-    }
-}
+                }
+                break;
+
+            case TextureProviderType.ProviderType.Video:
+                provider = new VideoTextureProvider(textureProvider as VideoTextureProvider, width, height);
+                break;
+            default:
+                throw new InvalidEnumArgumentException();
+        }
+        return provider;
+    }
+
+    protected Texture2D GetNextTexture()
+    {
+        return textureProvider.GetTexture();
+    }
+
+    void OnDisable()
+    {
+        nn.Dispose();
+        textureProvider.Stop();
+    }
+
+    protected void DrawResults(IEnumerable<ResultBox> results, Texture2D img)
+    {
+        results.ForEach(box => DrawBox(box, img));
+    }
+
+    protected virtual void DrawBox(ResultBox box, Texture2D img)
+    {
+        Color boxColor = colorArray[box.bestClassIndex % colorArray.Length];
+        int boxWidth = (int)(box.score / MinBoxConfidence);
+        TextureTools.DrawRectOutline(img, box.rect, boxColor, boxWidth, rectIsNormalized: false, revertY: true);
+    }
+
+    private void OnValidate()
+    {
+        Type t = TextureProviderType.GetProviderType(textureProviderType);
+        if (textureProvider == null || t != textureProvider.GetType())
+        {
+            if (nn == null)
+                textureProvider = RuntimeHelpers.GetUninitializedObject(t) as TextureProvider;
+            else
+            {
+                textureProvider = GetTextureProvider(nn.model);
+                textureProvider.Start();
+            }
+
+        }
+    }
+}