Merge pull request #8 from Chen-Yulin/yolo160

Yolo160

Assets/Scenes/SampleScene.unity

@@ -10962,14 +10962,13 @@ MonoBehaviour:
   m_Script: {fileID: 11500000, guid: aad934041a3d88c45b687ed811756753, type: 3}
   m_Name: 
   m_EditorClassIdentifier: 
-  ModelFile: {fileID: 5022602860645237092, guid: 28f8589cc7fa45140a5d20478601f140, type: 3}
+  ModelFile: {fileID: 5022602860645237092, guid: 4d9cc643a9bc424499c3f08397903820, type: 3}
   material_t: {fileID: 2100000, guid: d4c3396def303a84d81ff6a10a617c11, type: 2}
-  MinBoxConfidence: 0.2
+  MinBoxConfidence: 0.25
   textureProviderType: 0
   textureProvider:
     rid: 7844097659052228608
   holo: 1
-  cycle: 0.5
   references:
     version: 2
     RefIds:

Assets/Script/Detection/Detector.cs

@@ -42,20 +42,32 @@ public class Detector : MonoBehaviour
 
     YOLOv8 yolo;
 
-    float time = 0;
-    public float cycle = 0.2f;
 
-    bool detectionOn = false;
+    private bool detectionOn = false;
+
+    Texture2D CapturedTexture;
+    Texture2D DisplayTexture;
+    List<ResultBox> boxes = new List<ResultBox>();
 
     public void ToggleDetection()
     {
         detectionOn = !detectionOn;
+        if (detectionOn)
+        {
+            yolo.TurnOnYolo();
+        }
+        else
+        {
+            yolo.TurnOffYolo();
+        }
     }
 
     private void OnEnable()
     {
         nn = new NNHandler(ModelFile);
-        yolo = new YOLOv8Segmentation(nn);
+        //yolo = new YOLOv8Segmentation(nn);
+        yolo = gameObject.AddComponent<YOLOv8>();
+        yolo.InitYOLOv8(nn);
 
         textureProvider = GetTextureProvider(nn.model);
         textureProvider.Start();
@@ -65,22 +77,28 @@ private void Update()
     {
         if (detectionOn)
         {
-            time += Time.deltaTime;
-            if (time > cycle)
+            if (yolo.ResAvailable) // all data has been updated
             {
-                time = 0;
+                Debug.Log("update result "+yolo.Results.Count);
+                yolo.ResAvailable = false;
+                boxes = yolo.Results;
 
-                YOLOv8OutputReader.DiscardThreshold = MinBoxConfidence;
-                Texture2D texture = GetNextTexture();
+                DisplayTexture = new Texture2D(CapturedTexture.width, CapturedTexture.height, TextureFormat.RGB24, false);
+                Graphics.CopyTexture(CapturedTexture, DisplayTexture);
+                DrawResults(boxes, DisplayTexture);
+                material_t.mainTexture = DisplayTexture;
 
-                var boxes = yolo.Run(texture);
-                DrawResults(boxes, texture);
-                material_t.mainTexture = texture;
+                YOLOv8OutputReader.DiscardThreshold = MinBoxConfidence;
+                CapturedTexture = GetNextTexture();
+                yolo.SetSource(CapturedTexture);
+            }
+            else if (yolo.StartOfDetection)
+            {
+                yolo.StartOfDetection = false;
+                YOLOv8OutputReader.DiscardThreshold = MinBoxConfidence;
+                CapturedTexture = GetNextTexture();
+                yolo.SetSource(CapturedTexture);
             }
-        }
-        else
-        {
-            time = cycle / 2f;
         }
     }
 

Assets/Script/Detection/NN/YOLOv8.cs

@@ -1,68 +1,169 @@
-using System;
-using System.Collections.Generic;
-using System.Linq;
-using Unity.Barracuda;
-using UnityEngine;
-using UnityEngine.Profiling;
-
-namespace NN
-{
-    public class YOLOv8
-    {
-        protected YOLOv8OutputReader outputReader;
-
-        private NNHandler nn;
-
-        public YOLOv8(NNHandler nn)
-        {
-            this.nn = nn;
-            outputReader = new();
-        }
-
-        public List<ResultBox> Run(Texture2D image)
-        {
-            Profiler.BeginSample("YOLO.Run");
-            var outputs = ExecuteModel(image);
-            var results = Postprocess(outputs);
-            //var results = new List<ResultBox>();
-            Profiler.EndSample();
-            return results;
-        }
-
-        protected Tensor[] ExecuteModel(Texture2D image)
-        {
-            Tensor input = new Tensor(image);
-            ExecuteBlocking(input);
-            input.tensorOnDevice.Dispose();
-            return PeekOutputs().ToArray();
-        }
-
-        private void ExecuteBlocking(Tensor preprocessed)
-        {
-            Profiler.BeginSample("YOLO.Execute");
-            nn.worker.Execute(preprocessed);
-            nn.worker.FlushSchedule(blocking: true);
-            Profiler.EndSample();
-        }
-
-
-        private IEnumerable<Tensor> PeekOutputs()
-        {
-            foreach (string outputName in nn.model.outputs)
-            {
-                Tensor output = nn.worker.PeekOutput(outputName);
-                yield return output;
-            }
-        }
-
-        protected List<ResultBox> Postprocess(Tensor[] outputs)
-        {
-            Profiler.BeginSample("YOLOv8Postprocessor.Postprocess");
-            Tensor boxesOutput = outputs[0];
-            List<ResultBox> boxes = outputReader.ReadOutput(boxesOutput).ToList();
-            boxes = DuplicatesSupressor.RemoveDuplicats(boxes);
-            Profiler.EndSample();
-            return boxes;
-        }
-    }
+using System;
+using System.Collections;
+using System.Collections.Generic;
+using System.Linq;
+using System.Threading.Tasks;
+using Unity.Barracuda;
+using UnityEngine;
+using UnityEngine.Profiling;
+
+namespace NN
+{
+    public class YOLOv8 : MonoBehaviour
+    {
+        protected YOLOv8OutputReader outputReader;
+
+        protected NNHandler nn;
+
+        public bool workerbusy = false;
+
+        public bool jobcomplete = false;
+
+        public bool ResAvailable = false;
+
+        public bool StartOfDetection = false;
+
+        private List<ResultBox> results = new List<ResultBox>();
+
+        public List<ResultBox> Results
+        {
+            get { return results; }
+            set
+            {
+                if (value != null)
+                {
+                    results = value;
+                }
+            }
+        }
+
+        Tensor input;
+
+        Texture2D sourceImg;
+
+        bool on;
+
+        public void SetSource(Texture2D img)
+        {
+            sourceImg = img;
+        }
+
+        public void TurnOffYolo()
+        {
+            on = false;
+            workerbusy = false;
+            jobcomplete = false;
+            ResAvailable = false;
+            StartOfDetection = false;
+
+            StopAllCoroutines();
+            try
+            {
+                input.tensorOnDevice.Dispose();
+            }
+            catch { }
+        }
+
+        public void TurnOnYolo()
+        {
+            on = true;
+            StartOfDetection = true;
+        }
+
+        public YOLOv8(NNHandler nn)
+        {
+            this.nn = nn;
+            outputReader = new();
+        }
+        public virtual void InitYOLOv8(NNHandler nn)
+        {
+            this.nn = nn;
+            outputReader = new();
+        }
+
+        public void Run()
+        {
+            Results = Postprocess(ExecuteModel());
+        }
+        public void Update()
+        {
+            if (on)
+            {
+                Run();
+            }
+        }
+
+        protected Tensor[] ExecuteModel()
+        {
+            if (!workerbusy) // start of execution
+            {
+                if (jobcomplete) // complete
+                {
+                    jobcomplete = false;
+                    ResAvailable = true;
+                    input.tensorOnDevice.Dispose();
+                    return PeekOutputs().ToArray();
+                }
+                else // ready to start
+                {
+                    if (sourceImg)// start of job
+                    {
+                        input = new Tensor(sourceImg);
+                        // ExecuteBlocking(input);
+                        workerbusy = true;
+                        StartCoroutine(ExecuteUnblocking(input));
+                    }
+                    return null;
+                }
+            }
+            else //during execution
+            {
+                return null;
+            }
+        }
+
+        IEnumerator ExecuteUnblocking(Tensor preprocessed)
+        {
+            var it = nn.worker.StartManualSchedule(preprocessed);
+            workerbusy = true;
+            int cnt = 0;
+
+            while (it.MoveNext())
+            {
+                ++cnt;
+                if (cnt % 5 == 0)
+                {
+                    nn.worker.FlushSchedule(false);
+                    yield return null;
+                }
+            }
+
+            nn.worker.FlushSchedule(true);
+            workerbusy = false;
+            jobcomplete = true;
+        }
+
+        private IEnumerable<Tensor> PeekOutputs()
+        {
+            foreach (string outputName in nn.model.outputs)
+            {
+                Tensor output = nn.worker.PeekOutput(outputName);
+                yield return output;
+            }
+        }
+
+        protected List<ResultBox> Postprocess(Tensor[] outputs)
+        {
+            if (outputs == null)
+            {
+                return null;
+            }
+            Profiler.BeginSample("YOLOv8Postprocessor.Postprocess");
+            Tensor boxesOutput = outputs[0];
+            List<ResultBox> boxes = outputReader.ReadOutput(boxesOutput).ToList();
+            boxes = DuplicatesSupressor.RemoveDuplicats(boxes);
+            Profiler.EndSample();
+            return boxes;
+        }
+    }
 }

Assets/Script/Detection/NN/YOLOv8OutputReader.cs

@@ -11,9 +11,9 @@ public class YOLOv8OutputReader
     {
         public static float DiscardThreshold = 0.1f;
         protected const int ClassesNum = 80;
-        const int BoxesPerCell = 525;
-        const int InputWidth = 240;
-        const int InputHeight = 240;
+        const int BoxesPerCell = 2100;
+        const int InputWidth = 320;
+        const int InputHeight = 320;
 
         public IEnumerable<ResultBox> ReadOutput(Tensor output)
         {