Skip to main content

mimick/library/lightbox/ocr/
state.rs

1//! Selection state for the OCR boxes of one photo.
2
3use std::collections::BTreeSet;
4
5use super::geometry::{contain_rect, quad_contains, quad_in_area, reading_order};
6use crate::api_client::OcrBox;
7
8/// Boxes in reading order, plus what is selected and hovered; indices refer to that order.
9#[derive(Debug, Default)]
10pub(super) struct OcrState {
11    boxes: Vec<OcrBox>,
12    image: (f64, f64),
13    selected: BTreeSet<usize>,
14    /// Box a Shift-click or drag extends from.
15    anchor: Option<usize>,
16    hovered: Option<usize>,
17}
18
19impl OcrState {
20    /// `image` is the displayed texture's size, used to place the normalized boxes.
21    pub(super) fn new(boxes: Vec<OcrBox>, image: (f64, f64)) -> Self {
22        Self {
23            boxes: reading_order(boxes),
24            image,
25            ..Self::default()
26        }
27    }
28
29    pub(super) fn boxes(&self) -> &[OcrBox] {
30        &self.boxes
31    }
32
33    pub(super) fn image(&self) -> (f64, f64) {
34        self.image
35    }
36
37    pub(super) fn is_selected(&self, idx: usize) -> bool {
38        self.selected.contains(&idx)
39    }
40
41    pub(super) fn hovered(&self) -> Option<usize> {
42        self.hovered
43    }
44
45    /// Returns whether the hovered box changed, so callers redraw only when needed.
46    pub(super) fn set_hovered(&mut self, idx: Option<usize>) -> bool {
47        let changed = self.hovered != idx;
48        self.hovered = idx;
49        changed
50    }
51
52    /// Topmost box under `p`, in the coordinates of a widget of size `area`.
53    pub(super) fn box_at(&self, area: (f64, f64), p: (f64, f64)) -> Option<usize> {
54        let rect = contain_rect(area, self.image)?;
55        self.boxes
56            .iter()
57            .rposition(|ocr_box| quad_contains(&quad_in_area(rect, ocr_box.corners()), p))
58    }
59
60    /// Mouse press on a box; returns whether dragging should extend the selection.
61    pub(super) fn press(&mut self, idx: usize, shift: bool, ctrl: bool) -> bool {
62        if ctrl {
63            self.click(idx, true);
64            return false;
65        }
66        if shift {
67            self.extend_to(idx);
68        } else {
69            self.click(idx, false);
70        }
71        true
72    }
73
74    /// Selects only `idx`, or toggles it when `extend` is set.
75    pub(super) fn click(&mut self, idx: usize, extend: bool) {
76        if extend {
77            if !self.selected.remove(&idx) {
78                self.selected.insert(idx);
79            }
80        } else {
81            self.selected = BTreeSet::from([idx]);
82        }
83        self.anchor = Some(idx);
84    }
85
86    /// Selects every box between the anchor and `idx`.
87    pub(super) fn extend_to(&mut self, idx: usize) {
88        let anchor = *self.anchor.get_or_insert(idx);
89        self.selected = (anchor.min(idx)..=anchor.max(idx)).collect();
90    }
91
92    pub(super) fn select_all(&mut self) {
93        self.selected = (0..self.boxes.len()).collect();
94    }
95
96    pub(super) fn clear_selection(&mut self) {
97        self.selected.clear();
98        self.anchor = None;
99    }
100
101    pub(super) fn has_selection(&self) -> bool {
102        !self.selected.is_empty()
103    }
104
105    /// Selected boxes' text, one per line in reading order.
106    pub(super) fn selected_text(&self) -> Option<String> {
107        self.join_text(self.selected.iter().copied())
108    }
109
110    /// Every box's text, one per line in reading order.
111    pub(super) fn all_text(&self) -> Option<String> {
112        self.join_text(0..self.boxes.len())
113    }
114
115    fn join_text(&self, indices: impl Iterator<Item = usize>) -> Option<String> {
116        let lines: Vec<&str> = indices
117            .filter_map(|idx| self.boxes.get(idx))
118            .map(|ocr_box| ocr_box.text.as_str())
119            .collect();
120        (!lines.is_empty()).then(|| lines.join("\n"))
121    }
122}
123
124#[cfg(test)]
125mod tests {
126    use super::super::geometry::ob;
127    use super::*;
128
129    #[test]
130    fn click_replaces_selection_and_extend_toggles() {
131        let mut state = OcrState::new(
132            vec![ob("a", 0.0, 0.0), ob("b", 0.0, 0.25), ob("c", 0.0, 0.5)],
133            (100.0, 100.0),
134        );
135        state.click(0, false);
136        state.click(2, true);
137        assert_eq!(state.selected_text().as_deref(), Some("a\nc"));
138        state.click(0, true);
139        assert_eq!(state.selected_text().as_deref(), Some("c"));
140        state.click(1, false);
141        assert_eq!(state.selected_text().as_deref(), Some("b"));
142        state.clear_selection();
143        assert_eq!(state.selected_text(), None);
144    }
145
146    #[test]
147    fn select_all_copies_in_reading_order() {
148        let mut state = OcrState::new(
149            vec![
150                ob("right", 0.625, 0.5),
151                ob("left", 0.125, 0.51),
152                ob("top", 0.5, 0.25),
153            ],
154            (100.0, 100.0),
155        );
156        state.select_all();
157        assert_eq!(state.selected_text().as_deref(), Some("top\nleft\nright"));
158    }
159
160    #[test]
161    fn box_at_ignores_letterbox() {
162        let state = OcrState::new(vec![ob("a", 0.125, 0.25)], (2000.0, 1000.0));
163        assert_eq!(state.box_at((1000.0, 1000.0), (200.0, 390.0)), Some(0));
164        assert_eq!(state.box_at((1000.0, 1000.0), (200.0, 100.0)), None);
165        assert_eq!(state.box_at((0.0, 0.0), (0.0, 0.0)), None);
166    }
167
168    fn four_rows() -> OcrState {
169        OcrState::new(
170            vec![
171                ob("a", 0.0, 0.0),
172                ob("b", 0.0, 0.25),
173                ob("c", 0.0, 0.5),
174                ob("d", 0.0, 0.75),
175            ],
176            (100.0, 100.0),
177        )
178    }
179
180    #[test]
181    fn extend_to_selects_the_range_from_the_last_click() {
182        let mut state = four_rows();
183        state.click(1, false);
184        state.extend_to(3);
185        assert_eq!(state.selected_text().as_deref(), Some("b\nc\nd"));
186        state.extend_to(0);
187        assert_eq!(state.selected_text().as_deref(), Some("a\nb"));
188    }
189
190    #[test]
191    fn extend_to_without_a_click_selects_one_box() {
192        let mut state = four_rows();
193        state.extend_to(2);
194        assert_eq!(state.selected_text().as_deref(), Some("c"));
195    }
196
197    #[test]
198    fn all_text_joins_every_box_in_reading_order() {
199        let mut state = OcrState::new(
200            vec![ob("right", 0.625, 0.5), ob("top", 0.5, 0.25)],
201            (100.0, 100.0),
202        );
203        state.click(0, false);
204        assert_eq!(state.all_text().as_deref(), Some("top\nright"));
205        assert_eq!(OcrState::default().all_text(), None);
206    }
207
208    #[test]
209    fn has_selection_follows_clicks_and_clearing() {
210        let mut state = four_rows();
211        assert!(!state.has_selection());
212        state.click(2, false);
213        assert!(state.has_selection());
214        state.clear_selection();
215        assert!(!state.has_selection());
216    }
217
218    #[test]
219    fn press_selects_extends_or_toggles_and_reports_drag() {
220        let mut state = four_rows();
221        assert!(state.press(1, false, false));
222        assert_eq!(state.selected_text().as_deref(), Some("b"));
223        assert!(state.press(3, true, false));
224        assert_eq!(state.selected_text().as_deref(), Some("b\nc\nd"));
225        assert!(!state.press(0, false, true));
226        assert_eq!(state.selected_text().as_deref(), Some("a\nb\nc\nd"));
227    }
228}