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            is_shift_held: events.held_shift(),
323            released_key: map_released_key(events),
324        };
325        world = check_erase_mode::<DataAccessors>(params.released_key, set_visible, world);
326        (world, history) = on_key_released(world, history, events.mouse_pos_on_orig, &params);
327        (world, history)
328    }
329}
330
331impl Manipulate for Bbox {
332    fn new() -> Self {
333        Self {
334            prev_pos: PrevPos::default(),
335            mover: Mover::new(),
336            start_press_time: None,
337            points_after_held: None,
338            points_at_press: None,
339            last_proximal_circle_check: None,
340            prediction_receiver: None,
341        }
342    }
343
344    fn on_activate(&mut self, mut world: World) -> World {
345        self.prev_pos = PrevPos::default();
346        if let Some(data) = trace_ok_err(get_data_mut(&mut world)) {
347            data.menu_active = true;
348        }
349        set_visible(&mut world);
350        world
351    }
352
353    fn on_deactivate(&mut self, mut world: World) -> World {
354        self.prev_pos = PrevPos::default();
355        if let Some(td) = world.data.tools_data_map.get_mut(BBOX_NAME) {
356            td.menu_active = false;
357        }
358        world.request_redraw_annotations(BBOX_NAME, Visibility::None);
359        world
360    }
361    fn on_always_active_zoom(&mut self, mut world: World, history: History) -> (World, History) {
362        let visible = get_options(&world).map(|o| o.core.visible) == Some(true);
363        let vis = vis_from_lfoption(get_label_info(&world), visible);
364        world.request_redraw_annotations(BBOX_NAME, vis);
365        (world, history)
366    }
367    fn on_filechange(&mut self, mut world: World, mut history: History) -> (World, History) {
368        use_currentimageshape_for_annos(&mut world);
369
370        let bbox_data = get_specific_mut(&mut world);
371        if let Some(bbox_data) = bbox_data {
372            for (_, (anno, _)) in bbox_data.anno_iter_mut() {
373                anno.deselect_all();
374            }
375            let ild = get_instance_label_display(&world);
376            world = instance_label_display_sort::<_, DataAccessors, InstanceAnnoAccessors>(
377                world, ild, ACTOR_NAME,
378            );
379        }
380
381        let visible = get_options(&world).map(|o| o.core.visible) == Some(true);
382        let vis = vis_from_lfoption(get_label_info(&world), visible);
383        world.request_redraw_annotations(BBOX_NAME, vis);
384
385        (world, history) =
386            check_autopaste::<_, DataAccessors, InstanceAnnoAccessors>(world, history, ACTOR_NAME);
387
388        (world, history)
389    }
390
391    fn events_tf(
392        &mut self,
393        mut world: World,
394        mut history: History,
395        events: &Events,
396    ) -> (World, History) {
397        world = check_recolorboxes::<DataAccessors>(world, BBOX_NAME);
398
399        predictive_labeling::<DataAccessors>(
400            &mut world,
401            &mut history,
402            ACTOR_NAME,
403            &mut self.prediction_receiver,
404        );
405
406        (world, history) = check_trigger_history_update::<DataAccessors>(world, history, BBOX_NAME);
407
408        world = check_cocoexport(world);
409        let imported;
410        (world, imported) = check_cocoimport::<_, _, DataAccessors>(
411            world,
412            get_specific,
413            get_specific_mut,
414            import_coco,
415        );
416        if imported {
417            set_visible(&mut world);
418        }
419
420        let options = get_options(&world).copied();
421
422        self.last_proximal_circle_check = Some(show_grab_ball(
423            events.mouse_pos_on_orig,
424            &self.prev_pos,
425            &mut world,
426            self.last_proximal_circle_check,
427            options.as_ref(),
428        ));
429        if let Some(options) = options {
430            world = check_trigger_redraw::<DataAccessors>(world, BBOX_NAME);
431
432            let in_menu_selected_label = current_cat_idx(&world);
433            if let (Some(in_menu_selected_label), Some(mp), Some(pp_first)) = (
434                in_menu_selected_label,
435                events.mouse_pos_on_orig,
436                self.prev_pos.prev_pos.first(),
437            ) && !self.prev_pos.prev_pos.is_empty()
438            {
439                let geo = if self.prev_pos.prev_pos.len() == 1 {
440                    GeoFig::BB(BbF::from_points(mp, *pp_first))
441                } else {
442                    GeoFig::Poly(
443                        Polygon::from_vec(
444                            self.prev_pos
445                                .prev_pos
446                                .iter()
447                                .chain(iter::once(&mp))
448                                .copied()
449                                .collect::<Vec<_>>(),
450                        )
451                        .unwrap(),
452                    )
453                };
454                // animation
455                let circles = get_specific(&world).map(|d| d.highlight_circles.clone());
456                let label_info = get_specific(&world).map(|d| &d.label_info);
457
458                if let (Some(circles), Some(label_info)) = (circles, label_info)
459                    && let (Some(label), Some(color)) = (
460                        label_info.labels().get(in_menu_selected_label),
461                        label_info.colors().get(in_menu_selected_label),
462                    )
463                {
464                    let anno = BboxAnnotation {
465                        geofig: geo,
466                        label: Some(label.clone()),
467                        fill_color: Some(*color),
468                        fill_alpha: 0,
469                        outline: Stroke {
470                            color: *color,
471                            thickness: TPtF::from(options.outline_thickness) / 4.0,
472                        },
473                        outline_alpha: options.outline_alpha,
474                        is_selected: None,
475                        highlight_circles: circles,
476                        instance_label_display: options.core.instance_label_display,
477                    };
478                    let vis = get_visible(&world);
479                    world.request_redraw_annotations(BBOX_NAME, vis);
480                    world.request_redraw_tmp_anno(Annotation::Bbox(anno));
481                }
482            }
483        }
484        (world, history) = make_tool_transform!(
485            self,
486            world,
487            history,
488            events,
489            [
490                (pressed, KeyCode::MouseRight, mouse_pressed),
491                (pressed, KeyCode::MouseLeft, mouse_pressed),
492                (held, KeyCode::MouseRight, mouse_held),
493                (held, KeyCode::MouseLeft, mouse_held),
494                (released, KeyCode::MouseLeft, mouse_released),
495                (released, KeyCode::MouseRight, mouse_released),
496                (released, KeyCode::Delete, key_released),
497                (released, KeyCode::Back, key_released),
498                (released, KeyCode::H, key_released),
499                (released, KeyCode::A, key_released),
500                (released, KeyCode::D, key_released),
501                (released, KeyCode::E, key_released),
502                (released, KeyCode::C, key_released),
503                (released, KeyCode::V, key_released),
504                (released, KeyCode::L, key_released),
505                (released, KeyCode::Down, key_released),
506                (released, KeyCode::Up, key_released),
507                (released, KeyCode::Left, key_released),
508                (released, KeyCode::Right, key_released),
509                (released, KeyCode::Key1, key_released),
510                (released, KeyCode::Key2, key_released),
511                (released, KeyCode::Key3, key_released),
512                (released, KeyCode::Key4, key_released),
513                (released, KeyCode::Key5, key_released),
514                (released, KeyCode::Key6, key_released),
515                (released, KeyCode::Key7, key_released),
516                (released, KeyCode::Key8, key_released),
517                (released, KeyCode::Key9, key_released),
518                (held, KeyCode::Down, key_held),
519                (held, KeyCode::Up, key_held),
520                (held, KeyCode::Left, key_held),
521                (held, KeyCode::Right, key_held)
522            ]
523        );
524        (world, history)
525    }
526}
527
528#[cfg(test)]
529use {
530    super::on_events::test_data,
531    crate::Event,
532    crate::cfg::{ExportPath, ExportPathConnection},
533    std::{path::PathBuf, thread, time::Duration},
534};
535#[test]
536fn test_bbox_ctrl_h() {
537    let (_, mut world, mut history) = test_data();
538    let mut bbox = Bbox::new();
539    bbox.last_proximal_circle_check = Some(Instant::now());
540    thread::sleep(Duration::from_millis(3));
541    assert_eq!(get_visible(&world), Visibility::All);
542    let events = Events::default()
543        .events(vec![
544            Event::Held(KeyCode::Ctrl),
545            Event::Released(KeyCode::H),
546        ])
547        .mousepos_orig(Some((1.0, 1.0).into()));
548    (world, history) = bbox.events_tf(world, history, &events);
549    thread::sleep(Duration::from_millis(3));
550    (world, _) = bbox.events_tf(
551        world,
552        history,
553        &Events::default().mousepos_orig(Some((1.0, 1.0).into())),
554    );
555    assert_eq!(get_visible(&world), Visibility::None);
556}
557
558#[test]
559fn test_coco_import_label_info() {
560    const TEST_DATA_FOLDER: &str = "resources/test_data/";
561    let (_, mut world, history) = test_data();
562    let data = get_specific_mut(&mut world).unwrap();
563    data.coco_file = ExportPath {
564        path: PathBuf::from(format!("{}catids_12_coco.json", TEST_DATA_FOLDER)),
565        conn: ExportPathConnection::Local,
566    };
567    let label_info_before = data.label_info.clone();
568    data.options.core.import_export_trigger.trigger_import();
569    let mut bbox = Bbox::new();
570    let events = Events::default();
571    let (mut world, history) = bbox.events_tf(world, history, &events);
572    let data = get_specific(&world).unwrap();
573    assert_eq!(label_info_before.labels(), &["rvimage_fg", "label"]);
574    assert_eq!(label_info_before.cat_ids(), &[1, 2]);
575    assert_eq!(data.label_info.labels(), &["first label", "second label"]);
576    assert_eq!(data.label_info.cat_ids(), &[1, 2]);
577    assert!(!data.options.core.import_export_trigger.import_triggered());
578
579    // now we import another coco file with different labels
580    let data = get_specific_mut(&mut world).unwrap();
581    data.coco_file = ExportPath {
582        path: PathBuf::from(format!("{}catids_01_coco_3labels.json", TEST_DATA_FOLDER)),
583        conn: ExportPathConnection::Local,
584    };
585    data.options.core.import_export_trigger.trigger_import();
586    let (world, _) = bbox.events_tf(world, history, &events);
587    let data = get_specific(&world).unwrap();
588    assert_eq!(
589        data.label_info.labels(),
590        &["first label", "second label", "third label"]
591    );
592    assert_eq!(data.label_info.cat_ids(), &[0, 1, 2]);
593    let all_occurring_cats = data
594        .annotations_map
595        .iter()
596        .flat_map(|(_, (v, _))| v.cat_idxs().iter().copied())
597        .collect::<Vec<usize>>();
598    assert!(all_occurring_cats.contains(&0));
599    assert!(all_occurring_cats.contains(&1));
600    assert!(all_occurring_cats.contains(&2));
601}