rvlib/tools_data/
mod.rs

1pub use self::core::{
2    AccessInstanceData, Annotate, ExportAsCoco, ImportExportTrigger, ImportMode, InstanceAnnotate,
3    InstanceExportData, LabelInfo, OUTLINE_THICKNESS_CONVERSION, Options,
4    VisibleInactiveToolsState, vis_from_lfoption,
5};
6pub use self::{
7    attributes_data::AttributesToolData, bbox_data::BboxToolData, brush_data::BrushToolData,
8    coco_io::write_coco, label_map::LabelMap, plot_stats::PlotAnnotationStats,
9    rot90_data::Rot90ToolData,
10};
11use crate::tools::add_tools_initial_data;
12use crate::{
13    BrushAnnotation,
14    drawme::{Annotation, BboxAnnotation, Stroke},
15};
16use rvimage_domain::{RvResult, TPtF, rverr};
17use serde::{Deserialize, Serialize};
18use std::ops::Index;
19
20pub mod annotations;
21pub mod attributes_data;
22pub mod bbox_data;
23pub mod brush_data;
24pub mod coco_io;
25mod core;
26mod label_map;
27mod plot_stats;
28pub mod predictive_labeling;
29pub mod rot90_data;
30pub use core::{AnnotationsMap, InstanceLabelDisplay, Options as CoreOptions, merge};
31use std::collections::HashMap;
32
33macro_rules! variant_access {
34    ($variant:ident, $func_name:ident, $self:ty, $return_type:ty) => {
35        pub fn $func_name(self: $self) -> rvimage_domain::RvResult<$return_type> {
36            match self {
37                ToolSpecifics::$variant(x) => Ok(x),
38                _ => Err(rvimage_domain::rverr!(
39                    "this is not a {}",
40                    stringify!($variant)
41                )),
42            }
43        }
44    };
45}
46macro_rules! variant_access_free {
47    ($variant:ident, $func_name:ident, $lt:lifetime, $ToolsSpecific:ty, $return_type:ty) => {
48        pub fn $func_name<$lt>(x: $ToolsSpecific) -> rvimage_domain::RvResult<$return_type> {
49            match x {
50                ToolSpecifics::$variant(x) => Ok(x),
51                _ => Err(rvimage_domain::rverr!(
52                    "this is not a {}",
53                    stringify!($variant)
54                )),
55            }
56        }
57    };
58}
59
60variant_access_free!(Bbox, bbox, 'a, &'a ToolSpecifics, &'a BboxToolData);
61variant_access_free!(Bbox, bbox_mut, 'a, &'a mut ToolSpecifics, &'a mut BboxToolData);
62variant_access_free!(Brush, brush, 'a, &'a ToolSpecifics, &'a BrushToolData);
63variant_access_free!(Brush, brush_mut, 'a, &'a mut ToolSpecifics, &'a mut BrushToolData);
64variant_access_free!(Attributes, attributes, 'a, &'a ToolSpecifics, &'a AttributesToolData);
65variant_access_free!(Attributes, attributes_mut, 'a, &'a mut ToolSpecifics, &'a mut AttributesToolData);
66
67#[derive(Serialize, Deserialize, Clone, Debug, PartialEq)]
68#[allow(clippy::large_enum_variant)]
69pub enum ToolSpecifics {
70    Bbox(BboxToolData),
71    Brush(BrushToolData),
72    Rot90(Rot90ToolData),
73    Zoom(()),
74    AlwaysActiveZoom(()),
75    Attributes(AttributesToolData),
76}
77impl ToolSpecifics {
78    variant_access!(Bbox, bbox, &Self, &BboxToolData);
79    variant_access!(Brush, brush, &Self, &BrushToolData);
80    variant_access!(Rot90, rot90, &Self, &Rot90ToolData);
81    variant_access!(Attributes, attributes, &Self, &AttributesToolData);
82    variant_access!(Bbox, bbox_mut, &mut Self, &mut BboxToolData);
83    variant_access!(Brush, brush_mut, &mut Self, &mut BrushToolData);
84    variant_access!(Rot90, rot90_mut, &mut Self, &mut Rot90ToolData);
85    variant_access!(
86        Attributes,
87        attributes_mut,
88        &mut Self,
89        &mut AttributesToolData
90    );
91
92    pub fn apply_mut<T>(
93        &mut self,
94        f_bbox: impl FnOnce(&mut BboxToolData) -> RvResult<T>,
95        f_brush: impl FnOnce(&mut BrushToolData) -> RvResult<T>,
96        f_attr: impl FnOnce(&mut AttributesToolData) -> RvResult<T>,
97    ) -> RvResult<T> {
98        match self {
99            Self::Bbox(bbox_data) => f_bbox(bbox_data),
100            Self::Brush(brush_data) => f_brush(brush_data),
101            Self::Attributes(attr_data) => f_attr(attr_data),
102            _ => Err(rverr!("only brush tool and bbox tool can be used in apply")),
103        }
104    }
105    pub fn apply<T>(
106        &self,
107        f_bbox: impl FnOnce(&BboxToolData) -> RvResult<T>,
108        f_brush: impl FnOnce(&BrushToolData) -> RvResult<T>,
109    ) -> RvResult<T> {
110        match self {
111            Self::Bbox(bbox_data) => f_bbox(bbox_data),
112            Self::Brush(brush_data) => f_brush(brush_data),
113            _ => Err(rverr!("only brush tool and bbox tool can be used in apply")),
114        }
115    }
116
117    pub fn to_annotations_view(
118        &self,
119        file_path_relative: &str,
120        only_cat_idx: Option<usize>,
121    ) -> Option<Vec<Annotation>> {
122        match self {
123            ToolSpecifics::Bbox(bb_data) => {
124                if let Some(annos) = bb_data.get_annos(file_path_relative) {
125                    let geos = annos.elts();
126                    let cats = annos.cat_idxs();
127                    let selected_bbs = annos.selected_mask();
128                    let labels = bb_data.label_info.labels();
129                    let colors = bb_data.label_info.colors();
130                    let bbs_colored = geos
131                        .iter()
132                        .zip(cats.iter())
133                        .zip(selected_bbs.iter())
134                        .filter(|((_, cat_idx), _)| {
135                            if let Some(only_cat_idx) = only_cat_idx {
136                                **cat_idx == only_cat_idx
137                            } else {
138                                true
139                            }
140                        })
141                        .flat_map(|((geo, cat_idx), is_selected)| {
142                            colors.get(*cat_idx).map(|color| {
143                                Annotation::Bbox(BboxAnnotation {
144                                    geofig: geo.clone(),
145                                    fill_color: colors.get(*cat_idx).cloned(),
146                                    fill_alpha: bb_data.options.fill_alpha,
147                                    label: labels.get(*cat_idx).cloned(),
148                                    outline: Stroke {
149                                        thickness: TPtF::from(bb_data.options.outline_thickness)
150                                            / OUTLINE_THICKNESS_CONVERSION,
151                                        color: *color,
152                                    },
153                                    outline_alpha: bb_data.options.outline_alpha,
154                                    is_selected: Some(*is_selected),
155                                    highlight_circles: bb_data.highlight_circles.clone(),
156                                    instance_label_display: bb_data
157                                        .options
158                                        .core
159                                        .instance_label_display,
160                                })
161                            })
162                        })
163                        .collect::<Vec<Annotation>>();
164                    Some(bbs_colored)
165                } else {
166                    Some(vec![])
167                }
168            }
169            ToolSpecifics::Brush(br_data) => {
170                if let Some(annos) = br_data.get_annos(file_path_relative) {
171                    let colors = br_data.label_info.colors();
172                    let cats = annos.cat_idxs();
173                    let selected_mask = annos.selected_mask();
174                    let labels = br_data.label_info.labels();
175                    let annos = annos
176                        .elts()
177                        .iter()
178                        .zip(cats.iter())
179                        .zip(selected_mask.iter())
180                        .filter(|((_, cat_idx), _)| {
181                            if let Some(only_cat_idx) = only_cat_idx {
182                                **cat_idx == only_cat_idx
183                            } else {
184                                true
185                            }
186                        })
187                        .flat_map(|((brush_line, cat_idx), is_selected)| {
188                            colors.get(*cat_idx).map(|color| {
189                                Annotation::Brush(BrushAnnotation {
190                                    canvas: brush_line.clone(),
191                                    color: *color,
192                                    label: labels.get(*cat_idx).cloned(),
193                                    is_selected: Some(*is_selected),
194                                    fill_alpha: br_data.options.fill_alpha,
195                                    instance_display_label: br_data
196                                        .options
197                                        .core
198                                        .instance_label_display,
199                                })
200                            })
201                        })
202                        .collect::<Vec<Annotation>>();
203                    Some(annos)
204                } else {
205                    Some(vec![])
206                }
207            }
208            _ => None,
209        }
210    }
211}
212impl Default for ToolSpecifics {
213    fn default() -> Self {
214        ToolSpecifics::Bbox(BboxToolData::default())
215    }
216}
217
218#[derive(Serialize, Deserialize, Clone, Debug, Default, PartialEq)]
219pub struct ToolsData {
220    pub specifics: ToolSpecifics,
221    pub menu_active: bool,
222    #[serde(default)]
223    pub visible_inactive_tools: VisibleInactiveToolsState,
224}
225impl ToolsData {
226    pub fn new(
227        specifics: ToolSpecifics,
228        visible_inactive_tools: VisibleInactiveToolsState,
229    ) -> Self {
230        ToolsData {
231            specifics,
232            menu_active: false,
233            visible_inactive_tools,
234        }
235    }
236}
237
238#[macro_export]
239macro_rules! toolsdata_by_name {
240    ($name:expr, $acc:ident, $tdm:expr) => {
241        $tdm.get_mut($name)
242            .ok_or(rvimage_domain::rverr!("{} is not a tool", $name))?
243            .specifics
244            .$acc()?
245    };
246}
247
248#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
249pub struct ToolsDataMap {
250    // tool name -> tool's menu data type
251    #[serde(flatten)]
252    data: HashMap<String, ToolsData>,
253}
254impl ToolsDataMap {
255    pub fn new() -> Self {
256        let tdm = ToolsDataMap {
257            data: HashMap::new(),
258        };
259        add_tools_initial_data(tdm)
260    }
261    pub fn iter(&self) -> impl Iterator<Item = (&String, &ToolsData)> {
262        self.data.iter()
263    }
264    pub fn contains_key(&self, name: &str) -> bool {
265        self.data.contains_key(name)
266    }
267    pub fn get_specifics(&self, name: &str) -> Option<&ToolSpecifics> {
268        self.data.get(name).map(|d| &d.specifics)
269    }
270    pub fn get_specifics_mut(&mut self, name: &str) -> Option<&mut ToolSpecifics> {
271        self.data.get_mut(name).map(|d| &mut d.specifics)
272    }
273    pub fn get(&self, name: &str) -> Option<&ToolsData> {
274        self.data.get(name)
275    }
276    pub fn get_mut(&mut self, name: &str) -> Option<&mut ToolsData> {
277        self.data.get_mut(name)
278    }
279    pub fn len(&self) -> usize {
280        self.data.len()
281    }
282    pub fn is_empty(&self) -> bool {
283        self.data.is_empty()
284    }
285    pub fn values_mut(&mut self) -> impl Iterator<Item = &mut ToolsData> {
286        self.data.values_mut()
287    }
288
289    pub fn insert(&mut self, name: String, data: ToolsData) -> Option<ToolsData> {
290        self.data.insert(name, data)
291    }
292    pub fn set_tools_specific_data(&mut self, name: &str, specifics: ToolSpecifics) {
293        self.data.insert(
294            name.to_string(),
295            ToolsData::new(specifics, VisibleInactiveToolsState::default()),
296        );
297    }
298}
299impl Default for ToolsDataMap {
300    fn default() -> Self {
301        Self::new()
302    }
303}
304
305#[allow(clippy::indexing_slicing)]
306impl Index<&str> for ToolsDataMap {
307    type Output = ToolsData;
308    fn index(&self, index: &str) -> &Self::Output {
309        &self.data[index]
310    }
311}
312impl FromIterator<(String, ToolsData)> for ToolsDataMap {
313    fn from_iter<T: IntoIterator<Item = (std::string::String, ToolsData)>>(iter: T) -> Self {
314        let data = iter.into_iter().collect::<HashMap<String, ToolsData>>();
315        add_tools_initial_data(ToolsDataMap { data })
316    }
317}
318impl From<HashMap<String, ToolsData>> for ToolsDataMap {
319    fn from(data: HashMap<String, ToolsData>) -> Self {
320        add_tools_initial_data(ToolsDataMap { data })
321    }
322}
323
324#[macro_export]
325macro_rules! get_specifics_from_tdm {
326    ($actor_name:expr, $tdm:expr, $access_func:ident) => {
327        $tdm.get($actor_name)
328            .and_then(|x| x.specifics.$access_func().ok())
329    };
330}
331#[macro_export]
332macro_rules! get_annos_from_tdm {
333    ($actor_name:expr, $tdm:expr, $current_file_path:expr, $access_func:ident) => {
334        $crate::get_specifics_from_tdm!($actor_name, $tdm, $access_func)
335            .and_then(|d| d.get_annos($current_file_path))
336    };
337}
338
339#[macro_export]
340macro_rules! get_labelinfo_from_tdm {
341    ($actor_name:expr, $tdm:expr,  $access_func:ident) => {
342        $crate::get_specifics_from_tdm!($actor_name, $tdm, $access_func).map(|d| d.label_info())
343    };
344}
345
346#[cfg(test)]
347use crate::tools::{
348    ALWAYS_ACTIVE_ZOOM, ATTRIBUTES_NAME, BBOX_NAME, BRUSH_NAME, ROT90_NAME, ZOOM_NAME,
349};
350#[test]
351fn test_tools_data_map() {
352    let tdm = ToolsDataMap::new();
353    let tools = [
354        BBOX_NAME,
355        ROT90_NAME,
356        BRUSH_NAME,
357        ATTRIBUTES_NAME,
358        ZOOM_NAME,
359        ALWAYS_ACTIVE_ZOOM,
360    ];
361    for tool in tools.iter() {
362        assert!(tdm.contains_key(tool));
363    }
364    assert_eq!(tdm.len(), tools.len());
365
366    // test from hashmap
367    let data = HashMap::from([(BBOX_NAME.to_string(), ToolsData::default())]);
368    let tdm = ToolsDataMap::from(data);
369    for tool in tools.iter() {
370        assert!(tdm.contains_key(tool));
371    }
372    assert_eq!(tdm.len(), tools.len());
373
374    // test from iterator
375    let data = vec![
376        (BBOX_NAME.to_string(), ToolsData::default()),
377        (ROT90_NAME.to_string(), ToolsData::default()),
378    ];
379    let tdm = ToolsDataMap::from_iter(data);
380    for tool in tools.iter() {
381        assert!(tdm.contains_key(tool));
382    }
383    assert_eq!(tdm.len(), tools.len());
384}
rvlib/tools_data/mod.rs

rvlib/tools_data/
mod.rs
