Merge pull request #2 from PaddlePaddle/develop

mergepaddleocr

ppocr/postprocess/lanms/__init__.py

+import subprocess
+import os
+import numpy as np
+
+BASE_DIR = os.path.dirname(os.path.realpath(__file__))
+
+if subprocess.call(['make', '-C', BASE_DIR]) != 0:  # return value
+    raise RuntimeError('Cannot compile lanms: {}'.format(BASE_DIR))
+
+
+def merge_quadrangle_n9(polys, thres=0.3, precision=10000):
+    from .adaptor import merge_quadrangle_n9 as nms_impl
+    if len(polys) == 0:
+        return np.array([], dtype='float32')
+    p = polys.copy()
+    p[:,:8] *= precision
+    ret = np.array(nms_impl(p, thres), dtype='float32')
+    ret[:,:8] /= precision
+    return ret
+

ppocr/postprocess/lanms/__main__.py

+import numpy as np
+
+
+from . import merge_quadrangle_n9
+
+if __name__ == '__main__':
+    # unit square with confidence 1
+    q = np.array([0, 0, 0, 1, 1, 1, 1, 0, 1], dtype='float32')
+
+    print(merge_quadrangle_n9(np.array([q, q + 0.1, q + 2])))

ppocr/postprocess/lanms/adaptor.cpp

+#include "pybind11/pybind11.h"
+#include "pybind11/numpy.h"
+#include "pybind11/stl.h"
+#include "pybind11/stl_bind.h"
+
+#include "lanms.h"
+
+namespace py = pybind11;
+
+
+namespace lanms_adaptor {
+
+	std::vector<std::vector<float>> polys2floats(const std::vector<lanms::Polygon> &polys) {
+		std::vector<std::vector<float>> ret;
+		for (size_t i = 0; i < polys.size(); i ++) {
+			auto &p = polys[i];
+			auto &poly = p.poly;
+			ret.emplace_back(std::vector<float>{
+					float(poly[0].X), float(poly[0].Y),
+					float(poly[1].X), float(poly[1].Y),
+					float(poly[2].X), float(poly[2].Y),
+					float(poly[3].X), float(poly[3].Y),
+					float(p.score),
+					});
+		}
+
+		return ret;
+	}
+
+
+	/**
+	 *
+	 * \param quad_n9 an n-by-9 numpy array, where first 8 numbers denote the
+	 *		quadrangle, and the last one is the score
+	 * \param iou_threshold two quadrangles with iou score above this threshold
+	 *		will be merged
+	 *
+	 * \return an n-by-9 numpy array, the merged quadrangles
+	 */
+	std::vector<std::vector<float>> merge_quadrangle_n9(
+			py::array_t<float, py::array::c_style | py::array::forcecast> quad_n9,
+			float iou_threshold) {
+		auto pbuf = quad_n9.request();
+		if (pbuf.ndim != 2 || pbuf.shape[1] != 9)
+			throw std::runtime_error("quadrangles must have a shape of (n, 9)");
+		auto n = pbuf.shape[0];
+		auto ptr = static_cast<float *>(pbuf.ptr);
+		return polys2floats(lanms::merge_quadrangle_n9(ptr, n, iou_threshold));
+	}
+
+}
+
+PYBIND11_PLUGIN(adaptor) {
+	py::module m("adaptor", "NMS");
+
+	m.def("merge_quadrangle_n9", &lanms_adaptor::merge_quadrangle_n9,
+			"merge quadrangels");
+
+	return m.ptr();
+}
+