nms.hpp

Go to the documentation of this file.

// SPDX-FileCopyrightText: 2025 Povl Filip Sonne-Frederiksen
//
// SPDX-License-Identifier: GPL-3.0-or-later
 
#pragma once
#include <torch/torch.h>
 
namespace reusex::vision {

torch::Tensor xyxy_to_xywh(const torch::Tensor &x);

torch::Tensor xywh_to_xyxy(const torch::Tensor &x);

// TODO: Replace custom NMS with torchvision library implementation
// category=Vision estimate=4h
// Current implementation is custom-written. Consider using official torchvision
// NMS: Reference:
// https://github.com/pytorch/vision/blob/main/torchvision/csrc/ops/cpu/nms_kernel.cpp
// Benefits:
// 1. Optimized CPU/CUDA implementations available
// 2. Better maintained and tested by PyTorch team
// 3. Reduces custom code maintenance burden
// Trade-off: Adds torchvision as dependency (currently only use LibTorch)
torch::Tensor nms(const torch::Tensor &bboxes, const torch::Tensor &scores,
                  float iou_threshold = 0.45);

torch::Tensor non_max_suppression(torch::Tensor predictions,
                                  float confThreshold = 0.25,
                                  float iouThreshold = 0.45,
                                  int maxDetections = 300);
 
} // namespace reusex::vision