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