|
10 | 10 | #include <nanobind/stl/unique_ptr.h> |
11 | 11 | #include <nanobind/stl/vector.h> |
12 | 12 |
|
13 | | -#include "models/segmentation_model.h" |
14 | 13 | #include "models/results.h" |
| 14 | +#include "models/segmentation_model.h" |
15 | 15 | #include "py_utils.hpp" |
16 | 16 |
|
17 | 17 | namespace pyutils = vision::nanobind::utils; |
@@ -40,33 +40,34 @@ void init_segmentation(nb::module_& m) { |
40 | 40 | [](SegmentationModel& self, const nb::ndarray<>& input) { |
41 | 41 | return self.infer(pyutils::wrap_np_mat(input)); |
42 | 42 | }) |
43 | | - .def("infer_batch", [](SegmentationModel& self, const std::vector<nb::ndarray<>> inputs) { |
44 | | - std::vector<ImageInputData> input_mats; |
45 | | - input_mats.reserve(inputs.size()); |
| 43 | + .def("infer_batch", |
| 44 | + [](SegmentationModel& self, const std::vector<nb::ndarray<>> inputs) { |
| 45 | + std::vector<ImageInputData> input_mats; |
| 46 | + input_mats.reserve(inputs.size()); |
46 | 47 |
|
47 | | - for (const auto& input : inputs) { |
48 | | - input_mats.push_back(pyutils::wrap_np_mat(input)); |
49 | | - } |
| 48 | + for (const auto& input : inputs) { |
| 49 | + input_mats.push_back(pyutils::wrap_np_mat(input)); |
| 50 | + } |
50 | 51 |
|
51 | | - return self.inferBatch(input_mats); |
52 | | - }) |
53 | | - .def("postprocess", [](SegmentationModel& self, InferenceResult& infResult) { |
54 | | - return self.postprocess(infResult); |
55 | | - }) |
| 52 | + return self.inferBatch(input_mats); |
| 53 | + }) |
| 54 | + .def("postprocess", |
| 55 | + [](SegmentationModel& self, InferenceResult& infResult) { |
| 56 | + return self.postprocess(infResult); |
| 57 | + }) |
56 | 58 | .def_prop_ro_static("__model__", [](nb::object) { |
57 | 59 | return SegmentationModel::ModelType; |
58 | 60 | }); |
59 | 61 |
|
60 | 62 | nb::class_<ImageResult, ResultBase>(m, "ImageResult") |
61 | | - .def(nb::init<int64_t, std::shared_ptr<MetaData>>(), |
62 | | - nb::arg("frameId") = -1, |
63 | | - nb::arg("metaData") = nullptr) |
| 63 | + .def(nb::init<int64_t, std::shared_ptr<MetaData>>(), nb::arg("frameId") = -1, nb::arg("metaData") = nullptr) |
64 | 64 | .def_prop_ro( |
65 | 65 | "resultImage", |
66 | 66 | [](ImageResult& r) { |
67 | 67 | return nb::ndarray<uint8_t, nb::numpy, nb::c_contig>(r.resultImage.data, |
68 | | - {static_cast<size_t>(r.resultImage.rows), static_cast<size_t>(r.resultImage.cols), static_cast<size_t>(r.resultImage.channels())}); |
| 68 | + {static_cast<size_t>(r.resultImage.rows), |
| 69 | + static_cast<size_t>(r.resultImage.cols), |
| 70 | + static_cast<size_t>(r.resultImage.channels())}); |
69 | 71 | }, |
70 | | - nb::rv_policy::reference_internal |
71 | | - ); |
| 72 | + nb::rv_policy::reference_internal); |
72 | 73 | } |
0 commit comments