提交 762204e9 编写于 作者: Y Yang Zhang

Switch to "google" C++ style from customized "linux" style

上级 6f79f530
...@@ -4,36 +4,21 @@ ...@@ -4,36 +4,21 @@
#include <pybind11/numpy.h> #include <pybind11/numpy.h>
#include <pybind11/stl.h> #include <pybind11/stl.h>
namespace nms {
enum class Method : uint32_t { LINEAR = 0, GAUSSIAN, HARD };
namespace nms size_t soft_nms(float* boxes, int32_t* index, size_t count, Method method,
{ float Nt, float sigma, float threshold);
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);
} // namespace nms } // namespace nms
namespace binding namespace binding {
{
namespace py = pybind11; namespace py = pybind11;
using namespace pybind11::literals; using namespace pybind11::literals;
py::tuple py_soft_nms(py::array_t<float, py::array::c_style> boxes, py::tuple py_soft_nms(py::array_t<float, py::array::c_style> boxes,
nms::Method method = nms::Method::GAUSSIAN, nms::Method method = nms::Method::GAUSSIAN,
float Nt = 0.3, float Nt = 0.3, float sigma = 0.5,
float sigma = 0.5, float threshold = 0.001) {
float threshold = 0.001)
{
assert(boxes.ndim() == 2 && "Input should be 2-D NumPy array"); assert(boxes.ndim() == 2 && "Input should be 2-D NumPy array");
assert(boxes.shape()[1] == 5 && "Input should have size [N,5]"); assert(boxes.shape()[1] == 5 && "Input should have size [N,5]");
...@@ -44,8 +29,8 @@ py::tuple py_soft_nms(py::array_t<float, py::array::c_style> boxes, ...@@ -44,8 +29,8 @@ py::tuple py_soft_nms(py::array_t<float, py::array::c_style> boxes,
auto N = nms::soft_nms(b, i, count, method, Nt, sigma, threshold); auto N = nms::soft_nms(b, i, count, method, Nt, sigma, threshold);
std::vector<size_t> shape5 = {N, 5}; std::vector<size_t> shape5 = {N, 5};
std::vector<size_t> shape1 = {N}; std::vector<size_t> shape1 = {N};
std::vector<ssize_t> strides5 = {sizeof(float) * 5, sizeof(float)}; std::vector<ssize_t> strides5 = {sizeof(float) * 5, sizeof(float)};
std::vector<ssize_t> strides1 = {sizeof(float)}; std::vector<ssize_t> strides1 = {sizeof(float)};
...@@ -63,12 +48,12 @@ PYBIND11_MODULE(nms, m) { ...@@ -63,12 +48,12 @@ PYBIND11_MODULE(nms, m) {
m.doc() = "SoftNMS for object detection."; m.doc() = "SoftNMS for object detection.";
py::enum_<nms::Method>(m, "NMSMethod") py::enum_<nms::Method>(m, "NMSMethod")
.value("LINEAR", nms::Method::LINEAR) .value("LINEAR", nms::Method::LINEAR)
.value("GAUSSIAN", nms::Method::GAUSSIAN) .value("GAUSSIAN", nms::Method::GAUSSIAN)
.value("HARD", nms::Method::HARD) .value("HARD", nms::Method::HARD)
.export_values(); .export_values();
m.def("soft_nms", &py_soft_nms, "boxes"_a.noconvert(), m.def("soft_nms", &py_soft_nms, "boxes"_a.noconvert(),
"method"_a = nms::Method::GAUSSIAN, "method"_a = nms::Method::GAUSSIAN, "Nt"_a = 0.3, "sigma"_a = 0.5,
"Nt"_a = 0.3, "sigma"_a = 0.5, "threshold"_a = 0.001); "threshold"_a = 0.001);
} }
} /* namespace binding */ } /* namespace binding */
#include <cmath> #include <cmath>
#include <algorithm> #include <algorithm>
namespace nms namespace nms {
{ struct proposal {
struct proposal
{
float score, x1, y1, x2, y2; 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; 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; 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) { if (iw > 0) {
float ih = std::min(b.y2, a.y2) - std::max(b.y1, a.y1) + 1; float ih = std::min(b.y2, a.y2) - std::max(b.y1, a.y1) + 1;
if (ih > 0) { if (ih > 0) {
...@@ -31,21 +28,10 @@ static float iou(const proposal& a, const proposal& b) ...@@ -31,21 +28,10 @@ static float iou(const proposal& a, const proposal& b)
return overlap; return overlap;
} }
enum class Method : uint32_t enum class Method : uint32_t { LINEAR = 0, GAUSSIAN, HARD };
{
LINEAR = 0,
GAUSSIAN,
HARD
};
size_t soft_nms(float* boxes, size_t soft_nms(float* boxes, int32_t* index, size_t count, Method method,
int32_t* index, float Nt, float sigma, float threshold) {
size_t count,
Method method,
float Nt,
float sigma,
float threshold)
{
std::iota(index, index + count, 0); // np.arange() std::iota(index, index + count, 0); // np.arange()
auto p = reinterpret_cast<proposal*>(boxes); auto p = reinterpret_cast<proposal*>(boxes);
...@@ -55,7 +41,7 @@ size_t soft_nms(float* boxes, ...@@ -55,7 +41,7 @@ size_t soft_nms(float* boxes,
std::swap(p[i], *max); std::swap(p[i], *max);
std::swap(index[i], index[max - p]); std::swap(index[i], index[max - p]);
auto j = i + 1; auto j = i + 1;
auto weight = 0.f; auto weight = 0.f;
while (j < N) { while (j < N) {
auto ov = iou(p[i], p[j]); auto ov = iou(p[i], p[j]);
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册