rvlib/tools_data/
coco_io.rs

1use std::{
2    collections::HashMap,
3    fmt::Debug,
4    mem,
5    path::{Path, PathBuf},
6    thread::{self, JoinHandle},
7    vec,
8};
9
10use serde::{Deserialize, Serialize};
11use tracing::info;
12
13use crate::{
14    cfg::ExportPath,
15    file_util::{self, path_to_str, PathPair},
16    image_util,
17    meta_data::MetaData,
18    result::trace_ok_warn,
19    util::version_label,
20    GeoFig, Polygon,
21};
22use rvimage_domain::{rle_image_to_bb, rle_to_mask, BbF, Canvas, Point, ShapeI, TPtF};
23use rvimage_domain::{rverr, to_rv, RvError, RvResult};
24
25use super::{
26    annotations::InstanceAnnotations,
27    brush_data::BrushAnnoMap,
28    core::{new_random_colors, CocoSegmentation, ExportAsCoco},
29    BboxToolData, BrushToolData, InstanceAnnotate, InstanceExportData, Rot90ToolData,
30};
31
32#[derive(Serialize, Deserialize, Debug)]
33struct CocoInfo {
34    description: String,
35}
36
37#[derive(Serialize, Deserialize, Debug)]
38struct CocoImage {
39    id: u32,
40    width: u32,
41    height: u32,
42    file_name: String,
43}
44
45#[derive(Serialize, Deserialize, Debug)]
46struct CocoBboxCategory {
47    id: u32,
48    name: String,
49}
50
51#[derive(Serialize, Deserialize, Debug)]
52struct CocoAnnotation {
53    id: u32,
54    image_id: u32,
55    category_id: u32,
56    bbox: [TPtF; 4],
57    segmentation: Option<CocoSegmentation>,
58    area: Option<TPtF>,
59}
60
61fn colors_to_string(colors: &[[u8; 3]]) -> Option<String> {
62    colors
63        .iter()
64        .map(|[r, g, b]| format!("{r};{g};{b}"))
65        .reduce(|s1, s2| format!("{s1}_{s2}"))
66}
67
68fn string_to_colors(s: &str) -> RvResult<Vec<[u8; 3]>> {
69    let make_err = || rverr!("cannot convert str {} to rgb", s);
70    s.trim()
71        .split('_')
72        .map(|rgb_str| {
73            let mut rgb = [0; 3];
74            let mut it = rgb_str.split(';');
75            for c in &mut rgb {
76                *c = it
77                    .next()
78                    .and_then(|s| s.parse().ok())
79                    .ok_or_else(make_err)?;
80            }
81            Ok(rgb)
82        })
83        .collect::<RvResult<Vec<[u8; 3]>>>()
84}
85
86fn get_n_rotations(rotation_data: Option<&Rot90ToolData>, file_path: &str) -> u8 {
87    rotation_data
88        .and_then(|d| d.get_annos(file_path))
89        .map_or(0, |n_rot| n_rot.to_num())
90}
91
92fn insert_elt<A>(
93    elt: A,
94    annos: &mut HashMap<String, (Vec<A>, Vec<usize>, ShapeI)>,
95    cat_idx: usize,
96    n_rotations: u8,
97    path_as_key: String,
98    shape_coco: ShapeI,
99) where
100    A: InstanceAnnotate,
101{
102    let geo = trace_ok_warn(elt.rot90_with_image_ntimes(shape_coco, n_rotations));
103    if let Some(geo) = geo {
104        if let Some(annos_of_image) = annos.get_mut(&path_as_key) {
105            annos_of_image.0.push(geo);
106            annos_of_image.1.push(cat_idx);
107        } else {
108            annos.insert(
109                path_as_key,
110                (
111                    vec![geo],
112                    vec![cat_idx],
113                    ShapeI::new(shape_coco.w, shape_coco.h),
114                ),
115            );
116        }
117    }
118}
119
120fn instance_to_coco_anno<A>(
121    inst_anno: &A,
122    shape_im_unrotated: ShapeI,
123    n_rotations: u8,
124    is_export_coords_absolute: bool,
125    file_path: &str,
126) -> RvResult<([f64; 4], Option<CocoSegmentation>)>
127where
128    A: InstanceAnnotate,
129{
130    // to store data corresponding to the image on the disk, we need to invert the
131    // applied rotations
132    let n_rots_inverted = (4 - n_rotations) % 4;
133    let shape_rotated = shape_im_unrotated.rot90_with_image_ntimes(n_rotations);
134    let inst_anno = inst_anno
135        .clone()
136        .rot90_with_image_ntimes(shape_rotated, n_rots_inverted)?;
137
138    let bb = inst_anno.enclosing_bb();
139
140    let segmentation = inst_anno.to_cocoseg(shape_im_unrotated, is_export_coords_absolute)?;
141    let (imw, imh) = if is_export_coords_absolute {
142        (1.0, 1.0)
143    } else {
144        (
145            TPtF::from(shape_im_unrotated.w),
146            TPtF::from(shape_im_unrotated.h),
147        )
148    };
149
150    let bb_f = [bb.x / imw, bb.y / imh, bb.w / imw, bb.h / imh];
151    if bb_f[1] * bb_f[2] < 1e-6 {
152        tracing::warn!("annotation in {file_path} has no area {bb:?}.");
153    }
154    Ok((bb_f, segmentation))
155}
156
157struct WarnerCounting<'a> {
158    n_warnings: usize,
159    n_max_warnings: usize,
160    suppressing: bool,
161    suppress_str: &'a str,
162}
163impl<'a> WarnerCounting<'a> {
164    fn new(n_max_warnings: usize, suppress_str: &'a str) -> Self {
165        Self {
166            n_warnings: 0,
167            n_max_warnings,
168            suppressing: false,
169            suppress_str,
170        }
171    }
172    fn warn_str<'b>(&mut self, msg: &'b str) -> Option<&'b str>
173    where
174        'a: 'b,
175    {
176        if self.n_warnings < self.n_max_warnings {
177            self.n_warnings += 1;
178            Some(msg)
179        } else if !self.suppressing {
180            self.suppressing = true;
181            Some(self.suppress_str)
182        } else {
183            None
184        }
185    }
186    fn warn(&mut self, msg: &str) {
187        if let Some(msg) = self.warn_str(msg) {
188            tracing::warn!(msg);
189        }
190    }
191}
192
193#[derive(Serialize, Deserialize, Debug)]
194pub struct CocoExportData {
195    info: CocoInfo,
196    images: Vec<CocoImage>,
197    annotations: Vec<CocoAnnotation>,
198    categories: Vec<CocoBboxCategory>,
199}
200impl CocoExportData {
201    fn from_tools_data<T, A>(
202        tools_data: T,
203        rotation_data: Option<&Rot90ToolData>,
204        prj_path: Option<&Path>,
205    ) -> Self
206    where
207        T: ExportAsCoco<A>,
208        A: InstanceAnnotate + 'static,
209    {
210        type AnnoType<'a, A> = (usize, (&'a String, &'a (Vec<A>, Vec<usize>, ShapeI)));
211        type AnnotationMapValue<'a, A> = (&'a String, &'a (Vec<A>, Vec<usize>, ShapeI));
212
213        let (options, label_info, anno_map, coco_file) = tools_data.separate_data();
214        let color_str = if let Some(s) = colors_to_string(label_info.colors()) {
215            format!(", {s}")
216        } else {
217            String::new()
218        };
219        let info_str = format!(
220            "created with RV Image {}, https://github.com/bertiqwerty/rvimage{color_str}",
221            version_label()
222        );
223        let info = CocoInfo {
224            description: info_str,
225        };
226        let export_data =
227            InstanceExportData::from_tools_data(&options, label_info, coco_file, anno_map);
228
229        let make_image_map =
230            |(idx, (file_path, (_, _, shape))): (usize, AnnotationMapValue<A>)| CocoImage {
231                id: idx as u32,
232                width: shape.w,
233                height: shape.h,
234                file_name: file_path.clone(),
235            };
236        let images = export_data
237            .annotations
238            .iter()
239            .enumerate()
240            .map(make_image_map)
241            .collect::<Vec<_>>();
242
243        let categories = export_data
244            .labels
245            .iter()
246            .zip(export_data.cat_ids.iter())
247            .map(|(label, cat_id)| CocoBboxCategory {
248                id: *cat_id,
249                name: label.clone(),
250            })
251            .collect::<Vec<_>>();
252
253        let mut box_id = 0;
254        let make_anno_map = |(image_idx, (file_path, (bbs, cat_idxs, shape))): AnnoType<A>| {
255            let prj_path = if let Some(prj_path) = prj_path {
256                prj_path
257            } else {
258                Path::new("")
259            };
260            let p = PathPair::new(file_path.clone(), prj_path);
261            let p_abs = p.path_absolute();
262            let shape = if Path::new(p_abs).exists() {
263                let im = trace_ok_warn(image_util::read_image(file_path));
264                if let Some(im) = im {
265                    ShapeI::new(im.width(), im.height())
266                } else {
267                    *shape
268                }
269            } else {
270                *shape
271            };
272            let n_rotations = get_n_rotations(rotation_data, file_path);
273            bbs.iter()
274                .zip(cat_idxs.iter())
275                .filter_map(|(inst_anno, cat_idx): (&A, &usize)| {
276                    trace_ok_warn(instance_to_coco_anno(
277                        inst_anno,
278                        shape,
279                        n_rotations,
280                        options.is_export_absolute,
281                        file_path,
282                    ))
283                    .map(|(bb_f, segmentation)| {
284                        box_id += 1;
285                        CocoAnnotation {
286                            id: box_id - 1,
287                            image_id: image_idx as u32,
288                            category_id: export_data.cat_ids[*cat_idx],
289                            bbox: bb_f,
290                            segmentation,
291                            area: Some(bb_f[2] * bb_f[3]),
292                        }
293                    })
294                })
295                .collect::<Vec<_>>()
296        };
297        let annotations = export_data
298            .annotations
299            .iter()
300            .enumerate()
301            .flat_map(make_anno_map)
302            .collect::<Vec<_>>();
303
304        CocoExportData {
305            info,
306            images,
307            annotations,
308            categories,
309        }
310    }
311
312    #[allow(clippy::too_many_lines)]
313    fn convert_to_toolsdata(
314        self,
315        coco_file: ExportPath,
316        rotation_data: Option<&Rot90ToolData>,
317    ) -> RvResult<(BboxToolData, BrushToolData)> {
318        let cat_ids: Vec<u32> = self.categories.iter().map(|coco_cat| coco_cat.id).collect();
319        let labels: Vec<String> = self
320            .categories
321            .into_iter()
322            .map(|coco_cat| coco_cat.name)
323            .collect();
324        let color_str = self.info.description.split(',').next_back();
325        let colors: Vec<[u8; 3]> = if let Some(s) = color_str {
326            string_to_colors(s).unwrap_or_else(|_| new_random_colors(labels.len()))
327        } else {
328            new_random_colors(labels.len())
329        };
330        let id_image_map = self
331            .images
332            .iter()
333            .map(|coco_image: &CocoImage| {
334                Ok((
335                    coco_image.id,
336                    (
337                        coco_image.file_name.as_str(),
338                        coco_image.width,
339                        coco_image.height,
340                    ),
341                ))
342            })
343            .collect::<RvResult<HashMap<u32, (&str, u32, u32)>>>()?;
344
345        let mut annotations_bbox: HashMap<String, (Vec<GeoFig>, Vec<usize>, ShapeI)> =
346            HashMap::new();
347        let mut annotations_brush: HashMap<String, (Vec<Canvas>, Vec<usize>, ShapeI)> =
348            HashMap::new();
349
350        let n_annotations = self.annotations.len();
351        let mut warner = WarnerCounting::new(
352            n_annotations / 10,
353            "suppressing further warnings during coco import",
354        );
355        for coco_anno in self.annotations {
356            let (file_path, w_coco, h_coco) = id_image_map[&coco_anno.image_id];
357
358            let mut invalid_segmentation_exists = false;
359            // The annotations in the coco files created by RV Image are stored
360            // ignoring any orientation meta-data. Hence, if the image has been loaded
361            // and rotated with RV Image we correct the rotation.
362            let n_rotations = get_n_rotations(rotation_data, file_path);
363            let shape_coco = ShapeI::new(w_coco, h_coco);
364
365            let path_as_key = if file_path.starts_with("http") {
366                file_util::url_encode(file_path)
367            } else {
368                file_path.to_string()
369            };
370
371            let cat_idx = cat_ids
372                .iter()
373                .position(|cat_id| *cat_id == coco_anno.category_id)
374                .ok_or_else(|| {
375                    rverr!(
376                        "could not find cat id {}, we only have {:?}",
377                        coco_anno.category_id,
378                        cat_ids
379                    )
380                })?;
381            let coords_absolute = coco_anno.bbox.iter().any(|x| *x > 1.0);
382            let (w_factor, h_factor) = if coords_absolute {
383                (1.0, 1.0)
384            } else {
385                (f64::from(w_coco), f64::from(h_coco))
386            };
387            let bbox = [
388                (w_factor * coco_anno.bbox[0]),
389                (h_factor * coco_anno.bbox[1]),
390                (w_factor * coco_anno.bbox[2]),
391                (h_factor * coco_anno.bbox[3]),
392            ];
393
394            let mut insert_geo = |geo| {
395                insert_elt(
396                    geo,
397                    &mut annotations_bbox,
398                    cat_idx,
399                    n_rotations,
400                    path_as_key.clone(),
401                    shape_coco,
402                );
403            };
404
405            let bb = BbF::from(&bbox);
406            match coco_anno.segmentation {
407                Some(CocoSegmentation::Polygon(segmentation)) => {
408                    let geo = if segmentation.is_empty() {
409                        GeoFig::BB(bb)
410                    } else {
411                        if segmentation.len() > 1 {
412                            return Err(rverr!(
413                                "multiple polygons per box not supported. ignoring all but first."
414                            ));
415                        }
416                        let n_points = segmentation[0].len();
417                        let coco_data = &segmentation[0];
418
419                        let poly_points = (0..n_points)
420                            .step_by(2)
421                            .filter_map(|idx| {
422                                let p = Point {
423                                    x: (coco_data[idx] * w_factor),
424                                    y: (coco_data[idx + 1] * h_factor),
425                                };
426                                if bb.contains(p) {
427                                    Some(p)
428                                } else {
429                                    None
430                                }
431                            })
432                            .collect();
433                        let poly = Polygon::from_vec(poly_points);
434                        if let Ok(poly) = poly {
435                            let encl_bb = poly.enclosing_bb();
436                            if encl_bb.w * encl_bb.h < 1e-6 && bb.w * bb.h > 1e-6 {
437                                warner.warn(&format!("polygon has no area. using bb. bb: {bb:?}, poly: {encl_bb:?}, file: {file_path}"));
438                                GeoFig::BB(bb)
439                            } else {
440                                if !bb.all_corners_close(encl_bb) {
441                                    let msg = format!("bounding box and polygon enclosing box do not match. using bb. bb: {bb:?}, poly: {encl_bb:?}, file: {file_path}");
442                                    warner.warn(&msg);
443                                }
444                                // check if the poly is just a bounding box
445                                if poly.points().len() == 4
446                                // all points are bb corners
447                                && poly.points_iter().all(|p| {
448                                    encl_bb.points_iter().any(|p_encl| p == p_encl)})
449                                // all points are different
450                                && poly
451                                    .points_iter()
452                                    .all(|p| poly.points_iter().filter(|p_| p == *p_).count() == 1)
453                                {
454                                    GeoFig::BB(bb)
455                                } else {
456                                    GeoFig::Poly(poly)
457                                }
458                            }
459                        } else {
460                            if n_points > 0 {
461                                invalid_segmentation_exists = true;
462                            }
463                            // polygon might be empty, we continue with the BB
464                            GeoFig::BB(bb)
465                        }
466                    };
467                    insert_geo(geo);
468                }
469                Some(CocoSegmentation::Rle(rle)) => {
470                    let bb = bb.into();
471                    let rle_bb = rle_image_to_bb(&rle.counts, bb, ShapeI::from(rle.size))?;
472                    let mask = rle_to_mask(&rle_bb, bb.w, bb.h);
473                    let intensity = rle.intensity.unwrap_or(1.0);
474                    let canvas = Canvas {
475                        bb,
476                        mask,
477                        intensity,
478                    };
479                    insert_elt(
480                        canvas,
481                        &mut annotations_brush,
482                        cat_idx,
483                        n_rotations,
484                        path_as_key,
485                        shape_coco,
486                    );
487                }
488                _ => {
489                    let geo = GeoFig::BB(bb);
490                    insert_geo(geo);
491                }
492            }
493            if invalid_segmentation_exists {
494                warner.warn(&format!("invalid segmentation in coco file {file_path}"));
495            }
496        }
497        let bbox_data = BboxToolData::from_coco_export_data(InstanceExportData {
498            labels: labels.clone(),
499            colors: colors.clone(),
500            cat_ids: cat_ids.clone(),
501            annotations: annotations_bbox,
502            coco_file: coco_file.clone(),
503            is_export_absolute: false,
504        })?;
505        let brush_data = BrushToolData::from_coco_export_data(InstanceExportData {
506            labels,
507            colors,
508            cat_ids,
509            annotations: annotations_brush,
510            coco_file,
511            is_export_absolute: false,
512        })?;
513        Ok((bbox_data, brush_data))
514    }
515}
516
517fn meta_data_to_coco_path(meta_data: &MetaData) -> RvResult<PathBuf> {
518    let export_folder = Path::new(
519        meta_data
520            .export_folder
521            .as_ref()
522            .ok_or_else(|| RvError::new("no export folder given"))?,
523    );
524    let opened_folder = meta_data
525        .opened_folder
526        .as_ref()
527        .map(PathPair::path_absolute)
528        .ok_or_else(|| RvError::new("no folder open"))?;
529    let parent = Path::new(opened_folder)
530        .parent()
531        .and_then(|p| p.file_stem())
532        .and_then(|p| p.to_str());
533
534    let opened_folder_name = Path::new(opened_folder)
535        .file_stem()
536        .and_then(|of| of.to_str())
537        .ok_or_else(|| rverr!("cannot find folder name  of {}", opened_folder))?;
538    let file_name = if let Some(p) = parent {
539        format!("{p}_{opened_folder_name}_coco.json")
540    } else {
541        format!("{opened_folder_name}_coco.json")
542    };
543    Ok(export_folder.join(file_name))
544}
545fn get_cocofilepath(meta_data: &MetaData, coco_file: &ExportPath) -> RvResult<PathBuf> {
546    if path_to_str(&coco_file.path)?.is_empty() {
547        meta_data_to_coco_path(meta_data)
548    } else {
549        Ok(coco_file.path.clone())
550    }
551}
552
553pub fn to_per_file_crowd(brush_annotations_map: &mut BrushAnnoMap) {
554    for (i, (filename, (annos, _))) in brush_annotations_map.iter_mut().enumerate() {
555        if i % 10 == 0 {
556            info!("export - image #{i} converting {filename} to per-image-crowd");
557        }
558        if let Some(max_catidx) = annos.cat_idxs().iter().max() {
559            let mut canvas_idxes_of_cats = vec![vec![]; max_catidx + 1];
560            for i in 0..(annos.elts().len()) {
561                canvas_idxes_of_cats[annos.cat_idxs()[i]].push(i);
562            }
563            let mut merged_canvases = vec![None; max_catidx + 1];
564            for (cat_idx, canvas_idxes) in canvas_idxes_of_cats.iter().enumerate() {
565                let mut merged_canvas: Option<Canvas> = None;
566                for canvas_idx in canvas_idxes {
567                    let elt = &annos.elts()[*canvas_idx];
568                    if let Some(merged_canvas) = &mut merged_canvas {
569                        *merged_canvas = mem::take(merged_canvas).merge(elt);
570                    } else {
571                        merged_canvas = Some(elt.clone());
572                    }
573                }
574                merged_canvases[cat_idx] = merged_canvas;
575            }
576            let mut cat_idxes = vec![];
577            let elts = merged_canvases
578                .into_iter()
579                .enumerate()
580                .filter_map(|(i, cvs)| cvs.map(|cvs| (i, cvs)))
581                .map(|(i, cvs)| {
582                    cat_idxes.push(i);
583                    cvs
584                })
585                .collect();
586            let n_elts = cat_idxes.len();
587            let new_annos = trace_ok_warn(InstanceAnnotations::<Canvas>::new(
588                elts,
589                cat_idxes,
590                vec![false; n_elts],
591            ));
592            if let Some(new_annos) = new_annos {
593                *annos = new_annos;
594            }
595        }
596    }
597}
598
599/// Serialize annotations in Coco format. Any orientations changes applied with the rotation tool
600/// are reverted, since the rotation tool does not change the image file. Hence, the Coco file contains the annotation
601/// relative to the image as it is found in memory ignoring any meta-data.
602///
603/// # Errors
604/// - outpath name cannot be created due to weird characters or the like
605/// - serde write-to-json fails
606pub fn write_coco<T, A>(
607    meta_data: &MetaData,
608    tools_data: T,
609    rotation_data: Option<&Rot90ToolData>,
610    coco_file: &ExportPath,
611) -> RvResult<(PathBuf, JoinHandle<RvResult<()>>)>
612where
613    T: ExportAsCoco<A> + Send + 'static,
614    A: InstanceAnnotate + 'static,
615{
616    let meta_data = meta_data.clone();
617    let coco_out_path = get_cocofilepath(&meta_data, coco_file)?;
618    let coco_out_path_for_thr = coco_out_path.clone();
619    let rotation_data = rotation_data.cloned();
620    let conn = coco_file.conn.clone();
621    let handle = thread::spawn(move || {
622        let coco_data = CocoExportData::from_tools_data(
623            tools_data,
624            rotation_data.as_ref(),
625            meta_data.prj_path(),
626        );
627        let data_str = serde_json::to_string(&coco_data)
628            .map_err(to_rv)
629            .inspect_err(|e| tracing::error!("export failed due to {e:?}"))?;
630
631        conn.write(
632            &data_str,
633            &coco_out_path_for_thr,
634            meta_data.ssh_cfg.as_ref(),
635        )
636        .inspect_err(|e| tracing::error!("export failed due to {e:?}"))?;
637        tracing::info!("exported coco labels to {coco_out_path_for_thr:?}");
638        Ok(())
639    });
640    Ok((coco_out_path, handle))
641}
642
643/// Import annotations in Coco format. Any orientations changes applied with the rotation tool
644/// to images that have annotations in the Coco file are applied to the annotations before importing. We expect, that
645/// the Coco file contains the annotations relative to the image as it is found in memory ignoring any meta-data.
646///
647/// # Errors
648/// - not a coco file
649/// - ssh connection problems
650pub fn read_coco(
651    meta_data: &MetaData,
652    coco_file: &ExportPath,
653    rotation_data: Option<&Rot90ToolData>,
654) -> RvResult<(BboxToolData, BrushToolData)> {
655    let coco_inpath = get_cocofilepath(meta_data, coco_file)?;
656    let coco_str = coco_file
657        .conn
658        .read(&coco_inpath, meta_data.ssh_cfg.as_ref())?;
659    let read_data: CocoExportData = serde_json::from_str(coco_str.as_str()).map_err(to_rv)?;
660    read_data.convert_to_toolsdata(coco_file.clone(), rotation_data)
661}
662
663#[cfg(test)]
664use {
665    super::core::CocoRle,
666    crate::{
667        cfg::{ExportPathConnection, SshCfg},
668        defer_file_removal,
669        meta_data::{ConnectionData, MetaDataFlags},
670    },
671    file_util::DEFAULT_TMPDIR,
672    rvimage_domain::{make_test_bbs, BbI},
673    std::{fs, str::FromStr},
674};
675#[cfg(test)]
676fn make_meta_data(opened_folder: Option<&Path>) -> (MetaData, PathBuf) {
677    let opened_folder = if let Some(of) = opened_folder {
678        PathPair::new(of.to_str().unwrap().to_string(), Path::new(""))
679    } else {
680        PathPair::new("xi".to_string(), Path::new(""))
681    };
682    let test_export_folder = DEFAULT_TMPDIR.clone();
683
684    if !test_export_folder.exists() {
685        match fs::create_dir(&test_export_folder) {
686            Ok(_) => (),
687            Err(e) => {
688                println!("{e:?}");
689            }
690        }
691    }
692
693    let test_export_path = DEFAULT_TMPDIR.join(format!("{}.json", opened_folder.path_absolute()));
694    let mut meta = MetaData::from_filepath(
695        test_export_path
696            .with_extension("egal")
697            .to_str()
698            .unwrap()
699            .to_string(),
700        0,
701        Path::new("egal"),
702    );
703    meta.opened_folder = Some(opened_folder);
704    meta.export_folder = Some(test_export_folder.to_str().unwrap().to_string());
705    meta.connection_data = ConnectionData::Ssh(SshCfg::default());
706    (meta, test_export_path)
707}
708#[cfg(test)]
709fn make_data_brush(
710    image_file: &Path,
711    opened_folder: Option<&Path>,
712    export_absolute: bool,
713    n_boxes: Option<usize>,
714) -> (BrushToolData, MetaData, PathBuf, ShapeI) {
715    use super::InstanceLabelDisplay;
716
717    let shape = ShapeI::new(100, 40);
718    let mut bbox_data = BrushToolData::default();
719    bbox_data.options.core.is_export_absolute = export_absolute;
720    bbox_data.coco_file = ExportPath::default();
721    bbox_data
722        .label_info
723        .push("x".to_string(), None, None)
724        .unwrap();
725
726    bbox_data
727        .label_info
728        .remove_catidx(0, &mut bbox_data.annotations_map);
729
730    let mut bbs = make_test_bbs();
731    bbs.extend(bbs.clone());
732    bbs.extend(bbs.clone());
733    bbs.extend(bbs.clone());
734    bbs.extend(bbs.clone());
735    bbs.extend(bbs.clone());
736    bbs.extend(bbs.clone());
737    bbs.extend(bbs.clone());
738    if let Some(n) = n_boxes {
739        bbs = bbs[0..n].to_vec();
740    }
741
742    let annos = bbox_data.get_annos_mut(image_file.as_os_str().to_str().unwrap(), shape);
743    if let Some(a) = annos {
744        for bb in bbs {
745            let mut mask = vec![0; (bb.w * bb.h) as usize];
746            mask[4] = 1;
747            let c = Canvas {
748                bb: bb.into(),
749                mask,
750                intensity: 0.5,
751            };
752            a.add_elt(c, 0, InstanceLabelDisplay::None);
753        }
754    }
755
756    let (meta, test_export_path) = make_meta_data(opened_folder);
757    (bbox_data, meta, test_export_path, shape)
758}
759#[cfg(test)]
760pub fn make_data_bbox(
761    image_file: &Path,
762    opened_folder: Option<&Path>,
763    export_absolute: bool,
764    n_boxes: Option<usize>,
765) -> (BboxToolData, MetaData, PathBuf, ShapeI) {
766    let shape = ShapeI::new(20, 10);
767    let mut bbox_data = BboxToolData::new();
768    bbox_data.options.core.is_export_absolute = export_absolute;
769    bbox_data.coco_file = ExportPath::default();
770    bbox_data
771        .label_info
772        .push("x".to_string(), None, None)
773        .unwrap();
774
775    bbox_data
776        .label_info
777        .remove_catidx(0, &mut bbox_data.annotations_map);
778
779    let mut bbs = make_test_bbs();
780    bbs.extend(bbs.clone());
781    bbs.extend(bbs.clone());
782    bbs.extend(bbs.clone());
783    bbs.extend(bbs.clone());
784    bbs.extend(bbs.clone());
785    bbs.extend(bbs.clone());
786    bbs.extend(bbs.clone());
787    if let Some(n) = n_boxes {
788        bbs = bbs[0..n].to_vec();
789    }
790
791    let annos = bbox_data.get_annos_mut(image_file.as_os_str().to_str().unwrap(), shape);
792    if let Some(a) = annos {
793        for bb in bbs {
794            a.add_bb(bb, 0, super::InstanceLabelDisplay::IndexLr);
795        }
796    }
797    let (meta, test_export_path) = make_meta_data(opened_folder);
798    (bbox_data, meta, test_export_path, shape)
799}
800
801#[cfg(test)]
802fn is_image_duplicate_free(coco_data: &CocoExportData) -> bool {
803    let mut image_ids = coco_data.images.iter().map(|i| i.id).collect::<Vec<_>>();
804    image_ids.sort();
805    let len_prev = image_ids.len();
806    image_ids.dedup();
807    image_ids.len() == len_prev
808}
809
810#[cfg(test)]
811fn no_image_dups<P>(coco_file: P)
812where
813    P: AsRef<Path> + Debug,
814{
815    let s = file_util::read_to_string(&coco_file).unwrap();
816    let read_raw: CocoExportData = serde_json::from_str(s.as_str()).unwrap();
817
818    assert!(is_image_duplicate_free(&read_raw));
819}
820#[test]
821fn test_coco_export() {
822    fn assert_coco_eq<T, A>(data: T, read: T, coco_file: &PathBuf)
823    where
824        T: ExportAsCoco<A> + Send + 'static,
825        A: InstanceAnnotate + 'static + Debug,
826    {
827        assert_eq!(data.label_info().cat_ids(), read.label_info().cat_ids());
828        assert_eq!(data.label_info().labels(), read.label_info().labels());
829        for (brush_anno, read_anno) in data.anno_iter().zip(read.anno_iter()) {
830            let (name, (instance_annos, shape)) = brush_anno;
831            let (read_name, (read_instance_annos, read_shape)) = read_anno;
832            assert_eq!(instance_annos.cat_idxs(), read_instance_annos.cat_idxs());
833            assert_eq!(
834                instance_annos.elts().len(),
835                read_instance_annos.elts().len()
836            );
837            for (i, (a, b)) in instance_annos
838                .elts()
839                .iter()
840                .zip(read_instance_annos.elts().iter())
841                .enumerate()
842            {
843                assert_eq!(a, b, "annos at index {} differ", i);
844            }
845            assert_eq!(name, read_name);
846            assert_eq!(shape, read_shape);
847        }
848        no_image_dups(coco_file);
849    }
850    fn write_read<T, A>(meta: &MetaData, tools_data: T) -> ((BboxToolData, BrushToolData), PathBuf)
851    where
852        T: ExportAsCoco<A> + Send + 'static,
853        A: InstanceAnnotate + 'static,
854    {
855        let coco_file = tools_data.cocofile_conn();
856        let (coco_file, handle) = write_coco(meta, tools_data, None, &coco_file).unwrap();
857        handle.join().unwrap().unwrap();
858        (
859            read_coco(
860                meta,
861                &ExportPath {
862                    path: coco_file.clone(),
863                    conn: ExportPathConnection::Local,
864                },
865                None,
866            )
867            .unwrap(),
868            coco_file,
869        )
870    }
871    fn test_br(file_path: &Path, opened_folder: Option<&Path>, export_absolute: bool) {
872        let (brush_data, meta, _, _) =
873            make_data_brush(file_path, opened_folder, export_absolute, None);
874        let ((_, read), coco_file) = write_read(&meta, brush_data.clone());
875        defer_file_removal!(&coco_file);
876        assert_coco_eq(brush_data, read, &coco_file);
877    }
878    fn test_bb(file_path: &Path, opened_folder: Option<&Path>, export_absolute: bool) {
879        let (bbox_data, meta, _, _) =
880            make_data_bbox(file_path, opened_folder, export_absolute, None);
881        let ((read, _), coco_file) = write_read(&meta, bbox_data.clone());
882        defer_file_removal!(&coco_file);
883        assert_coco_eq(bbox_data, read, &coco_file);
884    }
885    let tmpdir = &DEFAULT_TMPDIR;
886    let file_path = tmpdir.join("test_image.png");
887    test_br(&file_path, None, true);
888    test_bb(&file_path, None, true);
889    let folder = Path::new("http://localhost:8000/some_path");
890    let file = Path::new("http://localhost:8000/some_path/xyz.png");
891    test_br(file, Some(folder), false);
892    test_bb(file, Some(folder), false);
893}
894
895#[cfg(test)]
896const TEST_DATA_FOLDER: &str = "resources/test_data/";
897
898#[test]
899fn test_coco_import_export() {
900    let meta = MetaData::new(
901        None,
902        None,
903        ConnectionData::None,
904        None,
905        Some(PathPair::new("ohm_somefolder".to_string(), Path::new(""))),
906        Some(TEST_DATA_FOLDER.to_string()),
907        MetaDataFlags::default(),
908        None,
909    );
910    let test_file_src = format!("{TEST_DATA_FOLDER}catids_12_coco_imwolab.json");
911    let test_file = "tmp_coco.json";
912    defer_file_removal!(&test_file);
913    fs::copy(test_file_src, test_file).unwrap();
914    let export_path = ExportPath {
915        path: PathBuf::from_str(test_file).unwrap(),
916        conn: ExportPathConnection::Local,
917    };
918
919    let (read, _) = read_coco(&meta, &export_path, None).unwrap();
920    let (_, handle) = write_coco(&meta, read.clone(), None, &export_path.clone()).unwrap();
921    handle.join().unwrap().unwrap();
922    no_image_dups(&read.coco_file.path);
923    let (read, _) = read_coco(&meta, &export_path, None).unwrap();
924    for anno in read.anno_iter() {
925        let (_, (annos, _)) = anno;
926        for a in annos.elts() {
927            println!("{a:?}");
928            assert!(a.enclosing_bb().w * a.enclosing_bb().h > 1e-3);
929        }
930    }
931}
932
933#[test]
934fn test_coco_import() -> RvResult<()> {
935    fn test(filename: &str, cat_ids: Vec<u32>, reference_bbs: &[(BbI, &str)]) {
936        let meta = MetaData::new(
937            None,
938            None,
939            ConnectionData::None,
940            None,
941            Some(PathPair::new(filename.to_string(), Path::new(""))),
942            Some(TEST_DATA_FOLDER.to_string()),
943            MetaDataFlags::default(),
944            None,
945        );
946        let (read, _) = read_coco(&meta, &ExportPath::default(), None).unwrap();
947        assert_eq!(read.label_info.cat_ids(), &cat_ids);
948        assert_eq!(
949            read.label_info.labels(),
950            &vec!["first label", "second label"]
951        );
952        for (bb, file_path) in reference_bbs {
953            let annos = read.get_annos(file_path);
954            println!();
955            println!("{file_path:?}");
956            println!("{annos:?}");
957            assert!(annos.unwrap().elts().contains(&GeoFig::BB((*bb).into())));
958        }
959    }
960
961    let bb_im_ref_abs1 = [
962        (
963            BbI::from_arr(&[1, 1, 5, 5]),
964            "http://localhost:5000/%2Bnowhere.png",
965        ),
966        (
967            BbI::from_arr(&[11, 11, 4, 7]),
968            "http://localhost:5000/%2Bnowhere.png",
969        ),
970        (
971            BbI::from_arr(&[1, 1, 5, 5]),
972            "http://localhost:5000/%2Bnowhere2.png",
973        ),
974    ];
975    let bb_im_ref_abs2 = [
976        (BbI::from_arr(&[1, 1, 5, 5]), "nowhere.png"),
977        (BbI::from_arr(&[11, 11, 4, 7]), "nowhere.png"),
978        (BbI::from_arr(&[1, 1, 5, 5]), "nowhere2.png"),
979    ];
980    let bb_im_ref_relative = [
981        (BbI::from_arr(&[10, 100, 50, 500]), "nowhere.png"),
982        (BbI::from_arr(&[91, 870, 15, 150]), "nowhere.png"),
983        (BbI::from_arr(&[10, 1, 50, 5]), "nowhere2.png"),
984    ];
985    test("catids_12", vec![1, 2], &bb_im_ref_abs1);
986    test("catids_01", vec![0, 1], &bb_im_ref_abs2);
987    test("catids_12_relative", vec![1, 2], &bb_im_ref_relative);
988    Ok(())
989}
990
991#[test]
992fn color_vs_str() {
993    let colors = vec![[0, 0, 7], [4, 0, 101], [210, 9, 0]];
994    let s = colors_to_string(&colors);
995    let colors_back = string_to_colors(&s.unwrap()).unwrap();
996    assert_eq!(colors, colors_back);
997}
998
999#[test]
1000fn test_rotation_export_import() {
1001    fn test<T, A>(
1002        coco_file: &PathBuf,
1003        bbox_specifics: T,
1004        meta_data: MetaData,
1005        shape: ShapeI,
1006        read_f: impl Fn(&MetaData, &ExportPath, Option<&Rot90ToolData>) -> T,
1007    ) where
1008        T: ExportAsCoco<A> + Send + 'static + Clone,
1009        A: InstanceAnnotate + 'static + Debug,
1010    {
1011        defer_file_removal!(&coco_file);
1012        let mut rotation_data = Rot90ToolData::default();
1013        let annos = rotation_data.get_annos_mut("some_path.png", shape);
1014        if let Some(annos) = annos {
1015            *annos = annos.increase();
1016        }
1017        let coco_file = bbox_specifics.cocofile_conn();
1018        let (out_path, handle) = write_coco(
1019            &meta_data,
1020            bbox_specifics.clone(),
1021            Some(&rotation_data),
1022            &coco_file,
1023        )
1024        .unwrap();
1025        handle.join().unwrap().unwrap();
1026        println!("write to {out_path:?}");
1027        let out_path = ExportPath {
1028            path: out_path,
1029            conn: ExportPathConnection::Local,
1030        };
1031        let read = read_f(&meta_data, &out_path, Some(&rotation_data));
1032
1033        for ((_, (anno_res, _)), (_, (anno_ref, _))) in
1034            bbox_specifics.anno_iter().zip(read.anno_iter())
1035        {
1036            for (read_elt, ref_elt) in anno_res.elts().iter().zip(anno_ref.elts().iter()) {
1037                assert_eq!(read_elt, ref_elt);
1038            }
1039        }
1040    }
1041    let (brush_specifics, meta_data, coco_file, shape) = make_data_brush(
1042        Path::new("some_path.png"),
1043        Some(Path::new("afolder")),
1044        false,
1045        None,
1046    );
1047    test(&coco_file, brush_specifics, meta_data, shape, |m, d, r| {
1048        read_coco(m, d, r).unwrap().1
1049    });
1050    let (bbox_specifics, meta_data, coco_file, shape) = make_data_bbox(
1051        Path::new("some_path.png"),
1052        Some(Path::new("afolder")),
1053        false,
1054        None,
1055    );
1056    test(&coco_file, bbox_specifics, meta_data, shape, |m, d, r| {
1057        read_coco(m, d, r).unwrap().0
1058    });
1059}
1060
1061#[test]
1062fn test_serialize_rle() {
1063    let rle = CocoRle {
1064        counts: vec![1, 2, 3, 4],
1065        size: (5, 6),
1066        intensity: None,
1067    };
1068    let rle = CocoSegmentation::Rle(rle);
1069    let s = serde_json::to_string(&rle).unwrap();
1070    println!("{s}");
1071    let rle2: CocoSegmentation = serde_json::from_str(&s).unwrap();
1072    assert_eq!(format!("{rle:?}"), format!("{rle2:?}"));
1073    let poly = CocoSegmentation::Polygon(vec![vec![1.0, 2.0]]);
1074    let s = serde_json::to_string(&poly).unwrap();
1075    println!("{s}");
1076    let poly2: CocoSegmentation = serde_json::from_str(&s).unwrap();
1077    assert_eq!(format!("{poly:?}"), format!("{poly2:?}"));
1078}
1079
1080#[test]
1081fn test_instance_to_coco() {
1082    let shape = ShapeI::new(2000, 2667);
1083    let bb = BbI::from_arr(&[1342, 1993, 8, 8]);
1084    let n_rot = 1;
1085    let canvas = Canvas {
1086        mask: vec![0; 64],
1087        bb,
1088        intensity: 0.5,
1089    };
1090    let coco_anno = instance_to_coco_anno(&canvas, shape, n_rot, false, "");
1091    assert!(coco_anno.is_err());
1092
1093    let shape_im = ShapeI::new(20, 40);
1094    let mut mask = vec![0; 4];
1095    mask[2] = 1;
1096    let canvas = Canvas {
1097        bb: BbI::from_arr(&[1, 1, 2, 2]),
1098        mask: mask.clone(),
1099        intensity: 0.5,
1100    };
1101    let n_rotations = 1;
1102
1103    let (_, segmentation) =
1104        instance_to_coco_anno(&canvas, shape_im, n_rotations, false, "").unwrap();
1105
1106    let coco_seg = canvas
1107        .rot90_with_image_ntimes(
1108            shape_im.rot90_with_image_ntimes(n_rotations),
1109            4 - n_rotations,
1110        )
1111        .unwrap()
1112        .to_cocoseg(shape_im, false)
1113        .unwrap();
1114    assert_ne!(coco_seg, None);
1115    assert_eq!(segmentation, coco_seg);
1116    let mut mask = [0; 4];
1117    mask[2] = 1;
1118    let geo = GeoFig::BB(BbF::from_arr(&[1.0, 1.0, 2.0, 8.0]));
1119
1120    let n_rotations = 1;
1121
1122    let (bb_rot, segmentation) =
1123        instance_to_coco_anno(&geo, shape_im, n_rotations, true, "").unwrap();
1124    println!("{bb_rot:?}");
1125    let coco_seg = geo
1126        .rot90_with_image_ntimes(
1127            shape_im.rot90_with_image_ntimes(n_rotations),
1128            4 - n_rotations,
1129        )
1130        .unwrap()
1131        .to_cocoseg(shape_im, true)
1132        .unwrap();
1133    assert_ne!(coco_seg, None);
1134    assert_eq!(segmentation, coco_seg);
1135}
1136
1137#[test]
1138fn test_warner() {
1139    let suppress_msg = "no further warnings";
1140    let mut warner = WarnerCounting::new(3, suppress_msg);
1141    assert_eq!(warner.warn_str("a"), Some("a"));
1142    assert_eq!(warner.warn_str("a"), Some("a"));
1143    assert_eq!(warner.warn_str("b"), Some("b"));
1144    assert_eq!(warner.warn_str("a"), Some(suppress_msg));
1145    assert_eq!(warner.warn_str("a"), None);
1146}