image classifier

package com.example.chesshelper;

import android.app.Activity;
import android.graphics.Bitmap;

import org.checkerframework.checker.nullness.qual.NonNull;
import org.tensorflow.lite.DataType;
import org.tensorflow.lite.Interpreter;
import org.tensorflow.lite.support.common.FileUtil;
import org.tensorflow.lite.support.common.TensorProcessor;
import org.tensorflow.lite.support.common.ops.NormalizeOp;
import org.tensorflow.lite.support.image.ImageProcessor;
import org.tensorflow.lite.support.image.TensorImage;
import org.tensorflow.lite.support.image.ops.ResizeOp;
import org.tensorflow.lite.support.image.ops.ResizeWithCropOrPadOp;
import org.tensorflow.lite.support.label.TensorLabel;
import org.tensorflow.lite.support.tensorbuffer.TensorBuffer;

import java.io.IOException;
import java.nio.MappedByteBuffer;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;

public class ImageIdentifier {
    private static final float PROBABILITY_MEAN = 0.0f;
    private static final float PROBABILITY_STD = 255.0f;
    private static final float IMAGE_MEAN = 0.0f;
    private static final float IMAGE_STD = 1.0f;
    private final int imageResizeX;
    private final int imageResizeY;
    private final List<String> labels;
    private TensorImage inputImageBuffer;
    private TensorBuffer probabilityImageBuffer;
    private TensorProcessor probabilityProcessor;
    private Interpreter imageClassifier;
    public ImageIdentifier(Activity activity) throws IOException
    {
        MappedByteBuffer classifierModel = FileUtil.loadMappedFile(activity, "detect.tflite");
        labels = FileUtil.loadLabels(activity, "labelMap.txt");
        imageClassifier = new Interpreter(classifierModel, null);
        int imageTensorIndex = 0;
        int outputImageTensorIndex = 0;

        int[] imageShape = imageClassifier.getInputTensor(imageTensorIndex).shape();
        DataType inputData = imageClassifier.getInputTensor(imageTensorIndex).dataType();
        int[] outputImageShape = imageClassifier.getOutputTensor(outputImageTensorIndex).shape();
        DataType outputData = imageClassifier.getOutputTensor(outputImageTensorIndex).dataType();

        imageResizeX = imageShape[1];
        imageResizeY = imageShape[2];

        inputImageBuffer = new TensorImage(inputData);
        probabilityImageBuffer = TensorBuffer.createFixedSize(outputImageShape, outputData);
        probabilityProcessor = new TensorProcessor.Builder().add(new NormalizeOp(PROBABILITY_MEAN, PROBABILITY_STD)).build();
    }

    public List<Recognition> recognisedImage(Bitmap bitmap, int sensorOrientation)
    {
        List<Recognition> recognitionsList = new ArrayList<>();
        inputImageBuffer = loadImage(bitmap, sensorOrientation);
        imageClassifier.run(inputImageBuffer.getBuffer(), probabilityImageBuffer.getBuffer().rewind());
        Map<String, Float> labelledProbability = new TensorLabel(labels,
                probabilityProcessor.process(probabilityImageBuffer)).getMapWithFloatValue();
        for (Map.Entry<String, Float> entry : labelledProbability.entrySet())
        {
            recognitionsList.add(new Recognition(entry.getKey(), entry.getValue()));
        }
        Collections.sort(recognitionsList);
        recognitionsList.subList(0, 5 > recognitionsList.size() ? recognitionsList.size() : 5).clear();
        return recognitionsList;
    }

    private TensorImage loadImage(Bitmap bitmap, int sensorOrientation) {
        inputImageBuffer.load(bitmap);
        int noOfRotations = sensorOrientation / 90;
        int cropSize = Math.min(bitmap.getWidth(), bitmap.getHeight());
        ImageProcessor imageProcessor = new ImageProcessor.Builder()
                .add(new ResizeWithCropOrPadOp(cropSize, cropSize))
                .add(new ResizeOp(imageResizeX, imageResizeY, ResizeOp.ResizeMethod.NEAREST_NEIGHBOR))
                .add(new NormalizeOp(IMAGE_MEAN, IMAGE_STD)).build();
        return imageProcessor.process(inputImageBuffer);
    }

    class Recognition implements Comparable
    {
        private String name;
        private float confidence;

        public Recognition() {
        }

        public Recognition(String name, float confidence) {
            this.name = name;
            this.confidence = confidence;
        }

        public String getName() {
            return name;
        }

        public void setName(String name) {
            this.name = name;
        }

        public float getConfidence() {
            return confidence;
        }

        public void setConfidence(float confidence) {
            this.confidence = confidence;
        }

        @Override
        public String toString() {
            return "Recognition{" +
                    "name='" + name + '\'' +
                    ", confidence=" + confidence +
                    '}';
        }

        @Override
        public int compareTo(Object object) {
            return Float.compare(((Recognition)object).confidence, this.confidence);
        }
    }
}