| 35 | } // namespace lambdas |
| 36 | |
| 37 | void bindUff(py::module& m) |
| 38 | { |
| 39 | py::enum_<UffInputOrder>(m, "UffInputOrder", UffInputOrderDoc::descr, py::module_local()) |
| 40 | .value("NCHW", UffInputOrder::kNCHW) |
| 41 | .value("NHWC", UffInputOrder::kNHWC) |
| 42 | .value("NC", UffInputOrder::kNC); |
| 43 | |
| 44 | py::enum_<FieldType>(m, "FieldType", FieldTypeDoc::descr, py::module_local()) |
| 45 | .value("FLOAT", FieldType::kFLOAT) |
| 46 | .value("INT32", FieldType::kINT32) |
| 47 | .value("CHAR", FieldType::kCHAR) |
| 48 | .value("DIMS", FieldType::kDIMS) |
| 49 | .value("DATATYPE", FieldType::kDATATYPE) |
| 50 | .value("UNKNOWN", FieldType::kUNKNOWN); |
| 51 | |
| 52 | py::class_<FieldMap>(m, "FieldMap", FieldMapDoc::descr, py::module_local()) |
| 53 | .def(py::init<const char*, const void*, const FieldType, int>(), "name"_a, "data"_a, "type"_a, "length"_a = 1) |
| 54 | .def_readwrite("name", &FieldMap::name) |
| 55 | .def_readwrite("data", &FieldMap::data) |
| 56 | .def_readwrite("type", &FieldMap::type) |
| 57 | .def_readwrite("length", &FieldMap::length); |
| 58 | |
| 59 | py::class_<FieldCollection>(m, "FieldCollection", FieldCollectionDoc::descr, py::module_local()) |
| 60 | .def_readwrite("num_fields", &FieldCollection::nbFields) |
| 61 | .def_readwrite("fields", &FieldCollection::fields); |
| 62 | |
| 63 | py::class_<IUffParser, std::unique_ptr<IUffParser, py::nodelete>>( |
| 64 | m, "UffParser", UffParserDoc::descr, py::module_local()) |
| 65 | .def(py::init(&createUffParser)) |
| 66 | .def_property_readonly("uff_required_version_major", &IUffParser::getUffRequiredVersionMajor) |
| 67 | .def_property_readonly("uff_required_version_minor", &IUffParser::getUffRequiredVersionMinor) |
| 68 | .def_property_readonly("uff_required_version_patch", &IUffParser::getUffRequiredVersionPatch) |
| 69 | .def_property( |
| 70 | "plugin_namespace", nullptr, py::cpp_function(&IUffParser::setPluginNamespace, py::keep_alive<1, 2>{})) |
| 71 | .def("register_input", &IUffParser::registerInput, "name"_a, "shape"_a, "order"_a = UffInputOrder::kNCHW, |
| 72 | UffParserDoc::register_input) |
| 73 | .def("register_output", &IUffParser::registerOutput, "name"_a, UffParserDoc::register_output) |
| 74 | .def("parse", &IUffParser::parse, "file"_a, "network"_a, "weights_type"_a = nvinfer1::DataType::kFLOAT, |
| 75 | UffParserDoc::parse, py::keep_alive<3, 1>{}) |
| 76 | .def("parse_buffer", lambdas::uff_parse_buffer, "buffer"_a, "network"_a, |
| 77 | "weights_type"_a = nvinfer1::DataType::kFLOAT, UffParserDoc::parse_buffer, py::keep_alive<3, 1>{}) |
| 78 | .def_property("error_recorder", &IUffParser::getErrorRecorder, |
| 79 | py::cpp_function(&IUffParser::setErrorRecorder, py::keep_alive<1, 2>{})) |
| 80 | .def("__del__", &utils::doNothingDel<IUffParser>); |
| 81 | } |
| 82 | } // namespace tensorrt |