Added code for image classification with OpenCV Java

Author: Anastasia

DNN-OpenCV-Classification-with-Java/DnnOpenCV.java

@@ -0,0 +1,122 @@
+import org.opencv.core.Core;
+import org.opencv.core.Mat;
+import org.opencv.core.Rect;
+import org.opencv.core.Scalar;
+import org.opencv.core.Size;
+import org.opencv.dnn.Net;
+import org.opencv.dnn.Dnn;
+import org.opencv.imgproc.Imgproc;
+import org.opencv.imgcodecs.Imgcodecs;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+import java.nio.file.Files;
+import java.nio.file.Paths;
+
+import org.opencv.core.CvType;
+
+
+public class DnnOpenCV {
+    private static final int TARGET_IMG_WIDTH = 224;
+    private static final int TARGET_IMG_HEIGHT = 224;
+
+    private static final double SCALE_FACTOR = 1 / 255.0;
+
+    private static final String IMAGENET_CLASSES = "imagenet_classes.txt";
+    private static final String MODEL_PATH = "models/pytorch_mobilenet.onnx";
+
+
+    public static ArrayList<String> getImgLabels(String imgLabelsFilePath) throws IOException {
+        ArrayList<String> imgLabels;
+        try (Stream<String> lines = Files.lines(Paths.get(imgLabelsFilePath))) {
+            imgLabels = lines.collect(Collectors.toCollection(ArrayList::new));
+        }
+        return imgLabels;
+    }
+
+    public static Mat centerCrop(Mat inputImage) {
+        int y1 = Math.round((inputImage.rows() - TARGET_IMG_HEIGHT) / 2);
+        int y2 = Math.round(y1 + TARGET_IMG_HEIGHT);
+        int x1 = Math.round((inputImage.cols() - TARGET_IMG_WIDTH) / 2);
+        int x2 = Math.round(x1 + TARGET_IMG_WIDTH);
+
+        Rect centerRect = new Rect(x1, y1, (x2 - x1), (y2 - y1));
+        Mat croppedImage = new Mat(inputImage, centerRect);
+
+        return croppedImage;
+    }
+
+    public static Mat getPreprocessedImage(String imagePath) {
+        // define mean and standard deviation
+        Scalar mean = new Scalar(0.485, 0.456, 0.406);
+        Scalar std = new Scalar(0.229, 0.224, 0.225);
+
+        // get the image from the internal resource folder
+        Mat image = Imgcodecs.imread(imagePath);
+
+        // resize input image
+        Imgproc.resize(image, image, new Size(256, 256));
+
+        // create empty Mat images for float conversions
+        Mat imgFloat = new Mat(image.rows(), image.cols(), CvType.CV_32FC3);
+
+        // convert input image to float type
+        image.convertTo(imgFloat, CvType.CV_32FC3, SCALE_FACTOR);
+
+        // crop input image
+        imgFloat = centerCrop(imgFloat);
+
+        // prepare DNN input
+        Mat blob = Dnn.blobFromImage(
+                imgFloat,
+                1.0, /* default scalefactor */
+                new Size(TARGET_IMG_WIDTH, TARGET_IMG_HEIGHT), /* target size */
+                mean,  /* mean */
+                true, /* swapRB */
+                false /* crop */
+        );
+
+        // divide on std
+        Core.divide(blob, std, blob);
+
+        return blob;
+    }
+
+    public static void getPredictedClass(Mat classificationResult) {
+        ArrayList<String> imgLabels = new ArrayList<String>();
+        try {
+            imgLabels = getImgLabels(IMAGENET_CLASSES);
+        } catch (IOException ex) {
+
+        }
+        // obtain max prediction result
+        Core.MinMaxLocResult mm = Core.minMaxLoc(classificationResult);
+        double maxValIndex = mm.maxLoc.x;
+        System.out.println("Predicted Class: " + imgLabels.get((int) maxValIndex));
+    }
+
+    public static void main(String[] args) {
+        String imageLocation = "images/coffee.jpg";
+
+        // load the OpenCV native library
+        System.loadLibrary(Core.NATIVE_LIBRARY_NAME);
+
+        // read and process the input image
+        Mat inputBlob = DnnOpenCV.getPreprocessedImage(imageLocation);
+
+        // read generated ONNX model into org.opencv.dnn.Net object
+        Net dnnNet = Dnn.readNetFromONNX(DnnOpenCV.MODEL_PATH);
+        System.out.println("DNN from ONNX was successfully loaded!");
+
+        // set OpenCV model input
+        dnnNet.setInput(inputBlob);
+
+        // provide inference
+        Mat classification = dnnNet.forward();
+
+        // decode classification results
+        DnnOpenCV.getPredictedClass(classification);
+    }
+}

