|
| 1 | +#include <pybind11/pybind11.h> |
| 2 | +#include <pybind11/stl.h> |
| 3 | +#include <pybind11/numpy.h> |
| 4 | + |
| 5 | +#include <pdal/pdal_config.hpp> |
| 6 | +#include <pdal/StageFactory.hpp> |
| 7 | + |
| 8 | +#include "PyArray.hpp" |
| 9 | +#include "PyDimension.hpp" |
| 10 | +#include "PyPipeline.hpp" |
| 11 | +#include "StreamableExecutor.hpp" |
| 12 | + |
| 13 | +namespace py = pybind11; |
| 14 | + |
| 15 | +namespace pdal { |
| 16 | + using namespace py::literals; |
| 17 | + |
| 18 | + py::object getInfo() { |
| 19 | + return py::module_::import("types").attr("SimpleNamespace")( |
| 20 | + "version"_a = pdal::Config::versionString(), |
| 21 | + "major"_a = pdal::Config::versionMajor(), |
| 22 | + "minor"_a = pdal::Config::versionMinor(), |
| 23 | + "patch"_a = pdal::Config::versionPatch(), |
| 24 | + "debug"_a = pdal::Config::debugInformation(), |
| 25 | + "sha1"_a = pdal::Config::sha1(), |
| 26 | + "plugin"_a = pdal::Config::pluginInstallPath() |
| 27 | + ); |
| 28 | + }; |
| 29 | + |
| 30 | + std::vector<py::dict> getDimensions() { |
| 31 | + py::object np = py::module_::import("numpy"); |
| 32 | + py::object dtype = np.attr("dtype"); |
| 33 | + std::vector<py::dict> dims; |
| 34 | + for (const auto& dim: getValidDimensions()) |
| 35 | + { |
| 36 | + py::dict d( |
| 37 | + "name"_a=dim.name, |
| 38 | + "description"_a=dim.description, |
| 39 | + "dtype"_a=dtype(dim.type + std::to_string(dim.size)) |
| 40 | + ); |
| 41 | + dims.push_back(std::move(d)); |
| 42 | + } |
| 43 | + return dims; |
| 44 | + }; |
| 45 | + |
| 46 | + using pdal::python::PipelineExecutor; |
| 47 | + using pdal::python::StreamableExecutor; |
| 48 | + |
| 49 | + class PipelineIterator : public StreamableExecutor { |
| 50 | + public: |
| 51 | + using StreamableExecutor::StreamableExecutor; |
| 52 | + |
| 53 | + py::object getSchema() { |
| 54 | + return py::module_::import("json").attr("loads")(StreamableExecutor::getSchema()); |
| 55 | + } |
| 56 | + |
| 57 | + py::array executeNext() { |
| 58 | + PyArrayObject* arr(StreamableExecutor::executeNext()); |
| 59 | + if (!arr) |
| 60 | + throw py::stop_iteration(); |
| 61 | + |
| 62 | + return py::reinterpret_steal<py::array>((PyObject*)arr); |
| 63 | + } |
| 64 | + |
| 65 | + }; |
| 66 | + |
| 67 | + class Pipeline { |
| 68 | + public: |
| 69 | + point_count_t execute() { return getExecutor()->execute(); } |
| 70 | + |
| 71 | + point_count_t executeStream(point_count_t streamLimit) { |
| 72 | + return getExecutor()->executeStream(streamLimit); |
| 73 | + } |
| 74 | + |
| 75 | + std::unique_ptr<PipelineIterator> iterator(int chunk_size, int prefetch) { |
| 76 | + return std::unique_ptr<PipelineIterator>( new PipelineIterator( |
| 77 | + getJson(), _inputs, _loglevel, chunk_size, prefetch |
| 78 | + )); |
| 79 | + } |
| 80 | + |
| 81 | + void setInputs(std::vector<py::array> ndarrays) { |
| 82 | + _inputs.clear(); |
| 83 | + for (const auto& ndarray: ndarrays) { |
| 84 | + PyArrayObject* ndarray_ptr = (PyArrayObject*)ndarray.ptr(); |
| 85 | + _inputs. push_back(std::make_shared<pdal::python::Array>(ndarray_ptr)); |
| 86 | + } |
| 87 | + delExecutor(); |
| 88 | + } |
| 89 | + |
| 90 | + int getLoglevel() { return _loglevel; } |
| 91 | + |
| 92 | + void setLogLevel(int level) { _loglevel = level; delExecutor(); } |
| 93 | + |
| 94 | + std::string getLog() { return getExecutor()->getLog(); } |
| 95 | + |
| 96 | + std::string getPipeline() { return getExecutor()->getPipeline(); } |
| 97 | + |
| 98 | + std::string getMetadata() { return getExecutor()->getMetadata(); } |
| 99 | + |
| 100 | + py::object getSchema() { |
| 101 | + return py::module_::import("json").attr("loads")(getExecutor()->getSchema()); |
| 102 | + } |
| 103 | + |
| 104 | + std::vector<py::array> getArrays() { |
| 105 | + std::vector<py::array> output; |
| 106 | + for (const auto &view: getExecutor()->views()) { |
| 107 | + PyArrayObject* arr(pdal::python::viewToNumpyArray(view)); |
| 108 | + output. push_back(py::reinterpret_steal<py::array>((PyObject*)arr)); |
| 109 | + } |
| 110 | + return output; |
| 111 | + } |
| 112 | + |
| 113 | + std::vector<py::array> getMeshes() { |
| 114 | + std::vector<py::array> output; |
| 115 | + for (const auto &view: getExecutor()->views()) { |
| 116 | + PyArrayObject* arr(pdal::python::meshToNumpyArray(view->mesh())); |
| 117 | + output. push_back(py::reinterpret_steal<py::array>((PyObject*)arr)); |
| 118 | + } |
| 119 | + return output; |
| 120 | + } |
| 121 | + |
| 122 | + std::string getJson() const { |
| 123 | + PYBIND11_OVERRIDE_PURE_NAME(std::string, Pipeline, "_get_json", getJson); |
| 124 | + } |
| 125 | + |
| 126 | + bool hasInputs() { return !_inputs.empty(); } |
| 127 | + |
| 128 | + void copyInputs(const Pipeline& other) { _inputs = other._inputs; } |
| 129 | + |
| 130 | + void delExecutor() { _executor.reset(); } |
| 131 | + |
| 132 | + PipelineExecutor* getExecutor() { |
| 133 | + if (!_executor) |
| 134 | + _executor.reset(new PipelineExecutor(getJson(), _inputs, _loglevel)); |
| 135 | + return _executor.get(); |
| 136 | + } |
| 137 | + |
| 138 | + private: |
| 139 | + std::unique_ptr<PipelineExecutor> _executor; |
| 140 | + std::vector<std::shared_ptr<pdal::python::Array>> _inputs; |
| 141 | + int _loglevel; |
| 142 | + }; |
| 143 | + |
| 144 | + PYBIND11_MODULE(libpdalpython, m) |
| 145 | + { |
| 146 | + py::class_<PipelineIterator>(m, "PipelineIterator") |
| 147 | + .def("__iter__", [](PipelineIterator &it) -> PipelineIterator& { return it; }) |
| 148 | + .def("__next__", &PipelineIterator::executeNext) |
| 149 | + .def_property_readonly("log", &PipelineIterator::getLog) |
| 150 | + .def_property_readonly("schema", &PipelineIterator::getSchema) |
| 151 | + .def_property_readonly("pipeline", &PipelineIterator::getPipeline) |
| 152 | + .def_property_readonly("metadata", &PipelineIterator::getMetadata); |
| 153 | + |
| 154 | + py::class_<Pipeline>(m, "Pipeline") |
| 155 | + .def(py::init<>()) |
| 156 | + .def("execute", &Pipeline::execute) |
| 157 | + .def("execute_streaming", &Pipeline::executeStream, "chunk_size"_a=10000) |
| 158 | + .def("iterator", &Pipeline::iterator, "chunk_size"_a=10000, "prefetch"_a=0) |
| 159 | + .def_property("inputs", nullptr, &Pipeline::setInputs) |
| 160 | + .def_property("loglevel", &Pipeline::getLoglevel, &Pipeline::setLogLevel) |
| 161 | + .def_property_readonly("log", &Pipeline::getLog) |
| 162 | + .def_property_readonly("schema", &Pipeline::getSchema) |
| 163 | + .def_property_readonly("pipeline", &Pipeline::getPipeline) |
| 164 | + .def_property_readonly("metadata", &Pipeline::getMetadata) |
| 165 | + .def_property_readonly("arrays", &Pipeline::getArrays) |
| 166 | + .def_property_readonly("meshes", &Pipeline::getMeshes) |
| 167 | + .def_property_readonly("_has_inputs", &Pipeline::hasInputs) |
| 168 | + .def("_copy_inputs", &Pipeline::copyInputs) |
| 169 | + .def("_get_json", &Pipeline::getJson) |
| 170 | + .def("_del_executor", &Pipeline::delExecutor); |
| 171 | + m.def("getInfo", &getInfo); |
| 172 | + m.def("getDimensions", &getDimensions); |
| 173 | + m.def("infer_reader_driver", &StageFactory::inferReaderDriver); |
| 174 | + m.def("infer_writer_driver", &StageFactory::inferWriterDriver); |
| 175 | + }; |
| 176 | + |
| 177 | +}; // namespace pdal |
0 commit comments