Folosirea unui model ONNX Resnet în Java

  • 6 minute de citit
  • 6 exemple de cod în Java și XML

Disponibil și în English

Pe această pagină

Salut tuturor! 👋

Acesta este un articol de continuare al celui precedent. Pe scurt, am început cursul de la fast.ai ca să învăț mai multe despre Deep Learning. Am construit un model simplu de etichetare a imaginilor folosind arhitectura Resnet (nu vă faceți griji în privința lucrării, nu am citit-o, GitHub Copilot a sugerat-o în timp ce scriam acest articol).

FastAI este o bibliotecă Python de nivel înalt care îți permite să antrenezi modele ML complexe foarte rapid și eficient. După antrenarea modelului, l-am exportat în formatul ONNX și în acest articol îl folosim într-o aplicație Java.

Modelul va face cuvinte cheie pentru imagini. Îi dai o imagine și îți va da câteva cuvinte cheie care descriu imaginea. Dacă ești fotograf, poți folosi acest model ca să îți etichetezi imaginile și să le faci mai ușor de găsit. 📸

Java ONNX Runtime

Creează un proiect Java nou cu Maven și adaugă următoarea dependență:

    <dependency>        <groupId>com.microsoft.onnxruntime</groupId>        <artifactId>onnxruntime</artifactId>        <version>1.15.1</version>    </dependency>

Aceasta este singura dependență de care avem nevoie.

Modelul

Dacă vrei să urmărești pas cu pas, poți lua modelul de aici

Clasele modelului sunt disponibile aici. Va trebui să le convertești într-un array sau o listă Java.

Am salvat categoriile într-o clasă Categories, în următorul format: Categories.java:

package main.java;import java.util.Arrays;import java.util.List;public class Categories {    public static List<String> CATEGORIES = Arrays.asList(        "1",        "2",        "3",    );}

Codul

Ca să obținem predicțiile pentru o imagine avem nevoie de modelul, categoriile și imaginea pe care vrem să o etichetăm cu cuvinte cheie. Trebuie să parcurgem următorii pași:

  1. Încărcăm modelul
  2. Creăm un tensor de intrare
  3. Rulăm modelul
  4. Obținem tensorul de ieșire
  5. Obținem predicțiile

