使用机器学习优化智能选区算法

qmy612 · Jul 3, 2019 · 5ece8ca · 5ece8ca
1 parent 9b70fca
commit 5ece8ca
Show file tree

Hide file tree

Showing 11 changed files with 167 additions and 11 deletions.
diff --git a/app/build.gradle b/app/build.gradle
@@ -21,6 +21,10 @@ android {
     lintOptions {
         abortOnError false
     }
+    aaptOptions {
+        noCompress "tflite"
+        noCompress "lite"
+    }
 }
 
 dependencies {

diff --git a/app/src/main/AndroidManifest.xml b/app/src/main/AndroidManifest.xml
@@ -4,6 +4,7 @@
     package="me.pqpo.smartcropper">
 
     <application
+        android:name=".App"
         android:allowBackup="true"
         android:icon="@mipmap/ic_launcher"
         android:label="@string/app_name"

diff --git a/app/src/main/java/me/pqpo/smartcropper/App.java b/app/src/main/java/me/pqpo/smartcropper/App.java
@@ -0,0 +1,18 @@
+package me.pqpo.smartcropper;
+
+import android.app.Application;
+
+import me.pqpo.smartcropperlib.SmartCropper;
+
+/**
+ * Created by [email protected] on 2019/7/3.
+ */
+public class App extends Application {
+
+    @Override
+    public void onCreate() {
+        super.onCreate();
+        // 如果使用机器学习代替 Canny 算子，请初始化 ImageDetector
+        SmartCropper.buildImageDetector(this);
+    }
+}
diff --git a/smartcropperlib/build.gradle b/smartcropperlib/build.gradle
@@ -38,4 +38,5 @@ android {
 
 dependencies {
     implementation fileTree(dir: 'libs', include: ['*.jar'])
+    implementation 'org.tensorflow:tensorflow-lite:0.0.0-nightly'
 }
diff --git a/smartcropperlib/src/main/assets/models/hed_lite_model_quantize.tflite b/smartcropperlib/src/main/assets/models/hed_lite_model_quantize.tflite
diff --git a/smartcropperlib/src/main/cpp/Scanner.cpp b/smartcropperlib/src/main/cpp/Scanner.cpp
@@ -12,8 +12,9 @@ static bool sortByArea(const vector<Point> &v1, const vector<Point> &v2) {
     return v1Area > v2Area;
 }
 
-Scanner::Scanner(cv::Mat& bitmap) {
+Scanner::Scanner(cv::Mat& bitmap, bool canny) {
     srcBitmap = bitmap;
+    Scanner::canny = canny;
 }
 
 Scanner::~Scanner() {
@@ -138,6 +139,9 @@ Mat Scanner::resizeImage() {
 Mat Scanner::preprocessedImage(Mat &image, int cannyValue, int blurValue) {
     Mat grayMat;
     cvtColor(image, grayMat, CV_BGR2GRAY);
+    if (!canny) {
+        return grayMat;
+    }
     if (isHisEqual){
         equalizeHist(grayMat, grayMat);
     }
@@ -147,7 +151,7 @@ Mat Scanner::preprocessedImage(Mat &image, int cannyValue, int blurValue) {
     Canny(blurMat, cannyMat, 50, cannyValue, 3);
     Mat thresholdMat;
     threshold(cannyMat, thresholdMat, 0, 255, CV_THRESH_OTSU);
-    return cannyMat;
+    return thresholdMat;
 }
 
 vector<Point> Scanner::selectPoints(vector<Point> points) {

diff --git a/smartcropperlib/src/main/cpp/include/Scanner.h b/smartcropperlib/src/main/cpp/include/Scanner.h
@@ -13,13 +13,15 @@ namespace scanner{
     public:
         int resizeThreshold = 500;
 
-        Scanner(cv::Mat& bitmap);
+        Scanner(cv::Mat& bitmap, bool canny);
         virtual ~Scanner();
         std::vector<cv::Point> scanPoint();
     private:
         cv::Mat srcBitmap;
         float resizeScale = 1.0f;
 
+        bool canny = true;
+
         bool isHisEqual = false;
 
         cv::Mat resizeImage();

diff --git a/smartcropperlib/src/main/cpp/smart_cropper.cpp b/smartcropperlib/src/main/cpp/smart_cropper.cpp
@@ -26,15 +26,15 @@ static jobject createJavaPoint(JNIEnv *env, Point point_) {
     return env -> NewObject(gPointInfo.jClassPoint, gPointInfo.jMethodInit, point_.x, point_.y);
 }
 
-static void native_scan(JNIEnv *env, jclass type, jobject srcBitmap, jobjectArray outPoint_) {
+static void native_scan(JNIEnv *env, jclass type, jobject srcBitmap, jobjectArray outPoint_, jboolean canny) {
     if (env -> GetArrayLength(outPoint_) != 4) {
         return;
     }
     Mat srcBitmapMat;
     bitmap_to_mat(env, srcBitmap, srcBitmapMat);
     Mat bgrData(srcBitmapMat.rows, srcBitmapMat.cols, CV_8UC3);
     cvtColor(srcBitmapMat, bgrData, CV_RGBA2BGR);
-    scanner::Scanner docScanner(bgrData);
+    scanner::Scanner docScanner(bgrData, canny);
     std::vector<Point> scanPoints = docScanner.scanPoint();
     if (scanPoints.size() == 4) {
         for (int i = 0; i < 4; ++i) {
@@ -98,7 +98,7 @@ static JNINativeMethod gMethods[] = {
 
         {
                 "nativeScan",
-                "(Landroid/graphics/Bitmap;[Landroid/graphics/Point;)V",
+                "(Landroid/graphics/Bitmap;[Landroid/graphics/Point;Z)V",
                 (void*)native_scan
         },
 

diff --git a/smartcropperlib/src/main/java/me/pqpo/smartcropperlib/ImageDetector.java b/smartcropperlib/src/main/java/me/pqpo/smartcropperlib/ImageDetector.java
@@ -0,0 +1,107 @@
+package me.pqpo.smartcropperlib;
+
+import android.content.Context;
+import android.content.res.AssetFileDescriptor;
+import android.graphics.Bitmap;
+import android.text.TextUtils;
+
+import org.tensorflow.lite.Interpreter;
+
+import java.io.FileInputStream;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.nio.MappedByteBuffer;
+import java.nio.channels.FileChannel;
+
+public class ImageDetector {
+
+    private static final String MODEL_FILE = "models/hed_lite_model_quantize.tflite";
+
+    private int desiredSize = 256;
+
+
+    private int[] intValues = new int[desiredSize * desiredSize];
+
+    protected ByteBuffer imgData = null;
+    protected ByteBuffer outImgData = null;
+
+    protected Interpreter tflite;
+
+    public ImageDetector(Context context) throws IOException {
+        this(context, MODEL_FILE);
+    }
+
+    public ImageDetector(Context context, String modelFile) throws IOException {
+        if (TextUtils.isEmpty(modelFile)) {
+            modelFile = MODEL_FILE;
+        }
+        MappedByteBuffer tfliteModel = loadModelFile(context, modelFile);
+        Interpreter.Options tfliteOptions = new Interpreter.Options();
+        tflite = new Interpreter(tfliteModel, tfliteOptions);
+        imgData = ByteBuffer.allocateDirect(desiredSize * desiredSize * 3 * Float.SIZE / Byte.SIZE);
+        imgData.order(ByteOrder.nativeOrder());
+
+        outImgData = ByteBuffer.allocateDirect(desiredSize * desiredSize * Float.SIZE / Byte.SIZE);
+        outImgData.order(ByteOrder.nativeOrder());
+    }
+
+    public synchronized Bitmap detectImage(Bitmap bitmap) {
+        if (bitmap == null) {
+            return null;
+        }
+        imgData.clear();
+        outImgData.clear();
+        bitmap = Bitmap.createScaledBitmap(bitmap, desiredSize, desiredSize, false);
+        convertBitmapToByteBuffer(bitmap);
+        tflite.run(imgData, outImgData);
+        return convertOutputBufferToBitmap(outImgData);
+    }
+
+    private void convertBitmapToByteBuffer(Bitmap bitmap) {
+        if (imgData == null) {
+            return;
+        }
+        bitmap.getPixels(intValues, 0, desiredSize, 0, 0, desiredSize, desiredSize);
+        imgData.rewind();
+        // Convert the image to floating point.
+        int pixel = 0;
+        for (int i = 0; i < desiredSize; ++i) {
+            for (int j = 0; j < desiredSize; ++j) {
+                final int pixelValue = intValues[pixel++];
+                imgData.putFloat(((pixelValue >> 16) & 0xFF));
+                imgData.putFloat(((pixelValue >> 8) & 0xFF));
+                imgData.putFloat((pixelValue & 0xFF));
+            }
+        }
+    }
+
+    private Bitmap convertOutputBufferToBitmap(ByteBuffer outImgData) {
+        if (outImgData == null) {
+            return null;
+        }
+        outImgData.rewind();
+        Bitmap bitmap_out = Bitmap.createBitmap(desiredSize , desiredSize, Bitmap.Config.ARGB_8888);
+        int[] pixels = new int[desiredSize * desiredSize];
+        for (int i = 0; i < desiredSize * desiredSize; i++) {
+            float val = outImgData.getFloat();
+            if (val > 0.2) {
+                pixels[i] = 0xFFFFFFFF;
+            } else {
+                pixels[i] = 0xFF000000;
+            }
+        }
+        bitmap_out.setPixels(pixels, 0, desiredSize, 0, 0, desiredSize, desiredSize);
+        return bitmap_out;
+    }
+
+    private MappedByteBuffer loadModelFile(Context activity, String modelFile) throws IOException {
+        AssetFileDescriptor fileDescriptor = activity.getAssets().openFd(modelFile);
+        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);
+    }
+
+}
diff --git a/smartcropperlib/src/main/java/me/pqpo/smartcropperlib/SmartCropper.java b/smartcropperlib/src/main/java/me/pqpo/smartcropperlib/SmartCropper.java
@@ -1,9 +1,10 @@
 package me.pqpo.smartcropperlib;
 
+import android.content.Context;
 import android.graphics.Bitmap;
 import android.graphics.Point;
 
-import java.util.ArrayList;
+import java.io.IOException;
 
 import me.pqpo.smartcropperlib.utils.CropUtils;
 
@@ -13,6 +14,20 @@
 
 public class SmartCropper {
 
+    private static ImageDetector sImageDetector = null;
+
+    public static void buildImageDetector(Context context) {
+        SmartCropper.buildImageDetector(context, null);
+    }
+
+    public static void buildImageDetector(Context context, String modelFile) {
+        try {
+            sImageDetector = new ImageDetector(context, modelFile);
+        } catch (IOException e) {
+            e.printStackTrace();
+        }
+    }
+
     /**
      *  输入图片扫描边框顶点
      * @param srcBmp 扫描图片
@@ -22,8 +37,14 @@ public static Point[] scan(Bitmap srcBmp) {
         if (srcBmp == null) {
             throw new IllegalArgumentException("srcBmp cannot be null");
         }
+        if (sImageDetector != null) {
+            Bitmap bitmap = sImageDetector.detectImage(srcBmp);
+            if (bitmap != null) {
+                srcBmp = Bitmap.createScaledBitmap(bitmap, srcBmp.getWidth(), srcBmp.getHeight(), false);
+            }
+        }
         Point[] outPoints = new Point[4];
-        nativeScan(srcBmp, outPoints);
+        nativeScan(srcBmp, outPoints, sImageDetector == null);
         return outPoints;
     }
 
@@ -57,7 +78,7 @@ public static Bitmap crop(Bitmap srcBmp, Point[] cropPoints) {
 
 
 
-    private static native void nativeScan(Bitmap srcBitmap, Point[] outPoints);
+    private static native void nativeScan(Bitmap srcBitmap, Point[] outPoints, boolean canny);
 
     private static native void nativeCrop(Bitmap srcBitmap, Point[] points, Bitmap outBitmap);
 

diff --git a/smartcropperlib/src/main/java/me/pqpo/smartcropperlib/view/CropImageView.java b/smartcropperlib/src/main/java/me/pqpo/smartcropperlib/view/CropImageView.java
@@ -25,8 +25,6 @@
 import android.view.MotionEvent;
 import android.widget.ImageView;
 
-import java.util.Arrays;
-
 import me.pqpo.smartcropperlib.R;
 import me.pqpo.smartcropperlib.SmartCropper;
 import me.pqpo.smartcropperlib.utils.CropUtils;