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    let shape = ShapeI::new(100, 40);
734    let mut bbox_data = BrushToolData::default();
735    bbox_data.options.core.is_export_absolute = export_absolute;
736    bbox_data.coco_file = ExportPath::default();
737    bbox_data
738        .label_info
739        .push("x".to_string(), None, None)
740        .unwrap();
741
742    bbox_data
743        .label_info
744        .remove_catidx(0, &mut bbox_data.annotations_map);
745
746    let mut bbs = make_test_bbs();
747    bbs.extend(bbs.clone());
748    bbs.extend(bbs.clone());
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    if let Some(n) = n_boxes {
755        bbs = bbs[0..n].to_vec();
756    }
757
758    let annos = bbox_data.get_annos_mut(image_file.as_os_str().to_str().unwrap(), shape);
759    if let Some(a) = annos {
760        for bb in bbs {
761            let mut mask = vec![0; (bb.w * bb.h) as usize];
762            mask[4] = 1;
763            let c = Canvas {
764                bb: bb.into(),
765                mask,
766                intensity: 0.5,
767            };
768            a.add_elt(c, 0);
769        }
770    }
771
772    let (meta, test_export_path) = make_meta_data(opened_folder);
773    (bbox_data, meta, test_export_path, shape)
774}
775#[cfg(test)]
776pub fn make_data_bbox(
777    image_file: &Path,
778    opened_folder: Option<&Path>,
779    export_absolute: bool,
780    n_boxes: Option<usize>,
781) -> (BboxToolData, MetaData, PathBuf, ShapeI) {
782    let shape = ShapeI::new(20, 10);
783    let mut bbox_data = BboxToolData::new();
784    bbox_data.options.core.is_export_absolute = export_absolute;
785    bbox_data.coco_file = ExportPath::default();
786    bbox_data
787        .label_info
788        .push("x".to_string(), None, None)
789        .unwrap();
790
791    bbox_data
792        .label_info
793        .remove_catidx(0, &mut bbox_data.annotations_map);
794
795    let mut bbs = make_test_bbs();
796    bbs.extend(bbs.clone());
797    bbs.extend(bbs.clone());
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    if let Some(n) = n_boxes {
804        bbs = bbs[0..n].to_vec();
805    }
806
807    let annos = bbox_data.get_annos_mut(image_file.as_os_str().to_str().unwrap(), shape);
808    if let Some(a) = annos {
809        for bb in bbs {
810            a.add_bb(bb, 0);
811        }
812    }
813    let (meta, test_export_path) = make_meta_data(opened_folder);
814    (bbox_data, meta, test_export_path, shape)
815}
816
817#[cfg(test)]
818fn is_image_duplicate_free(coco_data: &CocoExportData) -> bool {
819    let mut image_ids = coco_data.images.iter().map(|i| i.id).collect::<Vec<_>>();
820    image_ids.sort();
821    let len_prev = image_ids.len();
822    image_ids.dedup();
823    image_ids.len() == len_prev
824}
825
826#[cfg(test)]
827fn no_image_dups<P>(coco_file: P)
828where
829    P: AsRef<Path> + Debug,
830{
831    let s = file_util::read_to_string(&coco_file).unwrap();
832    let read_raw: CocoExportData = serde_json::from_str(s.as_str()).unwrap();
833
834    assert!(is_image_duplicate_free(&read_raw));
835}
836#[test]
837fn test_coco_export() {
838    fn assert_coco_eq<T, A>(data: T, read: T, coco_file: &PathBuf)
839    where
840        T: ExportAsCoco<A> + Send + 'static,
841        A: InstanceAnnotate + 'static + Debug,
842    {
843        assert_eq!(data.label_info().cat_ids(), read.label_info().cat_ids());
844        assert_eq!(data.label_info().labels(), read.label_info().labels());
845        for (brush_anno, read_anno) in data.anno_iter().zip(read.anno_iter()) {
846            let (name, (instance_annos, shape)) = brush_anno;
847            let (read_name, (read_instance_annos, read_shape)) = read_anno;
848            assert_eq!(instance_annos.cat_idxs(), read_instance_annos.cat_idxs());
849            assert_eq!(
850                instance_annos.elts().len(),
851                read_instance_annos.elts().len()
852            );
853            for (i, (a, b)) in instance_annos
854                .elts()
855                .iter()
856                .zip(read_instance_annos.elts().iter())
857                .enumerate()
858            {
859                assert_eq!(a, b, "annos at index {} differ", i);
860            }
861            assert_eq!(name, read_name);
862            assert_eq!(shape, read_shape);
863        }
864        no_image_dups(coco_file);
865    }
866    fn write_read<T, A>(meta: &MetaData, tools_data: T) -> ((BboxToolData, BrushToolData), PathBuf)
867    where
868        T: ExportAsCoco<A> + Send + 'static,
869        A: InstanceAnnotate + 'static,
870    {
871        let coco_file = tools_data.cocofile_conn();
872        let (coco_file, handle) = write_coco(meta, tools_data, None, &coco_file).unwrap();
873        handle.join().unwrap().unwrap();
874        (
875            read_coco(
876                meta,
877                &ExportPath {
878                    path: coco_file.clone(),
879                    conn: ExportPathConnection::Local,
880                },
881                None,
882            )
883            .unwrap(),
884            coco_file,
885        )
886    }
887    fn test_br(file_path: &Path, opened_folder: Option<&Path>, export_absolute: bool) {
888        let (brush_data, meta, _, _) =
889            make_data_brush(file_path, opened_folder, export_absolute, None);
890        let ((_, read), coco_file) = write_read(&meta, brush_data.clone());
891        defer_file_removal!(&coco_file);
892        assert_coco_eq(brush_data, read, &coco_file);
893    }
894    fn test_bb(file_path: &Path, opened_folder: Option<&Path>, export_absolute: bool) {
895        let (bbox_data, meta, _, _) =
896            make_data_bbox(file_path, opened_folder, export_absolute, None);
897        let ((read, _), coco_file) = write_read(&meta, bbox_data.clone());
898        defer_file_removal!(&coco_file);
899        assert_coco_eq(bbox_data, read, &coco_file);
900    }
901    let tmpdir = &DEFAULT_TMPDIR;
902    let file_path = tmpdir.join("test_image.png");
903    test_br(&file_path, None, true);
904    test_bb(&file_path, None, true);
905    let folder = Path::new("http://localhost:8000/some_path");
906    let file = Path::new("http://localhost:8000/some_path/xyz.png");
907    test_br(file, Some(folder), false);
908    test_bb(file, Some(folder), false);
909}
910
911#[cfg(test)]
912const TEST_DATA_FOLDER: &str = "resources/test_data/";
913
914#[test]
915fn test_coco_import_export() {
916    let meta = MetaData::new(
917        None,
918        None,
919        ConnectionData::None,
920        None,
921        Some(PathPair::new("ohm_somefolder".to_string(), Path::new(""))),
922        Some(TEST_DATA_FOLDER.to_string()),
923        MetaDataFlags::default(),
924        None,
925    );
926    let test_file_src = format!("{TEST_DATA_FOLDER}catids_12_coco_imwolab.json");
927    let test_file = "tmp_coco.json";
928    defer_file_removal!(&test_file);
929    fs::copy(test_file_src, test_file).unwrap();
930    let export_path = ExportPath {
931        path: PathBuf::from_str(test_file).unwrap(),
932        conn: ExportPathConnection::Local,
933    };
934
935    let (read, _) = read_coco(&meta, &export_path, None).unwrap();
936    let (_, handle) = write_coco(&meta, read.clone(), None, &export_path.clone()).unwrap();
937    handle.join().unwrap().unwrap();
938    no_image_dups(&read.coco_file.path);
939    let (read, _) = read_coco(&meta, &export_path, None).unwrap();
940    for anno in read.anno_iter() {
941        let (_, (annos, _)) = anno;
942        for a in annos.elts() {
943            println!("{a:?}");
944            assert!(a.enclosing_bb().w * a.enclosing_bb().h > 1e-3);
945        }
946    }
947}
948
949#[test]
950fn test_coco_import() -> RvResult<()> {
951    fn test(filename: &str, cat_ids: Vec<u32>, reference_bbs: &[(BbI, &str)]) {
952        let meta = MetaData::new(
953            None,
954            None,
955            ConnectionData::None,
956            None,
957            Some(PathPair::new(filename.to_string(), Path::new(""))),
958            Some(TEST_DATA_FOLDER.to_string()),
959            MetaDataFlags::default(),
960            None,
961        );
962        let (read, _) = read_coco(&meta, &ExportPath::default(), None).unwrap();
963        assert_eq!(read.label_info.cat_ids(), &cat_ids);
964        assert_eq!(
965            read.label_info.labels(),
966            &vec!["first label", "second label"]
967        );
968        for (bb, file_path) in reference_bbs {
969            let annos = read.get_annos(file_path);
970            println!();
971            println!("{file_path:?}");
972            println!("{annos:?}");
973            assert!(annos.unwrap().elts().contains(&GeoFig::BB((*bb).into())));
974        }
975    }
976
977    let bb_im_ref_abs1 = [
978        (
979            BbI::from_arr(&[1, 1, 5, 5]),
980            "http://localhost:5000/%2Bnowhere.png",
981        ),
982        (
983            BbI::from_arr(&[11, 11, 4, 7]),
984            "http://localhost:5000/%2Bnowhere.png",
985        ),
986        (
987            BbI::from_arr(&[1, 1, 5, 5]),
988            "http://localhost:5000/%2Bnowhere2.png",
989        ),
990    ];
991    let bb_im_ref_abs2 = [
992        (BbI::from_arr(&[1, 1, 5, 5]), "nowhere.png"),
993        (BbI::from_arr(&[11, 11, 4, 7]), "nowhere.png"),
994        (BbI::from_arr(&[1, 1, 5, 5]), "nowhere2.png"),
995    ];
996    let bb_im_ref_relative = [
997        (BbI::from_arr(&[10, 100, 50, 500]), "nowhere.png"),
998        (BbI::from_arr(&[91, 870, 15, 150]), "nowhere.png"),
999        (BbI::from_arr(&[10, 1, 50, 5]), "nowhere2.png"),
1000    ];
1001    test("catids_12", vec![1, 2], &bb_im_ref_abs1);
1002    test("catids_01", vec![0, 1], &bb_im_ref_abs2);
1003    test("catids_12_relative", vec![1, 2], &bb_im_ref_relative);
1004    Ok(())
1005}
1006
1007#[test]
1008fn color_vs_str() {
1009    let colors = vec![[0, 0, 7], [4, 0, 101], [210, 9, 0]];
1010    let s = colors_to_string(&colors);
1011    let colors_back = string_to_colors(&s.unwrap()).unwrap();
1012    assert_eq!(colors, colors_back);
1013}
1014
1015#[test]
1016fn test_rotation_export_import() {
1017    fn test<T, A>(
1018        coco_file: &PathBuf,
1019        bbox_specifics: T,
1020        meta_data: MetaData,
1021        shape: ShapeI,
1022        read_f: impl Fn(&MetaData, &ExportPath, Option<&Rot90ToolData>) -> T,
1023    ) where
1024        T: ExportAsCoco<A> + Send + 'static + Clone,
1025        A: InstanceAnnotate + 'static + Debug,
1026    {
1027        defer_file_removal!(&coco_file);
1028        let mut rotation_data = Rot90ToolData::default();
1029        let annos = rotation_data.get_annos_mut("some_path.png", shape);
1030        if let Some(annos) = annos {
1031            *annos = annos.increase();
1032        }
1033        let coco_file = bbox_specifics.cocofile_conn();
1034        let (out_path, handle) = write_coco(
1035            &meta_data,
1036            bbox_specifics.clone(),
1037            Some(&rotation_data),
1038            &coco_file,
1039        )
1040        .unwrap();
1041        handle.join().unwrap().unwrap();
1042        println!("write to {out_path:?}");
1043        let out_path = ExportPath {
1044            path: out_path,
1045            conn: ExportPathConnection::Local,
1046        };
1047        let read = read_f(&meta_data, &out_path, Some(&rotation_data));
1048
1049        for ((_, (anno_res, _)), (_, (anno_ref, _))) in
1050            bbox_specifics.anno_iter().zip(read.anno_iter())
1051        {
1052            for (read_elt, ref_elt) in anno_res.elts().iter().zip(anno_ref.elts().iter()) {
1053                assert_eq!(read_elt, ref_elt);
1054            }
1055        }
1056    }
1057    let (brush_specifics, meta_data, coco_file, shape) = make_data_brush(
1058        Path::new("some_path.png"),
1059        Some(Path::new("afolder")),
1060        false,
1061        None,
1062    );
1063    test(&coco_file, brush_specifics, meta_data, shape, |m, d, r| {
1064        read_coco(m, d, r).unwrap().1
1065    });
1066    let (bbox_specifics, meta_data, coco_file, shape) = make_data_bbox(
1067        Path::new("some_path.png"),
1068        Some(Path::new("afolder")),
1069        false,
1070        None,
1071    );
1072    test(&coco_file, bbox_specifics, meta_data, shape, |m, d, r| {
1073        read_coco(m, d, r).unwrap().0
1074    });
1075}
1076
1077#[test]
1078fn test_serialize_rle() {
1079    let rle = CocoRle {
1080        counts: vec![1, 2, 3, 4],
1081        size: (5, 6),
1082        intensity: None,
1083    };
1084    let rle = CocoSegmentation::Rle(rle);
1085    let s = serde_json::to_string(&rle).unwrap();
1086    println!("{s}");
1087    let rle2: CocoSegmentation = serde_json::from_str(&s).unwrap();
1088    assert_eq!(format!("{rle:?}"), format!("{rle2:?}"));
1089    let poly = CocoSegmentation::Polygon(vec![vec![1.0, 2.0]]);
1090    let s = serde_json::to_string(&poly).unwrap();
1091    println!("{s}");
1092    let poly2: CocoSegmentation = serde_json::from_str(&s).unwrap();
1093    assert_eq!(format!("{poly:?}"), format!("{poly2:?}"));
1094}
1095
1096#[test]
1097fn test_instance_to_coco() {
1098    let shape = ShapeI::new(2000, 2667);
1099    let bb = BbI::from_arr(&[1342, 1993, 8, 8]);
1100    let n_rot = 1;
1101    let canvas = Canvas {
1102        mask: vec![0; 64],
1103        bb,
1104        intensity: 0.5,
1105    };
1106    let coco_anno = instance_to_coco_anno(&canvas, shape, n_rot, false, "");
1107    assert!(coco_anno.is_err());
1108
1109    let shape_im = ShapeI::new(20, 40);
1110    let mut mask = vec![0; 4];
1111    mask[2] = 1;
1112    let canvas = Canvas {
1113        bb: BbI::from_arr(&[1, 1, 2, 2]),
1114        mask: mask.clone(),
1115        intensity: 0.5,
1116    };
1117    let n_rotations = 1;
1118
1119    let (_, segmentation) =
1120        instance_to_coco_anno(&canvas, shape_im, n_rotations, false, "").unwrap();
1121
1122    let coco_seg = canvas
1123        .rot90_with_image_ntimes(
1124            shape_im.rot90_with_image_ntimes(n_rotations),
1125            4 - n_rotations,
1126        )
1127        .unwrap()
1128        .to_cocoseg(shape_im, false)
1129        .unwrap();
1130    assert_ne!(coco_seg, None);
1131    assert_eq!(segmentation, coco_seg);
1132    let mut mask = [0; 4];
1133    mask[2] = 1;
1134    let geo = GeoFig::BB(BbF::from_arr(&[1.0, 1.0, 2.0, 8.0]));
1135
1136    let n_rotations = 1;
1137
1138    let (bb_rot, segmentation) =
1139        instance_to_coco_anno(&geo, shape_im, n_rotations, true, "").unwrap();
1140    println!("{bb_rot:?}");
1141    let coco_seg = geo
1142        .rot90_with_image_ntimes(
1143            shape_im.rot90_with_image_ntimes(n_rotations),
1144            4 - n_rotations,
1145        )
1146        .unwrap()
1147        .to_cocoseg(shape_im, true)
1148        .unwrap();
1149    assert_ne!(coco_seg, None);
1150    assert_eq!(segmentation, coco_seg);
1151}
1152
1153#[test]
1154fn test_warner() {
1155    let suppress_msg = "no further warnings";
1156    let mut warner = WarnerCounting::new(3, suppress_msg);
1157    assert_eq!(warner.warn_str("a"), Some("a"));
1158    assert_eq!(warner.warn_str("a"), Some("a"));
1159    assert_eq!(warner.warn_str("b"), Some("b"));
1160    assert_eq!(warner.warn_str("a"), Some(suppress_msg));
1161    assert_eq!(warner.warn_str("a"), None);
1162}