rvlib/tools_data/
mod.rs

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