rvlib/tools/bbox/
core.rs

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