1use serde::de::DeserializeOwned;
2use serde::{Deserialize, Serialize};
3use std::collections::HashMap;
4use std::fmt::Debug;
5use tracing::info;
6
7use crate::{cfg::ExportPath, util::Visibility, ShapeI};
8use rvimage_domain::{rverr, RvResult};
9use rvimage_domain::{BbF, PtF, TPtF, TPtI};
10
11use super::annotations::InstanceAnnotations;
12use super::label_map::LabelMap;
13
14pub const OUTLINE_THICKNESS_CONVERSION: TPtF = 10.0;
15
16const DEFAULT_LABEL: &str = "rvimage_fg";
17
18fn color_dist(c1: [u8; 3], c2: [u8; 3]) -> f32 {
19 let square_d = |i| (f32::from(c1[i]) - f32::from(c2[i])).powi(2);
20 (square_d(0) + square_d(1) + square_d(2)).sqrt()
21}
22
23#[derive(Default, Clone, Copy, Debug, PartialEq, Eq)]
24pub enum ImportMode {
25 Merge,
26 #[default]
27 Replace,
28}
29
30#[derive(Default, Clone, Copy, Debug, PartialEq, Eq)]
31pub struct ImportExportTrigger {
32 export_triggered: bool,
33 import_triggered: bool,
34 import_mode: ImportMode,
35}
36impl ImportExportTrigger {
37 pub fn import_triggered(self) -> bool {
38 self.import_triggered
39 }
40 pub fn import_mode(self) -> ImportMode {
41 self.import_mode
42 }
43 pub fn export_triggered(self) -> bool {
44 self.export_triggered
45 }
46 pub fn untrigger_export(&mut self) {
47 self.export_triggered = false;
48 }
49 pub fn untrigger_import(&mut self) {
50 self.import_triggered = false;
51 }
52 pub fn trigger_export(&mut self) {
53 self.export_triggered = true;
54 }
55 pub fn trigger_import(&mut self) {
56 self.import_triggered = true;
57 }
58 pub fn use_merge_import(&mut self) {
59 self.import_mode = ImportMode::Merge;
60 }
61 pub fn use_replace_import(&mut self) {
62 self.import_mode = ImportMode::Replace;
63 }
64 pub fn merge_mode(self) -> bool {
65 self.import_mode == ImportMode::Merge
66 }
67 pub fn from_export_triggered(export_triggered: bool) -> Self {
68 Self {
69 export_triggered,
70 ..Default::default()
71 }
72 }
73}
74
75pub type AnnotationsMap<T> = LabelMap<InstanceAnnotations<T>>;
76
77#[allow(clippy::struct_excessive_bools)]
78#[derive(Clone, Copy, Debug, PartialEq, Eq)]
79pub struct Options {
80 pub visible: bool,
81 pub is_colorchange_triggered: bool,
82 pub is_redraw_annos_triggered: bool,
83 pub is_export_absolute: bool,
84 pub import_export_trigger: ImportExportTrigger,
85 pub is_history_update_triggered: bool,
86 pub track_changes: bool,
87 pub erase: bool,
88 pub label_propagation: Option<usize>,
89 pub label_deletion: Option<usize>,
90 pub auto_paste: bool,
91}
92impl Default for Options {
93 fn default() -> Self {
94 Self {
95 visible: true,
96 is_colorchange_triggered: false,
97 is_redraw_annos_triggered: false,
98 is_export_absolute: false,
99 import_export_trigger: ImportExportTrigger::default(),
100 is_history_update_triggered: false,
101 track_changes: false,
102 erase: false,
103 label_propagation: None,
104 label_deletion: None,
105 auto_paste: false,
106 }
107 }
108}
109impl Options {
110 pub fn trigger_redraw_and_hist(mut self) -> Self {
111 self.is_history_update_triggered = true;
112 self.is_redraw_annos_triggered = true;
113 self
114 }
115}
116
117const N: usize = 1;
118#[derive(Serialize, Deserialize, Clone, Debug, Default, PartialEq)]
119pub struct VisibleInactiveToolsState {
120 show_mask: [bool; N],
122}
123impl VisibleInactiveToolsState {
124 pub fn new() -> Self {
125 Self::default()
126 }
127 #[allow(clippy::needless_lifetimes)]
128 pub fn iter<'a>(&'a self) -> impl Iterator<Item = bool> + 'a {
129 self.show_mask.iter().copied()
130 }
131 pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut bool> {
132 self.show_mask.iter_mut()
133 }
134 pub fn hide_all(&mut self) {
135 for show in &mut self.show_mask {
136 *show = false;
137 }
138 }
139 pub fn set_show(&mut self, idx: usize, is_visible: bool) {
140 self.show_mask[idx] = is_visible;
141 }
142}
143
144pub fn random_clr() -> [u8; 3] {
145 let r = rand::random::<u8>();
146 let g = rand::random::<u8>();
147 let b = rand::random::<u8>();
148 [r, g, b]
149}
150
151fn argmax_clr_dist(picklist: &[[u8; 3]], legacylist: &[[u8; 3]]) -> [u8; 3] {
152 let (idx, _) = picklist
153 .iter()
154 .enumerate()
155 .map(|(i, pickclr)| {
156 let min_dist = legacylist
157 .iter()
158 .map(|legclr| color_dist(*legclr, *pickclr))
159 .min_by(|a, b| a.partial_cmp(b).unwrap())
160 .unwrap_or(0.0);
161 (i, min_dist)
162 })
163 .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
164 .unwrap();
165 picklist[idx]
166}
167
168pub fn new_color(colors: &[[u8; 3]]) -> [u8; 3] {
169 let mut new_clr_proposals = [[0u8, 0u8, 0u8]; 10];
170 for new_clr in &mut new_clr_proposals {
171 *new_clr = random_clr();
172 }
173 argmax_clr_dist(&new_clr_proposals, colors)
174}
175
176pub fn new_random_colors(n: usize) -> Vec<[u8; 3]> {
177 let mut colors = vec![random_clr()];
178 for _ in 0..(n - 1) {
179 let color = new_color(&colors);
180 colors.push(color);
181 }
182 colors
183}
184
185fn get_visibility(visible: bool, show_only_current: bool, cat_idx_current: usize) -> Visibility {
186 if visible && show_only_current {
187 Visibility::Only(cat_idx_current)
188 } else if visible {
189 Visibility::All
190 } else {
191 Visibility::None
192 }
193}
194
195pub fn vis_from_lfoption(label_info: Option<&LabelInfo>, visible: bool) -> Visibility {
196 if let Some(label_info) = label_info {
197 label_info.visibility(visible)
198 } else if visible {
199 Visibility::All
200 } else {
201 Visibility::None
202 }
203}
204
205pub fn merge<T>(
206 annos1: AnnotationsMap<T>,
207 li1: LabelInfo,
208 annos2: AnnotationsMap<T>,
209 li2: LabelInfo,
210) -> (AnnotationsMap<T>, LabelInfo)
211where
212 T: InstanceAnnotate,
213{
214 let (li, idx_map) = li1.merge(li2);
215 let mut annotations_map = annos1;
216
217 for (k, (v2, s)) in annos2 {
218 if let Some((v1, _)) = annotations_map.get_mut(&k) {
219 let (elts, cat_idxs, _) = v2.separate_data();
220 v1.extend(
221 elts.into_iter(),
222 cat_idxs.into_iter().map(|old_idx| idx_map[old_idx]),
223 s,
224 );
225 v1.deselect_all();
226 } else {
227 let (elts, cat_idxs, _) = v2.separate_data();
228 let cat_idxs = cat_idxs
229 .into_iter()
230 .map(|old_idx| idx_map[old_idx])
231 .collect::<Vec<_>>();
232 let v2 = InstanceAnnotations::new_relaxed(elts, cat_idxs);
233 annotations_map.insert(k, (v2, s));
234 }
235 }
236 (annotations_map, li)
237}
238
239#[derive(Deserialize, Serialize, Clone, Debug, PartialEq, Eq)]
240pub struct LabelInfo {
241 pub new_label: String,
242 labels: Vec<String>,
243 colors: Vec<[u8; 3]>,
244 cat_ids: Vec<u32>,
245 pub cat_idx_current: usize,
246 pub show_only_current: bool,
247}
248impl LabelInfo {
249 pub fn merge(mut self, other: Self) -> (Self, Vec<usize>) {
252 let mut idx_map = vec![];
253 for other_label in other.labels {
254 let self_cat_idx = self.labels.iter().position(|slab| slab == &other_label);
255 if let Some(scidx) = self_cat_idx {
256 idx_map.push(scidx);
257 } else {
258 self.labels.push(other_label);
259 self.colors.push(new_color(&self.colors));
260 self.cat_ids.push(self.labels.len() as u32);
261 idx_map.push(self.labels.len() - 1);
262 }
263 }
264 (self, idx_map)
265 }
266
267 pub fn visibility(&self, visible: bool) -> Visibility {
268 get_visibility(visible, self.show_only_current, self.cat_idx_current)
269 }
270 pub fn new_random_colors(&mut self) {
271 info!("new random colors for annotations");
272 self.colors = new_random_colors(self.colors.len());
273 }
274 pub fn push(
275 &mut self,
276 label: String,
277 color: Option<[u8; 3]>,
278 cat_id: Option<u32>,
279 ) -> RvResult<()> {
280 if self.labels.contains(&label) {
281 Err(rverr!("label '{}' already exists", label))
282 } else {
283 info!("adding label '{label}'");
284 self.labels.push(label);
285 if let Some(clr) = color {
286 if self.colors.contains(&clr) {
287 return Err(rverr!("color '{:?}' already exists", clr));
288 }
289 self.colors.push(clr);
290 } else {
291 let new_clr = new_color(&self.colors);
292 self.colors.push(new_clr);
293 }
294 if let Some(cat_id) = cat_id {
295 if self.cat_ids.contains(&cat_id) {
296 return Err(rverr!("cat id '{:?}' already exists", cat_id));
297 }
298 self.cat_ids.push(cat_id);
299 } else if let Some(max_id) = self.cat_ids.iter().max() {
300 self.cat_ids.push(max_id + 1);
301 } else {
302 self.cat_ids.push(1);
303 }
304 Ok(())
305 }
306 }
307 pub fn rename_label(&mut self, idx: usize, label: String) -> RvResult<()> {
308 if self.labels.contains(&label) {
309 Err(rverr!("label '{label}' already exists"))
310 } else {
311 self.labels[idx] = label;
312 Ok(())
313 }
314 }
315 pub fn from_iter(it: impl Iterator<Item = ((String, [u8; 3]), u32)>) -> RvResult<Self> {
316 let mut info = Self::empty();
317 for ((label, color), cat_id) in it {
318 info.push(label, Some(color), Some(cat_id))?;
319 }
320 Ok(info)
321 }
322 pub fn is_empty(&self) -> bool {
323 self.labels.is_empty()
324 }
325 pub fn len(&self) -> usize {
326 self.labels.len()
327 }
328 pub fn remove(&mut self, idx: usize) -> (String, [u8; 3], u32) {
329 let removed_items = (
330 self.labels.remove(idx),
331 self.colors.remove(idx),
332 self.cat_ids.remove(idx),
333 );
334 info!("label '{}' removed", removed_items.0);
335 removed_items
336 }
337 pub fn find_default(&mut self) -> Option<&mut String> {
338 self.labels.iter_mut().find(|lab| lab == &DEFAULT_LABEL)
339 }
340 pub fn colors(&self) -> &Vec<[u8; 3]> {
341 &self.colors
342 }
343
344 pub fn labels(&self) -> &Vec<String> {
345 &self.labels
346 }
347
348 pub fn cat_ids(&self) -> &Vec<u32> {
349 &self.cat_ids
350 }
351
352 pub fn separate_data(self) -> (Vec<String>, Vec<[u8; 3]>, Vec<u32>) {
353 (self.labels, self.colors, self.cat_ids)
354 }
355
356 pub fn empty() -> Self {
357 Self {
358 new_label: DEFAULT_LABEL.to_string(),
359 labels: vec![],
360 colors: vec![],
361 cat_ids: vec![],
362 cat_idx_current: 0,
363 show_only_current: false,
364 }
365 }
366 pub fn remove_catidx<'a, T>(&mut self, cat_idx: usize, annotaions_map: &mut AnnotationsMap<T>)
367 where
368 T: InstanceAnnotate + PartialEq + Default + 'a,
369 {
370 if self.len() > 1 {
371 self.remove(cat_idx);
372 if self.cat_idx_current >= cat_idx.max(1) {
373 self.cat_idx_current -= 1;
374 }
375 for (anno, _) in annotaions_map.values_mut() {
376 let indices_for_rm = anno
377 .cat_idxs()
378 .iter()
379 .enumerate()
380 .filter(|(_, geo_cat_idx)| **geo_cat_idx == cat_idx)
381 .map(|(idx, _)| idx)
382 .collect::<Vec<_>>();
383 anno.remove_multiple(&indices_for_rm);
384 anno.reduce_cat_idxs(cat_idx);
385 }
386 }
387 }
388}
389
390impl Default for LabelInfo {
391 fn default() -> Self {
392 let new_label = DEFAULT_LABEL.to_string();
393 let new_color = [255, 255, 255];
394 let labels = vec![new_label.clone()];
395 let colors = vec![new_color];
396 let cat_ids = vec![1];
397 Self {
398 new_label,
399 labels,
400 colors,
401 cat_ids,
402 cat_idx_current: 0,
403 show_only_current: false,
404 }
405 }
406}
407
408#[derive(Deserialize, Serialize, Debug, Clone, PartialEq)]
409pub struct InstanceExportData<A> {
410 pub labels: Vec<String>,
411 pub colors: Vec<[u8; 3]>,
412 pub cat_ids: Vec<u32>,
413 pub annotations: HashMap<String, (Vec<A>, Vec<usize>, ShapeI)>,
415 pub coco_file: ExportPath,
416 pub is_export_absolute: bool,
417}
418
419impl<A> InstanceExportData<A>
420where
421 A: InstanceAnnotate,
422{
423 pub fn from_tools_data(
424 options: &Options,
425 label_info: LabelInfo,
426 coco_file: ExportPath,
427 annotations_map: AnnotationsMap<A>,
428 ) -> Self {
429 let is_export_absolute = options.is_export_absolute;
430 let annotations = annotations_map
431 .into_iter()
432 .map(|(filename, (annos, shape))| {
433 let (bbs, labels, _) = annos.separate_data();
434 (filename, (bbs, labels, shape))
435 })
436 .collect::<HashMap<_, _>>();
437 let (labels, colors, cat_ids) = label_info.separate_data();
438 InstanceExportData {
439 labels,
440 colors,
441 cat_ids,
442 annotations,
443 coco_file,
444 is_export_absolute,
445 }
446 }
447 pub fn label_info(&self) -> RvResult<LabelInfo> {
448 LabelInfo::from_iter(
449 self.labels
450 .clone()
451 .into_iter()
452 .zip(self.colors.clone())
453 .zip(self.cat_ids.clone()),
454 )
455 }
456}
457
458#[derive(Serialize, Deserialize, Debug, PartialEq)]
459pub struct CocoRle {
460 pub counts: Vec<TPtI>,
461 pub size: (TPtI, TPtI),
462 pub intensity: Option<TPtF>,
463}
464
465#[derive(Debug, Serialize, Deserialize, PartialEq)]
466#[serde(untagged)]
467pub enum CocoSegmentation {
468 Polygon(Vec<Vec<TPtF>>),
469 Rle(CocoRle),
470}
471
472#[macro_export]
473macro_rules! implement_annotate {
474 ($tooldata:ident) => {
475 impl $crate::tools_data::core::Annotate for $tooldata {
476 fn has_annos(&self, relative_path: &str) -> bool {
477 if let Some(v) = self.get_annos(relative_path) {
478 !v.is_empty()
479 } else {
480 false
481 }
482 }
483 }
484 };
485}
486
487pub trait Annotate {
488 fn has_annos(&self, relative_path: &str) -> bool;
491}
492
493pub trait InstanceAnnotate:
494 Clone + Default + Debug + PartialEq + Serialize + DeserializeOwned
495{
496 fn is_contained_in_image(&self, shape: ShapeI) -> bool;
497 fn contains<P>(&self, point: P) -> bool
498 where
499 P: Into<PtF>;
500 fn dist_to_boundary(&self, p: PtF) -> TPtF;
501 fn rot90_with_image_ntimes(self, shape: ShapeI, n: u8) -> RvResult<Self>;
504 fn enclosing_bb(&self) -> BbF;
505 fn to_cocoseg(
508 &self,
509 shape_im: ShapeI,
510 is_export_absolute: bool,
511 ) -> RvResult<Option<CocoSegmentation>>;
512}
513pub trait ExportAsCoco<A>
514where
515 A: InstanceAnnotate + 'static,
516{
517 fn separate_data(self) -> (Options, LabelInfo, AnnotationsMap<A>, ExportPath);
518 fn cocofile_conn(&self) -> ExportPath;
519 fn label_info(&self) -> &LabelInfo;
520 #[cfg(test)]
521 fn anno_iter(&self) -> impl Iterator<Item = (&String, &(InstanceAnnotations<A>, ShapeI))>;
522 fn set_annotations_map(&mut self, map: AnnotationsMap<A>) -> RvResult<()>;
523 fn set_labelinfo(&mut self, info: LabelInfo);
524 fn core_options_mut(&mut self) -> &mut Options;
525 fn new(
526 options: Options,
527 label_info: LabelInfo,
528 anno_map: AnnotationsMap<A>,
529 export_path: ExportPath,
530 ) -> Self;
531}
532
533#[cfg(test)]
534use crate::tools_data::brush_data;
535#[cfg(test)]
536use rvimage_domain::{BrushLine, Canvas, Line};
537#[test]
538fn test_argmax() {
539 let picklist = [
540 [200, 200, 200u8],
541 [1, 7, 3],
542 [0, 0, 1],
543 [45, 43, 52],
544 [1, 10, 15],
545 ];
546 let legacylist = [
547 [17, 16, 15],
548 [199, 199, 201u8],
549 [50, 50, 50u8],
550 [255, 255, 255u8],
551 ];
552 assert_eq!(argmax_clr_dist(&picklist, &legacylist), [0, 0, 1]);
553}
554
555#[test]
556fn test_labelinfo_merge() {
557 let li1 = LabelInfo::default();
558 let mut li2 = LabelInfo::default();
559 li2.new_random_colors();
560 let (mut li_merged, _) = li1.clone().merge(li2);
561 assert_eq!(li1, li_merged);
562 li_merged
563 .push("new_label".into(), Some([0, 0, 1]), None)
564 .unwrap();
565 let (li_merged, _) = li_merged.merge(li1);
566 let li_reference = LabelInfo {
567 new_label: DEFAULT_LABEL.to_string(),
568 labels: vec![DEFAULT_LABEL.to_string(), "new_label".to_string()],
569 colors: vec![[255, 255, 255], [0, 0, 1]],
570 cat_ids: vec![1, 2],
571 cat_idx_current: 0,
572 show_only_current: false,
573 };
574 assert_eq!(li_merged, li_reference);
575 assert_eq!(li_merged.clone().merge(li_merged.clone()).0, li_reference);
576 let li = LabelInfo {
577 new_label: DEFAULT_LABEL.to_string(),
578 labels: vec!["somelabel".to_string(), "new_label".to_string()],
579 colors: vec![[255, 255, 255], [0, 1, 1]],
580 cat_ids: vec![1, 2],
581 cat_idx_current: 0,
582 show_only_current: false,
583 };
584 let li_merged_ = li_merged.clone().merge(li.clone());
585 let li_reference = (
586 LabelInfo {
587 new_label: DEFAULT_LABEL.to_string(),
588 labels: vec![
589 DEFAULT_LABEL.to_string(),
590 "new_label".to_string(),
591 "somelabel".to_string(),
592 ],
593 colors: vec![[255, 255, 255], [0, 0, 1], li_merged_.0.colors[2]],
594 cat_ids: vec![1, 2, 3],
595 cat_idx_current: 0,
596 show_only_current: false,
597 },
598 vec![2, 1],
599 );
600 assert_ne!([255, 255, 255], li_merged_.0.colors[2]);
601 assert_eq!(li_merged_, li_reference);
602 let li_merged = li.merge(li_merged);
603 let li_reference = LabelInfo {
604 new_label: DEFAULT_LABEL.to_string(),
605 labels: vec![
606 "somelabel".to_string(),
607 "new_label".to_string(),
608 DEFAULT_LABEL.to_string(),
609 ],
610 colors: vec![[255, 255, 255], [0, 1, 1], li_merged.0.colors[2]],
611 cat_ids: vec![1, 2, 3],
612 cat_idx_current: 0,
613 show_only_current: false,
614 };
615 assert_eq!(li_merged.0, li_reference);
616}
617
618#[test]
619fn test_merge_annos() {
620 let orig_shape = ShapeI::new(100, 100);
621 let li1 = LabelInfo {
622 new_label: "x".to_string(),
623 labels: vec!["somelabel".to_string(), "x".to_string()],
624 colors: vec![[255, 255, 255], [0, 1, 1]],
625 cat_ids: vec![1, 2],
626 cat_idx_current: 0,
627 show_only_current: false,
628 };
629 let li2 = LabelInfo {
630 new_label: "x".to_string(),
631 labels: vec![
632 "somelabel".to_string(),
633 "new_label".to_string(),
634 "x".to_string(),
635 ],
636 colors: vec![[255, 255, 255], [0, 1, 2], [1, 1, 1]],
637 cat_ids: vec![1, 2, 3],
638 cat_idx_current: 0,
639 show_only_current: false,
640 };
641 let mut annos_map1: super::brush_data::BrushAnnoMap = AnnotationsMap::new();
642
643 let mut line = Line::new();
644 line.push(PtF { x: 5.0, y: 5.0 });
645 let anno1 = Canvas::new(
646 &BrushLine {
647 line: line.clone(),
648 thickness: 1.0,
649 intensity: 1.0,
650 },
651 orig_shape,
652 None,
653 )
654 .unwrap();
655 annos_map1.insert(
656 "file1".to_string(),
657 (
658 InstanceAnnotations::new(vec![anno1.clone()], vec![1], vec![true]).unwrap(),
659 orig_shape,
660 ),
661 );
662 let mut annos_map2: brush_data::BrushAnnoMap = AnnotationsMap::new();
663 let anno2 = Canvas::new(
664 &BrushLine {
665 line,
666 thickness: 2.0,
667 intensity: 2.0,
668 },
669 orig_shape,
670 None,
671 )
672 .unwrap();
673
674 annos_map2.insert(
675 "file1".to_string(),
676 (
677 InstanceAnnotations::new(vec![anno2.clone()], vec![1], vec![true]).unwrap(),
678 orig_shape,
679 ),
680 );
681 annos_map2.insert(
682 "file2".to_string(),
683 (
684 InstanceAnnotations::new(vec![anno2.clone()], vec![1], vec![true]).unwrap(),
685 orig_shape,
686 ),
687 );
688 let (merged_map, merged_li) = merge(annos_map1, li1, annos_map2, li2.clone());
689 let merged_li_ref = LabelInfo {
690 new_label: "x".to_string(),
691 labels: vec![
692 "somelabel".to_string(),
693 "x".to_string(),
694 "new_label".to_string(),
695 ],
696 colors: vec![[255, 255, 255], [0, 1, 1], merged_li.colors[2]],
697 cat_ids: vec![1, 2, 3],
698 cat_idx_current: 0,
699 show_only_current: false,
700 };
701
702 assert_eq!(merged_li, merged_li_ref);
703 let map_ref = [
704 (
705 "file1".to_string(),
706 (
707 InstanceAnnotations::new(
708 vec![anno1, anno2.clone()],
709 vec![1, 2],
710 vec![false, false],
711 )
712 .unwrap(),
713 orig_shape,
714 ),
715 ),
716 (
717 "file2".to_string(),
718 (
719 InstanceAnnotations::new(vec![anno2], vec![2], vec![false]).unwrap(),
720 orig_shape,
721 ),
722 ),
723 ]
724 .into_iter()
725 .collect::<AnnotationsMap<Canvas>>();
726 for (k, (v, s)) in merged_map.iter() {
727 assert_eq!(map_ref[k].0, *v);
728 assert_eq!(map_ref[k].1, *s);
729 }
730}