File size: 5,054 Bytes
3868539
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
using System;
using System.Globalization;
using Unity.Mathematics;
using Unity.Sentis;
using UnityEngine;

public static class BlazeUtils
{
    // matrix utility
    public static float2x3 mul(float2x3 a, float2x3 b)
    {
        return new float2x3(
            a[0][0] * b[0][0] + a[1][0] * b[0][1],
            a[0][0] * b[1][0] + a[1][0] * b[1][1],
            a[0][0] * b[2][0] + a[1][0] * b[2][1] + a[2][0],
            a[0][1] * b[0][0] + a[1][1] * b[0][1],
            a[0][1] * b[1][0] + a[1][1] * b[1][1],
            a[0][1] * b[2][0] + a[1][1] * b[2][1] + a[2][1]
        );
    }

    public static float2 mul(float2x3 a, float2 b)
    {
        return new float2(
            a[0][0] * b.x + a[1][0] * b.y + a[2][0],
            a[0][1] * b.x + a[1][1] * b.y + a[2][1]
        );
    }

    public static float2x3 RotationMatrix(float theta)
    {
        var sinTheta = math.sin(theta);
        var cosTheta = math.cos(theta);
        return new float2x3(
            cosTheta, -sinTheta, 0,
            sinTheta, cosTheta, 0
        );
    }

    public static float2x3 TranslationMatrix(float2 delta)
    {
        return new float2x3(
            1, 0, delta.x,
            0, 1, delta.y
        );
    }

    public static float2x3 ScaleMatrix(float2 scale)
    {
        return new float2x3(
            scale.x, 0, 0,
            0, scale.y, 0
        );
    }

    // model filtering utility
    static FunctionalTensor ScoreFiltering(FunctionalTensor rawScores, float scoreThreshold)
    {
        return Functional.Sigmoid(Functional.Clamp(rawScores, -scoreThreshold, scoreThreshold));
    }

    public static (FunctionalTensor, FunctionalTensor, FunctionalTensor) NMSFiltering(FunctionalTensor rawBoxes, FunctionalTensor rawScores, FunctionalTensor anchors, float inputSize, float iouThreshold, float scoreThreshold)
    {
        var xCenter = rawBoxes[0, .., 0] + anchors[.., 0] * inputSize;
        var yCenter = rawBoxes[0, .., 1] + anchors[.., 1] * inputSize;

        var widthHalf = 0.5f * rawBoxes[0, .., 2];
        var heightHalf = 0.5f * rawBoxes[0, .., 3];

        var nmsBoxes = Functional.Stack(new[]
        {
            yCenter - heightHalf,
            xCenter - widthHalf,
            yCenter + heightHalf,
            xCenter + widthHalf
        }, 1);

        var nmsScores = Functional.Squeeze(ScoreFiltering(rawScores, 100f));
        var selectedIndices = Functional.NMS(nmsBoxes, nmsScores, iouThreshold, scoreThreshold); // (N);

        var selectedBoxes = Functional.IndexSelect(rawBoxes, 1, selectedIndices).Unsqueeze(0); // (1, N, 16)
        var selectedScores = Functional.IndexSelect(rawScores, 1, selectedIndices).Unsqueeze(0); // (1, N, 1)

        return (selectedIndices, selectedScores, selectedBoxes);
    }

    // image transform utility
    static ComputeShader s_ImageTransformShader = Resources.Load<ComputeShader>("ComputeShaders/ImageTransform");
    static int s_ImageSample = s_ImageTransformShader.FindKernel("ImageSample");
    static int s_Optr = Shader.PropertyToID("Optr");
    static int s_X_tex2D = Shader.PropertyToID("X_tex2D");
    static int s_O_height = Shader.PropertyToID("O_height");
    static int s_O_width = Shader.PropertyToID("O_width");
    static int s_O_channels = Shader.PropertyToID("O_channels");
    static int s_X_height = Shader.PropertyToID("X_height");
    static int s_X_width = Shader.PropertyToID("X_width");
    static int s_affineMatrix = Shader.PropertyToID("affineMatrix");

    static int IDivC(int v, int div)
    {
        return (v + div - 1) / div;
    }

    public static void SampleImageAffine(Texture srcTexture, Tensor<float> dstTensor, float2x3 M)
    {
        var tensorData = ComputeTensorData.Pin(dstTensor, false);

        s_ImageTransformShader.SetTexture(s_ImageSample, s_X_tex2D, srcTexture);
        s_ImageTransformShader.SetBuffer(s_ImageSample, s_Optr, tensorData.buffer);

        s_ImageTransformShader.SetInt(s_O_height, dstTensor.shape[1]);
        s_ImageTransformShader.SetInt(s_O_width, dstTensor.shape[2]);
        s_ImageTransformShader.SetInt(s_O_channels, dstTensor.shape[3]);
        s_ImageTransformShader.SetInt(s_X_height, srcTexture.height);
        s_ImageTransformShader.SetInt(s_X_width, srcTexture.width);

        s_ImageTransformShader.SetMatrix(s_affineMatrix, new Matrix4x4(new Vector4(M[0][0], M[0][1]), new Vector4(M[1][0], M[1][1]), new Vector4(M[2][0], M[2][1]), Vector4.zero));

        s_ImageTransformShader.Dispatch(s_ImageSample, IDivC(dstTensor.shape[1], 8), IDivC(dstTensor.shape[1], 8), 1);
    }

    public static float[,] LoadAnchors(string csv, int numAnchors)
    {
        var anchors = new float[numAnchors, 4];
        var anchorLines = csv.Split('\n');

        for (var i = 0; i < numAnchors; i++)
        {
            var anchorValues = anchorLines[i].Split(',');
            for (var j = 0; j < 4; j++)
            {
                anchors[i, j] = float.Parse(anchorValues[j], CultureInfo.InvariantCulture);
            }
        }

        return anchors;
    }
}