mimick/library/lightbox/ocr/
state.rs1use std::collections::BTreeSet;
4
5use super::geometry::{contain_rect, quad_contains, quad_in_area, reading_order};
6use crate::api_client::OcrBox;
7
8#[derive(Debug, Default)]
10pub(super) struct OcrState {
11 boxes: Vec<OcrBox>,
12 image: (f64, f64),
13 selected: BTreeSet<usize>,
14 anchor: Option<usize>,
16 hovered: Option<usize>,
17}
18
19impl OcrState {
20 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 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 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 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 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 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 pub(super) fn selected_text(&self) -> Option<String> {
107 self.join_text(self.selected.iter().copied())
108 }
109
110 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}