Skip to main content

rvlib/tools_data/
coco_io.rs

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