content: package com.amitshekhar.tflite;
  
  import android.annotation.SuppressLint;
  import android.content.res.AssetFileDescriptor;
  import android.content.res.AssetManager;
  import android.graphics.Bitmap;
  
  import org.tensorflow.lite.Interpreter;
  
  import java.io.BufferedReader;
  import java.io.FileInputStream;
  import java.io.IOException;
  import java.io.InputStreamReader;
  import java.nio.ByteBuffer;
  import java.nio.ByteOrder;
  import java.nio.MappedByteBuffer;
  import java.nio.channels.FileChannel;
  import java.util.ArrayList;
  import java.util.Comparator;
  import java.util.List;
  import java.util.PriorityQueue;
  
  /**
   * Created by amitshekhar on 17/03/18.
   */
  
  public class TensorFlowImageClassifier implements Classifier {
  
      private static final int MAX_RESULTS = 3;
      private static final int BATCH_SIZE = 1;
      private static final int PIXEL_SIZE = 3;
      private static final float THRESHOLD = 0.1f;
  
      private static final int IMAGE_MEAN = 128;
      private static final float IMAGE_STD = 128.0f;
  
      private Interpreter interpreter;
      private int inputSize;
      private List<String> labelList;
      private boolean quant;
  
      private TensorFlowImageClassifier() {
  
      }
  
      static Classifier create(AssetManager assetManager,
                               String modelPath,
                               String labelPath,
                               int inputSize,
                               boolean quant) throws IOException {
  
          TensorFlowImageClassifier classifier = new TensorFlowImageClassifier();
          classifier.interpreter = new Interpreter(classifier.loadModelFile(assetManager, modelPath), new Interpreter.Options());
          classifier.labelList = classifier.loadLabelList(assetManager, labelPath);
          classifier.inputSize = inputSize;
          classifier.quant = quant;
  
          return classifier;
      }
  
      @Override
      public List<Recognition> recognizeImage(Bitmap bitmap) {
          ByteBuffer byteBuffer = convertBitmapToByteBuffer(bitmap);
          if(quant){
              byte[][] result = new byte[1][labelList.size()];
              interpreter.run(byteBuffer, result);
              return getSortedResultByte(result);
          } else {
              float [][] result = new float[1][labelList.size()];
              interpreter.run(byteBuffer, result);
              return getSortedResultFloat(result);
          }
  
      }
  
      @Override
      public void close() {
          interpreter.close();
          interpreter = null;
      }
  
      private MappedByteBuffer loadModelFile(AssetManager assetManager, String modelPath) throws IOException {
          AssetFileDescriptor fileDescriptor = assetManager.openFd(modelPath);
          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 List<String> loadLabelList(AssetManager assetManager, String labelPath) throws IOException {
          List<String> labelList = new ArrayList<>();
          BufferedReader reader = new BufferedReader(new InputStreamReader(assetManager.open(labelPath)));
          String line;
          while ((line = reader.readLine()) != null) {
              labelList.add(line);
          }
          reader.close();
          return labelList;
      }
  
      private ByteBuffer convertBitmapToByteBuffer(Bitmap bitmap) {
          ByteBuffer byteBuffer;
  
          if(quant) {
              byteBuffer = ByteBuffer.allocateDirect(BATCH_SIZE * inputSize * inputSize * PIXEL_SIZE);
          } else {
              byteBuffer = ByteBuffer.allocateDirect(4 * BATCH_SIZE * inputSize * inputSize * PIXEL_SIZE);
          }
  
          byteBuffer.order(ByteOrder.nativeOrder());
          int[] intValues = new int[inputSize * inputSize];
          bitmap.getPixels(intValues, 0, bitmap.getWidth(), 0, 0, bitmap.getWidth(), bitmap.getHeight());
          int pixel = 0;
          for (int i = 0; i < inputSize; ++i) {
              for (int j = 0; j < inputSize; ++j) {
                  final int val = intValues[pixel++];
                  if(quant){
                      byteBuffer.put((byte) ((val >> 16) & 0xFF));
                      byteBuffer.put((byte) ((val >> 8) & 0xFF));
                      byteBuffer.put((byte) (val & 0xFF));
                  } else {
                      byteBuffer.putFloat((((val >> 16) & 0xFF)-IMAGE_MEAN)/IMAGE_STD);
                      byteBuffer.putFloat((((val >> 8) & 0xFF)-IMAGE_MEAN)/IMAGE_STD);
                      byteBuffer.putFloat((((val) & 0xFF)-IMAGE_MEAN)/IMAGE_STD);
                  }
  
              }
          }
          return byteBuffer;
      }
  
      @SuppressLint("DefaultLocale")
      private List<Recognition> getSortedResultByte(byte[][] labelProbArray) {
  
          PriorityQueue<Recognition> pq =
                  new PriorityQueue<>(
                          MAX_RESULTS,
                          new Comparator<Recognition>() {
                              @Override
                              public int compare(Recognition lhs, Recognition rhs) {
                                  return Float.compare(rhs.getConfidence(), lhs.getConfidence());
                              }
                          });
  
          for (int i = 0; i < labelList.size(); ++i) {
              float confidence = (labelProbArray[0][i] & 0xff) / 255.0f;
              if (confidence > THRESHOLD) {
                  pq.add(new Recognition("" + i,
                          labelList.size() > i ? labelList.get(i) : "unknown",
                          confidence, quant));
              }
          }
  
          final ArrayList<Recognition> recognitions = new ArrayList<>();
          int recognitionsSize = Math.min(pq.size(), MAX_RESULTS);
          for (int i = 0; i < recognitionsSize; ++i) {
              recognitions.add(pq.poll());
          }
  
          return recognitions;
      }
  
      @SuppressLint("DefaultLocale")
      private List<Recognition> getSortedResultFloat(float[][] labelProbArray) {
  
          PriorityQueue<Recognition> pq =
                  new PriorityQueue<>(
                          MAX_RESULTS,
                          new Comparator<Recognition>() {
                              @Override
                              public int compare(Recognition lhs, Recognition rhs) {
                                  return Float.compare(rhs.getConfidence(), lhs.getConfidence());
                              }
                          });
  
          for (int i = 0; i < labelList.size(); ++i) {
              float confidence = labelProbArray[0][i];
              if (confidence > THRESHOLD) {
                  pq.add(new Recognition("" + i,
                          labelList.size() > i ? labelList.get(i) : "unknown",
                          confidence, quant));
              }
          }
  
          final ArrayList<Recognition> recognitions = new ArrayList<>();
          int recognitionsSize = Math.min(pq.size(), MAX_RESULTS);
          for (int i = 0; i < recognitionsSize; ++i) {
              recognitions.add(pq.poll());
          }
  
          return recognitions;
      }
  
  }
  