DNN-OpenCV-Classification-with-Java/Mobilenetv2ToOnnx.py

@@ -0,0 +1,127 @@
+import argparse
+import os
+
+import cv2
+import numpy as np
+import onnx
+import onnxruntime
+import torch
+from albumentations import (
+    CenterCrop,
+    Compose,
+    Normalize,
+    Resize,
+)
+from torchvision import models
+
+
+def compare_pytorch_onnx(
+    original_model_preds, onnx_model_path, input_image,
+):
+    # get onnx result
+    session = onnxruntime.InferenceSession(onnx_model_path)
+    input_name = session.get_inputs()[0].name
+    onnx_result = session.run([], {input_name: input_image})
+    onnx_result = np.squeeze(onnx_result, axis=0)
+
+    print("Checking PyTorch model and converted ONNX model outputs ... ")
+    for test_onnx_result, gold_result in zip(onnx_result, original_model_preds):
+        np.testing.assert_almost_equal(
+            gold_result, test_onnx_result, decimal=3,
+        )
+    print("PyTorch and ONNX output values are equal! \n")
+
+
+def get_onnx_model(
+    original_model, input_image, model_path="models", model_name="pytorch_mobilenet",
+):
+    # create model root dir
+    os.makedirs(model_path, exist_ok=True)
+
+    model_name = os.path.join(model_path, model_name + ".onnx")
+
+    torch.onnx.export(
+        original_model, torch.Tensor(input_image), model_name, verbose=True,
+    )
+    print("ONNX model was successfully generated: {} \n".format(model_name))
+
+    return model_name
+
+
+def get_preprocessed_image(image_name):
+    # read image
+    original_image = cv2.imread(image_name)
+
+    # convert original image to RGB format
+    image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)
+
+    # transform input image:
+    # 1. resize the image
+    # 2. crop the image
+    # 3. normalize: subtract mean and divide by standard deviation
+    transform = Compose(
+        [
+            Resize(height=256, width=256),
+            CenterCrop(224, 224),
+            Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ],
+    )
+    image = transform(image=image)["image"]
+
+    # change the order of channels
+    image = image.transpose(2, 0, 1)
+    return np.expand_dims(image, axis=0)
+
+
+def get_predicted_class(pytorch_preds):
+    # read ImageNet class id to name mapping
+    with open("imagenet_classes.txt") as f:
+        labels = [line.strip() for line in f.readlines()]
+
+    # find the class with the maximum score
+    pytorch_class_idx = np.argmax(pytorch_preds, axis=1)
+    predicted_pytorch_label = labels[pytorch_class_idx[0]]
+
+    # print top predicted class
+    print("Predicted class by PyTorch model: ", predicted_pytorch_label)
+
+
+def get_execution_arguments():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--input_image",
+        type=str,
+        help="Define the full input image path, including its name",
+        default="images/coffee.jpg",
+    )
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    # get the test case parameters
+    args = get_execution_arguments()
+
+    # read and process the input image
+    image = get_preprocessed_image(image_name=args.input_image)
+
+    # obtain original model
+    pytorch_model = models.mobilenet_v2(pretrained=True)
+
+    # provide inference of the original PyTorch model
+    pytorch_model.eval()
+    pytorch_predictions = pytorch_model(torch.Tensor(image)).detach().numpy()
+
+    # obtain OpenCV generated ONNX model
+    onnx_model_path = get_onnx_model(original_model=pytorch_model, input_image=image)
+
+    # check if conversion succeeded
+    onnx_model = onnx.load(onnx_model_path)
+    onnx.checker.check_model(onnx_model)
+
+    # check onnx model output
+    compare_pytorch_onnx(
+        pytorch_predictions, onnx_model_path, image,
+    )
+
+    # decode classification results
+    get_predicted_class(pytorch_preds=pytorch_predictions)