Skip to content

Commit e21a60a

Browse files
jcchrjcchr
committed
instance segmentation interface
1 parent f6126ba commit e21a60a

File tree

1 file changed

+163
-0
lines changed

1 file changed

+163
-0
lines changed

‎src/cpp/py_bindings/py_instance_segmentation.cpp‎

Lines changed: 163 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,163 @@
1+
/*
2+
* Copyright (C) 2025 Intel Corporation
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
#include <nanobind/ndarray.h>
7+
#include <nanobind/operators.h>
8+
#include <nanobind/stl/map.h>
9+
#include <nanobind/stl/string.h>
10+
#include <nanobind/stl/unique_ptr.h>
11+
#include <nanobind/stl/vector.h>
12+
13+
#include "models/instance_segmentation.h"
14+
#include "models/results.h"
15+
#include "py_utils.hpp"
16+
17+
namespace pyutils = vision::nanobind::utils;
18+
19+
using ScoresOutput = nb::ndarray<float, nb::numpy, nb::c_contig>;
20+
using LabelsOutput = nb::ndarray<size_t, nb::numpy, nb::c_contig>;
21+
22+
void init_instance_segmentation(nb::module_& m) {
23+
nb::class_<MaskRCNNModel, ImageModel>(m, "MaskRCNNModel")
24+
.def_static(
25+
"create_model",
26+
[](const std::string& model_path,
27+
const std::map<std::string, nb::object>& configuration,
28+
bool preload,
29+
const std::string& device) {
30+
auto ov_any_config = ov::AnyMap();
31+
for (const auto& item : configuration) {
32+
ov_any_config[item.first] = pyutils::py_object_to_any(item.second, item.first);
33+
}
34+
35+
36+
return MaskRCNNModel::create_model(model_path, ov_any_config, preload, device);
37+
},
38+
nb::arg("model_path"),
39+
nb::arg("configuration") = ov::AnyMap({}),
40+
nb::arg("preload") = true,
41+
nb::arg("device") = "AUTO")
42+
43+
.def("__call__",
44+
[](MaskRCNNModel& self, const nb::ndarray<>& input) {
45+
return self.infer(pyutils::wrap_np_mat(input));
46+
})
47+
.def("infer_batch",
48+
[](MaskRCNNModel& self, const std::vector<nb::ndarray<>> inputs) {
49+
std::vector<ImageInputData> input_mats;
50+
input_mats.reserve(inputs.size());
51+
52+
for (const auto& input : inputs) {
53+
input_mats.push_back(pyutils::wrap_np_mat(input));
54+
}
55+
56+
return self.inferBatch(input_mats);
57+
})
58+
.def_prop_ro_static("__model__", [](nb::object) {
59+
return MaskRCNNModel::ModelType;
60+
});
61+
62+
nb::class_<InstanceSegmentationResult, ResultBase>(m, "InstanceSegmentationResult")
63+
.def(nb::init<int64_t, std::shared_ptr<MetaData>>(), nb::arg("frameId") = -1, nb::arg("metaData") = nullptr)
64+
.def_prop_ro(
65+
"feature_vector",
66+
[](InstanceSegmentationResult& r) {
67+
if (!r.feature_vector) {
68+
return nb::ndarray<float, nb::numpy, nb::c_contig>();
69+
}
70+
71+
return nb::ndarray<float, nb::numpy, nb::c_contig>(r.feature_vector.data(),
72+
r.feature_vector.get_shape().size(),
73+
r.feature_vector.get_shape().data());
74+
},
75+
nb::rv_policy::reference_internal)
76+
.def_prop_ro("label_names",
77+
[](InstanceSegmentationResult& r) {
78+
size_t labels_count = static_cast<size_t>(r.segmentedObjects.size());
79+
std::vector<std::string> labels(labels_count);
80+
81+
for (size_t i = 0; i < labels_count; ++i) {
82+
labels[i] = r.segmentedObjects[i].label;
83+
}
84+
85+
return labels;
86+
})
87+
.def_prop_ro("labels",
88+
[](InstanceSegmentationResult& r) {
89+
size_t labels_count = static_cast<size_t>(r.segmentedObjects.size());
90+
std::vector<size_t> labels(labels_count);
91+
92+
for (size_t i = 0; i < labels_count; ++i) {
93+
labels[i] = r.segmentedObjects[i].labelID;
94+
}
95+
96+
return LabelsOutput(labels.data(), {labels_count}).cast();
97+
})
98+
.def_prop_ro("scores",
99+
[](InstanceSegmentationResult& r) {
100+
size_t scores_count = static_cast<size_t>(r.segmentedObjects.size());
101+
std::vector<float> scores(scores_count);
102+
103+
for (size_t i = 0; i < scores_count; ++i) {
104+
scores[i] = r.segmentedObjects[i].confidence;
105+
}
106+
107+
return ScoresOutput(scores.data(), {scores_count}).cast();
108+
})
109+
.def_prop_ro("bboxes",
110+
[](InstanceSegmentationResult& r) {
111+
size_t boxes_count = static_cast<size_t>(r.segmentedObjects.size());
112+
std::vector<std::vector<int>> boxes(boxes_count);
113+
114+
for (size_t i = 0; i < boxes_count; ++i) {
115+
std::vector<int> box(4);
116+
box[0] = r.segmentedObjects[i].tl().x;
117+
box[1] = r.segmentedObjects[i].tl().y;
118+
box[2] = r.segmentedObjects[i].br().x;
119+
box[3] = r.segmentedObjects[i].br().y;
120+
boxes[i] = box;
121+
}
122+
123+
return boxes;
124+
})
125+
.def_prop_ro("masks",
126+
[](InstanceSegmentationResult& r) {
127+
size_t elements_count = static_cast<size_t>(r.segmentedObjects.size());
128+
std::vector<std::vector<std::vector<int>>> masks(elements_count);
129+
130+
for (size_t i = 0; i < elements_count; ++i) {
131+
int rows = r.segmentedObjects[i].mask.rows;
132+
int cols = r.segmentedObjects[i].mask.cols;
133+
134+
std::vector<std::vector<int>> mask(rows, std::vector<int>(cols));
135+
136+
for (int row = 0; row < rows; ++row) {
137+
for (int col = 0; col < cols; ++col) {
138+
mask[row][col] = r.segmentedObjects[i].mask.at<uint8_t>(row, col);
139+
}
140+
}
141+
142+
masks[i] = mask;
143+
}
144+
145+
return masks;
146+
})
147+
.def_prop_ro(
148+
"saliency_map",
149+
[](InstanceSegmentationResult& r) {
150+
if (r.saliency_map.empty()) {
151+
return nb::ndarray<uint8_t, nb::numpy, nb::c_contig>();
152+
}
153+
int rows = r.saliency_map[0].rows;
154+
int cols = r.saliency_map[0].cols;
155+
int num_matrices = r.saliency_map.size();
156+
157+
return nb::ndarray<uint8_t, nb::numpy, nb::c_contig>(&r.saliency_map,
158+
{static_cast<size_t>(num_matrices),
159+
static_cast<size_t>(rows),
160+
static_cast<size_t>(cols)});
161+
},
162+
nb::rv_policy::reference_internal);
163+
}

0 commit comments

Comments
 (0)