A modern object detector inside fragment shaders
YOLOv4 Tiny is one of the fastest object detectors that exists currently. The goal of this project is to completely recreate it without any existing ML libraries such as Darknet, PyTorch, or TensorFlow in order to port it into a VR game called VRChat.
My naive implementation only runs around 30 FPS so it doesn't hog resources for VR. It is nowhere near as performant as the original.
This implementation is based on the TensorFlow version from https://github.com/hunglc007/tensorflow-yolov4-tflite
yolov4tiny.shader
nms.shader
The basic setup is: yolov4tiny.shader -> nms.shader -> output.shader. To read the bounding boxes information, we loop through the output of nms.shader
This is a basic setup of how output.shader works, refer to the file for more information.
Properties
{
_NMSout ("NMS Output", 2D) = "black" {}
}
#include "nms_include.cginc"
const float2 scale = 0.5.xx;
uint i;
uint j;
// Loop through the 26x26 grid output
for (i = 0; i < 26; i++) {
for (j = 0; j < 26; j++) {
// Only draw a box if the confidence is over 50%
uint4 buff = asuint(_NMSout[txL20nms.xy + uint2(i, j)]);
float conf = f16tof32(buff.a);
[branch]
if (conf > 0.5) {
// Class, 0 to 79
float c = f16tof32(buff.b >> 16);
// x, y is the center position of the bbox relative to 416, the initial image input size that goes into the network
float x = f16tof32(buff.r >> 16);
float y = f16tof32(buff.r);
// w, h are the width and height of the bbox relative to 416, the initial image input size that goes into the network
float w = f16tof32(buff.g >> 16);
float h = f16tof32(buff.g);
// Scale to camera resolution using UVs
float2 center = float2(x, y) / 416.0;
center.y = 1.0 - center.y;
float2 size = float2(w, h) / 416.0 * scale;
}
}
}
// YOLOv4 tiny has two outputs, remember to go through the second one too
// Loop through the 13x13 grid output
for (i = 0; i < 13; i++) {
for (j = 0; j < 13; j++) {
// Only draw a box if the confidence is over 50%
uint4 buff = asuint(_NMSout[txL17nms.xy + uint2(i, j)]);
float conf = f16tof32(buff.a);
[branch]
if (conf > 0.5) {
// Class, 0 to 79
float c = f16tof32(buff.b >> 16);
// x, y is the center position of the bbox relative to 416, the initial image input size that goes into the network
float x = f16tof32(buff.r >> 16);
float y = f16tof32(buff.r);
// w, h are the width and height of the bbox relative to 416, the initial image input size that goes into the network
float w = f16tof32(buff.g >> 16);
float h = f16tof32(buff.g);
// Scale to camera resolution using UVs
float2 center = float2(x, y) / 416.0;
center.y = 1.0 - center.y;
float2 size = float2(w, h) / 416.0 * scale;
}
}
}
The data is packed into two 16 bits per channel by the nms.shader and the layout is as follows:
R = X | Y
G = W | H
B = Best class | Best class probability
A = | Bounding box confidence
All the classes the network can detect, based on the COCO Dataset classes:
Index 0 starts at the top left, i.e. 0 = person, 1 = bicycle and so on.Since this is a direct implementation of a known architecture, you can refer to their original papers.
YOLOv4's paper is basically an ablation study on the different parameters and tuning to maximize speed and accuracy. I suggest reading the previous versions to have a better understanding of the actual architecture.
If you have questions or comments, you can reach me on Discord: SCRN#8008 or Twitter: https://twitter.com/SCRNinVR