saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[js/rn] Support UINT8 type for onnxruntime-react-native on Android (#12112)

Browse source 
* support uint8 for react native

* add test
This commit is contained in:
Yulong Wang 2022-07-08 14:07:46 -07:00 committed by GitHub
parent c04afae9a9
commit d45c1a144e
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 55 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

41
js/react_native/android/src/androidTest/java/ai/onnxruntime/reactnative/TensorHelperTest.java
View file

@ -6,6 +6,7 @@ package ai.onnxruntime.reactnative;
import static com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession;
import static org.mockito.Mockito.when;
import ai.onnxruntime.OnnxJavaType;
import ai.onnxruntime.OnnxTensor;
import ai.onnxruntime.OnnxValue;
import ai.onnxruntime.OrtEnvironment;
@ -448,6 +449,46 @@ public class TensorHelperTest {
}
}
@Test
public void createOutputTensor_uint8() throws Exception {
MockitoSession mockSession = mockitoSession().mockStatic(Arguments.class).startMocking();
try {
when(Arguments.createMap()).thenAnswer(i -> new JavaOnlyMap());
when(Arguments.createArray()).thenAnswer(i -> new JavaOnlyArray());
OrtSession.SessionOptions options = new OrtSession.SessionOptions();
byte[] modelData = readBytesFromResourceFile(ai.onnxruntime.reactnative.test.R.raw.test_types_uint8);
OrtSession session = ortEnvironment.createSession(modelData, options);
long[] dims = new long[] {1, 5};
byte[] inputData = new byte[] {1, 2, -3, Byte.MAX_VALUE, Byte.MAX_VALUE};
String inputName = session.getInputNames().iterator().next();
Map<String, OnnxTensor> container = new HashMap<>();
ByteBuffer inputBuffer = ByteBuffer.wrap(inputData);
OnnxTensor onnxTensor = OnnxTensor.createTensor(ortEnvironment, inputBuffer, dims, OnnxJavaType.UINT8);
container.put(inputName, onnxTensor);
OrtSession.Result result = session.run(container);
ReadableMap resultMap = TensorHelper.createOutputTensor(result);
ReadableMap outputMap = resultMap.getMap("output");
for (int i = 0; i < 2; ++i) {
Assert.assertEquals(outputMap.getArray("dims").getInt(i), dims[i]);
}
Assert.assertEquals(outputMap.getString("type"), TensorHelper.JsTensorTypeUnsignedByte);
String dataEncoded = outputMap.getString("data");
ByteBuffer buffer = ByteBuffer.wrap(Base64.decode(dataEncoded, Base64.DEFAULT));
for (int i = 0; i < 5; ++i) {
Assert.assertEquals(buffer.get(i), inputData[i]);
}
OnnxValue.close(container);
} finally {
mockSession.finishMocking();
}
}
private byte[] readBytesFromResourceFile(int resourceId) throws Exception {
Context context = InstrumentationRegistry.getInstrumentation().getContext();
InputStream inputStream = context.getResources().openRawResource(resourceId);

BIN
js/react_native/android/src/androidTest/res/raw/test_types_uint8.ort Normal file
View file

Binary file not shown.

16
js/react_native/android/src/main/java/ai/onnxruntime/reactnative/TensorHelper.java
View file

@ -3,6 +3,7 @@
package ai.onnxruntime.reactnative;
import ai.onnxruntime.OnnxJavaType;
import ai.onnxruntime.OnnxTensor;
import ai.onnxruntime.OnnxValue;
import ai.onnxruntime.OrtEnvironment;
@ -34,6 +35,7 @@ public class TensorHelper {
*/
public static final String JsTensorTypeBool = "bool";
public static final String JsTensorTypeByte = "int8";
public static final String JsTensorTypeUnsignedByte = "uint8";
public static final String JsTensorTypeShort = "int16";
public static final String JsTensorTypeInt = "int32";
public static final String JsTensorTypeLong = "int64";
@ -157,9 +159,13 @@ public class TensorHelper {
tensor = OnnxTensor.createTensor(ortEnvironment, buffer, dims);
break;
}
case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8: {
ByteBuffer buffer = values;
tensor = OnnxTensor.createTensor(ortEnvironment, buffer, dims, OnnxJavaType.UINT8);
break;
}
case ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL:
case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16:
case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8:
case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT16:
case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32:
case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64:
@ -204,9 +210,13 @@ public class TensorHelper {
buffer.asDoubleBuffer().put(onnxTensor.getDoubleBuffer());
break;
}
case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8: {
buffer = ByteBuffer.allocate(capacity).order(ByteOrder.nativeOrder());
buffer.put(onnxTensor.getByteBuffer());
break;
}
case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16:
case ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING:
case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8:
case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT16:
case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32:
case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64:
@ -223,6 +233,7 @@ public class TensorHelper {
.of(new Object[][] {
{JsTensorTypeFloat, TensorInfo.OnnxTensorType.ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT},
{JsTensorTypeByte, TensorInfo.OnnxTensorType.ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8},
{JsTensorTypeUnsignedByte, TensorInfo.OnnxTensorType.ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8},
{JsTensorTypeShort, TensorInfo.OnnxTensorType.ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16},
{JsTensorTypeInt, TensorInfo.OnnxTensorType.ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32},
{JsTensorTypeLong, TensorInfo.OnnxTensorType.ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64},
@ -245,6 +256,7 @@ public class TensorHelper {
.of(new Object[][] {
{TensorInfo.OnnxTensorType.ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, JsTensorTypeFloat},
{TensorInfo.OnnxTensorType.ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8, JsTensorTypeByte},
{TensorInfo.OnnxTensorType.ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8, JsTensorTypeUnsignedByte},
{TensorInfo.OnnxTensorType.ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16, JsTensorTypeShort},
{TensorInfo.OnnxTensorType.ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32, JsTensorTypeInt},
{TensorInfo.OnnxTensorType.ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64, JsTensorTypeLong},