Upload RunYOLO8n.cs

RunYOLO8n.cs

@@ -108,23 +108,28 @@ public class RunYOLO8n : MonoBehaviour
         model.AddConstant(new Lays.Constant("0", new int[] { 0 }));
         model.AddConstant(new Lays.Constant("1", new int[] { 1 }));
         model.AddConstant(new Lays.Constant("4", new int[] { 4 }));
+
+
         model.AddConstant(new Lays.Constant("classes_plus_4", new int[] { numClasses + 4 }));
         model.AddConstant(new Lays.Constant("maxOutputBoxes", new int[] { maxOutputBoxes }));
         model.AddConstant(new Lays.Constant("iouThreshold", new float[] { iouThreshold }));
         model.AddConstant(new Lays.Constant("scoreThreshold", new float[] { scoreThreshold }));
        
         //Add layers
-        model.AddLayer(new Lays.Slice("boxCoords0", "output0", "0", "4", "1"));
+        model.AddLayer(new Lays.Slice("boxCoords0", "output0", "0", "4", "1")); 
         model.AddLayer(new Lays.Transpose("boxCoords", "boxCoords0", new int[] { 0, 2, 1 }));
-        model.AddLayer(new Lays.Slice("scores0", "output0", "4", "classes_plus_4", "1"));
+        model.AddLayer(new Lays.Slice("scores0", "output0", "4", "classes_plus_4", "1")); 
+        model.AddLayer(new Lays.ReduceMax("scores", new[] { "scores0", "1" }));
+        model.AddLayer(new Lays.ArgMax("classIDs", "scores0", 1));
 
-        model.AddLayer(new Lays.NonMaxSuppression("NMS", "boxCoords", "scores0",
+        model.AddLayer(new Lays.NonMaxSuppression("NMS", "boxCoords", "scores",
             "maxOutputBoxes", "iouThreshold", "scoreThreshold",
             centerPointBox: Lays.CenterPointBox.Center
         ));
 
         model.outputs.Clear();
         model.AddOutput("boxCoords");
+        model.AddOutput("classIDs");
         model.AddOutput("NMS");
     }
 
@@ -165,12 +170,15 @@ public class RunYOLO8n : MonoBehaviour
 
         var boxCoords = engine.PeekOutput("boxCoords") as TensorFloat;
         var NMS = engine.PeekOutput("NMS") as TensorInt;
+        var classIDs = engine.PeekOutput("classIDs") as TensorInt;
 
         using var boxIDs = ops.Slice(NMS, new int[] { 2 }, new int[] { 3 }, new int[] { 1 }, new int[] { 1 });
         using var boxIDsFlat = boxIDs.ShallowReshape(new TensorShape(boxIDs.shape.length)) as TensorInt;
         using var output = ops.Gather(boxCoords, boxIDsFlat, 1);
+        using var labelIDs = ops.Gather(classIDs, boxIDsFlat, 2);
         
         output.MakeReadable();
+        labelIDs.MakeReadable();
         NMS.MakeReadable();
 
         float displayWidth = displayImage.rectTransform.rect.width;
@@ -188,7 +196,7 @@ public class RunYOLO8n : MonoBehaviour
                 centerY = output[0, n, 1] * scaleY - displayHeight / 2,
                 width = output[0, n, 2] * scaleX,
                 height = output[0, n, 3] * scaleY,
-                label = labels[(int)NMS[n, 1]]
+                label = labels[labelIDs[0, 0,n]],
             };
             DrawBox(box, n);
         }