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();
}