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

rvlib/tools_data/
mod.rs
