[YOLO] YoloTranslator: java.lang.IndexOutOfBoundsException

[momentNi]

When I use DJL ai.djl.modality.cv.translator.YoloTranslator to load a YOLO-V3 model and predict a picture, the program throws an Exception:

Exception in thread "main" ai.djl.translate.TranslateException: java.lang.IndexOutOfBoundsException: Index 25 out of bounds for length 1
at ai.djl.inference.Predictor.batchPredict(Predictor.java:186)
at ai.djl.inference.Predictor.predict(Predictor.java:123)
at modelTest.main(modelTest.java:99)

Caused by: java.lang.IndexOutOfBoundsException: Index 25 out of bounds for length 1
at java.base/jdk.internal.util.Preconditions.outOfBounds(Preconditions.java:64)
at java.base/jdk.internal.util.Preconditions.outOfBoundsCheckIndex(Preconditions.java:70)
at java.base/jdk.internal.util.Preconditions.checkIndex(Preconditions.java:248)
at java.base/java.util.Objects.checkIndex(Objects.java:372)
at java.base/java.util.ArrayList.get(ArrayList.java:459)
at ai.djl.modality.cv.translator.YoloTranslator.processOutput(YoloTranslator.java:63)
at ai.djl.modality.cv.translator.YoloTranslator.processOutput(YoloTranslator.java:27)
at ai.djl.inference.Predictor.processOutputs(Predictor.java:222)
at ai.djl.inference.Predictor.batchPredict(Predictor.java:180)
... 2 more

---

Enviroment: DJL 0.14.0 PyTorch 1.9.1 JDK11 IntelliJ IDEA Ultimate 2021.2.3

---

Here's the code to load and use the model:

    Pipeline pipeline = new Pipeline();
    pipeline.add(new Resize(640, 640));
    pipeline.add(new ToTensor());

    var builder = YoloTranslator.builder();
    builder.setPipeline(pipeline);
    builder.optSynsetArtifactName("names.txt");
    builder.optRescaleSize(640, 640);
    builder.optThreshold(0.8f);

    var translator = builder.build();

    Criteria<Image, DetectedObjects> criteria =
            Criteria.builder()
                    .setTypes(Image.class, DetectedObjects.class)
                    .optDevice(Device.cpu())
                    .optModelPath(Paths.get("ModelRepository\fire"))
                    .optModelName("best.torchscript.pt")
                    .optEngine("PyTorch")
                    .optTranslator(translator)
                    .optEngine(ModelEngineEnum.class.getEnumConstants()[0].name())
                    .optArgument("width", 640)
                    .optArgument("height", 640)
                    .build();

    ZooModel<Image, DetectedObjects> model = ai.djl.repository.zoo.ModelZoo.loadModel(criteria);

    DetectedObjects detectedObjects = model.newPredictor().predict(img);

---

The model can run by using Python, so the model itself has no mistakes. The model's input image size is [640*640]. I use similar code to run a YOLO-V5 model, and there is no problem.
When I debug the program, it seems that the problem occurs at the following position:

YoloTranslator.java

@Override
public DetectedObjects processOutput(TranslatorContext ctx, NDList list) {
    int[] classIndices = list.get(0).toType(DataType.INT32, true).flatten().toIntArray();
    double[] probs = list.get(1).toType(DataType.FLOAT64, true).flatten().toDoubleArray();
    NDArray boundingBoxes = list.get(2);
    int detected = Math.toIntExact(probs.length);

    NDArray xMin = boundingBoxes.get(":, 0").clip(0, imageWidth).div(imageWidth);
    NDArray yMin = boundingBoxes.get(":, 1").clip(0, imageHeight).div(imageHeight);
    NDArray xMax = boundingBoxes.get(":, 2").clip(0, imageWidth).div(imageWidth);
    NDArray yMax = boundingBoxes.get(":, 3").clip(0, imageHeight).div(imageHeight);

    float[] boxX = xMin.toFloatArray();
    float[] boxY = yMin.toFloatArray();
    float[] boxWidth = xMax.sub(xMin).toFloatArray();
    float[] boxHeight = yMax.sub(yMin).toFloatArray();

    List<String> retClasses = new ArrayList<>(detected);
    List<Double> retProbs = new ArrayList<>(detected);
    List<BoundingBox> retBB = new ArrayList<>(detected);
    for (int i = 0; i < detected; i++) {
        if (classIndices[i] < 0 || probs[i] < threshold) {
            continue;
        }
        retClasses.add(classes.get(classIndices[i]));                 // <--- Error: IndexOutOfBound
        retProbs.add(probs[i]);
        Rectangle rect = new Rectangle(boxX[i], boxY[i], boxWidth[i], boxHeight[i]);
        retBB.add(rect);
    }
    return new DetectedObjects(retClasses, retProbs, retBB);
}

---

If anyone can point out the mistakes in my code, I would be very grateful!!!

[frankfliu]

@momentNi Where is your yolov3 model comes from?

The YoloTranslator matches GluonCV yolov3 implementation. You might need to create your own translator if the model trained with other engine.

[zachgk]

In Yolo (and other object detection models), the output of the model should be some number of bounding box locations and probabilities that each location contains a particular class. However, this data must be represented using only NDArrays and there isn't just one way that it is possible to represent it. It is possible that your model is using some slightly different representation of the output data then the one that this implementation of YoloTranslator is expecting.

You should try to look into your model's output and especially it's output types and see if you can figure out how it represents the data. Then, you can use the DetectedObjects type as a standardized way of representing the output of object detection models. You may then need a slightly different processOutput in your translator to translate from your model's representation into the standard DetectedObjects representation of the output