| /* Copyright 2020 The TensorFlow Authors. All Rights Reserved. |
| |
| Licensed under the Apache License, Version 2.0 (the "License"); |
| you may not use this file except in compliance with the License. |
| You may obtain a copy of the License at |
| |
| http://www.apache.org/licenses/LICENSE-2.0 |
| |
| Unless required by applicable law or agreed to in writing, software |
| distributed under the License is distributed on an "AS IS" BASIS, |
| WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| See the License for the specific language governing permissions and |
| limitations under the License. |
| ==============================================================================*/ |
| |
| package org.tensorflow.lite.task.core; |
| |
| import android.content.Context; |
| import android.content.res.AssetFileDescriptor; |
| import android.util.Log; |
| import java.io.FileInputStream; |
| import java.io.IOException; |
| import java.nio.ByteBuffer; |
| import java.nio.MappedByteBuffer; |
| import java.nio.channels.FileChannel; |
| |
| /** JNI utils for Task API. */ |
| public class TaskJniUtils { |
| public static final long INVALID_POINTER = 0; |
| private static final String TAG = TaskJniUtils.class.getSimpleName(); |
| /** Syntax sugar to get nativeHandle from empty param list. */ |
| public interface EmptyHandleProvider { |
| long createHandle(); |
| } |
| |
| /** Syntax sugar to get nativeHandle from an array of {@link ByteBuffer}s. */ |
| public interface MultipleBuffersHandleProvider { |
| long createHandle(ByteBuffer... buffers); |
| } |
| |
| /** Syntax sugar to get nativeHandle from file descriptor and options. */ |
| public interface FdAndOptionsHandleProvider<T> { |
| long createHandle( |
| int fileDescriptor, long fileDescriptorLength, long fileDescriptorOffset, T options); |
| } |
| |
| /** |
| * Initializes the JNI and returns C++ handle with file descriptor and options for task API. |
| * |
| * @param context the Android app context |
| * @param provider provider to get C++ handle, usually returned from native call |
| * @param libName name of C++ lib to be loaded |
| * @param filePath path of the file to be loaded |
| * @param options options to set up the task API, used by the provider |
| * @return C++ handle as long |
| * @throws IOException If model file fails to load. |
| */ |
| public static <T> long createHandleFromFdAndOptions( |
| Context context, |
| final FdAndOptionsHandleProvider<T> provider, |
| String libName, |
| String filePath, |
| final T options) |
| throws IOException { |
| try (AssetFileDescriptor assetFileDescriptor = context.getAssets().openFd(filePath)) { |
| return createHandleFromLibrary( |
| new EmptyHandleProvider() { |
| @Override |
| public long createHandle() { |
| return provider.createHandle( |
| /*fileDescriptor=*/ assetFileDescriptor.getParcelFileDescriptor().getFd(), |
| /*fileDescriptorLength=*/ assetFileDescriptor.getLength(), |
| /*fileDescriptorOffset=*/ assetFileDescriptor.getStartOffset(), |
| options); |
| } |
| }, |
| libName); |
| } |
| } |
| |
| /** |
| * Initializes the JNI and returns C++ handle by first loading the C++ library and then invokes |
| * {@link EmptyHandleProvider#createHandle()}. |
| * |
| * @param provider provider to get C++ handle, usually returned from native call |
| * @return C++ handle as long |
| */ |
| public static long createHandleFromLibrary(EmptyHandleProvider provider, String libName) { |
| tryLoadLibrary(libName); |
| try { |
| return provider.createHandle(); |
| } catch (Exception e) { |
| String errorMessage = "Error getting native address of native library: " + libName; |
| Log.e(TAG, errorMessage, e); |
| throw new IllegalStateException(errorMessage, e); |
| } |
| } |
| |
| /** |
| * Initializes the JNI and returns C++ handle by first loading the C++ library and then invokes |
| * {@link MultipleBuffersHandleProvider#createHandle(ByteBuffer...)}. |
| * |
| * @param context app context |
| * @param provider provider to get C++ pointer, usually returned from native call |
| * @param libName name of C++ lib to load |
| * @param filePaths file paths to load |
| * @return C++ pointer as long |
| * @throws IOException If model file fails to load. |
| */ |
| public static long createHandleWithMultipleAssetFilesFromLibrary( |
| Context context, |
| final MultipleBuffersHandleProvider provider, |
| String libName, |
| String... filePaths) |
| throws IOException { |
| final MappedByteBuffer[] buffers = new MappedByteBuffer[filePaths.length]; |
| for (int i = 0; i < filePaths.length; i++) { |
| buffers[i] = loadMappedFile(context, filePaths[i]); |
| } |
| return createHandleFromLibrary( |
| new EmptyHandleProvider() { |
| @Override |
| public long createHandle() { |
| return provider.createHandle(buffers); |
| } |
| }, |
| libName); |
| } |
| |
| /** |
| * Loads a file from the asset folder through memory mapping. |
| * |
| * @param context Application context to access assets. |
| * @param filePath Asset path of the file. |
| * @return the loaded memory mapped file. |
| * @throws IOException If model file fails to load. |
| */ |
| public static MappedByteBuffer loadMappedFile(Context context, String filePath) |
| throws IOException { |
| try (AssetFileDescriptor fileDescriptor = context.getAssets().openFd(filePath); |
| FileInputStream inputStream = new FileInputStream(fileDescriptor.getFileDescriptor())) { |
| FileChannel fileChannel = inputStream.getChannel(); |
| long startOffset = fileDescriptor.getStartOffset(); |
| long declaredLength = fileDescriptor.getDeclaredLength(); |
| return fileChannel.map(FileChannel.MapMode.READ_ONLY, startOffset, declaredLength); |
| } |
| } |
| |
| private TaskJniUtils() {} |
| |
| /** |
| * Try load a native library, if it's already loaded return directly. |
| * |
| * @param libName name of the lib |
| */ |
| static void tryLoadLibrary(String libName) { |
| try { |
| System.loadLibrary(libName); |
| } catch (UnsatisfiedLinkError e) { |
| String errorMessage = "Error loading native library: " + libName; |
| Log.e(TAG, errorMessage, e); |
| throw new UnsatisfiedLinkError(errorMessage); |
| } |
| } |
| } |