Repository: deepjavalibrary/djl
Pull request #3305 (2024-07-06 01:36:27)

[api] Improve drawJoints behavior
- **[api] Reactor nms for yolo translator**
- **[api] Improve drawJoints behavior**

## Description ##

Brief description of what this PR is about

- If this change is a backward incompatible change, why must this change be made?
- Interesting edge cases to note here



Patch:
diff --git a/api/src/main/java/ai/djl/modality/cv/BufferedImageFactory.java b/api/src/main/java/ai/djl/modality/cv/BufferedImageFactory.java
index 40735ddec..63f3641c3 100644
--- a/api/src/main/java/ai/djl/modality/cv/BufferedImageFactory.java
+++ b/api/src/main/java/ai/djl/modality/cv/BufferedImageFactory.java
@@ -347,18 +347,37 @@ public class BufferedImageFactory extends ImageFactory {
             convertIdNeeded();
 
             Graphics2D g = (Graphics2D) image.getGraphics();
-            int stroke = 2;
-            g.setStroke(new BasicStroke(stroke));
+            g.setRenderingHint(RenderingHints.KEY_ANTIALIASING, RenderingHints.VALUE_ANTIALIAS_ON);
 
             int imageWidth = image.getWidth();
             int imageHeight = image.getHeight();
 
-            for (Joints.Joint joint : joints.getJoints()) {
-                g.setPaint(randomColor().darker());
+            List<Joints.Joint> list = joints.getJoints();
+            if (list.size() == 17) {
+                g.setColor(new Color(224, 255, 37));
+                g.setStroke(new BasicStroke(3));
+                drawLine(g, list.get(5), list.get(7), imageWidth, imageHeight);
+                drawLine(g, list.get(7), list.get(9), imageWidth, imageHeight);
+                drawLine(g, list.get(6), list.get(8), imageWidth, imageHeight);
+                drawLine(g, list.get(8), list.get(10), imageWidth, imageHeight);
+                drawLine(g, list.get(11), list.get(13), imageWidth, imageHeight);
+                drawLine(g, list.get(12), list.get(14), imageWidth, imageHeight);
+                drawLine(g, list.get(13), list.get(15), imageWidth, imageHeight);
+                drawLine(g, list.get(14), list.get(16), imageWidth, imageHeight);
+                drawLine(g, list.get(5), list.get(6), imageWidth, imageHeight);
+                drawLine(g, list.get(11), list.get(12), imageWidth, imageHeight);
+                drawLine(g, list.get(5), list.get(11), imageWidth, imageHeight);
+                drawLine(g, list.get(6), list.get(12), imageWidth, imageHeight);
+            }
+
+            g.setColor(new Color(37, 150, 190));
+            g.setStroke(new BasicStroke(2));
+            for (Joints.Joint joint : list) {
                 int x = (int) (joint.getX() * imageWidth);
                 int y = (int) (joint.getY() * imageHeight);
-                g.fillOval(x, y, 10, 10);
+                g.fillOval(x - 6, y - 6, 12, 12);
             }
+
             g.dispose();
         }
 
@@ -380,8 +399,13 @@ public class BufferedImageFactory extends ImageFactory {
             image = target;
         }
 
-        private Color randomColor() {
-            return new Color(RandomUtils.nextInt(255));
+        private void drawLine(
+                Graphics2D g, Joints.Joint from, Joints.Joint to, int width, int height) {
+            int x0 = (int) (from.getX() * width);
+            int y0 = (int) (from.getY() * height);
+            int x1 = (int) (to.getX() * width);
+            int y1 = (int) (to.getY() * height);
+            g.drawLine(x0, y0, x1, y1);
         }
 
         private void drawText(Graphics2D g, String text, int x, int y, int stroke, int padding) {
@@ -403,28 +427,38 @@ public class BufferedImageFactory extends ImageFactory {
             float b = RandomUtils.nextFloat();
             int imageWidth = image.getWidth();
             int imageHeight = image.getHeight();
-            int x = (int) (mask.getX() * imageWidth);
-            int y = (int) (mask.getY() * imageHeight);
-            float[][] probDist = mask.getProbDist();
-            // Correct some coordinates of box when going out of image
-            if (x < 0) {
-                x = 0;
-            }
-            if (y < 0) {
-                y = 0;
+            int x = 0;
+            int y = 0;
+            int w = imageWidth;
+            int h = imageHeight;
+            if (!mask.isFullImageMask()) {
+                x = (int) (mask.getX() * imageWidth);
+                y = (int) (mask.getY() * imageHeight);
+                w = (int) (mask.getWidth() * imageWidth);
+                h = (int) (mask.getHeight() * imageHeight);
+                // Correct some coordinates of box when going out of image
+                if (x < 0) {
+                    x = 0;
+                }
+                if (y < 0) {
+                    y = 0;
+                }
             }
+            float[][] probDist = mask.getProbDist();
 
             BufferedImage maskImage =
                     new BufferedImage(
-                            probDist.length, probDist[0].length, BufferedImage.TYPE_INT_ARGB);
-            for (int xCor = 0; xCor < probDist.length; xCor++) {
-                for (int yCor = 0; yCor < probDist[xCor].length; yCor++) {
-                    float opacity = probDist[xCor][yCor] * 0.8f;
+                            probDist[0].length, probDist.length, BufferedImage.TYPE_INT_ARGB);
+            for (int yCor = 0; yCor < probDist.length; yCor++) {
+                for (int xCor = 0; xCor < probDist[0].length; xCor++) {
+                    float opacity = probDist[yCor][xCor] * 0.8f;
                     maskImage.setRGB(xCor, yCor, new Color(r, g, b, opacity).darker().getRGB());
                 }
             }
+            java.awt.Image scaled = maskImage.getScaledInstance(w, h, java.awt.Image.SCALE_SMOOTH);
+
             Graphics2D gR = (Graphics2D) image.getGraphics();
-            gR.drawImage(maskImage, x, y, null);
+            gR.drawImage(scaled, x, y, null);
             gR.dispose();
         }
 
diff --git a/api/src/main/java/ai/djl/modality/cv/output/Mask.java b/api/src/main/java/ai/djl/modality/cv/output/Mask.java
index 0c4a9fd66..e9ada0710 100644
--- a/api/src/main/java/ai/djl/modality/cv/output/Mask.java
+++ b/api/src/main/java/ai/djl/modality/cv/output/Mask.java
@@ -21,6 +21,7 @@ public class Mask extends Rectangle {
 
     private static final long serialVersionUID = 1L;
     private float[][] probDist;
+    private boolean fullImageMask;
 
     /**
      * Constructs a Mask with the given data.
@@ -32,8 +33,29 @@ public class Mask extends Rectangle {
      * @param dist the probability distribution for each pixel in the rectangle
      */
     public Mask(double x, double y, double width, double height, float[][] dist) {
+        this(x, y, width, height, dist, false);
+    }
+
+    /**
+     * Constructs a Mask with the given data.
+     *
+     * @param x the left coordinate of the bounding rectangle
+     * @param y the top coordinate of the bounding rectangle
+     * @param width the width of the bounding rectangle
+     * @param height the height of the bounding rectangle
+     * @param dist the probability distribution for each pixel in the rectangle
+     * @param fullImageMask if the mask if for full image
+     */
+    public Mask(
+            double x,
+            double y,
+            double width,
+            double height,
+            float[][] dist,
+            boolean fullImageMask) {
         super(x, y, width, height);
         this.probDist = dist;
+        this.fullImageMask = fullImageMask;
     }
 
     /**
@@ -44,4 +66,13 @@ public class Mask extends Rectangle {
     public float[][] getProbDist() {
         return probDist;
     }
+
+    /**
+     * Returns if the mask is for full image.
+     *
+     * @return if the mask is for full image
+     */
+    public boolean isFullImageMask() {
+        return fullImageMask;
+    }
 }
diff --git a/api/src/main/java/ai/djl/modality/cv/output/Rectangle.java b/api/src/main/java/ai/djl/modality/cv/output/Rectangle.java
index 92afc6032..b7b16127a 100644
--- a/api/src/main/java/ai/djl/modality/cv/output/Rectangle.java
+++ b/api/src/main/java/ai/djl/modality/cv/output/Rectangle.java
@@ -14,6 +14,7 @@ package ai.djl.modality.cv.output;
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.PriorityQueue;
 
 /**
  * A {@code Rectangle} specifies an area in a coordinate space that is enclosed by the {@code
@@ -152,4 +153,81 @@ public class Rectangle implements BoundingBox {
         return String.format(
                 "{\"x\"=%.3f, \"y\"=%.3f, \"width\"=%.3f, \"height\"=%.3f}", x, y, width, height);
     }
+
+    /**
+     * Applies nms (non-maximum suppression) to the list of rectangles.
+     *
+     * @param boxes an list of {@code Rectangle}
+     * @param scores a list of scores
+     * @param nmsThreshold the nms threshold
+     * @return the filtered list with the index of the original list
+     */
+    public static List<Integer> nms(
+            List<Rectangle> boxes, List<Double> scores, float nmsThreshold) {
+        List<Integer> ret = new ArrayList<>();
+        PriorityQueue<Integer> pq =
+                new PriorityQueue<>(
+                        50,
+                        (lhs, rhs) -> {
+                            // Intentionally reversed to put high confidence at the head of the
+                            // queue.
+                            return Double.compare(scores.get(rhs), scores.get(lhs));
+                        });
+        for (int i = 0; i < boxes.size(); ++i) {
+            pq.add(i);
+        }
+
+        // do non maximum suppression
+        while (!pq.isEmpty()) {
+            // insert detection with max confidence
+            int[] detections = pq.stream().mapToInt(Integer::intValue).toArray();
+            ret.add(detections[0]);
+            Rectangle box = boxes.get(detections[0]);
+            pq.clear();
+            for (int i = 1; i < detections.length; i++) {
+                int detection = detections[i];
+                Rectangle location = boxes.get(detection);
+                if (box.boxIou(location) < nmsThreshold) {
+                    pq.add(detection);
+                }
+            }
+        }
+        return ret;
+    }
+
+    private double boxIou(Rectangle other) {
+        double intersection = intersection(other);
+        double union =
+                getWidth() * getHeight() + other.getWidth() * other.getHeight() - intersection;
+        return intersection / union;
+    }
+
+    private double intersection(Rectangle b) {
+        double w =
+                overlap(
+                        (getX() * 2 + getWidth()) / 2,
+                        getWidth(),
+                        (b.getX() * 2 + b.getWidth()) / 2,
+                        b.getWidth());
+        double h =
+                overlap(
+                        (getY() * 2 + getHeight()) / 2,
+                        getHeight(),
+                        (b.getY() * 2 + b.getHeight()) / 2,
+                        b.getHeight());
+        if (w < 0 || h < 0) {
+            return 0;
+        }
+        return w * h;
+    }
+
+    private double overlap(double x1, double w1, double x2, double w2) {
+        double l1 = x1 - w1 / 2;
+        double l2 = x2 - w2 / 2;
+        double left = Math.max(l1, l2);
+        double r1 = x1 + w1 / 2;
+        double r2 = x2 + w2 / 2;
+        double right = Math.min(r1, r2);
+        return right - left;
+    }
 }
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/InstanceSegmentationTranslator.java b/api/src/main/java/ai/djl/modality/cv/translator/InstanceSegmentationTranslator.java
index efd05d7be..458122bbc 100644
--- a/api/src/main/java/ai/djl/modality/cv/translator/InstanceSegmentationTranslator.java
+++ b/api/src/main/java/ai/djl/modality/cv/translator/InstanceSegmentationTranslator.java
@@ -12,7 +12,6 @@
  */
 package ai.djl.modality.cv.translator;
 
-import ai.djl.modality.cv.Image;
 import ai.djl.modality.cv.output.BoundingBox;
 import ai.djl.modality.cv.output.DetectedObjects;
 import ai.djl.modality.cv.output.Mask;
@@ -68,14 +67,6 @@ public class InstanceSegmentationTranslator extends BaseImageTranslator<Detected
         }
     }
 
-    /** {@inheritDoc} */
-    @Override
-    public NDList processInput(TranslatorContext ctx, Image image) {
-        ctx.setAttachment("originalHeight", image.getHeight());
-        ctx.setAttachment("originalWidth", image.getWidth());
-        return super.processInput(ctx, image);
-    }
-
     /** {@inheritDoc} */
     @Override
     public DetectedObjects processOutput(TranslatorContext ctx, NDList list) {
@@ -102,18 +93,15 @@ public class InstanceSegmentationTranslator extends BaseImageTranslator<Detected
                 double w = box[2] / rescaledWidth - x;
                 double h = box[3] / rescaledHeight - y;
 
-                int maskW = (int) (w * (int) ctx.getAttachment("originalWidth"));
-                int maskH = (int) (h * (int) ctx.getAttachment("originalHeight"));
-
                 // Reshape mask to actual image bounding box shape.
                 NDArray array = masks.get(i);
                 Shape maskShape = array.getShape();
-                array = array.reshape(maskShape.addAll(new Shape(1)));
-                NDArray maskArray = NDImageUtils.resize(array, maskW, maskH).transpose();
-                float[] flattened = maskArray.toFloatArray();
-                float[][] maskFloat = new float[maskW][maskH];
-                for (int j = 0; j < maskW; j++) {
-                    System.arraycopy(flattened, j * maskH, maskFloat[j], 0, maskH);
+                int maskH = (int) maskShape.get(0);
+                int maskW = (int) maskShape.get(1);
+                float[] flattened = array.toFloatArray();
+                float[][] maskFloat = new float[maskH][maskW];
+                for (int j = 0; j < maskH; j++) {
+                    System.arraycopy(flattened, j * maskW, maskFloat[j], 0, maskW);
                 }
                 Mask mask = new Mask(x, y, w, h, maskFloat);
 
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/YoloPoseTranslator.java b/api/src/main/java/ai/djl/modality/cv/translator/YoloPoseTranslator.java
new file mode 100644
index 000000000..b59da7020
--- /dev/null
+++ b/api/src/main/java/ai/djl/modality/cv/translator/YoloPoseTranslator.java
@@ -0,0 +1,188 @@
+/*
+ * Copyright 2024 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.Joints;
+import ai.djl.modality.cv.output.Joints.Joint;
+import ai.djl.modality.cv.output.Rectangle;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.translate.ArgumentsUtil;
+import ai.djl.translate.TranslatorContext;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+
+/** A translator for Yolov8 pose estimation models. */
+public class YoloPoseTranslator extends BaseImageTranslator<Joints[]> {
+
+    private static final int MAX_DETECTION = 300;
+
+    private int width;
+    private int height;
+    private float threshold;
+    private float nmsThreshold;
+
+    /**
+     * Creates the Pose Estimation translator from the given builder.
+     *
+     * @param builder the builder for the translator
+     */
+    public YoloPoseTranslator(Builder builder) {
+        super(builder);
+        this.width = builder.width;
+        this.height = builder.height;
+        this.threshold = builder.threshold;
+        this.nmsThreshold = builder.nmsThreshold;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Joints[] processOutput(TranslatorContext ctx, NDList list) {
+        NDArray pred = list.singletonOrThrow();
+        NDArray candidates = pred.get(4).gt(threshold);
+        pred = pred.transpose();
+        NDArray sub = pred.get("..., :4");
+        sub = xywh2xyxy(sub);
+        pred = sub.concat(pred.get("..., 4:"), -1);
+        pred = pred.get(candidates);
+
+        NDList split = pred.split(new long[] {4, 5}, 1);
+        NDArray box = split.get(0);
+
+        int numBox = Math.toIntExact(box.getShape().get(0));
+
+        float[] buf = box.toFloatArray();
+        float[] confidences = split.get(1).toFloatArray();
+        float[] mask = split.get(2).toFloatArray();
+
+        List<Rectangle> boxes = new ArrayList<>(numBox);
+        List<Double> scores = new ArrayList<>(numBox);
+
+        for (int i = 0; i < numBox; ++i) {
+            float xPos = buf[i * 4];
+            float yPos = buf[i * 4 + 1];
+            float w = buf[i * 4 + 2] - xPos;
+            float h = buf[i * 4 + 3] - yPos;
+            Rectangle rect = new Rectangle(xPos, yPos, w, h);
+            boxes.add(rect);
+            scores.add((double) confidences[i]);
+        }
+        List<Integer> nms = Rectangle.nms(boxes, scores, nmsThreshold);
+        if (nms.size() > MAX_DETECTION) {
+            nms = nms.subList(0, MAX_DETECTION);
+        }
+        Joints[] ret = new Joints[nms.size()];
+        for (int i = 0; i < ret.length; ++i) {
+            List<Joint> joints = new ArrayList<>();
+            ret[i] = new Joints(joints);
+
+            int index = nms.get(i);
+            int pos = index * 51;
+            for (int j = 0; j < 17; ++j) {
+                joints.add(
+                        new Joints.Joint(
+                                mask[pos + j * 3] / width,
+                                mask[pos + j * 3 + 1] / height,
+                                mask[pos + j * 3 + 2]));
+            }
+        }
+        return ret;
+    }
+
+    private NDArray xywh2xyxy(NDArray array) {
+        NDArray xy = array.get("..., :2");
+        NDArray wh = array.get("..., 2:").div(2);
+        return xy.sub(wh).concat(xy.add(wh), -1);
+    }
+
+    /**
+     * Creates a builder to build a {@code YoloPoseTranslator}.
+     *
+     * @return a new builder
+     */
+    public static Builder builder() {
+        return new Builder();
+    }
+
+    /**
+     * Creates a builder to build a {@code YoloPoseTranslator} with specified arguments.
+     *
+     * @param arguments arguments to specify builder options
+     * @return a new builder
+     */
+    public static Builder builder(Map<String, ?> arguments) {
+        Builder builder = new Builder();
+        builder.configPreProcess(arguments);
+        builder.configPostProcess(arguments);
+
+        return builder;
+    }
+
+    /** The builder for Pose Estimation translator. */
+    public static class Builder extends BaseBuilder<Builder> {
+
+        float threshold = 0.25f;
+        float nmsThreshold = 0.7f;
+
+        Builder() {}
+
+        /**
+         * Sets the threshold for prediction accuracy.
+         *
+         * <p>Predictions below the threshold will be dropped.
+         *
+         * @param threshold the threshold for prediction accuracy
+         * @return the builder
+         */
+        public Builder optThreshold(float threshold) {
+            this.threshold = threshold;
+            return self();
+        }
+
+        /**
+         * Sets the NMS threshold.
+         *
+         * @param nmsThreshold the NMS threshold
+         * @return this builder
+         */
+        public Builder optNmsThreshold(float nmsThreshold) {
+            this.nmsThreshold = nmsThreshold;
+            return this;
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        protected Builder self() {
+            return this;
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        protected void configPostProcess(Map<String, ?> arguments) {
+            optThreshold(ArgumentsUtil.floatValue(arguments, "threshold", threshold));
+            optNmsThreshold(ArgumentsUtil.floatValue(arguments, "nmsThreshold", nmsThreshold));
+        }
+
+        /**
+         * Builds the translator.
+         *
+         * @return the new translator
+         */
+        public YoloPoseTranslator build() {
+            validate();
+            return new YoloPoseTranslator(this);
+        }
+    }
+}
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/YoloPoseTranslatorFactory.java b/api/src/main/java/ai/djl/modality/cv/translator/YoloPoseTranslatorFactory.java
new file mode 100644
index 000000000..e27620acd
--- /dev/null
+++ b/api/src/main/java/ai/djl/modality/cv/translator/YoloPoseTranslatorFactory.java
@@ -0,0 +1,61 @@
+/*
+ * Copyright 2024 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.Model;
+import ai.djl.modality.Input;
+import ai.djl.modality.Output;
+import ai.djl.modality.cv.Image;
+import ai.djl.modality.cv.output.Joints;
+import ai.djl.translate.Translator;
+import ai.djl.translate.TranslatorFactory;
+import ai.djl.util.Pair;
+
+import java.io.Serializable;
+import java.lang.reflect.Type;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+
+/** An {@link TranslatorFactory} that creates a {@link YoloPoseTranslator} instance. */
+public class YoloPoseTranslatorFactory implements TranslatorFactory, Serializable {
+
+    private static final long serialVersionUID = 1L;
+
+    private static final Set<Pair<Type, Type>> SUPPORTED_TYPES = new HashSet<>();
+
+    static {
+        SUPPORTED_TYPES.add(new Pair<>(Image.class, Joints[].class));
+        SUPPORTED_TYPES.add(new Pair<>(Input.class, Output.class));
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Set<Pair<Type, Type>> getSupportedTypes() {
+        return SUPPORTED_TYPES;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    @SuppressWarnings("unchecked")
+    public <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
+        YoloPoseTranslator translator = YoloPoseTranslator.builder(arguments).build();
+        if (input == Image.class && output == Joints[].class) {
+            return (Translator<I, O>) translator;
+        } else if (input == Input.class && output == Output.class) {
+            return (Translator<I, O>) new ImageServingTranslator(translator);
+        }
+        throw new IllegalArgumentException("Unsupported input/output types.");
+    }
+}
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/YoloV5Translator.java b/api/src/main/java/ai/djl/modality/cv/translator/YoloV5Translator.java
index c31353766..67682d7c8 100644
--- a/api/src/main/java/ai/djl/modality/cv/translator/YoloV5Translator.java
+++ b/api/src/main/java/ai/djl/modality/cv/translator/YoloV5Translator.java
@@ -24,7 +24,6 @@ import java.util.ArrayList;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
-import java.util.PriorityQueue;
 
 /**
  * A translator for YoloV5 models. This was tested with ONNX exported Yolo models. For details check
@@ -68,104 +67,58 @@ public class YoloV5Translator extends ObjectDetectionTranslator {
         return builder;
     }
 
-    protected double boxIntersection(Rectangle a, Rectangle b) {
-        double w =
-                overlap(
-                        (a.getX() * 2 + a.getWidth()) / 2,
-                        a.getWidth(),
-                        (b.getX() * 2 + b.getWidth()) / 2,
-                        b.getWidth());
-        double h =
-                overlap(
-                        (a.getY() * 2 + a.getHeight()) / 2,
-                        a.getHeight(),
-                        (b.getY() * 2 + b.getHeight()) / 2,
-                        b.getHeight());
-        if (w < 0 || h < 0) {
-            return 0;
-        }
-        return w * h;
-    }
-
-    protected double boxIou(Rectangle a, Rectangle b) {
-        return boxIntersection(a, b) / boxUnion(a, b);
-    }
-
-    protected double boxUnion(Rectangle a, Rectangle b) {
-        double i = boxIntersection(a, b);
-        return (a.getWidth()) * (a.getHeight()) + (b.getWidth()) * (b.getHeight()) - i;
-    }
-
-    protected DetectedObjects nms(List<IntermediateResult> list) {
+    protected DetectedObjects nms(
+            List<Rectangle> boxes, List<Integer> classIds, List<Float> scores) {
         List<String> retClasses = new ArrayList<>();
         List<Double> retProbs = new ArrayList<>();
         List<BoundingBox> retBB = new ArrayList<>();
 
-        for (int k = 0; k < classes.size(); k++) {
-            // 1.find max confidence per class
-            PriorityQueue<IntermediateResult> pq =
-                    new PriorityQueue<>(
-                            50,
-                            (lhs, rhs) -> {
-                                // Intentionally reversed to put high confidence at the head of the
-                                // queue.
-                                return Double.compare(rhs.getConfidence(), lhs.getConfidence());
-                            });
-
-            for (IntermediateResult intermediateResult : list) {
-                if (intermediateResult.getDetectedClass() == k) {
-                    pq.add(intermediateResult);
+        for (int classId = 0; classId < classes.size(); classId++) {
+            List<Rectangle> r = new ArrayList<>();
+            List<Double> s = new ArrayList<>();
+            List<Integer> map = new ArrayList<>();
+            for (int j = 0; j < classIds.size(); ++j) {
+                if (classIds.get(j) == classId) {
+                    r.add(boxes.get(j));
+                    s.add(scores.get(j).doubleValue());
+                    map.add(j);
                 }
             }
-
-            // 2.do non maximum suppression
-            while (pq.size() > 0) {
-                // insert detection with max confidence
-                IntermediateResult[] a = new IntermediateResult[pq.size()];
-                IntermediateResult[] detections = pq.toArray(a);
-                Rectangle rec = detections[0].getLocation();
-                retClasses.add(detections[0].id);
-                retProbs.add(detections[0].confidence);
+            if (r.isEmpty()) {
+                continue;
+            }
+            List<Integer> nms = Rectangle.nms(r, s, nmsThreshold);
+            for (int index : nms) {
+                int pos = map.get(index);
+                int id = classIds.get(pos);
+                retClasses.add(classes.get(id));
+                retProbs.add(scores.get(pos).doubleValue());
+                Rectangle rect = boxes.get(pos);
                 if (applyRatio) {
                     retBB.add(
                             new Rectangle(
-                                    rec.getX() / imageWidth,
-                                    rec.getY() / imageHeight,
-                                    rec.getWidth() / imageWidth,
-                                    rec.getHeight() / imageHeight));
+                                    rect.getX() / imageWidth,
+                                    rect.getY() / imageHeight,
+                                    rect.getWidth() / imageWidth,
+                                    rect.getHeight() / imageHeight));
                 } else {
-                    retBB.add(
-                            new Rectangle(rec.getX(), rec.getY(), rec.getWidth(), rec.getHeight()));
-                }
-                pq.clear();
-                for (int j = 1; j < detections.length; j++) {
-                    IntermediateResult detection = detections[j];
-                    Rectangle location = detection.getLocation();
-                    if (boxIou(rec, location) < nmsThreshold) {
-                        pq.add(detection);
-                    }
+                    retBB.add(rect);
                 }
             }
         }
         return new DetectedObjects(retClasses, retProbs, retBB);
     }
 
-    protected double overlap(double x1, double w1, double x2, double w2) {
-        double l1 = x1 - w1 / 2;
-        double l2 = x2 - w2 / 2;
-        double left = Math.max(l1, l2);
-        double r1 = x1 + w1 / 2;
-        double r2 = x2 + w2 / 2;
-        double right = Math.min(r1, r2);
-        return right - left;
-    }
-
     protected DetectedObjects processFromBoxOutput(NDList list) {
         float[] flattened = list.get(0).toFloatArray();
-        ArrayList<IntermediateResult> intermediateResults = new ArrayList<>();
         int sizeClasses = classes.size();
         int stride = 5 + sizeClasses;
         int size = flattened.length / stride;
+
+        ArrayList<Rectangle> boxes = new ArrayList<>();
+        ArrayList<Float> scores = new ArrayList<>();
+        ArrayList<Integer> classIds = new ArrayList<>();
+
         for (int i = 0; i < size; i++) {
             int indexBase = i * stride;
             float maxClass = 0;
@@ -184,11 +137,12 @@ public class YoloV5Translator extends ObjectDetectionTranslator {
                 float h = flattened[indexBase + 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), score, maxIndex, rect));
+                boxes.add(rect);
+                scores.add(score);
+                classIds.add(maxIndex);
             }
         }
-        return nms(intermediateResults);
+        return nms(boxes, classIds, scores);
     }
 
     private DetectedObjects processFromDetectOutput() {
@@ -279,49 +233,4 @@ public class YoloV5Translator extends ObjectDetectionTranslator {
             return new YoloV5Translator(this);
         }
     }
-
-    protected static final class IntermediateResult {
-
-        /**
-         * A sortable score for how good the recognition is relative to others. Higher should be
-         * better.
-         */
-        private double confidence;
-
-        /** Display name for the recognition. */
-        private int detectedClass;
-
-        /**
-         * A unique identifier for what has been recognized. Specific to the class, not the instance
-         * of the object.
-         */
-        private String id;
-
-        /** Optional location within the source image for the location of the recognized object. */
-        private Rectangle location;
-
-        IntermediateResult(String id, double confidence, int detectedClass, Rectangle location) {
-            this.confidence = confidence;
-            this.id = id;
-            this.detectedClass = detectedClass;
-            this.location = location;
-        }
-
-        public double getConfidence() {
-            return confidence;
-        }
-
-        public int getDetectedClass() {
-            return detectedClass;
-        }
-
-        public String getId() {
-            return id;
-        }
-
-        public Rectangle getLocation() {
-            return new Rectangle(
-                    location.getX(), location.getY(), location.getWidth(), location.getHeight());
-        }
-    }
 }
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 7fb1fda62..7895ad15c 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
@@ -79,7 +79,10 @@ public class YoloV8Translator extends YoloV5Translator {
                     "Expected classes: " + (nClasses - 4) + ", got " + classes.size());
         }
 
-        ArrayList<IntermediateResult> intermediateResults = new ArrayList<>();
+        ArrayList<Rectangle> boxes = new ArrayList<>();
+        ArrayList<Float> scores = new ArrayList<>();
+        ArrayList<Integer> classIds = 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;
@@ -101,12 +104,13 @@ public class YoloV8Translator extends YoloV5Translator {
                 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));
+                boxes.add(rect);
+                scores.add(maxClassProb);
+                classIds.add(maxIndex);
             }
         }
-        return nms(intermediateResults);
+
+        return nms(boxes, classIds, scores);
     }
 
     /** The builder for {@link YoloV8Translator}. */
diff --git a/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/PtModelZoo.java b/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/PtModelZoo.java
index eacfe74c8..336e72e98 100644
--- a/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/PtModelZoo.java
+++ b/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/PtModelZoo.java
@@ -45,6 +45,7 @@ public class PtModelZoo extends ModelZoo {
         addModel(REPOSITORY.model(CV.OBJECT_DETECTION, GROUP_ID, "ssd", "0.0.1"));
         addModel(REPOSITORY.model(CV.OBJECT_DETECTION, GROUP_ID, "yolov5s", "0.0.1"));
         addModel(REPOSITORY.model(CV.OBJECT_DETECTION, GROUP_ID, "yolov8n", "0.0.1"));
+        addModel(REPOSITORY.model(CV.POSE_ESTIMATION, GROUP_ID, "yolov8n-pose", "0.0.1"));
         addModel(REPOSITORY.model(NLP.QUESTION_ANSWER, GROUP_ID, "bertqa", "0.0.1"));
         addModel(REPOSITORY.model(NLP.SENTIMENT_ANALYSIS, GROUP_ID, "distilbert", "0.0.1"));
         addModel(REPOSITORY.model(CV.IMAGE_GENERATION, GROUP_ID, "biggan-deep", "0.0.1"));
diff --git a/examples/src/main/java/ai/djl/examples/inference/cv/PoseEstimation.java b/examples/src/main/java/ai/djl/examples/inference/cv/PoseEstimation.java
index 3a5866c7b..9187c1f61 100644
--- a/examples/src/main/java/ai/djl/examples/inference/cv/PoseEstimation.java
+++ b/examples/src/main/java/ai/djl/examples/inference/cv/PoseEstimation.java
@@ -12,18 +12,14 @@
  */
 package ai.djl.examples.inference.cv;
 
-import ai.djl.MalformedModelException;
 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.DetectedObjects;
 import ai.djl.modality.cv.output.Joints;
-import ai.djl.modality.cv.output.Rectangle;
+import ai.djl.modality.cv.translator.YoloPoseTranslatorFactory;
 import ai.djl.repository.zoo.Criteria;
-import ai.djl.repository.zoo.ModelNotFoundException;
 import ai.djl.repository.zoo.ZooModel;
-import ai.djl.training.util.ProgressBar;
 import ai.djl.translate.TranslateException;
 
 import org.slf4j.Logger;
@@ -33,9 +29,7 @@ import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.List;
+import java.util.Arrays;
 
 /**
  * An example of inference using a pose estimation model.
@@ -51,98 +45,39 @@ public final class PoseEstimation {
     private PoseEstimation() {}
 
     public static void main(String[] args) throws IOException, ModelException, TranslateException {
-        List<Joints> joints = predict();
-        logger.info("{}", joints);
+        Joints[] joints = predict();
+        logger.info("{}", Arrays.toString(joints));
     }
 
-    public static List<Joints> predict() throws IOException, ModelException, TranslateException {
+    public static Joints[] predict() throws IOException, ModelException, TranslateException {
         Path imageFile = Paths.get("src/test/resources/pose_soccer.png");
         Image img = ImageFactory.getInstance().fromFile(imageFile);
 
-        List<Image> people = predictPeopleInImage(img);
-
-        if (people.isEmpty()) {
-            logger.warn("No people found in image.");
-            return Collections.emptyList();
-        }
-
-        return predictJointsForPeople(people);
-    }
-
-    private static List<Image> predictPeopleInImage(Image img)
-            throws MalformedModelException,
-                    ModelNotFoundException,
-                    IOException,
-                    TranslateException {
-
-        Criteria<Image, DetectedObjects> criteria =
-                Criteria.builder()
-                        .setTypes(Image.class, DetectedObjects.class)
-                        .optModelUrls("djl://ai.djl.mxnet/ssd/0.0.1/ssd_512_resnet50_v1_voc")
-                        .optEngine("MXNet")
-                        .optProgress(new ProgressBar())
-                        .build();
-
-        DetectedObjects detectedObjects;
-        try (ZooModel<Image, DetectedObjects> ssd = criteria.loadModel();
-                Predictor<Image, DetectedObjects> predictor = ssd.newPredictor()) {
-            detectedObjects = predictor.predict(img);
-        }
-
-        List<DetectedObjects.DetectedObject> items = detectedObjects.items();
-        List<Image> people = new ArrayList<>();
-        for (DetectedObjects.DetectedObject item : items) {
-            if ("person".equals(item.getClassName())) {
-                Rectangle rect = item.getBoundingBox().getBounds();
-                int width = img.getWidth();
-                int height = img.getHeight();
-                people.add(
-                        img.getSubImage(
-                                (int) (rect.getX() * width),
-                                (int) (rect.getY() * height),
-                                (int) (rect.getWidth() * width),
-                                (int) (rect.getHeight() * height)));
-            }
-        }
-        return people;
-    }
-
-    private static List<Joints> predictJointsForPeople(List<Image> people)
-            throws MalformedModelException,
-                    ModelNotFoundException,
-                    IOException,
-                    TranslateException {
-
-        // Use DJL MXNet model zoo model, model can be found:
-        // https://mlrepo.djl.ai/model/cv/pose_estimation/ai/djl/mxnet/simple_pose/0.0.1/simple_pose_resnet18_v1b-0000.params.gz
-        // https://mlrepo.djl.ai/model/cv/pose_estimation/ai/djl/mxnet/simple_pose/0.0.1/simple_pose_resnet18_v1b-symbol.json
-        Criteria<Image, Joints> criteria =
+        // Use DJL PyTorch model zoo model
+        Criteria<Image, Joints[]> criteria =
                 Criteria.builder()
-                        .setTypes(Image.class, Joints.class)
-                        .optModelUrls(
-                                "djl://ai.djl.mxnet/simple_pose/0.0.1/simple_pose_resnet18_v1b")
+                        .setTypes(Image.class, Joints[].class)
+                        .optModelUrls("djl://ai.djl.pytorch/yolov8n-pose")
+                        .optTranslatorFactory(new YoloPoseTranslatorFactory())
                         .build();
 
-        List<Joints> allJoints = new ArrayList<>();
-        try (ZooModel<Image, Joints> pose = criteria.loadModel();
-                Predictor<Image, Joints> predictor = pose.newPredictor()) {
-            int count = 0;
-            for (Image person : people) {
-                Joints joints = predictor.predict(person);
-                saveJointsImage(person, joints, count++);
-                allJoints.add(joints);
-            }
+        try (ZooModel<Image, Joints[]> pose = criteria.loadModel();
+                Predictor<Image, Joints[]> predictor = pose.newPredictor()) {
+            Joints[] allJoints = predictor.predict(img);
+            saveJointsImage(img, allJoints);
+            return allJoints;
         }
-        return allJoints;
     }
 
-    private static void saveJointsImage(Image img, Joints joints, int count) throws IOException {
+    private static void saveJointsImage(Image img, Joints[] allJoints) throws IOException {
         Path outputDir = Paths.get("build/output");
         Files.createDirectories(outputDir);
 
-        img.drawJoints(joints);
+        for (Joints joints : allJoints) {
+            img.drawJoints(joints);
+        }
 
-        Path imagePath = outputDir.resolve("joints-" + count + ".png");
+        Path imagePath = outputDir.resolve("joints.png");
         // Must use png format because you can't save as jpg with an alpha channel
         img.save(Files.newOutputStream(imagePath), "png");
         logger.info("Pose image has been saved in: {}", imagePath);
diff --git a/examples/src/main/java/ai/djl/examples/training/TrainTimeSeries.java b/examples/src/main/java/ai/djl/examples/training/TrainTimeSeries.java
index b9abb016e..66f734e9c 100644
--- a/examples/src/main/java/ai/djl/examples/training/TrainTimeSeries.java
+++ b/examples/src/main/java/ai/djl/examples/training/TrainTimeSeries.java
@@ -120,10 +120,10 @@ public final class TrainTimeSeries {
                 inputShapes[6] =
                         new Shape(
                                 1,
-                                predictionLength,
+                                contextLength,
                                 TimeFeature.timeFeaturesFromFreqStr(freq).size() + 1);
-                inputShapes[7] = new Shape(1, predictionLength);
-                inputShapes[8] = new Shape(1, predictionLength);
+                inputShapes[7] = new Shape(1, contextLength);
+                inputShapes[8] = new Shape(1, contextLength);
                 trainer.initialize(inputShapes);
 
                 EasyTrain.fit(trainer, arguments.getEpoch(), trainSet, null);
@@ -147,6 +147,7 @@ public final class TrainTimeSeries {
 
             Map<String, Object> arguments = new ConcurrentHashMap<>();
             arguments.put("prediction_length", predictionLength);
+            arguments.put("context_length", predictionNetwork.getContextLength());
             arguments.put("freq", freq);
             arguments.put("use_" + FieldName.FEAT_DYNAMIC_REAL.name().toLowerCase(), false);
             arguments.put("use_" + FieldName.FEAT_STATIC_CAT.name().toLowerCase(), true);
@@ -239,6 +240,7 @@ public final class TrainTimeSeries {
                         .setFreq(freq)
                         .setPredictionLength(predictionLength)
                         .optDistrOutput(distributionOutput)
+                        .optContextLength(8)
                         .optUseFeatStaticCat(true);
         return training ? builder.buildTrainingNetwork() : builder.buildPredictionNetwork();
     }
diff --git a/extensions/opencv/src/main/java/ai/djl/opencv/OpenCVImage.java b/extensions/opencv/src/main/java/ai/djl/opencv/OpenCVImage.java
index b6a229271..d4f29f32e 100644
--- a/extensions/opencv/src/main/java/ai/djl/opencv/OpenCVImage.java
+++ b/extensions/opencv/src/main/java/ai/djl/opencv/OpenCVImage.java
@@ -205,12 +205,25 @@ class OpenCVImage implements Image {
         int imageWidth = image.width();
         int imageHeight = image.height();
 
-        Scalar color =
-                new Scalar(
-                        RandomUtils.nextInt(178),
-                        RandomUtils.nextInt(178),
-                        RandomUtils.nextInt(178));
-        for (Joints.Joint joint : joints.getJoints()) {
+        List<Joints.Joint> list = joints.getJoints();
+        if (list.size() == 17) {
+            Scalar color = new Scalar(37, 255, 224);
+            drawLine(list.get(5), list.get(7), imageWidth, imageHeight, color);
+            drawLine(list.get(7), list.get(9), imageWidth, imageHeight, color);
+            drawLine(list.get(6), list.get(8), imageWidth, imageHeight, color);
+            drawLine(list.get(8), list.get(10), imageWidth, imageHeight, color);
+            drawLine(list.get(11), list.get(13), imageWidth, imageHeight, color);
+            drawLine(list.get(12), list.get(14), imageWidth, imageHeight, color);
+            drawLine(list.get(13), list.get(15), imageWidth, imageHeight, color);
+            drawLine(list.get(14), list.get(16), imageWidth, imageHeight, color);
+            drawLine(list.get(5), list.get(6), imageWidth, imageHeight, color);
+            drawLine(list.get(11), list.get(12), imageWidth, imageHeight, color);
+            drawLine(list.get(5), list.get(11), imageWidth, imageHeight, color);
+            drawLine(list.get(6), list.get(12), imageWidth, imageHeight, color);
+        }
+
+        Scalar color = new Scalar(190, 150, 37);
+        for (Joints.Joint joint : list) {
             int x = (int) (joint.getX() * imageWidth);
             int y = (int) (joint.getY() * imageHeight);
             Point point = new Point(x, y);
@@ -340,6 +353,14 @@ class OpenCVImage implements Image {
         return new OpenCVImage(result);
     }
 
+    private void drawLine(Joints.Joint from, Joints.Joint to, int width, int height, Scalar color) {
+        int x0 = (int) (from.getX() * width);
+        int y0 = (int) (from.getY() * height);
+        int x1 = (int) (to.getX() * width);
+        int y1 = (int) (to.getY() * height);
+        Imgproc.line(image, new Point(x0, y0), new Point(x1, y1), color, 2, Imgproc.LINE_AA);
+    }
+
     private void drawLandmarks(BoundingBox box) {
         Scalar color = new Scalar(0, 96, 246);
         for (ai.djl.modality.cv.output.Point point : box.getPath()) {
@@ -356,27 +377,37 @@ class OpenCVImage implements Image {
         float b = RandomUtils.nextFloat();
         int imageWidth = img.getWidth();
         int imageHeight = img.getHeight();
-        int x = (int) (mask.getX() * imageWidth);
-        int y = (int) (mask.getY() * imageHeight);
-        float[][] probDist = mask.getProbDist();
-        // Correct some coordinates of box when going out of image
-        if (x < 0) {
-            x = 0;
-        }
-        if (y < 0) {
-            y = 0;
+        int x = 0;
+        int y = 0;
+        int w = imageWidth;
+        int h = imageHeight;
+        if (!mask.isFullImageMask()) {
+            x = (int) (mask.getX() * imageWidth);
+            y = (int) (mask.getY() * imageHeight);
+            w = (int) (mask.getWidth() * imageWidth);
+            h = (int) (mask.getHeight() * imageHeight);
+            // Correct some coordinates of box when going out of image
+            if (x < 0) {
+                x = 0;
+            }
+            if (y < 0) {
+                y = 0;
+            }
         }
+        float[][] probDist = mask.getProbDist();
 
         BufferedImage maskImage =
-                new BufferedImage(probDist.length, probDist[0].length, BufferedImage.TYPE_INT_ARGB);
-        for (int xCor = 0; xCor < probDist.length; xCor++) {
-            for (int yCor = 0; yCor < probDist[xCor].length; yCor++) {
-                float opacity = probDist[xCor][yCor] * 0.8f;
-                maskImage.setRGB(xCor, yCor, new Color(r, g, b, opacity).getRGB());
+                new BufferedImage(probDist[0].length, probDist.length, BufferedImage.TYPE_INT_ARGB);
+        for (int yCor = 0; yCor < probDist.length; yCor++) {
+            for (int xCor = 0; xCor < probDist[0].length; xCor++) {
+                float opacity = probDist[yCor][xCor] * 0.8f;
+                maskImage.setRGB(xCor, yCor, new Color(r, g, b, opacity).darker().getRGB());
             }
         }
+        java.awt.Image scaled = maskImage.getScaledInstance(w, h, java.awt.Image.SCALE_SMOOTH);
+
         Graphics2D gR = (Graphics2D) img.getGraphics();
-        gR.drawImage(maskImage, x, y, null);
+        gR.drawImage(scaled, x, y, null);
         gR.dispose();
     }
 
diff --git a/extensions/timeseries/src/main/java/ai/djl/timeseries/transform/feature/Feature.java b/extensions/timeseries/src/main/java/ai/djl/timeseries/transform/feature/Feature.java
index 0ade03d3f..f7e1f6255 100644
--- a/extensions/timeseries/src/main/java/ai/djl/timeseries/transform/feature/Feature.java
+++ b/extensions/timeseries/src/main/java/ai/djl/timeseries/transform/feature/Feature.java
@@ -119,14 +119,14 @@ public final class Feature {
         if (!freq.matches("\\d+.*")) {
             sb.insert(0, 1);
         }
-        sb.insert(0, "P");
-        String formattedFreq = sb.toString();
 
         TemporalAmount timeFreq;
         if (freq.endsWith("H") || freq.endsWith("T") || freq.endsWith("S")) {
-            timeFreq = Duration.parse(formattedFreq);
+            sb.insert(0, "PT");
+            timeFreq = Duration.parse(sb.toString());
         } else {
-            timeFreq = Period.parse(formattedFreq);
+            sb.insert(0, "P");
+            timeFreq = Period.parse(sb.toString());
         }
 
         List<LocalDateTime> index = new ArrayList<>();
diff --git a/extensions/timeseries/src/main/java/ai/djl/timeseries/translator/BaseTimeSeriesTranslator.java b/extensions/timeseries/src/main/java/ai/djl/timeseries/translator/BaseTimeSeriesTranslator.java
index c84e874a3..d34bcfb6f 100644
--- a/extensions/timeseries/src/main/java/ai/djl/timeseries/translator/BaseTimeSeriesTranslator.java
+++ b/extensions/timeseries/src/main/java/ai/djl/timeseries/translator/BaseTimeSeriesTranslator.java
@@ -39,8 +39,7 @@ public abstract class BaseTimeSeriesTranslator implements Translator<TimeSeriesD
         this.batchifier = builder.batchifier;
         this.freq = builder.freq;
         this.predictionLength = builder.predictionLength;
-        // TODO: for inferring
-        this.contextLength = builder.predictionLength;
+        this.contextLength = builder.contextLength;
     }
 
     /** {@inheritDoc} */
@@ -57,6 +56,7 @@ public abstract class BaseTimeSeriesTranslator implements Translator<TimeSeriesD
     public abstract static class BaseBuilder<T extends BaseBuilder<T>> {
         protected Batchifier batchifier = Batchifier.STACK;
         protected int predictionLength;
+        protected int contextLength;
 
         protected String freq;
 
@@ -82,6 +82,8 @@ public abstract class BaseTimeSeriesTranslator implements Translator<TimeSeriesD
                 throw new IllegalArgumentException(
                         "The value of `prediction_length` should be > 0");
             }
+            this.contextLength =
+                    ArgumentsUtil.intValue(arguments, "context_length", predictionLength);
             if (arguments.containsKey("batchifier")) {
                 batchifier = Batchifier.fromString((String) arguments.get("batchifier"));
             }


Test patch:
diff --git a/engines/pytorch/pytorch-model-zoo/src/test/resources/mlrepo/model/cv/pose_estimation/ai/djl/pytorch/yolov8n-pose/metadata.json b/engines/pytorch/pytorch-model-zoo/src/test/resources/mlrepo/model/cv/pose_estimation/ai/djl/pytorch/yolov8n-pose/metadata.json
new file mode 100644
index 000000000..62741a61e
--- /dev/null
+++ b/engines/pytorch/pytorch-model-zoo/src/test/resources/mlrepo/model/cv/pose_estimation/ai/djl/pytorch/yolov8n-pose/metadata.json
@@ -0,0 +1,43 @@
+{
+  "metadataVersion": "0.2",
+  "resourceType": "model",
+  "application": "cv/pose_estimation",
+  "groupId": "ai.djl.pytorch",
+  "artifactId": "yolov8n-pose",
+  "name": "Yolov8n pose",
+  "description": "Yolov8n Pose Estimation",
+  "website": "http://www.djl.ai/engines/pytorch/pytorch-model-zoo",
+  "licenses": {
+    "license": {
+      "name": "The Apache License, Version 2.0",
+      "url": "https://www.apache.org/licenses/LICENSE-2.0"
+    }
+  },
+  "artifacts": [
+    {
+      "version": "0.0.1",
+      "snapshot": false,
+      "name": "yolov8n-pose",
+      "properties": {
+      },
+      "arguments": {
+        "width": 640,
+        "height": 640,
+        "resize": true,
+        "threshold": 0.25,
+        "translatorFactory": "ai.djl.modality.cv.translator.YoloPoseTranslatorFactory"
+      },
+      "options": {
+        "mapLocation": "true"
+      },
+      "files": {
+        "model": {
+          "uri": "0.0.1/yolov8n-pose.zip",
+          "sha1Hash": "b2358198586ea35189f30aa5f49285dc1842ec1b",
+          "name": "",
+          "size": 11684792
+        }
+      }
+    }
+  ]
+}
diff --git a/examples/src/test/java/ai/djl/examples/inference/PoseEstimationTest.java b/examples/src/test/java/ai/djl/examples/inference/PoseEstimationTest.java
index d0afdb90b..398c71ced 100644
--- a/examples/src/test/java/ai/djl/examples/inference/PoseEstimationTest.java
+++ b/examples/src/test/java/ai/djl/examples/inference/PoseEstimationTest.java
@@ -15,22 +15,19 @@ package ai.djl.examples.inference;
 import ai.djl.ModelException;
 import ai.djl.examples.inference.cv.PoseEstimation;
 import ai.djl.modality.cv.output.Joints;
-import ai.djl.testing.TestRequirements;
 import ai.djl.translate.TranslateException;
 
 import org.testng.Assert;
 import org.testng.annotations.Test;
 
 import java.io.IOException;
-import java.util.List;
 
 public class PoseEstimationTest {
 
     @Test
     public void testPoseEstimation() throws ModelException, TranslateException, IOException {
-        TestRequirements.linux();
-
-        List<Joints> result = PoseEstimation.predict();
-        Assert.assertTrue(result.get(0).getJoints().get(0).getConfidence() > 0.6d);
+        Joints[] result = PoseEstimation.predict();
+        Assert.assertEquals(result.length, 3);
+        Assert.assertEquals(result[0].getJoints().size(), 17);
     }
 }