Fix typo in YOLOv8-Libtorch-CPP-Inference (#9330)

Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
pull/9338/head^2
higor-melo 8 months ago committed by GitHub
parent 898cbcdc08
commit 1325889305
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
  1. 4
      examples/YOLOv8-LibTorch-CPP-Inference/main.cc

@ -139,7 +139,7 @@ torch::Tensor nms(const torch::Tensor& bboxes, const torch::Tensor& scores, floa
}
torch::Tensor non_max_supperession(torch::Tensor& prediction, float conf_thres = 0.25, float iou_thres = 0.45, int max_det = 300) {
torch::Tensor non_max_suppression(torch::Tensor& prediction, float conf_thres = 0.25, float iou_thres = 0.45, int max_det = 300) {
auto bs = prediction.size(0);
auto nc = prediction.size(1) - 4;
auto nm = prediction.size(1) - nc - 4;
@ -237,7 +237,7 @@ int main() {
torch::Tensor output = yolo_model.forward(inputs).toTensor().cpu();
// NMS
auto keep = non_max_supperession(output)[0];
auto keep = non_max_suppression(output)[0];
auto boxes = keep.index({Slice(), Slice(None, 4)});
keep.index_put_({Slice(), Slice(None, 4)}, scale_boxes({input_image.rows, input_image.cols}, boxes, {image.rows, image.cols}));

Loading…
Cancel
Save