Skip to main content

rvlib/tools/bbox/
core.rs

1use crate::{
2    GeoFig, Polygon, annotations_accessor_mut,
3    drawme::{Annotation, BboxAnnotation, Stroke},
4    events::{Events, KeyCode},
5    history::{History, Record},
6    instance_annotations_accessor, make_tool_transform,
7    result::trace_ok_err,
8    tools::{
9        BBOX_NAME, Manipulate,
10        core::{
11            Mover, check_autopaste, check_erase_mode, check_recolorboxes,
12            check_trigger_history_update, check_trigger_redraw, deselect_all,
13            instance_label_display_sort, map_released_key,
14        },
15        instance_anno_shared::{check_cocoimport, get_rot90_data, predictive_labeling},
16    },
17    tools_data::{
18        LabelInfo, OUTLINE_THICKNESS_CONVERSION, annotations::BboxAnnotations, bbox_data,
19        vis_from_lfoption,
20    },
21    tools_data_accessors, tools_data_accessors_objects,
22    util::Visibility,
23    world::World,
24    world_annotations_accessor,
25};
26use rvimage_domain::{BbF, Circle, PtF, TPtF, shape_unscaled};
27use std::{iter, mem, sync::mpsc::Receiver, time::Instant};
28
29use super::on_events::{
30    KeyReleasedParams, MouseHeldLeftParams, MouseReleaseParams, PrevPos, change_annos_bbox,
31    closest_corner, export_if_triggered, find_close_vertex, import_coco, move_corner_tol,
32    on_key_released, on_mouse_held_left, on_mouse_held_right, on_mouse_released_left,
33    on_mouse_released_right,
34};
35pub const ACTOR_NAME: &str = "Bbox";
36const MISSING_ANNO_MSG: &str = "bbox annotations have not yet been initialized";
37const MISSING_DATA_MSG: &str = "bbox tools data not available";
38annotations_accessor_mut!(ACTOR_NAME, bbox_mut, MISSING_ANNO_MSG, BboxAnnotations);
39world_annotations_accessor!(ACTOR_NAME, bbox, MISSING_ANNO_MSG, BboxAnnotations);
40instance_annotations_accessor!(GeoFig);
41tools_data_accessors!(
42    ACTOR_NAME,
43    MISSING_DATA_MSG,
44    bbox_data,
45    BboxToolData,
46    bbox,
47    bbox_mut
48);
49tools_data_accessors_objects!(
50    ACTOR_NAME,
51    MISSING_DATA_MSG,
52    bbox_data,
53    BboxSpecificData,
54    bbox,
55    bbox_mut
56);
57
58pub(super) fn current_cat_idx(world: &World) -> Option<usize> {
59    get_specific(world).map(|d| d.label_info.cat_idx_current)
60}
61
62fn check_cocoexport(mut world: World) -> World {
63    // export label file if demanded
64    let bbox_data = get_specific(&world);
65    if let Some(bbox_data) = bbox_data {
66        let rot90_data = get_rot90_data(&world);
67        export_if_triggered(&world.data.meta_data, bbox_data, rot90_data);
68        if let Some(o) = get_options_mut(&mut world) {
69            o.core.import_export_trigger.untrigger_export();
70        }
71    }
72    world
73}
74
75fn show_grab_ball(
76    mp: Option<PtF>,
77    prev_pos: &PrevPos,
78    world: &mut World,
79    last_proximal_circle_check: Option<Instant>,
80    options: Option<&bbox_data::Options>,
81) -> Instant {
82    if last_proximal_circle_check.map(|lc| lc.elapsed().as_millis()) > Some(2)
83        && let Some(mp) = mp
84    {
85        if prev_pos.prev_pos.is_empty() {
86            let label_info = get_label_info(world);
87            let geos = get_annos_if_some(world).map(|a| {
88                a.iter()
89                    .enumerate()
90                    .filter(|(elt_idx, _)| {
91                        let cur = label_info.map(|li| li.cat_idx_current);
92                        let show_only_current = label_info.map(|li| li.show_only_current);
93                        a.is_of_current_label(*elt_idx, cur, show_only_current)
94                    })
95                    .map(|(elt_idx, (geo, _, _))| (elt_idx, geo))
96            });
97            if let Some((bb_idx, c_idx)) = geos.and_then(|geos| {
98                let unscaled = shape_unscaled(world.zoom_box(), world.shape_orig());
99                let tolerance = move_corner_tol(unscaled);
100                find_close_vertex(mp, geos, tolerance)
101            }) {
102                let annos = get_annos(world);
103                let corner_point = annos.and_then(|a| a.elts().get(bb_idx).map(|a| a.point(c_idx)));
104                let data = get_specific_mut(world);
105                if let (Some(data), Some(corner_point), Some(options)) =
106                    (data, corner_point, options)
107                {
108                    data.highlight_circles = vec![Circle {
109                        center: corner_point,
110                        radius: TPtF::from(options.outline_thickness)
111                            / OUTLINE_THICKNESS_CONVERSION
112                            * 2.5,
113                    }];
114                    let vis = get_visible(world);
115                    world.request_redraw_annotations(BBOX_NAME, vis);
116                }
117            } else {
118                let data = get_specific_mut(world);
119                let n_circles = data.as_ref().map_or(0, |d| d.highlight_circles.len());
120                if let Some(data) = data {
121                    data.highlight_circles = vec![];
122                }
123                if n_circles > 0 {
124                    let vis = get_visible(world);
125                    world.request_redraw_annotations(BBOX_NAME, vis);
126                }
127            }
128        } else {
129            let (c_idx, c_dist) = closest_corner(mp, prev_pos.prev_pos.iter().copied());
130            let unscaled = shape_unscaled(world.zoom_box(), world.shape_orig());
131            let tolerance = move_corner_tol(unscaled);
132            if c_dist < tolerance
133                && let Some(center) = prev_pos.prev_pos.get(c_idx)
134            {
135                let data = get_specific_mut(world);
136                if let (Some(data), Some(options)) = (data, options) {
137                    data.highlight_circles = vec![Circle {
138                        center: *center,
139                        radius: TPtF::from(options.outline_thickness)
140                            / OUTLINE_THICKNESS_CONVERSION
141                            * 3.5,
142                    }];
143                    let vis = get_visible(world);
144                    world.request_redraw_annotations(BBOX_NAME, vis);
145                }
146            } else {
147                let data = get_specific_mut(world);
148                if let Some(data) = data {
149                    data.highlight_circles = vec![];
150                }
151                let vis = get_visible(world);
152                world.request_redraw_annotations(BBOX_NAME, vis);
153            }
154        }
155    }
156    Instant::now()
157}
158
159#[derive(Debug)]
160pub struct Bbox {
161    prev_pos: PrevPos,
162    mover: Mover,
163    start_press_time: Option<Instant>,
164    points_at_press: Option<usize>,
165    points_after_held: Option<usize>,
166    last_proximal_circle_check: Option<Instant>,
167    prediction_receiver: Option<Receiver<(World, History)>>,
168}
169impl Clone for Bbox {
170    fn clone(&self) -> Self {
171        Self {
172            prev_pos: self.prev_pos.clone(),
173            mover: self.mover,
174            start_press_time: self.start_press_time,
175            points_at_press: self.points_at_press,
176            points_after_held: self.points_after_held,
177            last_proximal_circle_check: self.last_proximal_circle_check,
178            prediction_receiver: None, // JoinHandle cannot be cloned
179        }
180    }
181}
182
183impl Bbox {
184    fn mouse_pressed(
185        &mut self,
186        event: &Events,
187        mut world: World,
188        history: History,
189    ) -> (World, History) {
190        if get_options(&world).map(|o| o.core.erase) != Some(true) {
191            if event.pressed(KeyCode::MouseRight) {
192                self.mover.move_mouse_pressed(event.mouse_pos_on_orig);
193            } else {
194                self.start_press_time = Some(Instant::now());
195                self.points_at_press = Some(self.prev_pos.prev_pos.len());
196                if !(event.held_alt() || event.held_ctrl() || event.held_shift()) {
197                    world =
198                        deselect_all::<_, DataAccessors, InstanceAnnoAccessors>(world, BBOX_NAME);
199                }
200            }
201        }
202        (world, history)
203    }
204
205    fn mouse_held(
206        &mut self,
207        event: &Events,
208        mut world: World,
209        mut history: History,
210    ) -> (World, History) {
211        if event.held(KeyCode::MouseRight) {
212            on_mouse_held_right(event.mouse_pos_on_orig, &mut self.mover, world, history)
213        } else {
214            let options = get_options(&world);
215            let params = MouseHeldLeftParams {
216                prev_pos: self.prev_pos.clone(),
217                is_alt_held: event.held_alt(),
218                is_shift_held: event.held_shift(),
219                is_ctrl_held: event.held_ctrl(),
220                distance: f64::from(options.map_or(2, |o| o.drawing_distance)),
221                elapsed_millis_since_press: self
222                    .start_press_time
223                    .map_or(0, |t| t.elapsed().as_millis()),
224            };
225            (world, history, self.prev_pos) =
226                on_mouse_held_left(event.mouse_pos_on_orig, params, world, history);
227            self.points_after_held = Some(self.prev_pos.prev_pos.len());
228            (world, history)
229        }
230    }
231
232    fn mouse_released(
233        &mut self,
234        event: &Events,
235        mut world: World,
236        mut history: History,
237    ) -> (World, History) {
238        // evaluate if a box or a polygon should be closed based on the number of points
239        // at the time of the press and the number of points after the held
240        let close_box_or_poly = self.points_at_press.map(|x| x + 4) < self.points_after_held;
241        self.points_at_press = None;
242        self.points_after_held = None;
243
244        let are_boxes_visible = get_visible(&world);
245        if event.released(KeyCode::MouseLeft) {
246            let params = MouseReleaseParams {
247                prev_pos: self.prev_pos.clone(),
248                visible: are_boxes_visible,
249                is_alt_held: event.held_alt(),
250                is_shift_held: event.held_shift(),
251                is_ctrl_held: event.held_ctrl(),
252                close_box_or_poly,
253            };
254            (world, history, self.prev_pos) =
255                on_mouse_released_left(event.mouse_pos_on_orig, params, world, history);
256        } else if event.released(KeyCode::MouseRight) {
257            (world, history, self.prev_pos) = on_mouse_released_right(
258                event.mouse_pos_on_orig,
259                self.prev_pos.clone(),
260                are_boxes_visible,
261                world,
262                history,
263            );
264        } else {
265            history.push(Record::new(world.clone(), ACTOR_NAME));
266        }
267        (world, history)
268    }
269
270    fn key_held(
271        &mut self,
272        events: &Events,
273        mut world: World,
274        history: History,
275    ) -> (World, History) {
276        // up, down, left, right
277        let shape_orig = world.data.shape();
278        let split_mode = get_options(&world).map(|o| o.split_mode);
279        let shift_annos = |annos: &mut BboxAnnotations| {
280            if let Some(split_mode) = split_mode {
281                if events.held(KeyCode::Up) && events.held_ctrl() {
282                    *annos = mem::take(annos).shift_min_bbs(0.0, -1.0, shape_orig, split_mode);
283                } else if events.held(KeyCode::Down) && events.held_ctrl() {
284                    *annos = mem::take(annos).shift_min_bbs(0.0, 1.0, shape_orig, split_mode);
285                } else if events.held(KeyCode::Right) && events.held_ctrl() {
286                    *annos = mem::take(annos).shift_min_bbs(1.0, 0.0, shape_orig, split_mode);
287                } else if events.held(KeyCode::Left) && events.held_ctrl() {
288                    *annos = mem::take(annos).shift_min_bbs(-1.0, 0.0, shape_orig, split_mode);
289                } else if events.held(KeyCode::Up) && events.held_alt() {
290                    *annos = mem::take(annos).shift(0.0, -1.0, shape_orig, split_mode);
291                } else if events.held(KeyCode::Down) && events.held_alt() {
292                    *annos = mem::take(annos).shift(0.0, 1.0, shape_orig, split_mode);
293                } else if events.held(KeyCode::Right) && events.held_alt() {
294                    *annos = mem::take(annos).shift(1.0, 0.0, shape_orig, split_mode);
295                } else if events.held(KeyCode::Left) && events.held_alt() {
296                    *annos = mem::take(annos).shift(-1.0, 0.0, shape_orig, split_mode);
297                } else if events.held(KeyCode::Up) {
298                    *annos = mem::take(annos).shift_max_bbs(0.0, -1.0, shape_orig, split_mode);
299                } else if events.held(KeyCode::Down) {
300                    *annos = mem::take(annos).shift_max_bbs(0.0, 1.0, shape_orig, split_mode);
301                } else if events.held(KeyCode::Right) {
302                    *annos = mem::take(annos).shift_max_bbs(1.0, 0.0, shape_orig, split_mode);
303                } else if events.held(KeyCode::Left) {
304                    *annos = mem::take(annos).shift_max_bbs(-1.0, 0.0, shape_orig, split_mode);
305                }
306            }
307        };
308        change_annos_bbox(&mut world, shift_annos);
309        let vis = get_visible(&world);
310        world.request_redraw_annotations(BBOX_NAME, vis);
311        (world, history)
312    }
313
314    fn key_released(
315        &mut self,
316        events: &Events,
317        mut world: World,
318        mut history: History,
319    ) -> (World, History) {
320        let params = KeyReleasedParams {
321            is_ctrl_held: events.held_ctrl(),
322            released_key: map_released_key(events),
323        };
324        world = check_erase_mode::<DataAccessors>(params.released_key, set_visible, world);
325        (world, history) = on_key_released(world, history, events.mouse_pos_on_orig, &params);
326        (world, history)
327    }
328}
329
330impl Manipulate for Bbox {
331    fn new() -> Self {
332        Self {
333            prev_pos: PrevPos::default(),
334            mover: Mover::new(),
335            start_press_time: None,
336            points_after_held: None,
337            points_at_press: None,
338            last_proximal_circle_check: None,
339            prediction_receiver: None,
340        }
341    }
342
343    fn on_activate(&mut self, mut world: World) -> World {
344        self.prev_pos = PrevPos::default();
345        if let Some(data) = trace_ok_err(get_data_mut(&mut world)) {
346            data.menu_active = true;
347        }
348        set_visible(&mut world);
349        world
350    }
351
352    fn on_deactivate(&mut self, mut world: World) -> World {
353        self.prev_pos = PrevPos::default();
354        if let Some(td) = world.data.tools_data_map.get_mut(BBOX_NAME) {
355            td.menu_active = false;
356        }
357        world.request_redraw_annotations(BBOX_NAME, Visibility::None);
358        world
359    }
360    fn on_always_active_zoom(&mut self, mut world: World, history: History) -> (World, History) {
361        let visible = get_options(&world).map(|o| o.core.visible) == Some(true);
362        let vis = vis_from_lfoption(get_label_info(&world), visible);
363        world.request_redraw_annotations(BBOX_NAME, vis);
364        (world, history)
365    }
366    fn on_filechange(&mut self, mut world: World, mut history: History) -> (World, History) {
367        use_currentimageshape_for_annos(&mut world);
368
369        let bbox_data = get_specific_mut(&mut world);
370        if let Some(bbox_data) = bbox_data {
371            for (_, (anno, _)) in bbox_data.anno_iter_mut() {
372                anno.deselect_all();
373            }
374            let ild = get_instance_label_display(&world);
375            world = instance_label_display_sort::<_, DataAccessors, InstanceAnnoAccessors>(
376                world, ild, ACTOR_NAME,
377            );
378        }
379
380        let visible = get_options(&world).map(|o| o.core.visible) == Some(true);
381        let vis = vis_from_lfoption(get_label_info(&world), visible);
382        world.request_redraw_annotations(BBOX_NAME, vis);
383
384        (world, history) =
385            check_autopaste::<_, DataAccessors, InstanceAnnoAccessors>(world, history, ACTOR_NAME);
386
387        (world, history)
388    }
389
390    fn events_tf(
391        &mut self,
392        mut world: World,
393        mut history: History,
394        events: &Events,
395    ) -> (World, History) {
396        world = check_recolorboxes::<DataAccessors>(world, BBOX_NAME);
397
398        predictive_labeling::<DataAccessors>(
399            &mut world,
400            &mut history,
401            ACTOR_NAME,
402            &mut self.prediction_receiver,
403        );
404
405        (world, history) = check_trigger_history_update::<DataAccessors>(world, history, BBOX_NAME);
406
407        world = check_cocoexport(world);
408        let imported;
409        (world, imported) = check_cocoimport::<_, _, DataAccessors>(
410            world,
411            get_specific,
412            get_specific_mut,
413            import_coco,
414        );
415        if imported {
416            set_visible(&mut world);
417        }
418
419        let options = get_options(&world).copied();
420
421        self.last_proximal_circle_check = Some(show_grab_ball(
422            events.mouse_pos_on_orig,
423            &self.prev_pos,
424            &mut world,
425            self.last_proximal_circle_check,
426            options.as_ref(),
427        ));
428        if let Some(options) = options {
429            world = check_trigger_redraw::<DataAccessors>(world, BBOX_NAME);
430
431            let in_menu_selected_label = current_cat_idx(&world);
432            if let (Some(in_menu_selected_label), Some(mp), Some(pp_first)) = (
433                in_menu_selected_label,
434                events.mouse_pos_on_orig,
435                self.prev_pos.prev_pos.first(),
436            ) && !self.prev_pos.prev_pos.is_empty()
437            {
438                let geo = if self.prev_pos.prev_pos.len() == 1 {
439                    GeoFig::BB(BbF::from_points(mp, *pp_first))
440                } else {
441                    GeoFig::Poly(
442                        Polygon::from_vec(
443                            self.prev_pos
444                                .prev_pos
445                                .iter()
446                                .chain(iter::once(&mp))
447                                .copied()
448                                .collect::<Vec<_>>(),
449                        )
450                        .unwrap(),
451                    )
452                };
453                // animation
454                let circles = get_specific(&world).map(|d| d.highlight_circles.clone());
455                let label_info = get_specific(&world).map(|d| &d.label_info);
456
457                if let (Some(circles), Some(label_info)) = (circles, label_info)
458                    && let (Some(label), Some(color)) = (
459                        label_info.labels().get(in_menu_selected_label),
460                        label_info.colors().get(in_menu_selected_label),
461                    )
462                {
463                    let anno = BboxAnnotation {
464                        geofig: geo,
465                        label: Some(label.clone()),
466                        fill_color: Some(*color),
467                        fill_alpha: 0,
468                        outline: Stroke {
469                            color: *color,
470                            thickness: TPtF::from(options.outline_thickness) / 4.0,
471                        },
472                        outline_alpha: options.outline_alpha,
473                        is_selected: None,
474                        highlight_circles: circles,
475                        instance_label_display: options.core.instance_label_display,
476                    };
477                    let vis = get_visible(&world);
478                    world.request_redraw_annotations(BBOX_NAME, vis);
479                    world.request_redraw_tmp_anno(Annotation::Bbox(anno));
480                }
481            }
482        }
483        (world, history) = make_tool_transform!(
484            self,
485            world,
486            history,
487            events,
488            [
489                (pressed, KeyCode::MouseRight, mouse_pressed),
490                (pressed, KeyCode::MouseLeft, mouse_pressed),
491                (held, KeyCode::MouseRight, mouse_held),
492                (held, KeyCode::MouseLeft, mouse_held),
493                (released, KeyCode::MouseLeft, mouse_released),
494                (released, KeyCode::MouseRight, mouse_released),
495                (released, KeyCode::Delete, key_released),
496                (released, KeyCode::Back, key_released),
497                (released, KeyCode::H, key_released),
498                (released, KeyCode::A, key_released),
499                (released, KeyCode::D, key_released),
500                (released, KeyCode::E, key_released),
501                (released, KeyCode::C, key_released),
502                (released, KeyCode::V, key_released),
503                (released, KeyCode::L, key_released),
504                (released, KeyCode::Down, key_released),
505                (released, KeyCode::Up, key_released),
506                (released, KeyCode::Left, key_released),
507                (released, KeyCode::Right, key_released),
508                (released, KeyCode::Key1, key_released),
509                (released, KeyCode::Key2, key_released),
510                (released, KeyCode::Key3, key_released),
511                (released, KeyCode::Key4, key_released),
512                (released, KeyCode::Key5, key_released),
513                (released, KeyCode::Key6, key_released),
514                (released, KeyCode::Key7, key_released),
515                (released, KeyCode::Key8, key_released),
516                (released, KeyCode::Key9, key_released),
517                (held, KeyCode::Down, key_held),
518                (held, KeyCode::Up, key_held),
519                (held, KeyCode::Left, key_held),
520                (held, KeyCode::Right, key_held)
521            ]
522        );
523        (world, history)
524    }
525}
526
527#[cfg(test)]
528use {
529    super::on_events::test_data,
530    crate::Event,
531    crate::cfg::{ExportPath, ExportPathConnection},
532    std::{path::PathBuf, thread, time::Duration},
533};
534#[test]
535fn test_bbox_ctrl_h() {
536    let (_, mut world, mut history) = test_data();
537    let mut bbox = Bbox::new();
538    bbox.last_proximal_circle_check = Some(Instant::now());
539    thread::sleep(Duration::from_millis(3));
540    assert_eq!(get_visible(&world), Visibility::All);
541    let events = Events::default()
542        .events(vec![
543            Event::Held(KeyCode::Ctrl),
544            Event::Released(KeyCode::H),
545        ])
546        .mousepos_orig(Some((1.0, 1.0).into()));
547    (world, history) = bbox.events_tf(world, history, &events);
548    thread::sleep(Duration::from_millis(3));
549    (world, _) = bbox.events_tf(
550        world,
551        history,
552        &Events::default().mousepos_orig(Some((1.0, 1.0).into())),
553    );
554    assert_eq!(get_visible(&world), Visibility::None);
555}
556
557#[test]
558fn test_coco_import_label_info() {
559    const TEST_DATA_FOLDER: &str = "resources/test_data/";
560    let (_, mut world, history) = test_data();
561    let data = get_specific_mut(&mut world).unwrap();
562    data.coco_file = ExportPath {
563        path: PathBuf::from(format!("{}catids_12_coco.json", TEST_DATA_FOLDER)),
564        conn: ExportPathConnection::Local,
565    };
566    let label_info_before = data.label_info.clone();
567    data.options.core.import_export_trigger.trigger_import();
568    let mut bbox = Bbox::new();
569    let events = Events::default();
570    let (mut world, history) = bbox.events_tf(world, history, &events);
571    let data = get_specific(&world).unwrap();
572    assert_eq!(label_info_before.labels(), &["rvimage_fg", "label"]);
573    assert_eq!(label_info_before.cat_ids(), &[1, 2]);
574    assert_eq!(data.label_info.labels(), &["first label", "second label"]);
575    assert_eq!(data.label_info.cat_ids(), &[1, 2]);
576    assert!(!data.options.core.import_export_trigger.import_triggered());
577
578    // now we import another coco file with different labels
579    let data = get_specific_mut(&mut world).unwrap();
580    data.coco_file = ExportPath {
581        path: PathBuf::from(format!("{}catids_01_coco_3labels.json", TEST_DATA_FOLDER)),
582        conn: ExportPathConnection::Local,
583    };
584    data.options.core.import_export_trigger.trigger_import();
585    let (world, _) = bbox.events_tf(world, history, &events);
586    let data = get_specific(&world).unwrap();
587    assert_eq!(
588        data.label_info.labels(),
589        &["first label", "second label", "third label"]
590    );
591    assert_eq!(data.label_info.cat_ids(), &[0, 1, 2]);
592    let all_occurring_cats = data
593        .annotations_map
594        .iter()
595        .flat_map(|(_, (v, _))| v.cat_idxs().iter().copied())
596        .collect::<Vec<usize>>();
597    assert!(all_occurring_cats.contains(&0));
598    assert!(all_occurring_cats.contains(&1));
599    assert!(all_occurring_cats.contains(&2));
600}