前言

TNN:由騰訊優圖實驗室打造,移動端高性能、輕量級推理框架,同時擁有跨平臺、高性能、模型壓縮、代碼裁剪等衆多突出優勢。TNN框架在原有Rapidnet、ncnn框架的基礎上進一步加強了移動端設備的支持以及性能優化,同時也借鑑了業界主流開源框架高性能和良好拓展性的優點。

教程源碼地址:https://github.com/yeyupiaoling/ClassificationForAndroid/tree/master/TNNClassification

編譯Android庫

  1. 安裝cmake 3.12
# 卸載舊的cmake
sudo apt-get autoremove cmake

# 下載cmake3.12
wget https://cmake.org/files/v3.12/cmake-3.12.2-Linux-x86_64.tar.gz
tar zxvf cmake-3.12.2-Linux-x86_64.tar.gz

# 移動目錄並添加軟連接
sudo mv cmake-3.12.2-Linux-x86_64 /opt/cmake-3.12.2
sudo ln -sf /opt/cmake-3.12.2/bin/*  /usr/bin/
  1. 添加Android NDK
wget https://dl.google.com/android/repository/android-ndk-r21b-linux-x86_64.zip
unzip android-ndk-r21b-linux-x86_64.zip
# 添加環境變量,留意你實際下載地址
export ANDROID_NDK=/mnt/d/android-ndk-r21b
  1. 安裝編譯環境
sudo apt-get install attr
  1. 開始編譯
git clone https://github.com/Tencent/TNN.git
cd TNN/scripts

vim build_android.sh
 ABIA32="armeabi-v7a"
 ABIA64="arm64-v8a"
 STL="c++_static"
 SHARED_LIB="ON"                # ON表示編譯動態庫,OFF表示編譯靜態庫
 ARM="ON"                       # ON表示編譯帶有Arm CPU版本的庫
 OPENMP="ON"                    # ON表示打開OpenMP
 OPENCL="ON"                    # ON表示編譯帶有Arm GPU版本的庫
 SHARING_MEM_WITH_OPENGL=0      # 1表示OpenGL的Texture可以與OpenCL共享

執行編譯

./build_android.sh

編譯完成後,會在當前目錄的release目錄下生成對應的armeabi-v7a庫,arm64-v8a庫和include頭文件,這些文件在下一步的Android開發都需要使用到。

模型轉換

接下來我們需要把Tensorflow,onnx等其他的模型轉換爲TNN的模型。目前 TNN 支持業界主流的模型文件格式,包括ONNX、PyTorch、TensorFlow 以及 Caffe 等。TNN 將 ONNX 作爲中間層,藉助於ONNX 開源社區的力量,來支持多種模型文件格式。如果要將PyTorch、TensorFlow 以及 Caffe 等模型文件格式轉換爲 TNN,首先需要使用對應的模型轉換工具,統一將各種模型格式轉換成爲 ONNX 模型格式,然後將 ONNX 模型轉換成 TNN 模型。

sudo docker pull turandotkay/tnn-convert
sudo docker tag turandotkay/tnn-convert:latest tnn-convert:latest
sudo docker rmi turandotkay/tnn-convert:latest

針對不同的模型轉換,有不同的命令,如onnx2tnn,caffe2tnn,tf2tnn。

docker run --volume=$(pwd):/workspace -it tnn-convert:latest  python3 ./converter.py tf2tnn \
    -tp /workspace/mobilenet_v1.pb \
    -in "input[1,224,224,3]" \
    -on MobilenetV1/Predictions/Reshape_1 \
    -v v1.0 \
    -optimize

通過上面的輸出,可以發現針對 TF 模型的轉換,convert2tnn 工具提供了很多參數,我們一次對下面的參數進行解釋:

  • tp 參數(必須)
    通過 “-tp” 參數指定需要轉換的模型的路徑。目前只支持單個 TF模型的轉換,不支持多個 TF 模型的一起轉換。
  • in 參數(必須)
    通過 “-in” 參數指定模型輸入的名稱,輸入的名稱需要放到“”中,例如,-in “name”。如果模型有多個輸入,請使用 “;”進行分割。有的 TensorFlow 模型沒有指定 batch 導致無法成功轉換爲 ONNX 模型,進而無法成功轉換爲 TNN 模型。你可以通過在名稱後添加輸入 shape 進行指定。shape 信息需要放在 [] 中。例如:-in “name[1,28,28,3]”。
  • on 參數(必須)
    通過 “-on” 參數指定模型輸入的名稱,如果模型有多個輸出,請使用 “;”進行分割
  • output_dir 參數:
    可以通過 “-o ” 參數指定輸出路徑,但是在 docker 中我們一般不使用這個參數,默認會將生成的 TNN 模型放在當前和 TF 模型相同的路徑下。
  • optimize 參數(可選)
    可以通過 “-optimize” 參數來對模型進行優化,我們強烈建議你開啓這個選項,只有在開啓這個選項模型轉換失敗時,我們才建議你去掉 “-optimize” 參數進行重新嘗試
  • v 參數(可選)
    可以通過 -v 來指定模型的版本號,以便於後期對模型進行追蹤和區分。
  • half 參數(可選)
    可以通過 -half 參數指定,模型數據通過 FP16 進行存儲,減少模型的大小,默認是通過 FP32 的方式進行存儲模型數據的。
  • align 參數(可選)
    可以通過 -align 參數指定,將 轉換得到的 TNN 模型和原模型進行對齊,確定 TNN 模型是否轉換成功。當前僅支持單輸入單輸出模型和單輸入多輸出模型。 align 只支持 FP32 模型的校驗,所以使用 align 的時候不能使用 half
  • input_file 參數(可選)
    可以通過 -input_file 參數指定模型對齊所需要的輸入文件的名稱,輸入需要遵循如下格式
  • ref_file 參數(可選)
    可以通過 -ref_file 參數指定待對齊的輸出文件的名稱,輸出需遵循如下格式。生成輸出的代碼可以參考

成功轉換會輸出以下的日誌。

----------  convert model, please wait a moment ----------

Converter Tensorflow to TNN model

Convert TensorFlow to ONNX model succeed!

Converter ONNX to TNN Model

Converter ONNX to TNN model succeed!

最終會得到這兩個模型文件,mobilenet_v1.opt.tnnmodel mobilenet_v1.opt.tnnproto

開發Android項目

  1. 將轉換的模型放在assets目錄下。
  2. 把上一步編譯得到的include目錄複製到Android項目的app目錄下。
  3. 把上一步編譯得到的armeabi-v7aarm64-v8a目錄複製到main/jniLibs下。
  4. app/src/main/cpp/目錄下編寫JNI的C++代碼。

TNN工具

編寫一個ImageClassifyUtil.java工具類,關於TNN的操作都在這裏完成,如加載模型、預測。

下面三個就是TNN的JNI接口,通過這個接口完成模型加載,預測,當不使用的時候和可以調用deinit()清空對象。

public native int init(String modelPath, String protoPath, int computeUnitType);

public native float[] predict(Bitmap image, int width, int height);

public native int deinit();

通過上面的JNI接口,下面就可以實現圖像識別了,WIDTHHEIGHT是模型輸入圖片的大小。爲了兼容圖片路徑和Bitmap格式的圖片預測,這裏創建了兩個重載方法。

private static final int WIDTH = 224;
private  static final int HEIGHT = 224;

public ImageClassifyUtil() {
    System.loadLibrary("TNN");
    System.loadLibrary("tnn_wrapper");
}

// 重載方法,根據圖片路徑轉Bitmap預測
public float[] predictImage(String image_path) throws Exception {
    if (!new File(image_path).exists()) {
        throw new Exception("image file is not exists!");
    }
    FileInputStream fis = new FileInputStream(image_path);
    Bitmap bitmap = BitmapFactory.decodeStream(fis);
    Bitmap scaleBitmap = Bitmap.createScaledBitmap(bitmap, WIDTH, HEIGHT, false);
    float[] result = predictImage(scaleBitmap);
    if (bitmap.isRecycled()) {
        bitmap.recycle();
    }
    return result;
}

// 重載方法,直接使用Bitmap預測
public float[] predictImage(Bitmap bitmap) {
    Bitmap scaleBitmap = Bitmap.createScaledBitmap(bitmap, WIDTH, HEIGHT, false);
    float[] results = predict(scaleBitmap, WIDTH, HEIGHT);
    int l = getMaxResult(results);
    return new float[]{l, results[l] * 0.01f};
}

這裏創建一個獲取最大概率值,並把下標返回的方法,其實就是獲取概率最大的預測標籤。

public static int getMaxResult(float[] result) {
    float probability = 0;
    int r = 0;
    for (int i = 0; i < result.length; i++) {
        if (probability < result[i]) {
            probability = result[i];
            r = i;
        }
    }
    return r;
}

不同的模型,訓練的預處理方式可能不一樣,TNN 的圖像預處理在C++中完成,代碼片段

TNN_NS::MatConvertParam input_cvt_param;
input_cvt_param.scale = {1.0 / (255 * 0.229), 1.0 / (255 * 0.224), 1.0 / (255 * 0.225), 0.0};
input_cvt_param.bias  = {-0.485 / 0.229, -0.456 / 0.224, -0.406 / 0.225, 0.0};
auto status = instance_->SetInputMat(input_mat, input_cvt_param);

選擇圖片預測

本教程會有兩個頁面,一個是選擇圖片進行預測的頁面,另一個是使用相機即時預測並顯示預測結果。以下爲activity_main.xml的代碼,通過按鈕選擇圖片,並在該頁面顯示圖片和預測結果。

<?xml version="1.0" encoding="utf-8"?>
<RelativeLayout xmlns:android="http://schemas.android.com/apk/res/android"
    xmlns:app="http://schemas.android.com/apk/res-auto"
    xmlns:tools="http://schemas.android.com/tools"
    android:layout_width="match_parent"
    android:layout_height="match_parent"
    android:orientation="vertical"
    tools:context=".MainActivity">

    <ImageView
        android:id="@+id/image_view"
        android:layout_width="match_parent"
        android:layout_height="400dp" />

    <TextView
        android:id="@+id/result_text"
        android:layout_width="match_parent"
        android:layout_height="wrap_content"
        android:layout_below="@id/image_view"
        android:text="識別結果"
        android:textSize="16sp" />


    <LinearLayout
        android:layout_width="match_parent"
        android:layout_height="wrap_content"
        android:layout_alignParentBottom="true"
        android:orientation="horizontal">

        <Button
            android:id="@+id/select_img_btn"
            android:layout_width="0dp"
            android:layout_height="wrap_content"
            android:layout_weight="1"
            android:text="選擇照片" />


        <Button
            android:id="@+id/open_camera"
            android:layout_width="0dp"
            android:layout_height="wrap_content"
            android:layout_weight="1"
            android:text="即時預測" />

    </LinearLayout>

</RelativeLayout>

MainActivity.java中,進入到頁面我們就要先加載模型,我們是把模型放在Android項目的assets目錄的,我們需要把模型複製到一個緩存目錄,然後再從緩存目錄加載模型,同時還有讀取標籤名,標籤名稱按照訓練的label順序存放在assets的label_list.txt,以下爲實現代碼。

classNames = Utils.ReadListFromFile(getAssets(), "label_list.txt");
String protoContent = getCacheDir().getAbsolutePath() + File.separator + "squeezenet_v1.1.tnnproto";
Utils.copyFileFromAsset(MainActivity.this, "squeezenet_v1.1.tnnproto", protoContent);
String modelContent = getCacheDir().getAbsolutePath() + File.separator + "squeezenet_v1.1.tnnmodel";
Utils.copyFileFromAsset(MainActivity.this, "squeezenet_v1.1.tnnmodel", modelContent);

imageClassifyUtil = new ImageClassifyUtil();
int status = imageClassifyUtil.init(modelContent, protoContent, USE_GPU ? 1 : 0);
if (status == 0){
    Toast.makeText(MainActivity.this, "模型加載成功!", Toast.LENGTH_SHORT).show();
}else {
    Toast.makeText(MainActivity.this, "模型加載失敗!", Toast.LENGTH_SHORT).show();
    finish();
}

添加兩個按鈕點擊事件,可以選擇打開相冊讀取圖片進行預測,或者打開另一個Activity進行調用攝像頭即時識別。

Button selectImgBtn = findViewById(R.id.select_img_btn);
Button openCamera = findViewById(R.id.open_camera);
imageView = findViewById(R.id.image_view);
textView = findViewById(R.id.result_text);
selectImgBtn.setOnClickListener(new View.OnClickListener() {
    @Override
    public void onClick(View v) {
        // 打開相冊
        Intent intent = new Intent(Intent.ACTION_PICK);
        intent.setType("image/*");
        startActivityForResult(intent, 1);
    }
});
openCamera.setOnClickListener(new View.OnClickListener() {
    @Override
    public void onClick(View v) {
        // 打開即時拍攝識別頁面
        Intent intent = new Intent(MainActivity.this, CameraActivity.class);
        startActivity(intent);
    }
});

當打開相冊選擇照片之後,回到原來的頁面,在下面這個回調方法中獲取選擇圖片的Uri,通過Uri可以獲取到圖片的絕對路徑。如果Android8以上的設備獲取不到圖片,需要在AndroidManifest.xml配置文件中的application添加android:requestLegacyExternalStorage="true"。拿到圖片路徑之後,調用TFLiteClassificationUtil類中的predictImage()方法預測並獲取預測值,在頁面上顯示預測的標籤、對應標籤的名稱、概率值和預測時間。

@Override
protected void onActivityResult(int requestCode, int resultCode, @Nullable Intent data) {
    super.onActivityResult(requestCode, resultCode, data);
    String image_path;
    if (resultCode == Activity.RESULT_OK) {
        if (requestCode == 1) {
            if (data == null) {
                Log.w("onActivityResult", "user photo data is null");
                return;
            }
            Uri image_uri = data.getData();
            image_path = getPathFromURI(MainActivity.this, image_uri);
            try {
                // 預測圖像
                FileInputStream fis = new FileInputStream(image_path);
                imageView.setImageBitmap(BitmapFactory.decodeStream(fis));
                long start = System.currentTimeMillis();
                float[] result = imageClassifyUtil.predictImage(image_path);
                long end = System.currentTimeMillis();
                String show_text = "預測結果標籤:" + (int) result[0] +
                        "\n名稱:" +  classNames[(int) result[0]] +
                        "\n概率:" + result[1] +
                        "\n時間:" + (end - start) + "ms";
                textView.setText(show_text);
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }
}

上面獲取的Uri可以通過下面這個方法把Url轉換成絕對路徑。

// get photo from Uri
public static String getPathFromURI(Context context, Uri uri) {
    String result;
    Cursor cursor = context.getContentResolver().query(uri, null, null, null, null);
    if (cursor == null) {
        result = uri.getPath();
    } else {
        cursor.moveToFirst();
        int idx = cursor.getColumnIndex(MediaStore.Images.ImageColumns.DATA);
        result = cursor.getString(idx);
        cursor.close();
    }
    return result;
}

攝像頭即時預測

在調用相機即時預測我就不再介紹了,原理都差不多,具體可以查看https://github.com/yeyupiaoling/ClassificationForAndroid/tree/master/TFLiteClassification中的源代碼。核心代碼如下,創建一個子線程,子線程中不斷從攝像頭預覽的AutoFitTextureView上獲取圖像,並執行預測,並在頁面上顯示預測的標籤、對應標籤的名稱、概率值和預測時間。每一次預測完成之後都立即獲取圖片繼續預測,只要預測速度夠快,就可以看成即時預測。

private Runnable periodicClassify =
        new Runnable() {
            @Override
            public void run() {
                synchronized (lock) {
                    if (runClassifier) {
                        // 開始預測前要判斷相機是否已經準備好
                        if (getApplicationContext() != null && mCameraDevice != null && mnnClassification != null) {
                            predict();
                        }
                    }
                }
                if (mInferThread != null && mInferHandler != null && mCaptureHandler != null && mCaptureThread != null) {
                    mInferHandler.post(periodicClassify);
                }
            }
        };

// 預測相機捕獲的圖像
private void predict() {
    // 獲取相機捕獲的圖像
    Bitmap bitmap = mTextureView.getBitmap();
    try {
        // 預測圖像
        long start = System.currentTimeMillis();
        float[] result = imageClassifyUtil.predictImage(bitmap);
        long end = System.currentTimeMillis();
        String show_text = "預測結果標籤:" + (int) result[0] +
                "\n名稱:" +  classNames[(int) result[0]] +
                "\n概率:" + result[1] +
                "\n時間:" + (end - start) + "ms";
        textView.setText(show_text);
    } catch (Exception e) {
        e.printStackTrace();
    }
}

本項目中使用的了讀取圖片的權限和打開相機的權限,所以不要忘記在AndroidManifest.xml添加以下權限申請。

<uses-permission android:name="android.permission.CAMERA"/>
<uses-permission android:name="android.permission.READ_EXTERNAL_STORAGE"/>
<uses-permission android:name="android.permission.WRITE_EXTERNAL_STORAGE"/>

如果是Android 6 以上的設備還要動態申請權限。

    // check had permission
    private boolean hasPermission() {
        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.M) {
            return checkSelfPermission(Manifest.permission.CAMERA) == PackageManager.PERMISSION_GRANTED &&
                    checkSelfPermission(Manifest.permission.READ_EXTERNAL_STORAGE) == PackageManager.PERMISSION_GRANTED &&
                    checkSelfPermission(Manifest.permission.WRITE_EXTERNAL_STORAGE) == PackageManager.PERMISSION_GRANTED;
        } else {
            return true;
        }
    }

    // request permission
    private void requestPermission() {
        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.M) {
            requestPermissions(new String[]{Manifest.permission.CAMERA,
                    Manifest.permission.READ_EXTERNAL_STORAGE,
                    Manifest.permission.WRITE_EXTERNAL_STORAGE}, 1);
        }
    }

效果圖:

小夜