File size: 5,220 Bytes
8fd532f |
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 |
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
{
const string modelName = "yolov7-tiny.sentis";
const string videoName = "giraffes.mp4";
// Link the classes.txt here:
public TextAsset labelsAsset;
// Create a Raw Image of size 640x640 and link it here:
public RawImage displayImage;
// Link to a bounding box texture here:
public Sprite boxTexture;
// 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;
private const int imageWidth = 640;
private const int imageHeight = 640;
private VideoPlayer video;
//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(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();
}
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.MakeReadable();
//Draw the bounding boxes
for (int n = 0; n < output.shape[0]; n++)
{
var box = new BoundingBox
{
centerX = (output[n, 1] + output[n, 3]) / 2 - 320,
centerY = (output[n, 2] + output[n, 4]) / 2 - 320,
width = output[n, 3] - output[n, 1],
height = output[n, 4] - output[n, 2],
label = labels[(int)output[n, 5]],
confidence = output[n, 6]
};
DrawBox(box);
}
input.Dispose();
}
public void DrawBox(BoundingBox box)
{
Color color = Color.yellow;
GameObject panel = new GameObject("ObjectBox");
panel.AddComponent<CanvasRenderer>();
Image img = panel.AddComponent<Image>();
img.color = color;
img.sprite = boxTexture;
panel.transform.SetParent(displayLocation, false);
panel.transform.localPosition = new Vector3(box.centerX, -box.centerY);
RectTransform rt = panel.GetComponent<RectTransform>();
rt.sizeDelta = new Vector2(box.width, box.height);
//add class label
var text = new GameObject("ObjectLabel");
text.AddComponent<CanvasRenderer>();
Text txt = text.AddComponent<Text>();
text.transform.SetParent(panel.transform, false);
txt.font = font;
txt.text = box.label;
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.offsetMax = new Vector2(rt2.offsetMax.x, 30);
rt2.offsetMin = new Vector2(rt2.offsetMin.x, 0);
rt2.anchorMin = new Vector2(0,0);
rt2.anchorMax = new Vector2(1, 1);
img.sprite = boxTexture;
img.type = Image.Type.Sliced;
}
public void ClearAnnotations()
{
foreach (Transform child in displayLocation) {
Destroy(child.gameObject);
}
}
private void OnDestroy()
{
engine?.Dispose();
}
}
|