/* Tencent is pleased to support the open source community by making ncnn available. * * Copyright (C) 2020 THL A29 Limited, a Tencent company. All rights reserved. * * Licensed under the BSD 3-Clause License (the "License"); you may not use this file except * in compliance with the License. You may obtain a copy of the License at * * https://opensource.org/licenses/BSD-3-Clause * * Unless required by applicable law or agreed to in writing, software distributed * under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR * CONDITIONS OF ANY KIND, either express or implied. See the License for the * specific language governing permissions and limitations under the License. */ #include #include #include #include #include #include #include #include #include #include #include "pybind11_mat.h" #include "pybind11_datareader.h" #include "pybind11_allocator.h" #include "pybind11_modelbin.h" #include "pybind11_layer.h" using namespace ncnn; namespace py = pybind11; struct LayerFactory { std::string name; int index; std::function creator; std::function destroyer; layer_creator_func creator_func; layer_destroyer_func destroyer_func; }; #define LayerFactoryDeclear(n) \ static ncnn::Layer* LayerCreator##n(void*); \ static void LayerDestroyer##n(ncnn::Layer*, void*); LayerFactoryDeclear(0); LayerFactoryDeclear(1); LayerFactoryDeclear(2); LayerFactoryDeclear(3); LayerFactoryDeclear(4); LayerFactoryDeclear(5); LayerFactoryDeclear(6); LayerFactoryDeclear(7); LayerFactoryDeclear(8); LayerFactoryDeclear(9); std::vector g_layer_factroys = { {"", -1, nullptr, nullptr, LayerCreator0, LayerDestroyer0}, {"", -1, nullptr, nullptr, LayerCreator1, LayerDestroyer1}, {"", -1, nullptr, nullptr, LayerCreator2, LayerDestroyer2}, {"", -1, nullptr, nullptr, LayerCreator3, LayerDestroyer3}, {"", -1, nullptr, nullptr, LayerCreator4, LayerDestroyer4}, {"", -1, nullptr, nullptr, LayerCreator5, LayerDestroyer5}, {"", -1, nullptr, nullptr, LayerCreator6, LayerDestroyer6}, {"", -1, nullptr, nullptr, LayerCreator7, LayerDestroyer7}, {"", -1, nullptr, nullptr, LayerCreator8, LayerDestroyer8}, {"", -1, nullptr, nullptr, LayerCreator9, LayerDestroyer9}, }; int g_layer_factroy_index = 0; #define LayerFactoryDefine(n) \ static ncnn::Layer* LayerCreator##n(void* p) \ { \ if (g_layer_factroys[n].creator != nullptr) \ { \ return g_layer_factroys[n].creator(); \ } \ return nullptr; \ } \ static void LayerDestroyer##n(ncnn::Layer* layer, void* p) \ { \ if (g_layer_factroys[n].destroyer) \ { \ g_layer_factroys[n].destroyer(layer); \ } \ } LayerFactoryDefine(0); LayerFactoryDefine(1); LayerFactoryDefine(2); LayerFactoryDefine(3); LayerFactoryDefine(4); LayerFactoryDefine(5); LayerFactoryDefine(6); LayerFactoryDefine(7); LayerFactoryDefine(8); LayerFactoryDefine(9); PYBIND11_MODULE(ncnn, m) { auto atexit = py::module_::import("atexit"); atexit.attr("register")(py::cpp_function([]() { for (int i = 0; i < g_layer_factroys.size(); i++) { g_layer_factroys[i].creator = nullptr; g_layer_factroys[i].destroyer = nullptr; } })); py::class_ >(m, "Allocator"); py::class_ >(m, "PoolAllocator") .def(py::init<>()) .def("set_size_compare_ratio", &PoolAllocator::set_size_compare_ratio, py::arg("src")) .def("clear", &PoolAllocator::clear) .def("fastMalloc", &PoolAllocator::fastMalloc, py::arg("size")) .def("fastFree", &PoolAllocator::fastFree, py::arg("ptr")); py::class_ >(m, "UnlockedPoolAllocator") .def(py::init<>()) .def("set_size_compare_ratio", &UnlockedPoolAllocator::set_size_compare_ratio, py::arg("src")) .def("clear", &UnlockedPoolAllocator::clear) .def("fastMalloc", &UnlockedPoolAllocator::fastMalloc, py::arg("size")) .def("fastFree", &UnlockedPoolAllocator::fastFree, py::arg("ptr")); py::class_ >(m, "DataReader") .def(py::init<>()) #if NCNN_STRING .def("scan", &DataReader::scan, py::arg("format"), py::arg("p")) #endif // NCNN_STRING .def("read", &DataReader::read, py::arg("buf"), py::arg("size")); py::class_ >(m, "DataReaderFromEmpty") .def(py::init<>()) #if NCNN_STRING .def("scan", &DataReaderFromEmpty::scan, py::arg("format"), py::arg("p")) #endif // NCNN_STRING .def("read", &DataReaderFromEmpty::read, py::arg("buf"), py::arg("size")); py::class_(m, "Blob") .def(py::init<>()) #if NCNN_STRING .def_readwrite("name", &Blob::name) #endif // NCNN_STRING .def_readwrite("producer", &Blob::producer) .def_readwrite("consumer", &Blob::consumer) .def_readwrite("shape", &Blob::shape); py::class_ >(m, "ModelBin") .def(py::init<>()) .def("load", (Mat(ModelBin::*)(int, int) const) & ModelBin::load, py::arg("w"), py::arg("type")) .def("load", (Mat(ModelBin::*)(int, int, int) const) & ModelBin::load, py::arg("w"), py::arg("h"), py::arg("type")) .def("load", (Mat(ModelBin::*)(int, int, int, int) const) & ModelBin::load, py::arg("w"), py::arg("h"), py::arg("c"), py::arg("type")) .def("load", (Mat(ModelBin::*)(int, int, int, int, int) const) & ModelBin::load, py::arg("w"), py::arg("h"), py::arg("d"), py::arg("c"), py::arg("type")); py::class_ >(m, "ModelBinFromDataReader") .def(py::init(), py::arg("dr")) .def("load", &ModelBinFromDataReader::load, py::arg("w"), py::arg("type")); py::class_ >(m, "ModelBinFromMatArray") .def(py::init(), py::arg("weights")) .def("load", &ModelBinFromMatArray::load, py::arg("w"), py::arg("type")); py::class_(m, "ParamDict") .def(py::init<>()) .def("type", &ParamDict::type, py::arg("id")) .def("get", (int (ParamDict::*)(int, int) const) & ParamDict::get, py::arg("id"), py::arg("def")) .def("get", (float (ParamDict::*)(int, float) const) & ParamDict::get, py::arg("id"), py::arg("def")) .def("get", (Mat(ParamDict::*)(int, const Mat&) const) & ParamDict::get, py::arg("id"), py::arg("def")) .def("set", (void (ParamDict::*)(int, int)) & ParamDict::set, py::arg("id"), py::arg("i")) .def("set", (void (ParamDict::*)(int, float)) & ParamDict::set, py::arg("id"), py::arg("f")) .def("set", (void (ParamDict::*)(int, const Mat&)) & ParamDict::set, py::arg("id"), py::arg("v")); py::class_