added support for gpu delegate on android

aazucena · Jun 4, 2020 · 5c1b8cd · 5c1b8cd
1 parent 4caff96
commit 5c1b8cd
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 6 deletions.
diff --git a/android/build.gradle b/android/build.gradle
@@ -34,5 +34,6 @@ android {
 
     dependencies {
         compile 'org.tensorflow:tensorflow-lite:+'
+        compile 'org.tensorflow:tensorflow-lite-gpu:+'
     }
 }
diff --git a/android/src/main/java/sq/flutter/tflite/TflitePlugin.java b/android/src/main/java/sq/flutter/tflite/TflitePlugin.java
@@ -26,6 +26,8 @@
 import org.tensorflow.lite.Interpreter;
 import org.tensorflow.lite.Tensor;
 
+import org.tensorflow.lite.gpu.GpuDelegate;
+
 import java.io.BufferedReader;
 import java.io.File;
 import java.io.FileInputStream;
@@ -195,6 +197,7 @@ public void onMethodCall(MethodCall call, Result result) {
   private String loadModel(HashMap args) throws IOException {
     String model = args.get("model").toString();
     Object isAssetObj = args.get("isAsset");
+    Object useGpuDelegateObj = args.get("useGpuDelegate");
     boolean isAsset = isAssetObj == null ? false : (boolean) isAssetObj;
     MappedByteBuffer buffer = null;
     String key = null;
@@ -216,8 +219,14 @@ private String loadModel(HashMap args) throws IOException {
     }
 
     int numThreads = (int) args.get("numThreads");
+    boolean useGpuDelegate = useGpuDelegateObj == null ? false : (boolean) useGpuDelegate;
+
     final Interpreter.Options tfliteOptions = new Interpreter.Options();
     tfliteOptions.setNumThreads(numThreads);
+    if (useGpuDelegate){
+      GpuDelegate delegate = new GpuDelegate();
+      tfliteOptions.addDelegate(delegate)
+    }
     tfLite = new Interpreter(buffer, tfliteOptions);
 
     String labels = args.get("labels").toString();

diff --git a/lib/tflite.dart b/lib/tflite.dart
@@ -7,19 +7,20 @@ import 'package:flutter/services.dart';
 class Tflite {
   static const MethodChannel _channel = const MethodChannel('tflite');
 
-  static Future<String> loadModel({
-    @required String model,
-    String labels = "",
-    int numThreads = 1,
-    bool isAsset = true,
-  }) async {
+  static Future<String> loadModel(
+      {@required String model,
+      String labels = "",
+      int numThreads = 1,
+      bool isAsset = true,
+      bool useGpuDelegate = false}) async {
     return await _channel.invokeMethod(
       'loadModel',
       {
         "model": model,
         "labels": labels,
         "numThreads": numThreads,
         "isAsset": isAsset,
+        'useGpuDelegate': useGpuDelegate
       },
     );
   }