rvlib/tools_data/
bbox_data.rs

1use std::iter;
2
3use serde::{Deserialize, Serialize};
4
5#[cfg(test)]
6use super::annotations::InstanceAnnotations;
7use super::{
8    annotations::{BboxAnnotations, ClipboardData},
9    core::{
10        AccessInstanceData, AnnotationsMap, CocoSegmentation, ExportAsCoco, InstanceExportData,
11        LabelInfo, OUTLINE_THICKNESS_CONVERSION,
12    },
13    predictive_labeling::PredictiveLabelingData,
14    InstanceAnnotate,
15};
16use crate::{
17    cfg::ExportPath,
18    file_util, implement_annotate, implement_annotations_getters,
19    tools_data::{annotations::SplitMode, core},
20    GeoFig,
21};
22use rvimage_domain::{rverr, BbF, Circle, PtF, RvResult, ShapeI, TPtF};
23
24/// filename -> (annotations per file, file dimensions)
25pub type BboxAnnoMap = AnnotationsMap<GeoFig>;
26
27#[derive(Clone, Copy, Deserialize, Serialize, Debug, PartialEq, Eq)]
28pub struct Options {
29    #[serde(skip)]
30    pub core: core::Options,
31    pub split_mode: SplitMode,
32    pub fill_alpha: u8,
33    pub outline_alpha: u8,
34    pub outline_thickness: u16,
35    pub drawing_distance: u8,
36}
37impl Default for Options {
38    fn default() -> Self {
39        Self {
40            core: core::Options::default(),
41            split_mode: SplitMode::default(),
42            fill_alpha: 30,
43            outline_alpha: 255,
44            outline_thickness: OUTLINE_THICKNESS_CONVERSION as u16,
45            drawing_distance: 10,
46        }
47    }
48}
49#[derive(Deserialize, Serialize, Clone, Debug, PartialEq)]
50pub struct BboxToolData {
51    pub label_info: LabelInfo,
52    pub annotations_map: BboxAnnoMap,
53    #[serde(skip)]
54    pub clipboard: Option<ClipboardData<GeoFig>>,
55    pub options: Options,
56    pub coco_file: ExportPath,
57    #[serde(skip)]
58    pub highlight_circles: Vec<Circle>,
59    #[serde(default)]
60    pub predictive_labeling_data: PredictiveLabelingData,
61}
62
63impl BboxToolData {
64    implement_annotations_getters!(BboxAnnotations);
65
66    pub fn separate_data(self) -> (LabelInfo, BboxAnnoMap, ExportPath) {
67        (self.label_info, self.annotations_map, self.coco_file)
68    }
69
70    pub fn from_coco_export_data(input_data: InstanceExportData<GeoFig>) -> RvResult<Self> {
71        let label_info = input_data.label_info()?;
72        let mut out_data = Self {
73            label_info,
74            annotations_map: AnnotationsMap::new(),
75            clipboard: None,
76            options: Options {
77                core: core::Options {
78                    visible: true,
79                    ..Default::default()
80                },
81                ..Default::default()
82            },
83            coco_file: input_data.coco_file,
84            highlight_circles: vec![],
85            predictive_labeling_data: PredictiveLabelingData::default(),
86        };
87        out_data.set_annotations_map(
88            input_data
89                .annotations
90                .into_iter()
91                .map(|(s, (geos, cat_ids, dims))| {
92                    (s, (BboxAnnotations::from_elts_cats(geos, cat_ids), dims))
93                })
94                .collect(),
95        )?;
96        Ok(out_data)
97    }
98
99    pub fn retain_fileannos_in_folder(&mut self, folder: &str) {
100        self.annotations_map
101            .retain(|f, _| file_util::url_encode(f).starts_with(folder));
102    }
103
104    pub fn new() -> Self {
105        let label_info = LabelInfo::default();
106
107        BboxToolData {
108            label_info,
109            annotations_map: AnnotationsMap::new(),
110            clipboard: None,
111            options: Options {
112                core: core::Options {
113                    visible: true,
114                    ..Default::default()
115                },
116                ..Default::default()
117            },
118            coco_file: ExportPath::default(),
119            highlight_circles: vec![],
120            predictive_labeling_data: PredictiveLabelingData::default(),
121        }
122    }
123
124    pub fn set_annotations_map(&mut self, map: BboxAnnoMap) -> RvResult<()> {
125        for (_, (annos, _)) in map.iter() {
126            for cat_idx in annos.cat_idxs() {
127                let len = self.label_info.len();
128                if *cat_idx >= len {
129                    return Err(rverr!(
130                        "cat idx {} does not have a label, out of bounds, {}",
131                        cat_idx,
132                        len
133                    ));
134                }
135            }
136        }
137        self.annotations_map = map;
138        Ok(())
139    }
140}
141
142implement_annotate!(BboxToolData);
143
144impl Default for BboxToolData {
145    fn default() -> Self {
146        Self::new()
147    }
148}
149
150impl AccessInstanceData<GeoFig> for BboxToolData {
151    fn annotations_map(&self) -> &AnnotationsMap<GeoFig> {
152        &self.annotations_map
153    }
154    fn label_info(&self) -> &LabelInfo {
155        &self.label_info
156    }
157}
158
159impl ExportAsCoco<GeoFig> for BboxToolData {
160    fn cocofile_conn(&self) -> ExportPath {
161        self.coco_file.clone()
162    }
163    fn separate_data(self) -> (core::Options, LabelInfo, AnnotationsMap<GeoFig>, ExportPath) {
164        (
165            self.options.core,
166            self.label_info,
167            self.annotations_map,
168            self.coco_file,
169        )
170    }
171    #[cfg(test)]
172    fn anno_iter(&self) -> impl Iterator<Item = (&String, &(InstanceAnnotations<GeoFig>, ShapeI))> {
173        self.anno_iter()
174    }
175    fn core_options_mut(&mut self) -> &mut core::Options {
176        &mut self.options.core
177    }
178    fn new(
179        options: core::Options,
180        label_info: LabelInfo,
181        anno_map: AnnotationsMap<GeoFig>,
182        export_path: ExportPath,
183    ) -> Self {
184        Self {
185            label_info,
186            annotations_map: anno_map,
187            clipboard: None,
188            options: Options {
189                core: options,
190                ..Default::default()
191            },
192            coco_file: export_path,
193            highlight_circles: vec![],
194            predictive_labeling_data: PredictiveLabelingData::default(),
195        }
196    }
197    fn set_annotations_map(&mut self, map: AnnotationsMap<GeoFig>) -> RvResult<()> {
198        self.set_annotations_map(map)
199    }
200    fn set_labelinfo(&mut self, info: LabelInfo) {
201        self.label_info = info;
202    }
203}
204
205impl InstanceAnnotate for GeoFig {
206    fn rot90_with_image_ntimes(self, shape: ShapeI, n: u8) -> RvResult<Self> {
207        Ok(match self {
208            Self::BB(bb) => Self::BB(bb.rot90_with_image_ntimes(shape, n)),
209            Self::Poly(poly) => Self::Poly(poly.rot90_with_image_ntimes(shape, n)),
210        })
211    }
212    fn is_contained_in_image(&self, shape: ShapeI) -> bool {
213        match self {
214            Self::BB(bb) => bb.is_contained_in_image(shape),
215            Self::Poly(poly) => poly.is_contained_in_image(shape),
216        }
217    }
218    fn enclosing_bb(&self) -> BbF {
219        match self {
220            Self::BB(bb) => *bb,
221            Self::Poly(poly) => poly.enclosing_bb(),
222        }
223    }
224    fn contains<P>(&self, point: P) -> bool
225    where
226        P: Into<PtF>,
227    {
228        match self {
229            Self::BB(bb) => bb.contains(point.into()),
230            Self::Poly(poly) => poly.contains(point),
231        }
232    }
233    fn dist_to_boundary(&self, point: PtF) -> TPtF {
234        match self {
235            Self::BB(bb) => bb.distance_to_boundary(point),
236            Self::Poly(poly) => poly.distance_to_boundary(point),
237        }
238    }
239    fn to_cocoseg(
240        &self,
241        shape_im: ShapeI,
242        is_export_absolute: bool,
243    ) -> RvResult<Option<core::CocoSegmentation>> {
244        Ok(Some(CocoSegmentation::Polygon(vec![
245            if is_export_absolute {
246                self.points()
247            } else {
248                self.points_normalized(TPtF::from(shape_im.w), TPtF::from(shape_im.h))
249            }
250            .iter()
251            .flat_map(|p| iter::once(p.x).chain(iter::once(p.y)))
252            .collect::<Vec<_>>(),
253        ])))
254    }
255}