diff --git a/python/src/main.cpp b/python/src/main.cpp index a7ed0528c6a..e5b1264264c 100644 --- a/python/src/main.cpp +++ b/python/src/main.cpp @@ -34,6 +34,20 @@ using namespace ncnn; namespace py = pybind11; +class DataReaderFromMemoryCopy : public DataReaderFromMemory +{ +public: + explicit DataReaderFromMemoryCopy(const unsigned char*& mem) + : DataReaderFromMemory(mem) + { + } + + virtual size_t reference(size_t size, const void** buf) const + { + return 0; + } +}; + struct LayerFactory { std::string name; @@ -956,6 +970,13 @@ PYBIND11_MODULE(ncnn, m) #endif // NCNN_STRING .def("load_param_bin", (int (Net::*)(const char*)) & Net::load_param_bin, py::arg("protopath")) .def("load_model", (int (Net::*)(const char*)) & Net::load_model, py::arg("modelpath")) + .def( + "load_model_mem", [](Net& net, const char* mem) { + const unsigned char* _mem = (const unsigned char*)mem; + DataReaderFromMemoryCopy dr(_mem); + net.load_model(dr); + }, + py::arg("mem")) #endif // NCNN_STDIO .def("clear", &Net::clear) diff --git a/python/tests/test_net.py b/python/tests/test_net.py index 03271aff462..362cc4791fb 100644 --- a/python/tests/test_net.py +++ b/python/tests/test_net.py @@ -42,6 +42,32 @@ def test_net(): assert len(net.blobs()) == 0 and len(net.layers()) == 0 +def test_net_mem(): + modelbin = bytearray(303940) + modelbin[0:4] = 71,107,48,1 + modelbin[180:184] = 71,107,48,1 + + with ncnn.Net() as net: + ret = net.load_param("tests/test.param") + net.load_model_mem(bytes(modelbin)) + assert ret == 0 and len(net.blobs()) == 3 and len(net.layers()) == 3 + + input_names = net.input_names() + output_names = net.output_names() + assert len(input_names) > 0 and len(output_names) > 0 + + in_mat = ncnn.Mat((227, 227, 3)) + + with net.create_extractor() as ex: + ex.input("data", in_mat) + ret, out_mat = ex.extract("output") + + assert ret == 0 and out_mat.dims == 1 and out_mat.w == 1 + + net.clear() + assert len(net.blobs()) == 0 and len(net.layers()) == 0 + + def test_net_vulkan(): if not hasattr(ncnn, "get_gpu_count"): return