Fix Pose val plotting bug (#3232)

single_channel
Glenn Jocher 2 years ago committed by GitHub
parent 1e5702a5b5
commit d3ca3e4b9b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -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,

Loading…
Cancel
Save