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 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 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
584pub 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
630pub 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}