Paul Bird commited on
Commit
3784e17
1 Parent(s): 80b735d

Upload RunYOLO.cs

Browse files
Files changed (1) hide show
  1. RunYOLO.cs +178 -0
RunYOLO.cs ADDED
@@ -0,0 +1,178 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ using System.Collections.Generic;
2
+ using Unity.Sentis;
3
+ using UnityEngine;
4
+ using UnityEngine.UI;
5
+ using UnityEngine.Video;
6
+
7
+ /*
8
+ * YOLO inference script
9
+ * =====================
10
+ *
11
+ * Place this script on the Main Camera.
12
+ *
13
+ * Place the yolov7-tiny.sentis file and a *.mp4 video file in the Assets/StreamingAssets folder
14
+ *
15
+ */
16
+
17
+
18
+ public class RunYOLO : MonoBehaviour
19
+ {
20
+ const string modelName = "yolov7-tiny.sentis";
21
+ const string videoName = "giraffes.mp4";
22
+ // Link the classes.txt here:
23
+ public TextAsset labelsAsset;
24
+ // Create a Raw Image of size 640x640 and link it here:
25
+ public RawImage displayImage;
26
+ // Link to a bounding box texture here:
27
+ public Sprite boxTexture;
28
+ // Link to the font for the labels:
29
+ public Font font;
30
+
31
+ private Transform displayLocation;
32
+ private Model model;
33
+ private IWorker engine;
34
+ private string[] labels;
35
+ private RenderTexture targetRT;
36
+ const BackendType backend = BackendType.GPUCompute;
37
+ private const int imageWidth = 640;
38
+ private const int imageHeight = 640;
39
+
40
+ private VideoPlayer video;
41
+
42
+ //bounding box data
43
+ public struct BoundingBox
44
+ {
45
+ public float centerX;
46
+ public float centerY;
47
+ public float width;
48
+ public float height;
49
+ public string label;
50
+ public float confidence;
51
+ }
52
+
53
+ void Start()
54
+ {
55
+ Application.targetFrameRate = 60;
56
+ Screen.orientation = ScreenOrientation.LandscapeLeft;
57
+
58
+ //Parse neural net labels
59
+ labels = labelsAsset.text.Split('\n');
60
+
61
+ //Load model
62
+ model = ModelLoader.Load(Application.streamingAssetsPath + "/" + modelName);
63
+
64
+ targetRT = new RenderTexture(imageWidth, imageHeight, 0);
65
+
66
+ //Create image to display video
67
+ displayLocation = displayImage.transform;
68
+
69
+ //Create engine to run model
70
+ engine = WorkerFactory.CreateWorker(backend, model);
71
+
72
+ SetupInput();
73
+ }
74
+ void SetupInput()
75
+ {
76
+ video = gameObject.AddComponent<VideoPlayer>();
77
+ video.renderMode = VideoRenderMode.APIOnly;
78
+ video.source = VideoSource.Url;
79
+ video.url = Application.streamingAssetsPath + "/" + videoName;
80
+ video.isLooping = true;
81
+ video.Play();
82
+ }
83
+
84
+ private void Update()
85
+ {
86
+ ExecuteML();
87
+
88
+ if (Input.GetKeyDown(KeyCode.Escape))
89
+ {
90
+ Application.Quit();
91
+ }
92
+ }
93
+
94
+ public void ExecuteML()
95
+ {
96
+ ClearAnnotations();
97
+
98
+ if (video && video.texture)
99
+ {
100
+ float aspect = video.width * 1f / video.height;
101
+ Graphics.Blit(video.texture, targetRT, new Vector2(1f / aspect, 1), new Vector2(0, 0));
102
+ displayImage.texture = targetRT;
103
+ }
104
+ else return;
105
+
106
+ using var input = TextureConverter.ToTensor(targetRT, imageWidth, imageHeight, 3);
107
+ engine.Execute(input);
108
+
109
+ //Read output tensors
110
+ var output = engine.PeekOutput() as TensorFloat;
111
+ output.MakeReadable();
112
+
113
+ //Draw the bounding boxes
114
+ for (int n = 0; n < output.shape[0]; n++)
115
+ {
116
+ var box = new BoundingBox
117
+ {
118
+ centerX = (output[n, 1] + output[n, 3]) / 2 - 320,
119
+ centerY = (output[n, 2] + output[n, 4]) / 2 - 320,
120
+ width = output[n, 3] - output[n, 1],
121
+ height = output[n, 4] - output[n, 2],
122
+ label = labels[(int)output[n, 5]],
123
+ confidence = Mathf.FloorToInt(output[n, 6] * 100 + 0.5f)
124
+ };
125
+ DrawBox(box);
126
+ }
127
+ }
128
+
129
+ public void DrawBox(BoundingBox box)
130
+ {
131
+ Color color = Color.yellow;
132
+
133
+ GameObject panel = new GameObject("ObjectBox");
134
+ panel.AddComponent<CanvasRenderer>();
135
+ Image img = panel.AddComponent<Image>();
136
+ img.color = color;
137
+ img.sprite = boxTexture;
138
+
139
+ panel.transform.SetParent(displayLocation, false);
140
+ panel.transform.localPosition = new Vector3(box.centerX, -box.centerY);
141
+
142
+ RectTransform rt = panel.GetComponent<RectTransform>();
143
+ rt.sizeDelta = new Vector2(box.width, box.height);
144
+
145
+ //add class label
146
+ var text = new GameObject("ObjectLabel");
147
+ text.AddComponent<CanvasRenderer>();
148
+ Text txt = text.AddComponent<Text>();
149
+ text.transform.SetParent(panel.transform, false);
150
+ txt.font = font;
151
+ txt.text = box.label + " (" + box.confidence + "%)";
152
+ txt.color = color;
153
+ txt.fontSize = 40;
154
+ txt.horizontalOverflow = HorizontalWrapMode.Overflow;
155
+ RectTransform rt2 = text.GetComponent<RectTransform>();
156
+ rt2.offsetMin = new Vector2(20, rt2.offsetMin.y);
157
+ rt2.offsetMax = new Vector2(0, rt2.offsetMax.y);
158
+ rt2.offsetMax = new Vector2(rt2.offsetMax.x, 30);
159
+ rt2.offsetMin = new Vector2(rt2.offsetMin.x, 0);
160
+ rt2.anchorMin = new Vector2(0,0);
161
+ rt2.anchorMax = new Vector2(1, 1);
162
+
163
+ img.sprite = boxTexture;
164
+ img.type = Image.Type.Sliced;
165
+ }
166
+
167
+ public void ClearAnnotations()
168
+ {
169
+ foreach (Transform child in displayLocation) {
170
+ Destroy(child.gameObject);
171
+ }
172
+ }
173
+
174
+ private void OnDestroy()
175
+ {
176
+ engine?.Dispose();
177
+ }
178
+ }