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;
26pub mod parameters;
27mod plot_stats;
28pub mod predictive_labeling;
29pub mod rot90_data;
30pub use core::{merge, AnnotationsMap, InstanceLabelDisplay, Options as CoreOptions};
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                        .map(|((geo, cat_idx), is_selected)| {
142                            Annotation::Bbox(BboxAnnotation {
143                                geofig: geo.clone(),
144                                fill_color: Some(colors[*cat_idx]),
145                                fill_alpha: bb_data.options.fill_alpha,
146                                label: Some(labels[*cat_idx].clone()),
147                                outline: Stroke {
148                                    thickness: TPtF::from(bb_data.options.outline_thickness)
149                                        / OUTLINE_THICKNESS_CONVERSION,
150                                    color: colors[*cat_idx],
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.options.core.instance_label_display,
156                            })
157                        })
158                        .collect::<Vec<Annotation>>();
159                    Some(bbs_colored)
160                } else {
161                    Some(vec![])
162                }
163            }
164            ToolSpecifics::Brush(br_data) => {
165                if let Some(annos) = br_data.get_annos(file_path_relative) {
166                    let colors = br_data.label_info.colors();
167                    let cats = annos.cat_idxs();
168                    let selected_mask = annos.selected_mask();
169                    let labels = br_data.label_info.labels();
170                    let annos = annos
171                        .elts()
172                        .iter()
173                        .zip(cats.iter())
174                        .zip(selected_mask.iter())
175                        .filter(|((_, cat_idx), _)| {
176                            if let Some(only_cat_idx) = only_cat_idx {
177                                **cat_idx == only_cat_idx
178                            } else {
179                                true
180                            }
181                        })
182                        .map(|((brush_line, cat_idx), is_selected)| {
183                            Annotation::Brush(BrushAnnotation {
184                                canvas: brush_line.clone(),
185                                color: colors[*cat_idx],
186                                label: Some(labels[*cat_idx].clone()),
187                                is_selected: Some(*is_selected),
188                                fill_alpha: br_data.options.fill_alpha,
189                                instance_display_label: br_data.options.core.instance_label_display,
190                            })
191                        })
192                        .collect::<Vec<Annotation>>();
193                    Some(annos)
194                } else {
195                    Some(vec![])
196                }
197            }
198            _ => None,
199        }
200    }
201}
202impl Default for ToolSpecifics {
203    fn default() -> Self {
204        ToolSpecifics::Bbox(BboxToolData::default())
205    }
206}
207
208#[derive(Serialize, Deserialize, Clone, Debug, Default, PartialEq)]
209pub struct ToolsData {
210    pub specifics: ToolSpecifics,
211    pub menu_active: bool,
212    #[serde(default)]
213    pub visible_inactive_tools: VisibleInactiveToolsState,
214}
215impl ToolsData {
216    pub fn new(
217        specifics: ToolSpecifics,
218        visible_inactive_tools: VisibleInactiveToolsState,
219    ) -> Self {
220        ToolsData {
221            specifics,
222            menu_active: false,
223            visible_inactive_tools,
224        }
225    }
226}
227
228#[macro_export]
229macro_rules! toolsdata_by_name {
230    ($name:expr, $acc:ident, $tdm:expr) => {
231        $tdm.get_mut($name)
232            .ok_or(rvimage_domain::rverr!("{} is not a tool", $name))?
233            .specifics
234            .$acc()?
235    };
236}
237
238#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
239pub struct ToolsDataMap {
240    // tool name -> tool's menu data type
241    #[serde(flatten)]
242    data: HashMap<String, ToolsData>,
243}
244impl ToolsDataMap {
245    pub fn new() -> Self {
246        let tdm = ToolsDataMap {
247            data: HashMap::new(),
248        };
249        add_tools_initial_data(tdm)
250    }
251    pub fn iter(&self) -> impl Iterator<Item = (&String, &ToolsData)> {
252        self.data.iter()
253    }
254    pub fn contains_key(&self, name: &str) -> bool {
255        self.data.contains_key(name)
256    }
257    pub fn get_specifics(&self, name: &str) -> Option<&ToolSpecifics> {
258        self.data.get(name).map(|d| &d.specifics)
259    }
260    pub fn get_specifics_mut(&mut self, name: &str) -> Option<&mut ToolSpecifics> {
261        self.data.get_mut(name).map(|d| &mut d.specifics)
262    }
263    pub fn get(&self, name: &str) -> Option<&ToolsData> {
264        self.data.get(name)
265    }
266    pub fn get_mut(&mut self, name: &str) -> Option<&mut ToolsData> {
267        self.data.get_mut(name)
268    }
269    pub fn len(&self) -> usize {
270        self.data.len()
271    }
272    pub fn is_empty(&self) -> bool {
273        self.data.is_empty()
274    }
275    pub fn values_mut(&mut self) -> impl Iterator<Item = &mut ToolsData> {
276        self.data.values_mut()
277    }
278
279    pub fn insert(&mut self, name: String, data: ToolsData) -> Option<ToolsData> {
280        self.data.insert(name, data)
281    }
282    pub fn set_tools_specific_data(&mut self, name: &str, specifics: ToolSpecifics) {
283        self.data.insert(
284            name.to_string(),
285            ToolsData::new(specifics, VisibleInactiveToolsState::default()),
286        );
287    }
288}
289impl Default for ToolsDataMap {
290    fn default() -> Self {
291        Self::new()
292    }
293}
294impl Index<&str> for ToolsDataMap {
295    type Output = ToolsData;
296    fn index(&self, index: &str) -> &Self::Output {
297        &self.data[index]
298    }
299}
300impl FromIterator<(String, ToolsData)> for ToolsDataMap {
301    fn from_iter<T: IntoIterator<Item = (std::string::String, ToolsData)>>(iter: T) -> Self {
302        let data = iter.into_iter().collect::<HashMap<String, ToolsData>>();
303        add_tools_initial_data(ToolsDataMap { data })
304    }
305}
306impl From<HashMap<String, ToolsData>> for ToolsDataMap {
307    fn from(data: HashMap<String, ToolsData>) -> Self {
308        add_tools_initial_data(ToolsDataMap { data })
309    }
310}
311
312#[macro_export]
313macro_rules! get_specifics_from_tdm {
314    ($actor_name:expr, $tdm:expr, $access_func:ident) => {
315        $tdm.get($actor_name)
316            .and_then(|x| x.specifics.$access_func().ok())
317    };
318}
319#[macro_export]
320macro_rules! get_annos_from_tdm {
321    ($actor_name:expr, $tdm:expr, $current_file_path:expr, $access_func:ident) => {
322        $crate::get_specifics_from_tdm!($actor_name, $tdm, $access_func)
323            .and_then(|d| d.get_annos($current_file_path))
324    };
325}
326
327#[macro_export]
328macro_rules! get_labelinfo_from_tdm {
329    ($actor_name:expr, $tdm:expr,  $access_func:ident) => {
330        $crate::get_specifics_from_tdm!($actor_name, $tdm, $access_func).map(|d| d.label_info())
331    };
332}
333
334#[cfg(test)]
335use crate::tools::{
336    ALWAYS_ACTIVE_ZOOM, ATTRIBUTES_NAME, BBOX_NAME, BRUSH_NAME, ROT90_NAME, ZOOM_NAME,
337};
338#[test]
339fn test_tools_data_map() {
340    let tdm = ToolsDataMap::new();
341    let tools = [
342        BBOX_NAME,
343        ROT90_NAME,
344        BRUSH_NAME,
345        ATTRIBUTES_NAME,
346        ZOOM_NAME,
347        ALWAYS_ACTIVE_ZOOM,
348    ];
349    for tool in tools.iter() {
350        assert!(tdm.contains_key(tool));
351    }
352    assert_eq!(tdm.len(), tools.len());
353
354    // test from hashmap
355    let data = HashMap::from([(BBOX_NAME.to_string(), ToolsData::default())]);
356    let tdm = ToolsDataMap::from(data);
357    for tool in tools.iter() {
358        assert!(tdm.contains_key(tool));
359    }
360    assert_eq!(tdm.len(), tools.len());
361
362    // test from iterator
363    let data = vec![
364        (BBOX_NAME.to_string(), ToolsData::default()),
365        (ROT90_NAME.to_string(), ToolsData::default()),
366    ];
367    let tdm = ToolsDataMap::from_iter(data);
368    for tool in tools.iter() {
369        assert!(tdm.contains_key(tool));
370    }
371    assert_eq!(tdm.len(), tools.len());
372}