add nms c function,compile with python setuptools

ops/clean.sh

@@ -7,5 +7,6 @@ python setup.py clean --all
 cd ..
 
 cd pybind11
+python setup.py clean --all
 rm -f *.so
 

ops/make.sh

@@ -6,5 +6,7 @@ python setup.py build_ext --inplace
 cd ..
 
 cd pybind11
-include=`python -m pybind11 --includes`
-g++ -std=c++11 -shared -o box_ops_cc.so box_ops.cc -fPIC -O3 ${include}
+python setup.py build_ext --inplace
+
+# include=`python -m pybind11 --includes`
+# g++ -std=c++11 -shared -o box_ops_cc.so box_ops.cc -fPIC -O3 ${include}

ops/pybind11/compile.sh

@@ -5,4 +5,4 @@ set -e
 include=`python -m pybind11 --includes`
 
 g++ -std=c++11 -shared -o box_ops_cc.so box_ops.cc -fPIC -O3 ${include}
-
+g++ -std=c++11 -shared -o nms.so nms.cc -fPIC -O3 ${include}

ops/pybind11/nms.cc

@@ -0,0 +1,20 @@
+#include "nms_cpu.h"
+PYBIND11_MODULE(nms, m)
+{
+    m.doc() = "non_max_suppression asd";
+
+    m.def("non_max_suppression_cpu", &non_max_suppression_cpu<double>, py::return_value_policy::reference_internal, "bbox iou", 
+          "boxes"_a = 1, "order"_a = 2, "nms_overlap_thresh"_a = 3, "eps"_a = 4);
+    m.def("non_max_suppression_cpu", &non_max_suppression_cpu<float>, py::return_value_policy::reference_internal, "bbox iou", 
+          "boxes"_a = 1, "order"_a = 2, "nms_overlap_thresh"_a = 3, "eps"_a = 4);
+
+    m.def("rotate_non_max_suppression_cpu", &rotate_non_max_suppression_cpu<float>, py::return_value_policy::reference_internal, "bbox iou", 
+          "box_corners"_a = 1, "order"_a = 2, "standup_iou"_a = 3, "thresh"_a = 4);
+    m.def("rotate_non_max_suppression_cpu", &rotate_non_max_suppression_cpu<double>, py::return_value_policy::reference_internal, "bbox iou", 
+          "box_corners"_a = 1, "order"_a = 2, "standup_iou"_a = 3, "thresh"_a = 4);
+
+    m.def("rotate_non_max_suppression_3d_cpu", &rotate_non_max_suppression_3d_cpu<float>, py::return_value_policy::reference_internal, "bbox iou", 
+          "box_corners"_a = 1, "order"_a = 2, "standup_iou"_a = 3, "thresh"_a = 4);
+    m.def("rotate_non_max_suppression_3d_cpu", &rotate_non_max_suppression_3d_cpu<double>, py::return_value_policy::reference_internal, "bbox iou", 
+          "box_corners"_a = 1, "order"_a = 2, "standup_iou"_a = 3, "thresh"_a = 4);
+}

ops/pybind11/nms_cpu.h

