refactor

Author: MrKnyaz

src/main/java/Classificator.java

@@ -5,6 +5,7 @@
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.util.ArrayList;
 import java.util.List;
 
 public class Classificator {
@@ -23,37 +24,39 @@ public Classificator() {
             modelGraph = new Graph();
             modelGraph.importGraphDef(graphData);
             session = new Session(modelGraph);
-
-            //Just print two main operations to look at shapes
-            System.out.println(modelGraph.operation("input").output(0));
-            System.out.println(modelGraph.operation("output").output(0));
         } catch(Exception e) {e.printStackTrace(); throw new RuntimeException(e);}
     }
 
-    public String classify(float[][][][] imageData) {
+    public List<String> classify(float[][][][] imageData) {
         Tensor imageTensor = Tensor.create(imageData, Float.class);
-        float[][] output = predict(imageTensor);
-        return findPredictedLabel(output);
+        float[][] prediction = predict(imageTensor);
+        return findPredictedLabel(prediction);
     }
 
     private float[][] predict(Tensor imageTensor) {
         Tensor result = session.runner()
                 .feed("input", imageTensor)
                 .fetch("output").run().get(0);
+        int batchSize = (int)result.shape()[0];
         //create prediction buffer
-        float[][] prediction = new float[1][1008];
+        float[][] prediction = new float[batchSize][1008];
         result.copyTo(prediction);
         return prediction;
     }
 
-    private String findPredictedLabel(float[][] prediction) {
-        int maxValueIndex = 0;
-        for (int i = 1; i < prediction[0].length; i++) {
-            if (prediction[0][maxValueIndex] < prediction[0][i]) {
-                maxValueIndex = i;
+    private List<String> findPredictedLabel(float[][] prediction) {
+        List<String> result = new ArrayList<>();
+        int batchSize = prediction.length;
+        for (int i = 0; i < batchSize; i++) {
+            //Finding maximum value for each predicted image
+            int maxValueIndex = 0;
+            for (int j = 1; j < prediction[i].length; j++) {
+                if (prediction[i][maxValueIndex] < prediction[i][j]) {
+                    maxValueIndex = j;
+                }
             }
+            result.add(labels.get(maxValueIndex) + ": " + (prediction[i][maxValueIndex] * 100) + "%");
         }
-        System.out.println(prediction[0][maxValueIndex]);
-        return labels.get(maxValueIndex);
+        return result;
     }
 }

src/main/java/ImageProcessor.java

@@ -25,22 +25,27 @@ public class ImageProcessor {
     public ImageProcessor() {
     }
 
-    public float[][][][] loadAndNormalizeImage(String path) {
-        IplImage origImg = cvLoadImage(getFullPath(path));
-        //Creating image placeholder to put resized image data
-        IplImage resizedImg = IplImage.create(width, height, origImg.depth(), origImg.nChannels());
-        cvResize(origImg, resizedImg);
-        return getRGBArray(resizedImg);
+    public float[][][][] loadAndNormalizeImages(String... path) {
+        //First dimension of the result is a number of images, because we may accept multiple paths
+        float[][][][] result = new float[path.length][height][width][3];
+        for (int i = 0; i < path.length; i++) {
+            IplImage origImg = cvLoadImage(getFullPath(path[i]));
+            //Creating image placeholder to put resized image data
+            IplImage resizedImg = IplImage.create(width, height, origImg.depth(), origImg.nChannels());
+            cvResize(origImg, resizedImg);
+            result[i] = getRGBArray(resizedImg);
+        }
+        return result;
     }
 
-    private float[][][][] getRGBArray(IplImage image) {
-        float[][][][] result = new float[1][image.height()][image.width()][3];
+    private float[][][] getRGBArray(IplImage image) {
+        float[][][] result = new float[image.height()][image.width()][3];
         for (int i = 0; i < image.height(); i++) {
             for (int j = 0; j < image.width(); j++) {
                 CvScalar pixel = cvGet2D(image, i, j);
-                result[0][i][j][0] = (float)(pixel.val(2) - mean) / scale; //R
-                result[0][i][j][1] = (float)(pixel.val(1) - mean) / scale; //G
-                result[0][i][j][2] = (float)(pixel.val(0) - mean) / scale; //B
+                result[i][j][0] = (float)(pixel.val(2) - mean) / scale; //R
+                result[i][j][1] = (float)(pixel.val(1) - mean) / scale; //G
+                result[i][j][2] = (float)(pixel.val(0) - mean) / scale; //B
             }
         }
         return result;

src/main/java/Main.java

@@ -1,24 +1,17 @@
-import org.tensorflow.*;
-
-import java.io.PrintStream;
-import java.nio.ByteBuffer;
-import java.nio.file.Files;
-import java.nio.file.Path;
-import java.nio.file.Paths;
-import java.time.LocalDateTime;
-import java.util.Iterator;
 import java.util.List;
 
 public class Main {
 
     public static void main(String[] args) throws Exception {
-        /* ImageProcessor returns 4-dimensional float array -> float[1][224][224][3]
+        /* ImageProcessor returns 4-dimensional float array -> float[?][224][224][3]
          * First dimension of the array is a batch for neural network,
-         * for this simple example I use just one image at a time for classification
-         * The code can be easily modified to load and classify multiple images.
+         * You may pass multiple image paths to ImageProcessor
          * */
         ImageProcessor imageProcessor = new ImageProcessor();
         Classificator classificator = new Classificator();
-        System.out.println(classificator.classify(imageProcessor.loadAndNormalizeImage("images/hyndai.jpg")));
+        List<String> result = classificator.classify(imageProcessor.loadAndNormalizeImages("images/ship.jpg", "images/hyndai.jpg"));
+        for(String label: result) {
+            System.out.println(label);
+        }
     }
 }