Repository: deepjavalibrary/djl
Pull request #2908 (2023-12-30 23:49:40)

Yelov8 Translator optimization
## Description ##

Brief description of what this PR is about

- Improved post-processing performance up to 40x by reducing expensive native calls
- Additional argument 'maxBox' added to improve post-processing performance by reducing number of considered bounding boxes
- Sunset file fixed,, previous version ignored first 4 rows, so recognized classes were 4 off. Adding 4 rows header fixes the problem. New headers pointing to ultralytics doc pages and original coco dataset page.

- Change on yolov8_synset.txt is a backward incompatible, example sunset file was incorrectly formatted and predicted class categories were off by 4. 



codecov-commenter: ## [Codecov](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2908?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
Attention: `1365 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`bb5073f`)](https://app.codecov.io/gh/deepjavalibrary/djl/commit/bb5073f7ffc33f1883621fc1783dbb01e70b418f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) 72.08% compared to head [(`83c7925`)](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2908?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) 72.21%.
> Report is 940 commits behind head on master.

| [Files](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2908?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | Patch % | Lines |
|---|---|---|
| [...va/ai/djl/modality/nlp/generate/TextGenerator.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2908?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L25scC9nZW5lcmF0ZS9UZXh0R2VuZXJhdG9yLmphdmE=) | 2.81% | [276 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2908?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [.../java/ai/djl/modality/nlp/generate/SeqBatcher.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2908?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L25scC9nZW5lcmF0ZS9TZXFCYXRjaGVyLmphdmE=) | 0.75% | [132 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2908?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [...ity/nlp/generate/ContrastiveSeqBatchScheduler.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2908?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L25scC9nZW5lcmF0ZS9Db250cmFzdGl2ZVNlcUJhdGNoU2NoZWR1bGVyLmphdmE=) | 2.97% | [98 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2908?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [...i/djl/modality/nlp/generate/SeqBatchScheduler.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2908?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L25scC9nZW5lcmF0ZS9TZXFCYXRjaFNjaGVkdWxlci5qYXZh) | 9.83% | [55 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2908?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [.../java/ai/djl/modality/cv/BufferedImageFactory.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2908?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L0J1ZmZlcmVkSW1hZ2VGYWN0b3J5LmphdmE=) | 40.96% | [47 Missing and 2 partials :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2908?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [...a/ai/djl/modality/nlp/generate/StepGeneration.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2908?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L25scC9nZW5lcmF0ZS9TdGVwR2VuZXJhdGlvbi5qYXZh) | 2.04% | [48 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2908?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [api/src/main/java/ai/djl/ndarray/NDArray.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2908?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL25kYXJyYXkvTkRBcnJheS5qYXZh) | 43.42% | [39 Missing and 4 partials :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2908?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [...n/java/ai/djl/modality/cv/output/CategoryMask.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2908?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9DYXRlZ29yeU1hc2suamF2YQ==) | 22.00% | [39 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2908?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [...i/src/main/java/ai/djl/ndarray/NDArrayAdapter.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2908?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL25kYXJyYXkvTkRBcnJheUFkYXB0ZXIuamF2YQ==) | 71.21% | [31 Missing and 7 partials :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2908?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [.../cv/translator/SemanticSegmentationTranslator.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2908?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvU2VtYW50aWNTZWdtZW50YXRpb25UcmFuc2xhdG9yLmphdmE=) | 37.50% | [35 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2908?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| ... and [75 more](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2908?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | |

:exclamation: Your organization needs to install the [Codecov GitHub app](https://github.com/apps/codecov/installations/select_target) to enable full functionality.

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff              @@
##             master    #2908      +/-   ##
============================================
+ Coverage     72.08%   72.21%   +0.12%     
- Complexity     5126     7183    +2057     
============================================
  Files           473      708     +235     
  Lines         21970    32021   +10051     
  Branches       2351     3337     +986     
============================================
+ Hits          15838    23123    +7285     
- Misses         4925     7310    +2385     
- Partials       1207     1588     +381     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2908?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).


Patch:
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/YoloV8Translator.java b/api/src/main/java/ai/djl/modality/cv/translator/YoloV8Translator.java
index dc160ba75..d47f7a4a1 100644
--- a/api/src/main/java/ai/djl/modality/cv/translator/YoloV8Translator.java
+++ b/api/src/main/java/ai/djl/modality/cv/translator/YoloV8Translator.java
@@ -1,103 +1,124 @@
-/*
- * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
- *
- * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
- * with the License. A copy of the License is located at
- *
- * http://aws.amazon.com/apache2.0/
- *
- * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
- * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
- * and limitations under the License.
- */
-package ai.djl.modality.cv.translator;
-
-import ai.djl.modality.cv.output.DetectedObjects;
-import ai.djl.modality.cv.output.Rectangle;
-import ai.djl.ndarray.NDArray;
-import ai.djl.ndarray.NDList;
-import ai.djl.ndarray.types.DataType;
-
-import java.util.ArrayList;
-import java.util.Map;
-
-/**
- * A translator for YoloV8 models. This was tested with ONNX exported Yolo models. For details check
- * here: https://github.com/ultralytics/ultralytics
- */
-public class YoloV8Translator extends YoloV5Translator {
-
-    /**
-     * Constructs an ImageTranslator with the provided builder.
-     *
-     * @param builder the data to build with
-     */
-    protected YoloV8Translator(Builder builder) {
-        super(builder);
-    }
-
-    /**
-     * Creates a builder to build a {@code YoloV8Translator} with specified arguments.
-     *
-     * @param arguments arguments to specify builder options
-     * @return a new builder
-     */
-    public static YoloV8Translator.Builder builder(Map<String, ?> arguments) {
-        YoloV8Translator.Builder builder = new YoloV8Translator.Builder();
-        builder.configPreProcess(arguments);
-        builder.configPostProcess(arguments);
-
-        return builder;
-    }
-
-    @Override
-    protected DetectedObjects processFromBoxOutput(NDList list) {
-        NDArray features4OneImg = list.get(0);
-        int sizeClasses = classes.size();
-        long sizeBoxes = features4OneImg.size(1);
-        ArrayList<IntermediateResult> intermediateResults = new ArrayList<>();
-
-        for (long b = 0; b < sizeBoxes; b++) {
-            float maxClass = 0;
-            int maxIndex = 0;
-            for (int c = 4; c < sizeClasses; c++) {
-                float classProb = features4OneImg.getFloat(c, b);
-                if (classProb > maxClass) {
-                    maxClass = classProb;
-                    maxIndex = c;
-                }
-            }
-
-            if (maxClass > threshold) {
-                float xPos = features4OneImg.getFloat(0, b); // center x
-                float yPos = features4OneImg.getFloat(1, b); // center y
-                float w = features4OneImg.getFloat(2, b);
-                float h = features4OneImg.getFloat(3, b);
-                Rectangle rect =
-                        new Rectangle(Math.max(0, xPos - w / 2), Math.max(0, yPos - h / 2), w, h);
-                intermediateResults.add(
-                        new IntermediateResult(classes.get(maxIndex), maxClass, maxIndex, rect));
-            }
-        }
-
-        return nms(intermediateResults);
-    }
-
-    /** The builder for {@link YoloV8Translator}. */
-    public static class Builder extends YoloV5Translator.Builder {
-        /**
-         * Builds the translator.
-         *
-         * @return the new translator
-         */
-        @Override
-        public YoloV8Translator build() {
-            if (pipeline == null) {
-                addTransform(
-                        array -> array.transpose(2, 0, 1).toType(DataType.FLOAT32, false).div(255));
-            }
-            validate();
-            return new YoloV8Translator(this);
-        }
-    }
-}
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.modality.cv.translator;
+
+import ai.djl.modality.cv.output.DetectedObjects;
+import ai.djl.modality.cv.output.Rectangle;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.types.DataType;
+import ai.djl.ndarray.types.Shape;
+import ai.djl.translate.ArgumentsUtil;
+
+import java.util.ArrayList;
+import java.util.Map;
+
+/**
+ * A translator for YoloV8 models. This was tested with ONNX exported Yolo models. For details check
+ * here: https://github.com/ultralytics/ultralytics
+ */
+public class YoloV8Translator extends YoloV5Translator {
+
+    private int maxBoxes;
+
+    /**
+     * Constructs an ImageTranslator with the provided builder.
+     *
+     * @param builder the data to build with
+     */
+    protected YoloV8Translator(Builder builder) {
+        super(builder);
+        maxBoxes = builder.maxBox;
+    }
+
+    /**
+     * Creates a builder to build a {@code YoloV8Translator} with specified arguments.
+     *
+     * @param arguments arguments to specify builder options
+     * @return a new builder
+     */
+    public static YoloV8Translator.Builder builder(Map<String, ?> arguments) {
+        YoloV8Translator.Builder builder = new YoloV8Translator.Builder();
+        builder.configPreProcess(arguments);
+        builder.configPostProcess(arguments);
+
+        return builder;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    protected DetectedObjects processFromBoxOutput(NDList list) {
+        NDArray rawResult = list.get(0);
+        NDArray reshapedResult = rawResult.transpose();
+        Shape shape = reshapedResult.getShape();
+        float[] buf = reshapedResult.toFloatArray();
+        int numberRows = Math.toIntExact(shape.get(0));
+        int nClasses = Math.toIntExact(shape.get(1));
+
+        ArrayList<IntermediateResult> intermediateResults = new ArrayList<>();
+        // reverse order search in heap; searches through #maxBoxes for optimization when set
+        for (int i = numberRows - 1; i > numberRows - maxBoxes; --i) {
+            int index = i * nClasses;
+            float maxClassProb = -1f;
+            int maxIndex = -1;
+            for (int c = 4; c < nClasses; c++) {
+                float classProb = buf[index + c];
+                if (classProb > maxClassProb) {
+                    maxClassProb = classProb;
+                    maxIndex = c;
+                }
+            }
+
+            if (maxClassProb > threshold) {
+                float xPos = buf[index]; // center x
+                float yPos = buf[index + 1]; // center y
+                float w = buf[index + 2];
+                float h = buf[index + 3];
+                Rectangle rect =
+                        new Rectangle(Math.max(0, xPos - w / 2), Math.max(0, yPos - h / 2), w, h);
+                intermediateResults.add(
+                        new IntermediateResult(
+                                classes.get(maxIndex), maxClassProb, maxIndex, rect));
+            }
+        }
+        return nms(intermediateResults);
+    }
+
+    /** The builder for {@link YoloV8Translator}. */
+    public static class Builder extends YoloV5Translator.Builder {
+
+        private int maxBox = 8400;
+
+        /**
+         * Builds the translator.
+         *
+         * @return the new translator
+         */
+        @Override
+        public YoloV8Translator build() {
+            if (pipeline == null) {
+                addTransform(
+                        array -> array.transpose(2, 0, 1).toType(DataType.FLOAT32, false).div(255));
+            }
+            validate();
+            return new YoloV8Translator(this);
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        protected void configPostProcess(Map<String, ?> arguments) {
+            super.configPostProcess(arguments);
+            maxBox = ArgumentsUtil.intValue(arguments, "maxBox", 8400);
+        }
+    }
+}
diff --git a/examples/src/main/java/ai/djl/examples/inference/Yolov8Detection.java b/examples/src/main/java/ai/djl/examples/inference/Yolov8Detection.java
index 5a474c1da..27ee1211d 100644
--- a/examples/src/main/java/ai/djl/examples/inference/Yolov8Detection.java
+++ b/examples/src/main/java/ai/djl/examples/inference/Yolov8Detection.java
@@ -16,27 +16,21 @@ import ai.djl.ModelException;
 import ai.djl.inference.Predictor;
 import ai.djl.modality.cv.Image;
 import ai.djl.modality.cv.ImageFactory;
-import ai.djl.modality.cv.output.BoundingBox;
 import ai.djl.modality.cv.output.DetectedObjects;
-import ai.djl.modality.cv.output.DetectedObjects.DetectedObject;
-import ai.djl.modality.cv.output.Rectangle;
 import ai.djl.modality.cv.translator.YoloV8TranslatorFactory;
 import ai.djl.repository.zoo.Criteria;
 import ai.djl.repository.zoo.ZooModel;
 import ai.djl.training.util.ProgressBar;
 import ai.djl.translate.TranslateException;
-import ai.djl.translate.Translator;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.io.IOException;
+import java.io.OutputStream;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
-import java.util.List;
-import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
 
 /** An example of inference using an yolov8 model. */
 public final class Yolov8Detection {
@@ -51,82 +45,44 @@ public final class Yolov8Detection {
     }
 
     public static DetectedObjects predict() throws IOException, ModelException, TranslateException {
-        String classPath = System.getProperty("java.class.path");
-        String pathSeparator = System.getProperty("path.separator");
-        classPath = classPath.split(pathSeparator)[0];
-        Path modelPath = Paths.get(classPath + "/yolov8n.onnx");
-        Path imgPath = Paths.get(classPath + "/yolov8_test.jpg");
+        Path modelPath = Paths.get("src/test/resources/yolov8n.onnx");
+        Path imgPath = Paths.get("src/test/resources/yolov8_test.jpg");
         Image img = ImageFactory.getInstance().fromFile(imgPath);
 
-        Map<String, Object> arguments = new ConcurrentHashMap<>();
-        arguments.put("width", 640);
-        arguments.put("height", 640);
-        arguments.put("resize", "true");
-        arguments.put("toTensor", true);
-        arguments.put("applyRatio", true);
-        arguments.put("threshold", 0.6f);
-        arguments.put("synsetFileName", "yolov8_synset.txt");
-
-        YoloV8TranslatorFactory yoloV8TranslatorFactory = new YoloV8TranslatorFactory();
-        Translator<Image, DetectedObjects> translator =
-                yoloV8TranslatorFactory.newInstance(
-                        Image.class, DetectedObjects.class, null, arguments);
-
         Criteria<Image, DetectedObjects> criteria =
                 Criteria.builder()
                         .setTypes(Image.class, DetectedObjects.class)
                         .optModelPath(modelPath)
                         .optEngine("OnnxRuntime")
-                        .optTranslator(translator)
+                        .optArgument("width", 640)
+                        .optArgument("height", 640)
+                        .optArgument("resize", true)
+                        .optArgument("toTensor", true)
+                        .optArgument("applyRatio", true)
+                        .optArgument("threshold", 0.6f)
+                        // for performance optimization maxBox parameter can reduce number of
+                        // considered boxes from 8400
+                        .optArgument("maxBox", 1000)
+                        .optArgument("synsetFileName", "yolov8_synset.txt")
+                        .optTranslatorFactory(new YoloV8TranslatorFactory())
                         .optProgress(new ProgressBar())
                         .build();
 
-        DetectedObjects detectedObjects;
-        DetectedObject detectedObject;
         try (ZooModel<Image, DetectedObjects> model = criteria.loadModel();
                 Predictor<Image, DetectedObjects> predictor = model.newPredictor()) {
-            Path outputPath = Paths.get(classPath + "/output");
+            Path outputPath = Paths.get("build/output");
             Files.createDirectories(outputPath);
 
-            detectedObjects = predictor.predict(img);
-
-            if (detectedObjects.getNumberOfObjects() > 0) {
-                List<DetectedObject> detectedObjectList = detectedObjects.items();
-                for (DetectedObject object : detectedObjectList) {
-                    detectedObject = object;
-                    BoundingBox boundingBox = detectedObject.getBoundingBox();
-                    Rectangle tectangle = boundingBox.getBounds();
-                    logger.info(
-                            detectedObject.getClassName()
-                                    + " "
-                                    + detectedObject.getProbability()
-                                    + " "
-                                    + tectangle.getX()
-                                    + " "
-                                    + tectangle.getY()
-                                    + " "
-                                    + tectangle.getWidth()
-                                    + " "
-                                    + tectangle.getHeight());
+            DetectedObjects detection = predictor.predict(img);
+            if (detection.getNumberOfObjects() > 0) {
+                img.drawBoundingBoxes(detection);
+                Path output = outputPath.resolve("yolov8_detected.png");
+                try (OutputStream os = Files.newOutputStream(output)) {
+                    img.save(os, "png");
                 }
-
-                saveBoundingBoxImage(
-                        img.resize(640, 640, false),
-                        detectedObjects,
-                        outputPath,
-                        imgPath.toFile().getName());
+                logger.info("Detected object saved in: {}", output);
             }
-
-            return detectedObjects;
+            return detection;
         }
     }
-
-    private static void saveBoundingBoxImage(
-            Image img, DetectedObjects detectedObjects, Path outputPath, String outputFileName)
-            throws IOException {
-        img.drawBoundingBoxes(detectedObjects);
-
-        Path imagePath = outputPath.resolve(outputFileName);
-        img.save(Files.newOutputStream(imagePath), "png");
-    }
 }


Test patch:
diff --git a/examples/src/test/java/ai/djl/examples/inference/Yolov8DetectionTest.java b/examples/src/test/java/ai/djl/examples/inference/Yolov8DetectionTest.java
new file mode 100644
index 000000000..35e3fc434
--- /dev/null
+++ b/examples/src/test/java/ai/djl/examples/inference/Yolov8DetectionTest.java
@@ -0,0 +1,40 @@
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.examples.inference;
+
+import ai.djl.ModelException;
+import ai.djl.modality.Classifications;
+import ai.djl.modality.cv.output.DetectedObjects;
+import ai.djl.testing.TestRequirements;
+import ai.djl.translate.TranslateException;
+
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+import java.io.IOException;
+
+public class Yolov8DetectionTest {
+
+    @Test
+    public void testYolov8Detection() throws ModelException, TranslateException, IOException {
+        TestRequirements.engine("MXNet", "PyTorch");
+
+        DetectedObjects result = Yolov8Detection.predict();
+
+        Assert.assertTrue(result.getNumberOfObjects() >= 1);
+        Classifications.Classification obj = result.best();
+        String className = obj.getClassName();
+        Assert.assertEquals(className, "dog");
+        Assert.assertTrue(obj.getProbability() > 0.6);
+    }
+}
diff --git a/examples/src/test/resources/yolov8_synset.txt b/examples/src/test/resources/yolov8_synset.txt
index 7139f0cc6..ffba20649 100644
--- a/examples/src/test/resources/yolov8_synset.txt
+++ b/examples/src/test/resources/yolov8_synset.txt
@@ -1,3 +1,7 @@
+# Classes for coco dataset on which yelov8 is trained
+# source config https://github.com/ultralytics/ultralytics/blob/main/ultralytics/cfg/datasets/coco.yaml.
+# COCO dataset website: https://cocodataset.org/#home
+# Ultralytics Coco doc page: https://docs.ultralytics.com/datasets/detect/coco/
 person
 bicycle
 car