HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

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

Browse source 
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
higor-melo 2024-03-27 15:31:42 -03:00 • committed by GitHub
parent 898cbcdc08
commit 1325889305
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 2 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

4
examples/YOLOv8-LibTorch-CPP-Inference/main.cc
View file

@ -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 bs = prediction.size(0);
auto nc = prediction.size(1) - 4; auto nc = prediction.size(1) - 4;
auto nm = prediction.size(1) - nc - 4; auto nm = prediction.size(1) - nc - 4;
@ -237,7 +237,7 @@ int main() {
torch::Tensor output = yolo_model.forward(inputs).toTensor().cpu(); torch::Tensor output = yolo_model.forward(inputs).toTensor().cpu();
// NMS // NMS
auto keep = non_max_supperession(output)[0]; auto keep = non_max_suppression(output)[0];
auto boxes = keep.index({Slice(), Slice(None, 4)}); 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})); keep.index_put_({Slice(), Slice(None, 4)}, scale_boxes({input_image.rows, input_image.cols}, boxes, {image.rows, image.cols}));