diff --git a/src/binding.cc b/src/binding.cc index 445f853f328cf6183e98cc2f86c631ba1056890c..13186dd7177a1e0ecb37cb85d9dfe5cbbba00490 100644 --- a/src/binding.cc +++ b/src/binding.cc @@ -4,36 +4,21 @@ #include #include +namespace nms { +enum class Method : uint32_t { LINEAR = 0, GAUSSIAN, HARD }; -namespace nms -{ -enum class Method : uint32_t -{ - LINEAR = 0, - GAUSSIAN, - HARD -}; - -size_t soft_nms(float* boxes, - int32_t* index, - size_t count, - Method method, - float Nt, - float sigma, - float threshold); +size_t soft_nms(float* boxes, int32_t* index, size_t count, Method method, + float Nt, float sigma, float threshold); } // namespace nms -namespace binding -{ +namespace binding { namespace py = pybind11; using namespace pybind11::literals; py::tuple py_soft_nms(py::array_t boxes, nms::Method method = nms::Method::GAUSSIAN, - float Nt = 0.3, - float sigma = 0.5, - float threshold = 0.001) -{ + float Nt = 0.3, float sigma = 0.5, + float threshold = 0.001) { assert(boxes.ndim() == 2 && "Input should be 2-D NumPy array"); assert(boxes.shape()[1] == 5 && "Input should have size [N,5]"); @@ -44,8 +29,8 @@ py::tuple py_soft_nms(py::array_t boxes, auto N = nms::soft_nms(b, i, count, method, Nt, sigma, threshold); - std::vector shape5 = {N, 5}; - std::vector shape1 = {N}; + std::vector shape5 = {N, 5}; + std::vector shape1 = {N}; std::vector strides5 = {sizeof(float) * 5, sizeof(float)}; std::vector strides1 = {sizeof(float)}; @@ -63,12 +48,12 @@ PYBIND11_MODULE(nms, m) { m.doc() = "SoftNMS for object detection."; py::enum_(m, "NMSMethod") - .value("LINEAR", nms::Method::LINEAR) - .value("GAUSSIAN", nms::Method::GAUSSIAN) - .value("HARD", nms::Method::HARD) - .export_values(); + .value("LINEAR", nms::Method::LINEAR) + .value("GAUSSIAN", nms::Method::GAUSSIAN) + .value("HARD", nms::Method::HARD) + .export_values(); m.def("soft_nms", &py_soft_nms, "boxes"_a.noconvert(), - "method"_a = nms::Method::GAUSSIAN, - "Nt"_a = 0.3, "sigma"_a = 0.5, "threshold"_a = 0.001); + "method"_a = nms::Method::GAUSSIAN, "Nt"_a = 0.3, "sigma"_a = 0.5, + "threshold"_a = 0.001); } } /* namespace binding */ diff --git a/src/nms.cc b/src/nms.cc index 84122685c129483d605d33344b49cf5cb27672c2..cbf7ae6207c0a5e3bd0334806cd0e6f95c73095d 100644 --- a/src/nms.cc +++ b/src/nms.cc @@ -1,24 +1,21 @@ #include #include -namespace nms -{ -struct proposal -{ +namespace nms { +struct proposal { float score, x1, y1, x2, y2; }; -inline static bool cmp(const proposal& a, const proposal& b) -{ +inline static bool cmp(const proposal& a, const proposal& b) { return a.score < b.score; } -inline static float iou(const proposal&, const proposal&) __attribute__((always_inline)); +inline static float iou(const proposal&, const proposal&) + __attribute__((always_inline)); -static float iou(const proposal& a, const proposal& b) -{ +static float iou(const proposal& a, const proposal& b) { auto overlap = 0.f; - float iw = std::min(b.x2, a.x2) - std::max(b.x1, a.x1) + 1; + float iw = std::min(b.x2, a.x2) - std::max(b.x1, a.x1) + 1; if (iw > 0) { float ih = std::min(b.y2, a.y2) - std::max(b.y1, a.y1) + 1; if (ih > 0) { @@ -31,21 +28,10 @@ static float iou(const proposal& a, const proposal& b) return overlap; } -enum class Method : uint32_t -{ - LINEAR = 0, - GAUSSIAN, - HARD -}; +enum class Method : uint32_t { LINEAR = 0, GAUSSIAN, HARD }; -size_t soft_nms(float* boxes, - int32_t* index, - size_t count, - Method method, - float Nt, - float sigma, - float threshold) -{ +size_t soft_nms(float* boxes, int32_t* index, size_t count, Method method, + float Nt, float sigma, float threshold) { std::iota(index, index + count, 0); // np.arange() auto p = reinterpret_cast(boxes); @@ -55,7 +41,7 @@ size_t soft_nms(float* boxes, std::swap(p[i], *max); std::swap(index[i], index[max - p]); - auto j = i + 1; + auto j = i + 1; auto weight = 0.f; while (j < N) { auto ov = iou(p[i], p[j]);