First commit

pkg/rknn/rknn.go

@@ -0,0 +1,145 @@
+package rknn
+
+/*
+#cgo linux LDFLAGS: -lrknnrt
+#cgo linux,arm64 LDFLAGS: -lrknnrt
+
+#include <stdint.h>
+#include <stdlib.h>
+
+typedef void* rknn_context;
+
+typedef enum {
+	RKNN_QUERY_IN_OUT_NUM = 0,
+	RKNN_QUERY_INPUT_ATTR = 1,
+	RKNN_QUERY_OUTPUT_ATTR = 2,
+	RKNN_QUERY_SDK_VERSION = 5,
+} rknn_query_cmd;
+
+typedef enum {
+	RKNN_TENSOR_UINT8 = 0,
+	RKNN_TENSOR_FLOAT32 = 2,
+} rknn_tensor_type;
+
+typedef enum {
+	RKNN_TENSOR_NCHW = 0,
+	RKNN_TENSOR_NHWC = 1,
+} rknn_tensor_format;
+
+typedef struct {
+	uint32_t index;
+	int32_t  type;
+	int32_t  fmt;
+	void*    buf;
+	uint32_t size;
+	uint8_t  want_float;
+	uint8_t  is_prealloc;
+} rknn_output;
+
+typedef struct {
+	uint32_t n_input;
+	uint32_t n_output;
+} rknn_input_output_num;
+
+typedef struct {
+	void*    buf;
+	uint32_t size;
+	uint8_t  pass_through;
+	uint32_t type;
+	uint32_t fmt;
+} rknn_input;
+
+extern int rknn_init(rknn_context* ctx, void* model, uint32_t size, uint32_t flag, rknn_input_output_num* io_num);
+extern int rknn_query(rknn_context ctx, rknn_query_cmd cmd, void* info, uint32_t size);
+extern int rknn_inputs_set(rknn_context ctx, uint32_t n_inputs, rknn_input inputs[]);
+extern int rknn_run(rknn_context ctx, void* ext);
+extern int rknn_outputs_get(rknn_context ctx, uint32_t n_outputs, rknn_output outputs[], void* ext);
+extern int rknn_outputs_release(rknn_context ctx, uint32_t n_outputs, rknn_output outputs[]);
+extern int rknn_destroy(rknn_context ctx);
+*/
+import "C"
+import (
+	"fmt"
+	"os"
+	"unsafe"
+)
+
+type Context struct {
+	ctx     C.rknn_context
+	nInput  uint32
+	nOutput uint32
+}
+
+func LoadModel(modelPath string) (*Context, error) {
+	data, err := os.ReadFile(modelPath)
+	if err != nil {
+		return nil, fmt.Errorf("rknn: read model file: %w", err)
+	}
+
+	var ctx C.rknn_context
+	var ioNum C.rknn_input_output_num
+
+	ret := C.rknn_init(
+		&ctx,
+		unsafe.Pointer(&data[0]),
+		C.uint32_t(len(data)),
+		0,
+		&ioNum,
+	)
+	if ret != 0 {
+		return nil, fmt.Errorf("rknn_init: error %d", int(ret))
+	}
+
+	return &Context{
+		ctx:     ctx,
+		nInput:  uint32(ioNum.n_input),
+		nOutput: uint32(ioNum.n_output),
+	}, nil
+}
+
+func (c *Context) InferenceRGB(inputData []uint8, height, width int) ([][]float32, error) {
+	rknnIn := C.rknn_input{
+		buf:          unsafe.Pointer(&inputData[0]),
+		size:         C.uint32_t(len(inputData)),
+		pass_through: 0,
+		_type:        C.RKNN_TENSOR_UINT8,
+		fmt:          C.RKNN_TENSOR_NHWC,
+	}
+
+	ret := C.rknn_inputs_set(c.ctx, 1, &rknnIn)
+	if ret < 0 {
+		return nil, fmt.Errorf("rknn_inputs_set: error %d", int(ret))
+	}
+
+	ret = C.rknn_run(c.ctx, nil)
+	if ret < 0 {
+		return nil, fmt.Errorf("rknn_run: error %d", int(ret))
+	}
+
+	outputs := make([]C.rknn_output, c.nOutput)
+	for i := range outputs {
+		outputs[i].want_float = 1
+	}
+
+	ret = C.rknn_outputs_get(c.ctx, C.uint32_t(c.nOutput), &outputs[0], nil)
+	if ret < 0 {
+		return nil, fmt.Errorf("rknn_outputs_get: error %d", int(ret))
+	}
+
+	results := make([][]float32, c.nOutput)
+	for i := range outputs {
+		nVals := int(outputs[i].size) / 4
+		vals := make([]float32, nVals)
+		src := unsafe.Slice((*float32)(outputs[i].buf), nVals)
+		copy(vals, src)
+		results[i] = vals
+	}
+
+	C.rknn_outputs_release(c.ctx, C.uint32_t(c.nOutput), &outputs[0])
+
+	return results, nil
+}
+
+func (c *Context) Release() {
+	C.rknn_destroy(c.ctx)
+}