File size: 6,608 Bytes
93c60d5 8fd532f 93c60d5 8fd532f 8a51653 8fd532f 905e832 8fd532f 93c60d5 8fd532f 905e832 8fd532f 92c260e 8fd532f 93c60d5 8fd532f 93c60d5 8fd532f 93c60d5 8fd532f 93c60d5 8fd532f 93c60d5 8fd532f 905e832 93c60d5 8fd532f 93c60d5 8fd532f 93c60d5 8fd532f 8284277 8fd532f 92c260e 8fd532f 93c60d5 8fd532f 92c260e 93c60d5 92c260e 8fd532f 8284277 8fd532f 93c60d5 92c260e 8fd532f 92c260e 8fd532f 92c260e 8fd532f 92c260e 8fd532f 92c260e 8fd532f 92c260e 8fd532f 93c60d5 92c260e 8fd532f 93c60d5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 |
using System.Collections.Generic;
using Unity.Sentis;
using UnityEngine;
using UnityEngine.UI;
using UnityEngine.Video;
/*
* YOLO inference script
* =====================
*
* Place this script on the Main Camera.
*
* Place the yolov7-tiny.sentis file and a *.mp4 video file in the Assets/StreamingAssets folder
*
*/
public class RunYOLO : MonoBehaviour
{
public ModelAsset modelAsset;
const string modelName = "yolov7-tiny.sentis";
// Change this to the name of the video you put in StreamingAssets folder:
const string videoName = "giraffes.mp4";
// Link the classes.txt here:
public TextAsset labelsAsset;
// Create a Raw Image in the scene and link it here:
public RawImage displayImage;
// Link to a bounding box texture here:
public Sprite borderSprite;
public Texture2D borderTexture;
// Link to the font for the labels:
public Font font;
private Transform displayLocation;
private Model model;
private IWorker engine;
private string[] labels;
private RenderTexture targetRT;
const BackendType backend = BackendType.GPUCompute;
//Image size for the model
private const int imageWidth = 640;
private const int imageHeight = 640;
private VideoPlayer video;
List<GameObject> boxPool = new List<GameObject>();
//bounding box data
public struct BoundingBox
{
public float centerX;
public float centerY;
public float width;
public float height;
public string label;
public float confidence;
}
void Start()
{
Application.targetFrameRate = 60;
Screen.orientation = ScreenOrientation.LandscapeLeft;
//Parse neural net labels
labels = labelsAsset.text.Split('\n');
//Load model
model = ModelLoader.Load(modelAsset);
//model = ModelLoader.Load(Application.streamingAssetsPath + "/" + modelName);
targetRT = new RenderTexture(imageWidth, imageHeight, 0);
//Create image to display video
displayLocation = displayImage.transform;
//Create engine to run model
engine = WorkerFactory.CreateWorker(backend, model);
SetupInput();
if (borderSprite == null)
{
borderSprite = Sprite.Create(borderTexture, new Rect(0, 0, borderTexture.width, borderTexture.height), new Vector2(borderTexture.width / 2, borderTexture.height / 2));
}
}
void SetupInput()
{
video = gameObject.AddComponent<VideoPlayer>();
video.renderMode = VideoRenderMode.APIOnly;
video.source = VideoSource.Url;
video.url = Application.streamingAssetsPath + "/" + videoName;
video.isLooping = true;
video.Play();
}
private void Update()
{
ExecuteML();
if (Input.GetKeyDown(KeyCode.Escape))
{
Application.Quit();
}
}
public void ExecuteML()
{
ClearAnnotations();
if (video && video.texture)
{
float aspect = video.width * 1f / video.height;
Graphics.Blit(video.texture, targetRT, new Vector2(1f / aspect, 1), new Vector2(0, 0));
displayImage.texture = targetRT;
}
else return;
using var input = TextureConverter.ToTensor(targetRT, imageWidth, imageHeight, 3);
engine.Execute(input);
//Read output tensors
var output = engine.PeekOutput() as TensorFloat;
output.CompleteOperationsAndDownload();
float displayWidth = displayImage.rectTransform.rect.width;
float displayHeight = displayImage.rectTransform.rect.height;
float scaleX = displayWidth / imageWidth;
float scaleY = displayHeight / imageHeight;
int foundBoxes = output.shape[0];
//Draw the bounding boxes
for (int n = 0; n < foundBoxes; n++)
{
var box = new BoundingBox
{
centerX = ((output[n, 1] + output[n, 3]) * scaleX - displayWidth) / 2,
centerY = ((output[n, 2] + output[n, 4]) * scaleY - displayHeight) / 2,
width = (output[n, 3] - output[n, 1]) * scaleX,
height = (output[n, 4] - output[n, 2]) * scaleY,
label = labels[(int)output[n, 5]],
confidence = Mathf.FloorToInt(output[n, 6] * 100 + 0.5f)
};
DrawBox(box, n);
}
}
public void DrawBox(BoundingBox box, int id)
{
//Create the bounding box graphic or get from pool
GameObject panel;
if (id < boxPool.Count)
{
panel = boxPool[id];
panel.SetActive(true);
}
else
{
panel = CreateNewBox(Color.yellow);
}
//Set box position
panel.transform.localPosition = new Vector3(box.centerX, -box.centerY);
//Set box size
RectTransform rt = panel.GetComponent<RectTransform>();
rt.sizeDelta = new Vector2(box.width, box.height);
//Set label text
var label = panel.GetComponentInChildren<Text>();
label.text = box.label + " (" + box.confidence + "%)";
}
public GameObject CreateNewBox(Color color)
{
//Create the box and set image
var panel = new GameObject("ObjectBox");
panel.AddComponent<CanvasRenderer>();
Image img = panel.AddComponent<Image>();
img.color = color;
img.sprite = borderSprite;
img.type = Image.Type.Sliced;
panel.transform.SetParent(displayLocation, false);
//Create the label
var text = new GameObject("ObjectLabel");
text.AddComponent<CanvasRenderer>();
text.transform.SetParent(panel.transform, false);
Text txt = text.AddComponent<Text>();
txt.font = font;
txt.color = color;
txt.fontSize = 40;
txt.horizontalOverflow = HorizontalWrapMode.Overflow;
RectTransform rt2 = text.GetComponent<RectTransform>();
rt2.offsetMin = new Vector2(20, rt2.offsetMin.y);
rt2.offsetMax = new Vector2(0, rt2.offsetMax.y);
rt2.offsetMin = new Vector2(rt2.offsetMin.x, 0);
rt2.offsetMax = new Vector2(rt2.offsetMax.x, 30);
rt2.anchorMin = new Vector2(0, 0);
rt2.anchorMax = new Vector2(1, 1);
boxPool.Add(panel);
return panel;
}
public void ClearAnnotations()
{
foreach (var box in boxPool)
{
box.SetActive(false);
}
}
private void OnDestroy()
{
engine?.Dispose();
}
} |