Fix Pose val plotting bug (#3232)

2023-06-17 19:35:19 +02:00 · 2023-06-17 19:35:19 +02:00 · d3ca3e4b9b
commit d3ca3e4b9b
parent 1e5702a5b5
1 changed files with 1 additions and 1 deletions
--- a/ultralytics/yolo/v8/pose/val.py
+++ b/ultralytics/yolo/v8/pose/val.py
@ -154,7 +154,7 @@ class PoseValidator(DetectionValidator):

    def plot_predictions(self, batch, preds, ni):
        """Plots predictions for YOLO model."""
-        pred_kpts = torch.cat([p[:, 6:].view(-1, *self.kpt_shape)[:15] for p in preds], 0)
+        pred_kpts = torch.cat([p[:, 6:].view(-1, *self.kpt_shape) for p in preds], 0)
        plot_images(batch['img'],
                    *output_to_target(preds, max_det=self.args.max_det),
                    kpts=pred_kpts,