Skip to content

Commit b6cf566

Browse files
author
Gin
committed
refactor ML label program
1 parent 9086b0f commit b6cf566

File tree

2 files changed

+71
-65
lines changed

2 files changed

+71
-65
lines changed

‎SerialPrograms/Source/ML/Programs/ML_LabelImages.cpp‎

Lines changed: 69 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -79,72 +79,9 @@ void DrawnBoundingBox::on_mouse_release(double, double){
7979
auto& m_program = m_widget.m_program;
8080
auto& m_overlay_set = m_widget.m_overlay_set;
8181

82-
const size_t source_width = m_program.source_image_width;
83-
const size_t source_height = m_program.source_image_height;
84-
85-
const int box_x = int(m_program.X * source_width + 0.5);
86-
const int box_y = int(m_program.Y * source_height + 0.5);
87-
const int box_width = int(m_program.WIDTH * source_width + 0.5);
88-
const int box_height = int(m_program.HEIGHT * source_height + 0.5);
89-
if (box_width == 0 || box_height == 0){
90-
return;
91-
}
92-
93-
if (m_program.m_image_embedding.size() == 0){
94-
// no embedding file loaded
95-
return;
96-
}
97-
m_program.m_sam_session.run(
98-
m_program.m_image_embedding,
99-
(int)source_height, (int)source_width, {}, {},
100-
{box_x, box_y, box_x + box_width, box_y + box_height},
101-
m_program.m_output_boolean_mask
102-
);
103-
104-
size_t min_mask_x = INT_MAX, max_mask_x = 0;
105-
size_t min_mask_y = INT_MAX, max_mask_y = 0;
106-
for (size_t y = 0; y < source_height; y++){
107-
for (size_t x = 0; x < source_width; x++){
108-
bool mask = m_program.m_output_boolean_mask[y*source_width + x];
109-
uint32_t& pixel = m_program.m_mask_image.pixel(x, y);
110-
// if the pixel's mask value is true, set a semi-transparent 45-degree blue strip color
111-
// otherwise: fully transparent (alpha = 0)
112-
uint32_t color = 0;
113-
if (mask){
114-
color = (std::abs(int(x) - int(y)) % 4 <= 1) ? combine_argb(150, 30, 144, 255) : combine_argb(150, 0, 0, 60);
115-
min_mask_x = std::min(x, min_mask_x);
116-
max_mask_x = std::max(x, max_mask_x);
117-
min_mask_y = std::min(y, min_mask_y);
118-
max_mask_y = std::max(y, max_mask_y);
119-
}
120-
pixel = color;
121-
}
122-
}
123-
if (min_mask_x < INT_MAX && max_mask_x > min_mask_x && min_mask_y < INT_MAX && max_mask_y > min_mask_y){
124-
const size_t mask_width = max_mask_x - min_mask_x + 1;
125-
const size_t mask_height = max_mask_y - min_mask_y + 1;
126-
ImageFloatBox mask_box(
127-
min_mask_x/double(source_width), min_mask_y/double(source_height),
128-
mask_width/double(source_width), mask_height/double(source_height));
129-
const std::string label = m_program.FORM_LABEL.slug();
130-
131-
132-
ObjectAnnotation annotation;
133-
annotation.user_box = ImagePixelBox(box_x, box_y, box_x + box_width + 1, box_y + box_height + 1);
134-
annotation.mask_box = ImagePixelBox(min_mask_x, min_mask_y, max_mask_x+1, max_mask_y+1);
135-
annotation.mask.resize(mask_width * mask_height);
136-
for(size_t row = 0; row < mask_height; row++){
137-
auto it = m_program.m_output_boolean_mask.begin() + (min_mask_y + row) * source_width + min_mask_x;
138-
auto it2 = annotation.mask.begin() + row * mask_width;
139-
std::copy(it, it + mask_width, it2);
140-
}
141-
142-
annotation.label = label;
143-
m_program.m_annotated_objects.emplace_back(std::move(annotation));
144-
145-
m_program.set_rendered_objects(m_overlay_set);
146-
}
82+
m_program.compute_mask(m_overlay_set);
14783
}
84+
14885
void DrawnBoundingBox::on_mouse_move(double x, double y){
14986
auto& program = m_widget.m_program;
15087
if (!m_mouse_start){
@@ -268,6 +205,73 @@ void LabelImages::set_rendered_objects(VideoOverlaySet& overlay_set){
268205
}
269206
}
270207

208+
void LabelImages::compute_mask(VideoOverlaySet& overlay_set){
209+
const size_t source_width = source_image_width;
210+
const size_t source_height = source_image_height;
211+
212+
const int box_x = int(X * source_width + 0.5);
213+
const int box_y = int(Y * source_height + 0.5);
214+
const int box_width = int(WIDTH * source_width + 0.5);
215+
const int box_height = int(HEIGHT * source_height + 0.5);
216+
if (box_width == 0 || box_height == 0){
217+
return;
218+
}
219+
220+
if (m_image_embedding.size() == 0){
221+
// no embedding file loaded
222+
return;
223+
}
224+
m_sam_session.run(
225+
m_image_embedding,
226+
(int)source_height, (int)source_width, {}, {},
227+
{box_x, box_y, box_x + box_width, box_y + box_height},
228+
m_output_boolean_mask
229+
);
230+
231+
size_t min_mask_x = INT_MAX, max_mask_x = 0;
232+
size_t min_mask_y = INT_MAX, max_mask_y = 0;
233+
for (size_t y = 0; y < source_height; y++){
234+
for (size_t x = 0; x < source_width; x++){
235+
bool mask = m_output_boolean_mask[y*source_width + x];
236+
uint32_t& pixel = m_mask_image.pixel(x, y);
237+
// if the pixel's mask value is true, set a semi-transparent 45-degree blue strip color
238+
// otherwise: fully transparent (alpha = 0)
239+
uint32_t color = 0;
240+
if (mask){
241+
color = (std::abs(int(x) - int(y)) % 4 <= 1) ? combine_argb(150, 30, 144, 255) : combine_argb(150, 0, 0, 60);
242+
min_mask_x = std::min(x, min_mask_x);
243+
max_mask_x = std::max(x, max_mask_x);
244+
min_mask_y = std::min(y, min_mask_y);
245+
max_mask_y = std::max(y, max_mask_y);
246+
}
247+
pixel = color;
248+
}
249+
}
250+
if (min_mask_x < INT_MAX && max_mask_x > min_mask_x && min_mask_y < INT_MAX && max_mask_y > min_mask_y){
251+
const size_t mask_width = max_mask_x - min_mask_x + 1;
252+
const size_t mask_height = max_mask_y - min_mask_y + 1;
253+
ImageFloatBox mask_box(
254+
min_mask_x/double(source_width), min_mask_y/double(source_height),
255+
mask_width/double(source_width), mask_height/double(source_height));
256+
const std::string label = FORM_LABEL.slug();
257+
258+
259+
ObjectAnnotation annotation;
260+
annotation.user_box = ImagePixelBox(box_x, box_y, box_x + box_width + 1, box_y + box_height + 1);
261+
annotation.mask_box = ImagePixelBox(min_mask_x, min_mask_y, max_mask_x+1, max_mask_y+1);
262+
annotation.mask.resize(mask_width * mask_height);
263+
for(size_t row = 0; row < mask_height; row++){
264+
auto it = m_output_boolean_mask.begin() + (min_mask_y + row) * source_width + min_mask_x;
265+
auto it2 = annotation.mask.begin() + row * mask_width;
266+
std::copy(it, it + mask_width, it2);
267+
}
268+
269+
annotation.label = label;
270+
m_annotated_objects.emplace_back(std::move(annotation));
271+
272+
set_rendered_objects(overlay_set);
273+
}
274+
}
271275

272276
LabelImages_Widget::~LabelImages_Widget(){
273277
m_program.FORM_LABEL.remove_listener(*this);

‎SerialPrograms/Source/ML/Programs/ML_LabelImages.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,8 @@ class LabelImages : public PanelInstance{
7272

7373
void set_rendered_objects(VideoOverlaySet& overlayset);
7474

75+
void compute_mask(VideoOverlaySet& overlay_set);
76+
7577
private:
7678
friend class LabelImages_Widget;
7779
friend class DrawnBoundingBox;

0 commit comments

Comments
 (0)