Update RunYOLO.cs

93c60d5 verified about 2 months ago

No virus

6.61 kB

	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();
	}
	}