Codul care face asta este următorul:

    public void run() throws OrtException {        System.out.println("Onnx Runtime Java Image Prediction");        // 1. Încarcă modelul.        var env = OrtEnvironment.getEnvironment();        var session = env.createSession(this.modelPath, new OrtSession.SessionOptions());        // Obține numele intrărilor și ieșirilor        var inputName = session.getInputNames().iterator().next();        var outputName = session.getOutputNames().iterator().next();        // 2. Creează tensorul de intrare        OnnxTensor inputTensor = OnnxTensor.createTensor(env, processImage(this.imagePath));        // 3. Rulează modelul.        var inputs = Map.of(inputName, inputTensor);        var results = session.run(inputs);        // 4. Obține tensorul de ieșire        var outputTensor = results.get(outputName);        if (outputTensor.isPresent()) {            // 5. Obține rezultatele predicției            float[][] floatBuffer = (float[][]) outputTensor.get().getValue();            ArrayList<String> predictions = new ArrayList<String>();            // filtrează bufferul după prag            for (int i = 0; i < floatBuffer[0].length; i++) {                if (floatBuffer[0][i] > -0.5) {                    predictions.add(String.format("%s: %f - %s", i, floatBuffer[0][i], Categories.CATEGORIES.get(i)));                }            }            // Afișează rezultatele            System.out.println("Predictions: ");            for (String prediction : predictions) {                System.out.println(prediction);            }        } else {            System.out.println("Failed to predict!");        }    }

☕

Iar pasul doi: 2. Creează tensorul de intrare nu este inclus în codul de mai sus pentru că este puțin mai complicat. Codul se află în metoda processImage(this.imagePath).

Dacă rulezi programul, vei obține următorul rezultat, în funcție de imaginea pe care o folosești:

Onnx Runtime Java Image PredictionPredictions: 25: -0.032844 - architecture150: 0.288728 - day381: -0.130801 - nature387: 0.343386 - no people404: 1.527045 - outdoors

Procesarea imaginii

Această secțiune descrie codul din metoda processImage și pasul doi: 2. Creează tensorul de intrare.

Tensorul de intrare este un array de 4 dimensiuni cu valori float. Dimensiunile sunt: 1 x 3 x 224 x 224.

Prima dimensiune este dimensiunea batch-ului. În cazul nostru, vom prezice câte o imagine pe rând, deci dimensiunea batch-ului este 1.

A doua dimensiune este numărul de canale. Folosim imagini RGB, deci avem 3 canale.

Ultimele două dimensiuni sunt lățimea și înălțimea imaginii. În cazul nostru, folosim imagini de 224x224.

Notă: Nu aveam nicio idee că prima dimensiune este dimensiunea batch-ului, paragraful acela a fost generat de GitHub Copilot 😲.

Codul de procesare a imaginii face următoarele lucruri:

  1. Citește imaginea
  2. Decupează imaginea
  3. Redimensionează imaginea la 224x224
  4. Calculează valorile tensorului din RGB.
    public float[][][][] processImage(String imagePath) {        try {            float[][][][] tensorData = new float[1][3][224][224];            var mean = new float[] { 0.485f, 0.456f, 0.406f };            var standardDeviation = new float[] { 0.229f, 0.224f, 0.225f };            // Citește imaginea            File imageFile = new File(imagePath);            var image = ImageIO.read(imageFile);            // decupează imaginea            int width = image.getWidth();            int height = image.getHeight();            int startX = 0;            int startY = 0;            if (width > height) {                startX = (width - height) / 2;                width = height;            } else {                startY = (height - width) / 2;                height = width;            }            image = image.getSubimage(startX, startY, width, height);            // DEBUG salvează imaginea pe disc            // ImageIO.write(image, "jpg", new File("C:\\Users\\NUCULABS\\IdeaProjects\\untitled\\src\\test\\java\\main\\resources\\debug.jpg"));            // Redimensionează imaginea            var resizedImage = image.getScaledInstance(224, 224, 4);            // Procesează imaginea            BufferedImage scaledImage = new BufferedImage(224, 224, BufferedImage.TYPE_INT_ARGB);            scaledImage.getGraphics().drawImage(resizedImage, 0, 0, null);            for (var y = 0; y < scaledImage.getHeight(); y++) {                for (var x = 0; x < scaledImage.getWidth(); x++) {                    int pixel = scaledImage.getRGB(x,y);                    // Obține valorile RGB                    tensorData[0][0][y][x] = (((pixel >> 16) & 0xFF) / 255f - mean[0]) / standardDeviation[0];                    tensorData[0][1][y][x] = (((pixel >> 16) & 0xFF) / 255f - mean[1]) / standardDeviation[1];                    tensorData[0][2][y][x] = (((pixel >> 16) & 0xFF) / 255f - mean[2]) / standardDeviation[2];                }            }            return tensorData;        } catch (IOException e) {            throw new RuntimeException(e);        }    }

Concluzie

Așa poți folosi API-ul ONNX Runtime pentru Java ca să faci predicții pe imagini.

Mulțumesc pentru lectură! Dacă ai întrebări, lasă un comentariu mai jos. ❤️

Cod complet:

Codul complet pentru fișierul Program.java:

package main.java;import ai.onnxruntime.OnnxTensor;import ai.onnxruntime.OrtEnvironment;import ai.onnxruntime.OrtException;import ai.onnxruntime.OrtSession;import javax.imageio.ImageIO;import java.awt.image.BufferedImage;import java.io.File;import java.io.IOException;import java.util.ArrayList;import java.util.Map;public class Program {    String modelPath = "C:\\Users\\nutiu\\IdeaProjects\\untitled\\src\\test\\java\\main\\resources\\resnet34_10_epochs.onnx";    String imagePath = "C:\\Users\\nutiu\\Downloads\\4.jpg";    public float[][][][] processImage(String imagePath) {        try {            float[][][][] tensorData = new float[1][3][224][224];            var mean = new float[] { 0.485f, 0.456f, 0.406f };            var standardDeviation = new float[] { 0.229f, 0.224f, 0.225f };            // Citește imaginea            File imageFile = new File(imagePath);            BufferedImage bufferedImage = ImageIO.read(imageFile);            // Decupează imaginea            int width = image.getWidth();            int height = image.getHeight();            int startX = 0;            int startY = 0;            if (width > height) {                startX = (width - height) / 2;                width = height;            } else {                startY = (height - width) / 2;                height = width;            }            var image = image.getSubimage(startX, startY, width, height);            // ImageIO.write(image, "jpg", new File("C:\\Users\\nutiu\\IdeaProjects\\untitled\\src\\test\\java\\main\\resources\\debug.jpg"));            // Redimensionează imaginea            var resizedImage = image.getScaledInstance(224, 224, 4);            // Procesează imaginea            BufferedImage scaledImage = new BufferedImage(224, 224, BufferedImage.TYPE_4BYTE_ABGR);            scaledImage.getGraphics().drawImage(resizedImage, 0, 0, null);            for (var y = 0; y < scaledImage.getHeight(); y++) {                for (var x = 0; x < scaledImage.getWidth(); x++) {                    int pixel = scaledImage.getRGB(x,y);                    // Obține valorile RGB                    tensorData[0][0][y][x] = (((pixel >> 16) & 0xFF) / 255f - mean[0]) / standardDeviation[0];                    tensorData[0][1][y][x] = (((pixel >> 16) & 0xFF) / 255f - mean[1]) / standardDeviation[1];                    tensorData[0][2][y][x] = (((pixel >> 16) & 0xFF) / 255f - mean[2]) / standardDeviation[2];                }            }            return tensorData;        } catch (IOException e) {            throw new RuntimeException(e);        }    }    public void run() throws OrtException {        System.out.println("Onnx Runtime Java Image Prediction");        // 1. Încarcă modelul.        var env = OrtEnvironment.getEnvironment();        var session = env.createSession(this.modelPath, new OrtSession.SessionOptions());        // Obține numele intrărilor și ieșirilor        var inputName = session.getInputNames().iterator().next();        var outputName = session.getOutputNames().iterator().next();        // 2. Creează tensorul de intrare        OnnxTensor inputTensor = OnnxTensor.createTensor(env, processImage(this.imagePath));        // 3. Rulează modelul.        var inputs = Map.of(inputName, inputTensor);        var results = session.run(inputs);        // 4. Obține tensorul de ieșire        var outputTensor = results.get(outputName);        if (outputTensor.isPresent()) {            // 5. Obține rezultatele predicției            float[][] floatBuffer = (float[][]) outputTensor.get().getValue();            ArrayList<String> predictions = new ArrayList<String>();            // filtrează bufferul după prag            for (int i = 0; i < floatBuffer[0].length; i++) {                if (floatBuffer[0][i] > -0.5) {                    predictions.add(String.format("%s: %f - %s", i, floatBuffer[0][i], Categories.CATEGORIES.get(i)));                }            }            // Afișează rezultatele            System.out.println("Predictions: ");            for (String prediction : predictions) {                System.out.println(prediction);            }        } else {            System.out.println("Failed to predict!");        }    }}