DJL ResNet with JNI: Accelerating Image Preprocessing in Java

package org.estech;

import java.nio.ByteBuffer;

public final class NativeImageOps {
    static { System.loadLibrary("imageops"); } // pass -Djava.library.path at runtime
    public static native ByteBuffer preprocessToCHW(
            ByteBuffer encoded, int outW, int outH,
            float m0, float m1, float m2,
            float s0, float s1, float s2);
    public static native void freeBuffer(ByteBuffer buf);
}

public class ResNetJniTranslator implements Translator<byte[], Classifications> {

    private static final int SIZE = 224;
    private static final float M0 = 0.485f, M1 = 0.456f, M2 = 0.406f;
    private static final float S0 = 0.229f, S1 = 0.224f, S2 = 0.225f;

    @Override
    public NDList processInput(TranslatorContext ctx, byte[] input) throws Exception {
        ByteBuffer encoded = ByteBuffer.allocateDirect(input.length);
        encoded.put(input).flip();

        long t0 = System.nanoTime();
        ByteBuffer chw = NativeImageOps.preprocessToCHW(
                encoded, SIZE, SIZE, M0, M1, M2, S0, S1, S2);
        long t1 = System.nanoTime();

        if (chw == null) {
            throw new IOException("JNI preprocessToCHW returned null");
        }
        System.out.printf("[JNI] preprocess: %.2f ms, bytes=%d%n",
                (t1 - t0) / 1e6, chw.remaining());

        NDManager manager = ctx.getNDManager();
        chw.rewind();
        NDArray x = manager.create(chw, new Shape(3, SIZE, SIZE), DataType.FLOAT32);

        NativeImageOps.freeBuffer(chw);

        return new NDList(x);
    }

    @Override
    public Classifications processOutput(TranslatorContext ctx, NDList list) throws Exception {
        NDArray logits = list.singletonOrThrow();

        NDArray probs;
        Shape sh = logits.getShape();
        if (sh.dimension() == 2 && sh.get(0) == 1) {
            probs = logits.softmax(1).squeeze(0);   // [1,1000] -> [1000]
        } else if (sh.dimension() == 1) {
            probs = logits.softmax(0);              // [1000]
        } else {
            int last = (int) sh.dimension() - 1;
            probs = logits.softmax(last);
            if (probs.getShape().dimension() > 1) {
                probs = probs.flatten();           
            }
        }

        List<String> synset = loadSynset(ctx, (int) probs.getShape().get(0));
        return new Classifications(synset, probs);
    }

    @Override
    public Batchifier getBatchifier() {
        return Batchifier.STACK;
    }

    private List<String> loadSynset(TranslatorContext ctx, int nClasses) throws IOException {
        try {
            List<String> labels = ctx.getModel().getArtifact("synset.txt", Utils::readLines);
            if (labels != null && !labels.isEmpty()) return labels;
        } catch (Exception ignore) { /* fallthrough */ }

        Path local = Paths.get("path/to/your/synset.txt");
        if (Files.exists(local)) {
            return Files.readAllLines(local);
        }

        try (InputStream is = ResNetJniTranslator.class.getResourceAsStream("path/to/your.txt")) {
            if (is != null) {
                return Utils.readLines(is);
            }
        }

        List<String> fallback = new ArrayList<>(nClasses);
        for (int i = 0; i < nClasses; i++) fallback.add(String.valueOf(i));
        return fallback;
    }
}

if (useJni) {
    // JNI Translator 
    byte[] imgBytes = Files.readAllBytes(imgPath);

    Translator<byte[], Classifications> translator = new ResNetJniTranslator();

    Criteria<byte[], Classifications> criteria = Criteria.builder()
        .optEngine("PyTorch")
        .optApplication(Application.CV.IMAGE_CLASSIFICATION)
        .setTypes(byte[].class, Classifications.class)
        .optModelPath(modelDir)
        .optModelName("traced_resnet18")
        .optTranslator(translator)
        .optProgress(new ProgressBar())
        .build();

    try (ZooModel<byte[], Classifications> model = ModelZoo.loadModel(criteria);
        Predictor<byte[], Classifications> predictor = model.newPredictor()) {

        long t0 = System.nanoTime();
        Classifications result = predictor.predict(imgBytes);
        long t1 = System.nanoTime();

        System.out.printf("Predict total (JNI): %.2f ms%n", (t1 - t0) / 1e6);
        System.out.println("== Top-5 ==");
        result.topK(5).forEach(c ->
            System.out.printf("%-25s %.5f%n", c.getClassName(), c.getProbability()));
    }
} else {
    // Java DJL Translator
    Image img = ImageFactory.getInstance().fromFile(imgPath);

    Translator<Image, Classifications> translator =
        ai.djl.modality.cv.translator.ImageClassificationTranslator.builder()
            .addTransform(new Resize(256))
            .addTransform(new CenterCrop(224, 224))
            .addTransform(new ToTensor())
            .addTransform(new Normalize(
                new float[]{0.485f, 0.456f, 0.406f},
                new float[]{0.229f, 0.224f, 0.225f}))
            .optApplySoftmax(true)
            .optSynsetUrl(Paths.get("path/to/your/synset.txt")
            .toUri().toString())
            .build();

    Criteria<Image, Classifications> criteria = Criteria.builder()
            .optEngine("PyTorch")
            .optApplication(Application.CV.IMAGE_CLASSIFICATION)
            .setTypes(Image.class, Classifications.class)
            .optModelPath(modelDir)
            .optModelName("traced_resnet18")
            .optTranslator(translator)
            .optProgress(new ProgressBar())
            .build();

    try (ZooModel<Image, Classifications> model = ModelZoo.loadModel(criteria);
        Predictor<Image, Classifications> predictor = model.newPredictor()) {

        long t0 = System.nanoTime();
        Classifications result = predictor.predict(img);
        long t1 = System.nanoTime();

        System.out.printf("Predict total (DJL): %.2f ms%n", (t1 - t0) / 1e6);
        System.out.println("== Top-5 ==");
        result.topK(5).forEach(c ->
            System.out.printf("%-25s %.5f%n", c.getClassName(), c.getProbability()));
    }
}