@@ -0,0 +1,243 @@
+#ifndef NMS_CPU_H
+#define NMS_CPU_H
+#include <pybind11/pybind11.h>
+// must include pybind11/stl.h if using containers in STL in arguments.
+#include <pybind11/stl.h>
+#include <pybind11/numpy.h>
+#include <vector>
+#include <algorithm>
+#include <boost/geometry.hpp>
+
+namespace py = pybind11;
+using namespace pybind11::literals;
+
+
+template<typename DType, typename ShapeContainer>
+inline py::array_t<DType> constant(ShapeContainer shape, DType value){
+    // create ROWMAJOR array.
+    py::array_t<DType> array(shape);
+    std::fill(array.mutable_data(), array.mutable_data() + array.size(), value);
+    return array;
+}
+
+template<typename DType>
+inline py::array_t<DType> zeros(std::vector<long int> shape){
+    return constant<DType, std::vector<long int>>(shape, 0);
+}
+
+template <typename DType>
+std::vector<int> non_max_suppression_cpu(
+    py::array_t<DType> boxes,
+    py::array_t<int> order,
+    DType thresh,
+    DType eps=0)
+{
+    auto ndets = boxes.shape(0);
+    auto boxes_r = boxes.template unchecked<2>();
+    auto order_r = order.template unchecked<1>();
+    auto suppressed = zeros<int>({ndets});
+    auto suppressed_rw = suppressed.template mutable_unchecked<1>();
+    auto area = zeros<DType>({ndets});
+    auto area_rw = area.template mutable_unchecked<1>();
+    // get areas
+    for(int i = 0; i < ndets; ++i){
+        area_rw(i) = (boxes_r(i, 2) - boxes_r(i, 0) + eps) * (boxes_r(i, 3) - boxes_r(i, 1) + eps);
+    }
+    std::vector<int> keep;
+    int i, j;
+    DType xx1, xx2, w, h, inter, ovr;
+    for(int _i = 0; _i < ndets; ++_i){
+        i = order_r(_i);
+        if(suppressed_rw(i) == 1)
+            continue;
+        keep.push_back(i);
+        for(int _j = _i + 1; _j < ndets; ++_j){
+            j = order_r(_j);
+            if(suppressed_rw(j) == 1)
+                continue;
+            xx2 = std::min(boxes_r(i, 2), boxes_r(j, 2));
+            xx1 = std::max(boxes_r(i, 0), boxes_r(j, 0));
+            w = xx2 - xx1 + eps;
+            if (w > 0){
+                xx2 = std::min(boxes_r(i, 3), boxes_r(j, 3));
+                xx1 = std::max(boxes_r(i, 1), boxes_r(j, 1));
+                h = xx2 - xx1 + eps;
+                if (h > 0){
+                    inter = w * h;
+                    ovr = inter / (area_rw(i) + area_rw(j) - inter);
+                    if(ovr >= thresh)
+                        suppressed_rw(j) = 1;
+                }
+            }
+        }
+    }
+    return keep;
+}
+
+template <typename DType>
+std::vector<int> rotate_non_max_suppression_cpu(
+    py::array_t<DType> box_corners,
+    py::array_t<int> order,
+    py::array_t<DType> standup_iou,
+    DType thresh)
+{
+    auto ndets = box_corners.shape(0);
+    auto box_corners_r = box_corners.template unchecked<3>();
+    auto order_r = order.template unchecked<1>();
+    auto suppressed = zeros<int>({ndets});
+    auto suppressed_rw = suppressed.template mutable_unchecked<1>();
+    auto standup_iou_r = standup_iou.template unchecked<2>();
+    std::vector<int> keep;
+    int i, j;
+
+    namespace bg = boost::geometry;
+    typedef bg::model::point<DType, 2, bg::cs::cartesian> point_t;
+    typedef bg::model::polygon<point_t> polygon_t;
+    polygon_t poly, qpoly;
+    std::vector<polygon_t> poly_inter, poly_union;
+    DType inter_area, union_area, overlap;
+
+    for(int _i = 0; _i < ndets; ++_i){
+        i = order_r(_i);
+        if(suppressed_rw(i) == 1)
+            continue;
+        keep.push_back(i);
+        for(int _j = _i + 1; _j < ndets; ++_j){
+            j = order_r(_j);
+            if(suppressed_rw(j) == 1)
+                continue;
+            if (standup_iou_r(i, j) <= 0.0)
+                continue;
+
+            bg::append(poly, point_t(box_corners_r(i, 0, 0), box_corners_r(i, 0, 1)));
+            bg::append(poly, point_t(box_corners_r(i, 1, 0), box_corners_r(i, 1, 1)));
+            bg::append(poly, point_t(box_corners_r(i, 2, 0), box_corners_r(i, 2, 1)));
+            bg::append(poly, point_t(box_corners_r(i, 3, 0), box_corners_r(i, 3, 1)));
+            bg::append(poly, point_t(box_corners_r(i, 0, 0), box_corners_r(i, 0, 1)));
+            bg::append(qpoly, point_t(box_corners_r(j, 0, 0), box_corners_r(j, 0, 1)));
+            bg::append(qpoly, point_t(box_corners_r(j, 1, 0), box_corners_r(j, 1, 1)));
+            bg::append(qpoly, point_t(box_corners_r(j, 2, 0), box_corners_r(j, 2, 1)));
+            bg::append(qpoly, point_t(box_corners_r(j, 3, 0), box_corners_r(j, 3, 1)));
+            bg::append(qpoly, point_t(box_corners_r(j, 0, 0), box_corners_r(j, 0, 1)));
+            bg::intersection(poly, qpoly, poly_inter);
+
+            if (!poly_inter.empty())
+            {
+                inter_area = bg::area(poly_inter.front());
+                bg::union_(poly, qpoly, poly_union);
+
+                if (!poly_union.empty()){ // ignore invalid box
+                    union_area = bg::area(poly_union.front());
+                    // std::cout << "post union area" << std::endl;
+                    // std::cout << union_area << "debug" << std::endl;
+                    overlap = inter_area / union_area;
+                    if(overlap >= thresh)
+                        suppressed_rw(j) = 1;
+                    poly_union.clear();
+                }
+            }
+            poly.clear();
+            qpoly.clear();
+            poly_inter.clear();
+
+        }
+    }
+    return keep;
+}
+
+template <typename DType>
+std::vector<int> rotate_non_max_suppression_3d_cpu(
+    py::array_t<DType> box_corners,
+    py::array_t<int> order,
+    py::array_t<DType> standup_iou,
+    DType thresh)  
+ /*
+  camera coordinate
+            7 -------- 4
+           /|         /|
+          6 -------- 5 .
+          | |        | |
+          . 3 -------- 0
+          |/         |/
+          2 -------- 1
+  */
+{
+    auto ndets = box_corners.shape(0);
+    auto box_corners_r = box_corners.template unchecked<3>();
+    auto order_r = order.template unchecked<1>();
+    auto suppressed = zeros<int>({ndets});
+    auto suppressed_rw = suppressed.template mutable_unchecked<1>();
+    auto standup_iou_r = standup_iou.template unchecked<2>();
+    std::vector<int> keep;
+    int i, j;
+
+    namespace bg = boost::geometry;
+    typedef bg::model::point<DType, 2, bg::cs::cartesian> point_t;
+    typedef bg::model::polygon<point_t> polygon_t;
+    polygon_t poly, qpoly;
+    std::vector<polygon_t> poly_inter, poly_union;
+    DType inter_area, inter_vol, overlap;
+    DType ymin, ymax;
+    DType area, qarea, h, qh;
+    DType vol, qvol;
+    DType zero = 0.0;
+    for(int _i = 0; _i < ndets; ++_i){
+        i = order_r(_i);
+        if(suppressed_rw(i) == 1)
+            continue;
+        keep.push_back(i);
+        for(int _j = _i + 1; _j < ndets; ++_j){
+            j = order_r(_j);
+            if(suppressed_rw(j) == 1)
+                continue;
+            if (standup_iou_r(i, j) <= 0.0)
+                continue;
+            bg::append(poly, point_t(box_corners_r(i, 6, 0), box_corners_r(i, 6, 2)));
+            bg::append(poly, point_t(box_corners_r(i, 7, 0), box_corners_r(i, 7, 2)));
+            bg::append(poly, point_t(box_corners_r(i, 4, 0), box_corners_r(i, 4, 2)));
+            bg::append(poly, point_t(box_corners_r(i, 5, 0), box_corners_r(i, 5, 2)));
+            bg::append(poly, point_t(box_corners_r(i, 6, 0), box_corners_r(i, 6, 2)));
+            bg::append(qpoly, point_t(box_corners_r(j, 6, 0), box_corners_r(j, 6, 2)));
+            bg::append(qpoly, point_t(box_corners_r(j, 7, 0), box_corners_r(j, 7, 2)));
+            bg::append(qpoly, point_t(box_corners_r(j, 4, 0), box_corners_r(j, 4, 2)));
+            bg::append(qpoly, point_t(box_corners_r(j, 5, 0), box_corners_r(j, 5, 2)));
+            bg::append(qpoly, point_t(box_corners_r(j, 6, 0), box_corners_r(j, 6, 2)));
+            bg::intersection(poly, qpoly, poly_inter);
+
+            if (!poly_inter.empty())
+            {
+                inter_area = bg::area(poly_inter.front());
+                bg::union_(poly, qpoly, poly_union);
+                if (!poly_union.empty()){ 
+                    ymax = std::min(box_corners_r(i, 0, 1), box_corners_r(j, 0, 1));
+                    ymin = std::max(box_corners_r(i, 4, 1), box_corners_r(j, 4, 1));
+
+                    h = box_corners_r(i, 0, 1) - box_corners_r(i, 4, 1);
+                    qh = box_corners_r(j, 0, 1) - box_corners_r(j, 4, 1);
+
+                    area = bg::area(poly);
+                    qarea = bg::area(qpoly);
+
+                    inter_vol = inter_area * std::max(zero, ymax - ymin);;
+
+                    vol = std::max(zero, area * h);
+                    qvol = std::max(zero, qarea * qh);
+
+                    overlap = inter_vol / (vol + qvol - inter_vol);
+
+                    if(overlap >= thresh)
+                        suppressed_rw(j) = 1;
+                    poly_union.clear();
+                }
+            }
+            poly.clear();
+            qpoly.clear();
+            poly_inter.clear();
+
+        }
+    }
+    return keep;
+}
+
+
+#endif

