1use super::types::{CocoCompressedRle, CocoRle, CocoSegmentation};
14use crate::{Box2d, Error, MaskData, Polygon};
15
16pub fn coco_bbox_to_box2d(bbox: &[f64; 4], image_width: u32, image_height: u32) -> Box2d {
45 let [x, y, w, h] = *bbox;
46 let img_w = image_width as f64;
47 let img_h = image_height as f64;
48
49 Box2d::new(
50 (x / img_w) as f32,
51 (y / img_h) as f32,
52 (w / img_w) as f32,
53 (h / img_h) as f32,
54 )
55}
56
57pub fn box2d_to_coco_bbox(box2d: &Box2d, image_width: u32, image_height: u32) -> [f64; 4] {
67 let img_w = image_width as f64;
68 let img_h = image_height as f64;
69
70 [
71 (box2d.left() as f64) * img_w,
72 (box2d.top() as f64) * img_h,
73 (box2d.width() as f64) * img_w,
74 (box2d.height() as f64) * img_h,
75 ]
76}
77
78pub fn validate_coco_bbox(
88 bbox: &[f64; 4],
89 image_width: u32,
90 image_height: u32,
91) -> Result<(), Error> {
92 let [x, y, w, h] = *bbox;
93
94 if w <= 0.0 || h <= 0.0 {
95 return Err(Error::CocoError(format!(
96 "Width and height must be positive: w={}, h={}",
97 w, h
98 )));
99 }
100
101 let epsilon = 1.0;
103 if x < -epsilon || y < -epsilon {
104 return Err(Error::CocoError(format!(
105 "Bbox has negative coordinates: x={}, y={}",
106 x, y
107 )));
108 }
109
110 if x + w > (image_width as f64) + epsilon || y + h > (image_height as f64) + epsilon {
111 return Err(Error::CocoError(format!(
112 "Bbox exceeds image bounds: [{}, {}, {}, {}] for {}x{} image",
113 x, y, w, h, image_width, image_height
114 )));
115 }
116
117 Ok(())
118}
119
120pub fn coco_polygon_to_polygon(
138 polygons: &[Vec<f64>],
139 image_width: u32,
140 image_height: u32,
141) -> Polygon {
142 let img_w = image_width as f64;
143 let img_h = image_height as f64;
144
145 let converted: Vec<Vec<(f32, f32)>> = polygons
146 .iter()
147 .filter(|poly| poly.len() >= 6) .map(|polygon| {
149 polygon
150 .chunks(2)
151 .filter_map(|chunk| {
152 if chunk.len() == 2 {
153 Some(((chunk[0] / img_w) as f32, (chunk[1] / img_h) as f32))
154 } else {
155 None
156 }
157 })
158 .collect()
159 })
160 .filter(|poly: &Vec<(f32, f32)>| poly.len() >= 3) .collect();
162
163 Polygon::new(converted)
164}
165
166pub fn polygon_to_coco_polygon(
176 polygon: &Polygon,
177 image_width: u32,
178 image_height: u32,
179) -> Vec<Vec<f64>> {
180 let img_w = image_width as f64;
181 let img_h = image_height as f64;
182
183 polygon
184 .rings
185 .iter()
186 .filter(|poly| poly.len() >= 3) .map(|ring| {
188 ring.iter()
189 .flat_map(|(x, y)| [(*x as f64) * img_w, (*y as f64) * img_h])
190 .collect()
191 })
192 .collect()
193}
194
195pub fn decode_rle(rle: &CocoRle) -> Result<(Vec<u8>, u32, u32), Error> {
210 let [height, width] = rle.size;
211 let total_pixels = (width as usize) * (height as usize);
212
213 let counts_sum: u64 = rle.counts.iter().map(|&c| c as u64).sum();
215 if counts_sum != total_pixels as u64 {
216 return Err(Error::CocoError(format!(
217 "RLE counts sum {} does not match image size {}x{} = {}",
218 counts_sum, width, height, total_pixels
219 )));
220 }
221
222 let mut column_major = vec![0u8; total_pixels];
224 let mut pos = 0usize;
225 let mut is_foreground = false; for &count in &rle.counts {
228 let count = count as usize;
229 if is_foreground {
230 for i in pos..(pos + count).min(column_major.len()) {
231 column_major[i] = 1;
232 }
233 }
234 pos += count;
235 is_foreground = !is_foreground;
236 }
237
238 let mut row_major = vec![0u8; total_pixels];
240 for col in 0..width as usize {
241 for row in 0..height as usize {
242 let col_idx = col * (height as usize) + row;
243 let row_idx = row * (width as usize) + col;
244 if col_idx < column_major.len() && row_idx < row_major.len() {
245 row_major[row_idx] = column_major[col_idx];
246 }
247 }
248 }
249
250 Ok((row_major, height, width))
251}
252
253fn decode_leb128(s: &str) -> Result<Vec<u32>, Error> {
257 let bytes = s.as_bytes();
258 let mut counts = Vec::new();
259 let mut i = 0;
260
261 while i < bytes.len() {
262 let mut value: i64 = 0;
263 let mut shift = 0;
264 let mut more = true;
265
266 while more && i < bytes.len() {
267 let byte = bytes[i] as i64;
268 i += 1;
269
270 let decoded = if (48..96).contains(&byte) {
272 byte - 48 } else if byte >= 96 {
274 byte - 96 + 48 } else {
276 return Err(Error::CocoError(format!(
277 "Invalid LEB128 character: {}",
278 byte as u8 as char
279 )));
280 };
281
282 value |= (decoded & 0x1F) << shift;
283 more = decoded >= 32;
284 shift += 5;
285 }
286
287 if shift < 32 && (value & (1 << (shift - 1))) != 0 {
289 value |= (-1i64) << shift;
290 }
291
292 counts.push(value);
293 }
294
295 let mut result = Vec::with_capacity(counts.len());
297 let mut prev: i64 = 0;
298 for diff in counts {
299 prev += diff;
300 result.push(prev.max(0) as u32);
301 }
302
303 Ok(result)
304}
305
306pub fn decode_compressed_rle(compressed: &CocoCompressedRle) -> Result<(Vec<u8>, u32, u32), Error> {
308 let counts = decode_leb128(&compressed.counts)?;
309
310 let rle = CocoRle {
311 counts,
312 size: compressed.size,
313 };
314
315 decode_rle(&rle)
316}
317
318pub fn mask_to_contours(mask: &[u8], width: u32, height: u32) -> Vec<Vec<(f64, f64)>> {
335 let mut contours = Vec::new();
336 let mut visited = vec![false; mask.len()];
337
338 let w = width as usize;
339 let h = height as usize;
340
341 for start_y in 0..h {
342 for start_x in 0..w {
343 let idx = start_y * w + start_x;
344 if mask[idx] == 1 && !visited[idx] {
345 let is_boundary = start_x == 0
348 || start_x == w - 1
349 || start_y == 0
350 || start_y == h - 1
351 || (start_x > 0 && mask[idx - 1] == 0)
352 || (start_x < w - 1 && mask[idx + 1] == 0)
353 || (start_y > 0 && mask[idx - w] == 0)
354 || (start_y < h - 1 && mask[idx + w] == 0);
355
356 if is_boundary
357 && let Some(contour) = trace_contour(mask, w, h, start_x, start_y, &mut visited)
358 && contour.len() >= 3
359 {
360 contours.push(contour);
361 }
362 }
363 }
364 }
365
366 contours
367}
368
369fn trace_contour(
371 mask: &[u8],
372 width: usize,
373 height: usize,
374 start_x: usize,
375 start_y: usize,
376 visited: &mut [bool],
377) -> Option<Vec<(f64, f64)>> {
378 let mut contour = Vec::new();
379 let mut x = start_x;
380 let mut y = start_y;
381
382 let dx: [i32; 8] = [1, 1, 0, -1, -1, -1, 0, 1];
384 let dy: [i32; 8] = [0, 1, 1, 1, 0, -1, -1, -1];
385
386 let mut dir = 0usize; let max_steps = width * height;
388 let mut steps = 0;
389
390 loop {
391 let idx = y * width + x;
392 if !visited[idx] {
393 contour.push((x as f64, y as f64));
394 visited[idx] = true;
395 }
396
397 let mut found = false;
399 for i in 0..8 {
400 let new_dir = (dir + i) % 8;
401 let nx = x as i32 + dx[new_dir];
402 let ny = y as i32 + dy[new_dir];
403
404 if nx >= 0 && nx < width as i32 && ny >= 0 && ny < height as i32 {
405 let nidx = (ny as usize) * width + (nx as usize);
406 if mask[nidx] == 1 {
407 x = nx as usize;
408 y = ny as usize;
409 dir = (new_dir + 5) % 8; found = true;
411 break;
412 }
413 }
414 }
415
416 if !found || (x == start_x && y == start_y && contour.len() > 2) {
417 break;
418 }
419
420 steps += 1;
421 if steps > max_steps {
422 break; }
424 }
425
426 if contour.len() >= 3 {
427 Some(contour)
428 } else {
429 None
430 }
431}
432
433pub fn coco_rle_to_polygon(
437 rle: &CocoRle,
438 image_width: u32,
439 image_height: u32,
440) -> Result<Polygon, Error> {
441 let (binary_mask, height, width) = decode_rle(rle)?;
442 let contours = mask_to_contours(&binary_mask, width, height);
443
444 let normalized: Vec<Vec<(f32, f32)>> = contours
446 .iter()
447 .map(|contour| {
448 contour
449 .iter()
450 .map(|(x, y)| {
451 (
452 (*x / image_width as f64) as f32,
453 (*y / image_height as f64) as f32,
454 )
455 })
456 .collect()
457 })
458 .collect();
459
460 Ok(Polygon::new(normalized))
461}
462
463pub fn coco_segmentation_to_polygon(
467 segmentation: &CocoSegmentation,
468 image_width: u32,
469 image_height: u32,
470) -> Result<Polygon, Error> {
471 match segmentation {
472 CocoSegmentation::Polygon(polygons) => {
473 Ok(coco_polygon_to_polygon(polygons, image_width, image_height))
474 }
475 CocoSegmentation::Rle(rle) => coco_rle_to_polygon(rle, image_width, image_height),
476 CocoSegmentation::CompressedRle(compressed) => {
477 let counts = decode_leb128(&compressed.counts)?;
478 let rle = CocoRle {
479 counts,
480 size: compressed.size,
481 };
482 coco_rle_to_polygon(&rle, image_width, image_height)
483 }
484 }
485}
486
487pub fn rle_to_mask_data(rle: &CocoRle) -> Result<MaskData, Error> {
493 let (pixels, height, width) = decode_rle(rle)?;
494 MaskData::encode(&pixels, width, height, 1)
495}
496
497pub fn coco_segmentation_to_mask_data(seg: &CocoSegmentation) -> Result<Option<MaskData>, Error> {
502 match seg {
503 CocoSegmentation::Polygon(_) => Ok(None),
504 CocoSegmentation::Rle(rle) => Ok(Some(rle_to_mask_data(rle)?)),
505 CocoSegmentation::CompressedRle(crle) => {
506 let (pixels, height, width) = decode_compressed_rle(crle)?;
507 Ok(Some(MaskData::encode(&pixels, width, height, 1)?))
508 }
509 }
510}
511
512pub fn encode_rle(mask: &[u8], width: u32, height: u32) -> Result<CocoRle, Error> {
530 let total = (width as usize) * (height as usize);
531
532 if mask.len() != total {
533 return Err(Error::CocoError(format!(
534 "mask length {} does not match {}x{} = {}",
535 mask.len(),
536 width,
537 height,
538 total
539 )));
540 }
541
542 let mut column_major = vec![0u8; total];
544 for row in 0..height as usize {
545 for col in 0..width as usize {
546 column_major[col * height as usize + row] = mask[row * width as usize + col];
547 }
548 }
549
550 let mut counts = Vec::new();
552 let mut current = 0u8; let mut run = 0u32;
554 for &pixel in &column_major {
555 let val = if pixel != 0 { 1 } else { 0 };
556 if val == current {
557 run += 1;
558 } else {
559 counts.push(run);
560 current = val;
561 run = 1;
562 }
563 }
564 counts.push(run);
565
566 Ok(CocoRle {
567 counts,
568 size: [height, width],
569 })
570}
571
572pub fn calculate_coco_area(segmentation: &CocoSegmentation) -> Result<f64, Error> {
578 match segmentation {
579 CocoSegmentation::Polygon(polygons) => {
580 let mut total_area = 0.0;
582 for polygon in polygons {
583 total_area += shoelace_area(polygon);
584 }
585 Ok(total_area)
586 }
587 CocoSegmentation::Rle(rle) => {
588 let (mask, _, _) = decode_rle(rle)?;
589 let area = mask.iter().filter(|&&v| v == 1).count() as f64;
590 Ok(area)
591 }
592 CocoSegmentation::CompressedRle(compressed) => {
593 let (mask, _, _) = decode_compressed_rle(compressed)?;
594 let area = mask.iter().filter(|&&v| v == 1).count() as f64;
595 Ok(area)
596 }
597 }
598}
599
600fn shoelace_area(polygon: &[f64]) -> f64 {
602 if polygon.len() < 6 {
603 return 0.0;
604 }
605
606 let n = polygon.len() / 2;
607 let mut area = 0.0;
608
609 for i in 0..n {
610 let j = (i + 1) % n;
611 let x1 = polygon[i * 2];
612 let y1 = polygon[i * 2 + 1];
613 let x2 = polygon[j * 2];
614 let y2 = polygon[j * 2 + 1];
615 area += x1 * y2 - x2 * y1;
616 }
617
618 (area / 2.0).abs()
619}
620
621#[cfg(test)]
622mod tests {
623 use super::*;
624
625 #[test]
630 fn test_coco_bbox_to_box2d() {
631 let bbox = [100.0, 50.0, 200.0, 150.0];
632 let box2d = coco_bbox_to_box2d(&bbox, 640, 480);
633
634 assert!((box2d.left() - 100.0 / 640.0).abs() < 1e-6);
635 assert!((box2d.top() - 50.0 / 480.0).abs() < 1e-6);
636 assert!((box2d.width() - 200.0 / 640.0).abs() < 1e-6);
637 assert!((box2d.height() - 150.0 / 480.0).abs() < 1e-6);
638 }
639
640 #[test]
641 fn test_box2d_to_coco_bbox() {
642 let box2d = Box2d::new(0.15625, 0.104167, 0.3125, 0.3125);
643 let bbox = box2d_to_coco_bbox(&box2d, 640, 480);
644
645 assert!((bbox[0] - 100.0).abs() < 1.0);
646 assert!((bbox[1] - 50.0).abs() < 1.0);
647 assert!((bbox[2] - 200.0).abs() < 1.0);
648 assert!((bbox[3] - 150.0).abs() < 1.0);
649 }
650
651 #[test]
652 fn test_bbox_roundtrip() {
653 let original = [123.5, 456.7, 89.1, 234.5];
654 let image_w = 1920;
655 let image_h = 1080;
656
657 let box2d = coco_bbox_to_box2d(&original, image_w, image_h);
658 let restored = box2d_to_coco_bbox(&box2d, image_w, image_h);
659
660 for i in 0..4 {
661 assert!(
662 (original[i] - restored[i]).abs() < 1.0,
663 "Mismatch at index {}: {} vs {}",
664 i,
665 original[i],
666 restored[i]
667 );
668 }
669 }
670
671 #[test]
672 fn test_validate_coco_bbox_valid() {
673 assert!(validate_coco_bbox(&[10.0, 20.0, 100.0, 80.0], 640, 480).is_ok());
674 assert!(validate_coco_bbox(&[0.0, 0.0, 640.0, 480.0], 640, 480).is_ok());
675 }
676
677 #[test]
678 fn test_validate_coco_bbox_invalid() {
679 assert!(validate_coco_bbox(&[10.0, 20.0, -100.0, 80.0], 640, 480).is_err());
681 assert!(validate_coco_bbox(&[10.0, 20.0, 0.0, 80.0], 640, 480).is_err());
683 assert!(validate_coco_bbox(&[600.0, 400.0, 100.0, 100.0], 640, 480).is_err());
685 }
686
687 #[test]
692 fn test_coco_polygon_to_polygon() {
693 let polygons = vec![vec![100.0, 100.0, 200.0, 100.0, 200.0, 200.0, 100.0, 200.0]];
694 let polygon = coco_polygon_to_polygon(&polygons, 400, 400);
695
696 assert_eq!(polygon.rings.len(), 1);
697 assert_eq!(polygon.rings[0].len(), 4);
698
699 assert!((polygon.rings[0][0].0 - 0.25).abs() < 1e-6);
701 assert!((polygon.rings[0][0].1 - 0.25).abs() < 1e-6);
702 }
703
704 #[test]
705 fn test_polygon_to_coco_polygon() {
706 let polygon = Polygon::new(vec![vec![
707 (0.25, 0.25),
708 (0.5, 0.25),
709 (0.5, 0.5),
710 (0.25, 0.5),
711 ]]);
712
713 let polygons = polygon_to_coco_polygon(&polygon, 400, 400);
714
715 assert_eq!(polygons.len(), 1);
716 assert_eq!(polygons[0].len(), 8); assert!((polygons[0][0] - 100.0).abs() < 1e-6);
719 assert!((polygons[0][1] - 100.0).abs() < 1e-6);
720 }
721
722 #[test]
723 fn test_polygon_roundtrip() {
724 let original = vec![vec![
725 50.0, 60.0, 150.0, 60.0, 180.0, 120.0, 150.0, 180.0, 50.0, 180.0, 20.0, 120.0,
726 ]];
727
728 let image_w = 300;
729 let image_h = 300;
730
731 let polygon = coco_polygon_to_polygon(&original, image_w, image_h);
732 let restored = polygon_to_coco_polygon(&polygon, image_w, image_h);
733
734 assert_eq!(original.len(), restored.len());
735 assert_eq!(original[0].len(), restored[0].len());
736
737 for i in 0..original[0].len() {
738 assert!(
739 (original[0][i] - restored[0][i]).abs() < 1.0,
740 "Mismatch at index {}: {} vs {}",
741 i,
742 original[0][i],
743 restored[0][i]
744 );
745 }
746 }
747
748 #[test]
749 fn test_polygon_multiple_regions() {
750 let polygons = vec![
751 vec![10.0, 10.0, 50.0, 10.0, 50.0, 50.0, 10.0, 50.0],
752 vec![60.0, 60.0, 90.0, 60.0, 90.0, 90.0, 60.0, 90.0],
753 ];
754
755 let polygon = coco_polygon_to_polygon(&polygons, 100, 100);
756
757 assert_eq!(polygon.rings.len(), 2);
758 assert_eq!(polygon.rings[0].len(), 4);
759 assert_eq!(polygon.rings[1].len(), 4);
760 }
761
762 #[test]
763 fn test_polygon_filters_too_small() {
764 let polygons = vec![
765 vec![10.0, 10.0], vec![10.0, 10.0, 50.0, 50.0], vec![10.0, 10.0, 50.0, 10.0, 50.0, 50.0], ];
769
770 let polygon = coco_polygon_to_polygon(&polygons, 100, 100);
771
772 assert_eq!(polygon.rings.len(), 1);
773 }
774
775 #[test]
776 fn test_polygon_empty_ring_handled() {
777 let polygons: Vec<Vec<f64>> = vec![vec![]];
779 let polygon = coco_polygon_to_polygon(&polygons, 100, 100);
780 assert!(
781 polygon.rings.is_empty(),
782 "Empty polygon ring should be filtered out"
783 );
784 }
785
786 #[test]
791 fn test_decode_rle_simple() {
792 let rle = CocoRle {
799 counts: vec![1, 2, 1, 2], size: [3, 2], };
803
804 let result = decode_rle(&rle);
810 assert!(result.is_ok());
811
812 let (mask, height, width) = result.unwrap();
813 assert_eq!(height, 3);
814 assert_eq!(width, 2);
815 assert_eq!(mask.len(), 6);
816
817 }
822
823 #[test]
824 fn test_decode_rle_all_background() {
825 let rle = CocoRle {
826 counts: vec![100], size: [10, 10],
828 };
829
830 let (mask, _, _) = decode_rle(&rle).unwrap();
831 assert!(mask.iter().all(|&v| v == 0));
832 }
833
834 #[test]
835 fn test_decode_rle_all_foreground() {
836 let rle = CocoRle {
837 counts: vec![0, 100], size: [10, 10],
839 };
840
841 let (mask, _, _) = decode_rle(&rle).unwrap();
842 assert!(mask.iter().all(|&v| v == 1));
843 }
844
845 #[test]
846 fn test_decode_rle_invalid_counts() {
847 let rle = CocoRle {
848 counts: vec![50], size: [10, 10],
850 };
851
852 let result = decode_rle(&rle);
853 assert!(result.is_err());
854 }
855
856 #[test]
861 fn test_shoelace_area_square() {
862 let polygon = vec![0.0, 0.0, 100.0, 0.0, 100.0, 100.0, 0.0, 100.0];
864 let area = shoelace_area(&polygon);
865 assert!((area - 10000.0).abs() < 1e-6);
866 }
867
868 #[test]
869 fn test_shoelace_area_triangle() {
870 let polygon = vec![0.0, 0.0, 100.0, 0.0, 50.0, 100.0];
873 let area = shoelace_area(&polygon);
874 assert!((area - 5000.0).abs() < 1e-6);
875 }
876
877 #[test]
878 fn test_calculate_coco_area_polygon() {
879 let seg =
880 CocoSegmentation::Polygon(vec![vec![0.0, 0.0, 100.0, 0.0, 100.0, 100.0, 0.0, 100.0]]);
881 let area = calculate_coco_area(&seg).unwrap();
882 assert!((area - 10000.0).abs() < 1e-6);
883 }
884
885 #[test]
890 fn test_encode_rle_all_background() {
891 let mask = vec![0u8; 100];
892 let rle = encode_rle(&mask, 10, 10).unwrap();
893 assert_eq!(rle.size, [10, 10]);
894 assert_eq!(rle.counts, vec![100]);
896 }
897
898 #[test]
899 fn test_encode_rle_all_foreground() {
900 let mask = vec![1u8; 100];
901 let rle = encode_rle(&mask, 10, 10).unwrap();
902 assert_eq!(rle.size, [10, 10]);
903 assert_eq!(rle.counts, vec![0, 100]);
905 }
906
907 #[test]
908 fn test_encode_decode_rle_roundtrip() {
909 #[rustfmt::skip]
911 let mask = vec![
912 1, 1, 0, 0,
913 1, 1, 0, 0,
914 0, 0, 0, 0,
915 0, 0, 0, 0,
916 ];
917
918 let rle = encode_rle(&mask, 4, 4).unwrap();
919 assert_eq!(rle.size, [4, 4]);
920 let counts_sum: u32 = rle.counts.iter().sum();
921 assert_eq!(counts_sum, 16, "RLE counts should sum to total pixels");
922
923 let (decoded, height, width) = decode_rle(&rle).unwrap();
925 assert_eq!(height, 4);
926 assert_eq!(width, 4);
927 assert_eq!(decoded, mask);
928 }
929
930 #[test]
931 fn test_encode_rle_single_pixel_foreground() {
932 #[rustfmt::skip]
934 let mask = vec![
935 1, 0, 0,
936 0, 0, 0,
937 0, 0, 0,
938 ];
939
940 let rle = encode_rle(&mask, 3, 3).unwrap();
941 let (decoded, _, _) = decode_rle(&rle).unwrap();
942 assert_eq!(decoded, mask);
943 }
944
945 #[test]
946 fn test_encode_rle_size_mismatch() {
947 let mask = vec![0u8; 50];
948 assert!(
949 encode_rle(&mask, 10, 10).is_err(),
950 "Should reject mask length != width * height"
951 );
952 }
953}