Classify images with TensorFlow and Deep Java Library

Playing with image classification is great. However, it is also good to know how to put such a model into operation, at least as a console application.

In this tutorial I will not deal with optimizing the model with TFLite or ONNX. We will look at that some other time.

I use Maven for dependencies, so let's add them right away.

<dependencies>
    <dependency>
        <groupId>ai.djl</groupId>
        <artifactId>api</artifactId>
        <version>0.15.0</version>
    </dependency>

    <dependency>
        <groupId>ai.djl.tensorflow</groupId>
        <artifactId>tensorflow-model-zoo</artifactId>
        <version>0.15.0</version>
    </dependency>

    <dependency>
        <groupId>ai.djl.tensorflow</groupId>
        <artifactId>tensorflow-engine</artifactId>
        <version>0.15.0</version>
    </dependency>
</dependencies>

These are the basic Deep Java Library and TensorFlow packages that let the Java code run inference on our model. api is the DJL core, tensorflow-engine runs the model and downloads the native TensorFlow library on the first run, and tensorflow-model-zoo adds the TensorFlow model loader.

Run the installation with the following command:

mvn clean install

I won't be a Picasso, but we have to edit the picture

I'm sure you know this from Python: you install Pillow and NumPy and combine it all to make it work. Here you also need to prepare the image for the classifier, but it is much easier. You implement a Translator, which makes sure all inputs are uniform, so the model understands them.

import ai.djl.modality.cv.transform.Resize;
import ai.djl.modality.cv.translator.ImageClassificationTranslator;

ImageClassificationTranslator translator = ImageClassificationTranslator.builder()
        .addTransform(new Resize(224, 224))
        .addTransform(array -> array.div(255f))
        .build();

This Translator does nothing more than resize the input image to 224x224 and scale the pixel values to the range 0 to 1. The result is a tensor the model can take as input.

The original version of this article used ToTensor in the second step. That is a mistake with TensorFlow: ToTensor also transposes the image from HWC to CHW (channels first), while TensorFlow models expect channels last (NHWC). Use the same scaling the model was trained with. Many Keras models expect values from -1 to 1, which is array.div(127.5f).sub(1f).

Are we playing hide-and-seek?

The first thing we need to do is initialize the builder in which we describe our model.

import ai.djl.modality.Classifications;
import ai.djl.modality.cv.Image;
import ai.djl.repository.zoo.Criteria;
import java.nio.file.Paths;

Criteria<Image, Classifications> criteria = Criteria.builder()
        .setTypes(Image.class, Classifications.class)
        .optModelPath(Paths.get("/model/path"))
        .optTranslator(translator)
        .optEngine("TensorFlow")
        .build();

setTypes says that the input is an image and the output is a classification. optModelPath points to the SavedModel directory, the one with saved_model.pb and the variables folder.

The important thing is to set the engine. If you don't, DJL picks the default engine itself, which is fine only when TensorFlow is the only engine installed. If you have OnnxRuntime installed alongside it, DJL may pick a different engine than you expect and loading ends with an exception. optEngine("TensorFlow") removes the guessing.

Since we have not set the class names in the Translator, we need a synset.txt file, where each line is one class name. DJL pairs the first model output with the first line, the second with the second, and so on. The file goes into the same folder as the model, the one set in optModelPath(...).

cat
dog
bird

If you don't want a separate file, you can pass the names directly to the Translator builder with optSynset(List.of("cat", "dog", "bird")).

Let's pack our bags and go

Now that we have everything configured, we can load the model and create a predictor, which brings us to the final step.

import ai.djl.inference.Predictor;
import ai.djl.repository.zoo.ZooModel;

try (ZooModel<Image, Classifications> model = criteria.loadModel();
     Predictor<Image, Classifications> predictor = model.newPredictor()) {
    // prediction goes here
}

ZooModel and Predictor hold native memory, so close them with try-with-resources.

Now you just need to load the image and pass it to the input.

import ai.djl.modality.cv.ImageFactory;
import java.nio.file.Path;

Path imagePath = Paths.get("/Users/xyz/Image.jpg");
Image image = ImageFactory.getInstance().fromFile(imagePath);
Classifications classifications = predictor.predict(image);
System.out.println(classifications.topK(3));

topK(3) prints the three classes with the highest probability.

This article was first published on Medium on March 21, 2022. I am its author, the text was taken over from there with minor technical edits.