1use crate::{
2 drawme::{Annotation, BboxAnnotation, Stroke},
3 BrushAnnotation,
4};
5
6pub use self::core::{
7 vis_from_lfoption, Annotate, ExportAsCoco, ImportExportTrigger, ImportMode, InstanceAnnotate,
8 InstanceExportData, LabelInfo, Options, VisibleInactiveToolsState,
9 OUTLINE_THICKNESS_CONVERSION,
10};
11pub use self::{
12 attributes_data::AttributesToolData, bbox_data::BboxToolData, brush_data::BrushToolData,
13 coco_io::write_coco, rot90_data::Rot90ToolData,
14};
15use rvimage_domain::{rverr, RvResult, TPtF};
16use serde::{Deserialize, Serialize};
17pub mod annotations;
18pub mod attributes_data;
19pub mod bbox_data;
20pub mod brush_data;
21pub mod coco_io;
22mod core;
23mod label_map;
24pub mod rot90_data;
25pub use core::{merge, AnnotationsMap, InstanceLabelDisplay, Options as CoreOptions};
26use std::collections::HashMap;
27
28macro_rules! variant_access {
29 ($variant:ident, $func_name:ident, $self:ty, $return_type:ty) => {
30 pub fn $func_name(self: $self) -> rvimage_domain::RvResult<$return_type> {
31 match self {
32 ToolSpecifics::$variant(x) => Ok(x),
33 _ => Err(rvimage_domain::rverr!(
34 "this is not a {}",
35 stringify!($variant)
36 )),
37 }
38 }
39 };
40}
41macro_rules! variant_access_free {
42 ($variant:ident, $func_name:ident, $lt:lifetime, $ToolsSpecific:ty, $return_type:ty) => {
43 pub fn $func_name<$lt>(x: $ToolsSpecific) -> rvimage_domain::RvResult<$return_type> {
44 match x {
45 ToolSpecifics::$variant(x) => Ok(x),
46 _ => Err(rvimage_domain::rverr!(
47 "this is not a {}",
48 stringify!($variant)
49 )),
50 }
51 }
52 };
53}
54
55variant_access_free!(Bbox, bbox, 'a, &'a ToolSpecifics, &'a BboxToolData);
56variant_access_free!(Bbox, bbox_mut, 'a, &'a mut ToolSpecifics, &'a mut BboxToolData);
57variant_access_free!(Brush, brush, 'a, &'a ToolSpecifics, &'a BrushToolData);
58variant_access_free!(Brush, brush_mut, 'a, &'a mut ToolSpecifics, &'a mut BrushToolData);
59variant_access_free!(Attributes, attributes, 'a, &'a ToolSpecifics, &'a AttributesToolData);
60variant_access_free!(Attributes, attributes_mut, 'a, &'a mut ToolSpecifics, &'a mut AttributesToolData);
61
62#[derive(Serialize, Deserialize, Clone, Debug, PartialEq)]
63#[allow(clippy::large_enum_variant)]
64pub enum ToolSpecifics {
65 Bbox(BboxToolData),
66 Brush(BrushToolData),
67 Rot90(Rot90ToolData),
68 Zoom(()),
69 AlwaysActiveZoom(()),
70 Attributes(AttributesToolData),
71}
72impl ToolSpecifics {
73 variant_access!(Bbox, bbox, &Self, &BboxToolData);
74 variant_access!(Brush, brush, &Self, &BrushToolData);
75 variant_access!(Rot90, rot90, &Self, &Rot90ToolData);
76 variant_access!(Attributes, attributes, &Self, &AttributesToolData);
77 variant_access!(Bbox, bbox_mut, &mut Self, &mut BboxToolData);
78 variant_access!(Brush, brush_mut, &mut Self, &mut BrushToolData);
79 variant_access!(Rot90, rot90_mut, &mut Self, &mut Rot90ToolData);
80 variant_access!(
81 Attributes,
82 attributes_mut,
83 &mut Self,
84 &mut AttributesToolData
85 );
86
87 pub fn apply_mut<T>(
88 &mut self,
89 f_bbox: impl FnOnce(&mut BboxToolData) -> RvResult<T>,
90 f_brush: impl FnOnce(&mut BrushToolData) -> RvResult<T>,
91 f_attr: impl FnOnce(&mut AttributesToolData) -> RvResult<T>,
92 ) -> RvResult<T> {
93 match self {
94 Self::Bbox(bbox_data) => f_bbox(bbox_data),
95 Self::Brush(brush_data) => f_brush(brush_data),
96 Self::Attributes(attr_data) => f_attr(attr_data),
97 _ => Err(rverr!("only brush tool and bbox tool can be used in apply")),
98 }
99 }
100 pub fn apply<T>(
101 &self,
102 f_bbox: impl FnOnce(&BboxToolData) -> RvResult<T>,
103 f_brush: impl FnOnce(&BrushToolData) -> RvResult<T>,
104 ) -> RvResult<T> {
105 match self {
106 Self::Bbox(bbox_data) => f_bbox(bbox_data),
107 Self::Brush(brush_data) => f_brush(brush_data),
108 _ => Err(rverr!("only brush tool and bbox tool can be used in apply")),
109 }
110 }
111
112 pub fn to_annotations_view(
113 &self,
114 file_path_relative: &str,
115 only_cat_idx: Option<usize>,
116 ) -> Option<Vec<Annotation>> {
117 match self {
118 ToolSpecifics::Bbox(bb_data) => {
119 if let Some(annos) = bb_data.get_annos(file_path_relative) {
120 let geos = annos.elts();
121 let cats = annos.cat_idxs();
122 let selected_bbs = annos.selected_mask();
123 let labels = bb_data.label_info.labels();
124 let colors = bb_data.label_info.colors();
125 let bbs_colored = geos
126 .iter()
127 .zip(cats.iter())
128 .zip(selected_bbs.iter())
129 .filter(|((_, cat_idx), _)| {
130 if let Some(only_cat_idx) = only_cat_idx {
131 **cat_idx == only_cat_idx
132 } else {
133 true
134 }
135 })
136 .map(|((geo, cat_idx), is_selected)| {
137 Annotation::Bbox(BboxAnnotation {
138 geofig: geo.clone(),
139 fill_color: Some(colors[*cat_idx]),
140 fill_alpha: bb_data.options.fill_alpha,
141 label: Some(labels[*cat_idx].clone()),
142 outline: Stroke {
143 thickness: TPtF::from(bb_data.options.outline_thickness)
144 / OUTLINE_THICKNESS_CONVERSION,
145 color: colors[*cat_idx],
146 },
147 outline_alpha: bb_data.options.outline_alpha,
148 is_selected: Some(*is_selected),
149 highlight_circles: bb_data.highlight_circles.clone(),
150 instance_label_display: bb_data.options.core.instance_label_display,
151 })
152 })
153 .collect::<Vec<Annotation>>();
154 Some(bbs_colored)
155 } else {
156 Some(vec![])
157 }
158 }
159 ToolSpecifics::Brush(br_data) => {
160 if let Some(annos) = br_data.get_annos(file_path_relative) {
161 let colors = br_data.label_info.colors();
162 let cats = annos.cat_idxs();
163 let selected_mask = annos.selected_mask();
164 let labels = br_data.label_info.labels();
165 let annos = annos
166 .elts()
167 .iter()
168 .zip(cats.iter())
169 .zip(selected_mask.iter())
170 .filter(|((_, cat_idx), _)| {
171 if let Some(only_cat_idx) = only_cat_idx {
172 **cat_idx == only_cat_idx
173 } else {
174 true
175 }
176 })
177 .map(|((brush_line, cat_idx), is_selected)| {
178 Annotation::Brush(BrushAnnotation {
179 canvas: brush_line.clone(),
180 color: colors[*cat_idx],
181 label: Some(labels[*cat_idx].clone()),
182 is_selected: Some(*is_selected),
183 fill_alpha: br_data.options.fill_alpha,
184 instance_display_label: br_data.options.core.instance_label_display,
185 })
186 })
187 .collect::<Vec<Annotation>>();
188 Some(annos)
189 } else {
190 Some(vec![])
191 }
192 }
193 _ => None,
194 }
195 }
196}
197impl Default for ToolSpecifics {
198 fn default() -> Self {
199 ToolSpecifics::Bbox(BboxToolData::default())
200 }
201}
202
203#[derive(Serialize, Deserialize, Clone, Debug, Default, PartialEq)]
204pub struct ToolsData {
205 pub specifics: ToolSpecifics,
206 pub menu_active: bool,
207 #[serde(default)]
208 pub visible_inactive_tools: VisibleInactiveToolsState,
209}
210impl ToolsData {
211 pub fn new(
212 specifics: ToolSpecifics,
213 visible_inactive_tools: VisibleInactiveToolsState,
214 ) -> Self {
215 ToolsData {
216 specifics,
217 menu_active: false,
218 visible_inactive_tools,
219 }
220 }
221}
222
223#[macro_export]
224macro_rules! toolsdata_by_name {
225 ($name:expr, $acc:ident, $tdm:expr) => {
226 $tdm.get_mut($name)
227 .ok_or(rvimage_domain::rverr!("{} is not a tool", $name))?
228 .specifics
229 .$acc()?
230 };
231}
232
233pub type ToolsDataMap = HashMap<String, ToolsData>;
235
236#[macro_export]
237macro_rules! get_specifics_mut_from_tdm {
238 ($actor_name:expr, $tdm:expr, $access_func:ident) => {
239 $tdm.get_mut($actor_name)
240 .and_then(|x| x.specifics.$access_func().ok())
241 };
242}
243#[macro_export]
244macro_rules! get_specifics_from_tdm {
245 ($actor_name:expr, $tdm:expr, $access_func:ident) => {
246 $tdm.get($actor_name)
247 .and_then(|x| x.specifics.$access_func().ok())
248 };
249}
250#[macro_export]
251macro_rules! get_annos_from_tdm {
252 ($actor_name:expr, $tdm:expr, $current_file_path:expr, $access_func:ident) => {
253 $crate::get_specifics_from_tdm!($actor_name, $tdm, $access_func)
254 .and_then(|d| d.get_annos($current_file_path))
255 };
256}
257
258#[macro_export]
259macro_rules! get_labelinfo_from_tdm {
260 ($actor_name:expr, $tdm:expr, $access_func:ident) => {
261 $crate::get_specifics_from_tdm!($actor_name, $tdm, $access_func).map(|d| d.label_info())
262 };
263}