ops/pybind11/rbbox_iou.py

@@ -1,10 +1,7 @@
 import numpy as np
 
 from . import box_ops_cc
-
-
-# import nms
-
+from . import nms
 
 def bbox_overlaps_1d(ex, gt):
     '''
@@ -275,6 +272,45 @@ def bev_nms_np(dets, nms_thresh, top_k=300):
     return keep
 
 
+def rotate_nms_bev_cc(dets, thresh, top_k=300):
+    '''
+    dets: [[cx, cz, l, w, ry, score]] n,6
+    '''
+    assert dets.shape[1] == 6
+    scores = dets[:, 5]
+
+    order = scores.argsort()[::-1].astype(np.int32)  # highest->lowest
+
+    boxes_corners = rbbox2corner(dets[:, :5])
+    boxes_standup = corner2standup(boxes_corners)
+
+    standup_iou = bbox_overlaps_2d(boxes_standup, boxes_standup)
+
+    keep = nms.rotate_non_max_suppression_cpu(boxes_corners, order, standup_iou, thresh)
+
+    return keep[:top_k]
+
+
+def rotate_nms_3d_cc(dets, thresh, top_k=300):
+    '''
+    dets: [[cx, cy, cz, l, w, h, ry, score]] n,8
+
+    '''
+    assert dets.shape[1] == 8
+    scores = dets[:, 7]
+
+    order = scores.argsort()[::-1].astype(np.int32)  # highest->lowest
+
+    bbox_corner_3d = boxes3d2corners(dets)  # n, 8, 3
+    bbox_standup = np.concatenate([np.min(bbox_corner_3d, 1), np.max(bbox_corner_3d, 1)], 1)  # n, 6
+
+    standup_iou = bbox_overlaps_3d(bbox_standup, bbox_standup)
+
+    keep = nms.rotate_non_max_suppression_3d_cpu(bbox_corner_3d, order, standup_iou, thresh)
+
+    return keep[:top_k]
+
+
 if __name__ == '__main__':
     # import matplotlib.pyplot as plt