1use super::types::{CocoCompressedRle, CocoRle, CocoSegmentation};
14use crate::{Box2d, Error, Mask};
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_mask(polygons: &[Vec<f64>], image_width: u32, image_height: u32) -> Mask {
137 let img_w = image_width as f64;
138 let img_h = image_height as f64;
139
140 let converted: Vec<Vec<(f32, f32)>> = polygons
141 .iter()
142 .filter(|poly| poly.len() >= 6) .map(|polygon| {
144 polygon
145 .chunks(2)
146 .filter_map(|chunk| {
147 if chunk.len() == 2 {
148 Some(((chunk[0] / img_w) as f32, (chunk[1] / img_h) as f32))
149 } else {
150 None
151 }
152 })
153 .collect()
154 })
155 .filter(|poly: &Vec<(f32, f32)>| poly.len() >= 3) .collect();
157
158 Mask::new(converted)
159}
160
161pub fn mask_to_coco_polygon(mask: &Mask, image_width: u32, image_height: u32) -> Vec<Vec<f64>> {
171 let img_w = image_width as f64;
172 let img_h = image_height as f64;
173
174 mask.polygon
175 .iter()
176 .filter(|poly| poly.len() >= 3) .map(|polygon| {
178 polygon
179 .iter()
180 .flat_map(|(x, y)| vec![(*x as f64) * img_w, (*y as f64) * img_h])
181 .collect()
182 })
183 .collect()
184}
185
186pub fn decode_rle(rle: &CocoRle) -> Result<(Vec<u8>, u32, u32), Error> {
201 let [height, width] = rle.size;
202 let total_pixels = (width as usize) * (height as usize);
203
204 let counts_sum: u64 = rle.counts.iter().map(|&c| c as u64).sum();
206 if counts_sum != total_pixels as u64 {
207 return Err(Error::CocoError(format!(
208 "RLE counts sum {} does not match image size {}x{} = {}",
209 counts_sum, width, height, total_pixels
210 )));
211 }
212
213 let mut column_major = vec![0u8; total_pixels];
215 let mut pos = 0usize;
216 let mut is_foreground = false; for &count in &rle.counts {
219 let count = count as usize;
220 if is_foreground {
221 for i in pos..(pos + count).min(column_major.len()) {
222 column_major[i] = 1;
223 }
224 }
225 pos += count;
226 is_foreground = !is_foreground;
227 }
228
229 let mut row_major = vec![0u8; total_pixels];
231 for col in 0..width as usize {
232 for row in 0..height as usize {
233 let col_idx = col * (height as usize) + row;
234 let row_idx = row * (width as usize) + col;
235 if col_idx < column_major.len() && row_idx < row_major.len() {
236 row_major[row_idx] = column_major[col_idx];
237 }
238 }
239 }
240
241 Ok((row_major, height, width))
242}
243
244fn decode_leb128(s: &str) -> Result<Vec<u32>, Error> {
248 let bytes = s.as_bytes();
249 let mut counts = Vec::new();
250 let mut i = 0;
251
252 while i < bytes.len() {
253 let mut value: i64 = 0;
254 let mut shift = 0;
255 let mut more = true;
256
257 while more && i < bytes.len() {
258 let byte = bytes[i] as i64;
259 i += 1;
260
261 let decoded = if (48..96).contains(&byte) {
263 byte - 48 } else if byte >= 96 {
265 byte - 96 + 48 } else {
267 return Err(Error::CocoError(format!(
268 "Invalid LEB128 character: {}",
269 byte as u8 as char
270 )));
271 };
272
273 value |= (decoded & 0x1F) << shift;
274 more = decoded >= 32;
275 shift += 5;
276 }
277
278 if shift < 32 && (value & (1 << (shift - 1))) != 0 {
280 value |= (-1i64) << shift;
281 }
282
283 counts.push(value);
284 }
285
286 let mut result = Vec::with_capacity(counts.len());
288 let mut prev: i64 = 0;
289 for diff in counts {
290 prev += diff;
291 result.push(prev.max(0) as u32);
292 }
293
294 Ok(result)
295}
296
297pub fn decode_compressed_rle(compressed: &CocoCompressedRle) -> Result<(Vec<u8>, u32, u32), Error> {
299 let counts = decode_leb128(&compressed.counts)?;
300
301 let rle = CocoRle {
302 counts,
303 size: compressed.size,
304 };
305
306 decode_rle(&rle)
307}
308
309pub fn mask_to_contours(mask: &[u8], width: u32, height: u32) -> Vec<Vec<(f64, f64)>> {
326 let mut contours = Vec::new();
327 let mut visited = vec![false; mask.len()];
328
329 let w = width as usize;
330 let h = height as usize;
331
332 for start_y in 0..h {
333 for start_x in 0..w {
334 let idx = start_y * w + start_x;
335 if mask[idx] == 1 && !visited[idx] {
336 let is_boundary = start_x == 0
339 || start_x == w - 1
340 || start_y == 0
341 || start_y == h - 1
342 || (start_x > 0 && mask[idx - 1] == 0)
343 || (start_x < w - 1 && mask[idx + 1] == 0)
344 || (start_y > 0 && mask[idx - w] == 0)
345 || (start_y < h - 1 && mask[idx + w] == 0);
346
347 if is_boundary
348 && let Some(contour) = trace_contour(mask, w, h, start_x, start_y, &mut visited)
349 && contour.len() >= 3
350 {
351 contours.push(contour);
352 }
353 }
354 }
355 }
356
357 contours
358}
359
360fn trace_contour(
362 mask: &[u8],
363 width: usize,
364 height: usize,
365 start_x: usize,
366 start_y: usize,
367 visited: &mut [bool],
368) -> Option<Vec<(f64, f64)>> {
369 let mut contour = Vec::new();
370 let mut x = start_x;
371 let mut y = start_y;
372
373 let dx: [i32; 8] = [1, 1, 0, -1, -1, -1, 0, 1];
375 let dy: [i32; 8] = [0, 1, 1, 1, 0, -1, -1, -1];
376
377 let mut dir = 0usize; let max_steps = width * height;
379 let mut steps = 0;
380
381 loop {
382 let idx = y * width + x;
383 if !visited[idx] {
384 contour.push((x as f64, y as f64));
385 visited[idx] = true;
386 }
387
388 let mut found = false;
390 for i in 0..8 {
391 let new_dir = (dir + i) % 8;
392 let nx = x as i32 + dx[new_dir];
393 let ny = y as i32 + dy[new_dir];
394
395 if nx >= 0 && nx < width as i32 && ny >= 0 && ny < height as i32 {
396 let nidx = (ny as usize) * width + (nx as usize);
397 if mask[nidx] == 1 {
398 x = nx as usize;
399 y = ny as usize;
400 dir = (new_dir + 5) % 8; found = true;
402 break;
403 }
404 }
405 }
406
407 if !found || (x == start_x && y == start_y && contour.len() > 2) {
408 break;
409 }
410
411 steps += 1;
412 if steps > max_steps {
413 break; }
415 }
416
417 if contour.len() >= 3 {
418 Some(contour)
419 } else {
420 None
421 }
422}
423
424pub fn coco_rle_to_mask(rle: &CocoRle, image_width: u32, image_height: u32) -> Result<Mask, Error> {
428 let (binary_mask, height, width) = decode_rle(rle)?;
429 let contours = mask_to_contours(&binary_mask, width, height);
430
431 let normalized: Vec<Vec<(f32, f32)>> = contours
433 .iter()
434 .map(|contour| {
435 contour
436 .iter()
437 .map(|(x, y)| {
438 (
439 (*x / image_width as f64) as f32,
440 (*y / image_height as f64) as f32,
441 )
442 })
443 .collect()
444 })
445 .collect();
446
447 Ok(Mask::new(normalized))
448}
449
450pub fn coco_segmentation_to_mask(
454 segmentation: &CocoSegmentation,
455 image_width: u32,
456 image_height: u32,
457) -> Result<Mask, Error> {
458 match segmentation {
459 CocoSegmentation::Polygon(polygons) => {
460 Ok(coco_polygon_to_mask(polygons, image_width, image_height))
461 }
462 CocoSegmentation::Rle(rle) => coco_rle_to_mask(rle, image_width, image_height),
463 CocoSegmentation::CompressedRle(compressed) => {
464 let counts = decode_leb128(&compressed.counts)?;
465 let rle = CocoRle {
466 counts,
467 size: compressed.size,
468 };
469 coco_rle_to_mask(&rle, image_width, image_height)
470 }
471 }
472}
473
474pub fn calculate_coco_area(segmentation: &CocoSegmentation) -> Result<f64, Error> {
480 match segmentation {
481 CocoSegmentation::Polygon(polygons) => {
482 let mut total_area = 0.0;
484 for polygon in polygons {
485 total_area += shoelace_area(polygon);
486 }
487 Ok(total_area)
488 }
489 CocoSegmentation::Rle(rle) => {
490 let (mask, _, _) = decode_rle(rle)?;
491 let area = mask.iter().filter(|&&v| v == 1).count() as f64;
492 Ok(area)
493 }
494 CocoSegmentation::CompressedRle(compressed) => {
495 let (mask, _, _) = decode_compressed_rle(compressed)?;
496 let area = mask.iter().filter(|&&v| v == 1).count() as f64;
497 Ok(area)
498 }
499 }
500}
501
502fn shoelace_area(polygon: &[f64]) -> f64 {
504 if polygon.len() < 6 {
505 return 0.0;
506 }
507
508 let n = polygon.len() / 2;
509 let mut area = 0.0;
510
511 for i in 0..n {
512 let j = (i + 1) % n;
513 let x1 = polygon[i * 2];
514 let y1 = polygon[i * 2 + 1];
515 let x2 = polygon[j * 2];
516 let y2 = polygon[j * 2 + 1];
517 area += x1 * y2 - x2 * y1;
518 }
519
520 (area / 2.0).abs()
521}
522
523#[cfg(test)]
524mod tests {
525 use super::*;
526
527 #[test]
532 fn test_coco_bbox_to_box2d() {
533 let bbox = [100.0, 50.0, 200.0, 150.0];
534 let box2d = coco_bbox_to_box2d(&bbox, 640, 480);
535
536 assert!((box2d.left() - 100.0 / 640.0).abs() < 1e-6);
537 assert!((box2d.top() - 50.0 / 480.0).abs() < 1e-6);
538 assert!((box2d.width() - 200.0 / 640.0).abs() < 1e-6);
539 assert!((box2d.height() - 150.0 / 480.0).abs() < 1e-6);
540 }
541
542 #[test]
543 fn test_box2d_to_coco_bbox() {
544 let box2d = Box2d::new(0.15625, 0.104167, 0.3125, 0.3125);
545 let bbox = box2d_to_coco_bbox(&box2d, 640, 480);
546
547 assert!((bbox[0] - 100.0).abs() < 1.0);
548 assert!((bbox[1] - 50.0).abs() < 1.0);
549 assert!((bbox[2] - 200.0).abs() < 1.0);
550 assert!((bbox[3] - 150.0).abs() < 1.0);
551 }
552
553 #[test]
554 fn test_bbox_roundtrip() {
555 let original = [123.5, 456.7, 89.1, 234.5];
556 let image_w = 1920;
557 let image_h = 1080;
558
559 let box2d = coco_bbox_to_box2d(&original, image_w, image_h);
560 let restored = box2d_to_coco_bbox(&box2d, image_w, image_h);
561
562 for i in 0..4 {
563 assert!(
564 (original[i] - restored[i]).abs() < 1.0,
565 "Mismatch at index {}: {} vs {}",
566 i,
567 original[i],
568 restored[i]
569 );
570 }
571 }
572
573 #[test]
574 fn test_validate_coco_bbox_valid() {
575 assert!(validate_coco_bbox(&[10.0, 20.0, 100.0, 80.0], 640, 480).is_ok());
576 assert!(validate_coco_bbox(&[0.0, 0.0, 640.0, 480.0], 640, 480).is_ok());
577 }
578
579 #[test]
580 fn test_validate_coco_bbox_invalid() {
581 assert!(validate_coco_bbox(&[10.0, 20.0, -100.0, 80.0], 640, 480).is_err());
583 assert!(validate_coco_bbox(&[10.0, 20.0, 0.0, 80.0], 640, 480).is_err());
585 assert!(validate_coco_bbox(&[600.0, 400.0, 100.0, 100.0], 640, 480).is_err());
587 }
588
589 #[test]
594 fn test_coco_polygon_to_mask() {
595 let polygons = vec![vec![100.0, 100.0, 200.0, 100.0, 200.0, 200.0, 100.0, 200.0]];
596 let mask = coco_polygon_to_mask(&polygons, 400, 400);
597
598 assert_eq!(mask.polygon.len(), 1);
599 assert_eq!(mask.polygon[0].len(), 4);
600
601 assert!((mask.polygon[0][0].0 - 0.25).abs() < 1e-6);
603 assert!((mask.polygon[0][0].1 - 0.25).abs() < 1e-6);
604 }
605
606 #[test]
607 fn test_mask_to_coco_polygon() {
608 let mask = Mask::new(vec![vec![
609 (0.25, 0.25),
610 (0.5, 0.25),
611 (0.5, 0.5),
612 (0.25, 0.5),
613 ]]);
614
615 let polygons = mask_to_coco_polygon(&mask, 400, 400);
616
617 assert_eq!(polygons.len(), 1);
618 assert_eq!(polygons[0].len(), 8); assert!((polygons[0][0] - 100.0).abs() < 1e-6);
621 assert!((polygons[0][1] - 100.0).abs() < 1e-6);
622 }
623
624 #[test]
625 fn test_polygon_roundtrip() {
626 let original = vec![vec![
627 50.0, 60.0, 150.0, 60.0, 180.0, 120.0, 150.0, 180.0, 50.0, 180.0, 20.0, 120.0,
628 ]];
629
630 let image_w = 300;
631 let image_h = 300;
632
633 let mask = coco_polygon_to_mask(&original, image_w, image_h);
634 let restored = mask_to_coco_polygon(&mask, image_w, image_h);
635
636 assert_eq!(original.len(), restored.len());
637 assert_eq!(original[0].len(), restored[0].len());
638
639 for i in 0..original[0].len() {
640 assert!(
641 (original[0][i] - restored[0][i]).abs() < 1.0,
642 "Mismatch at index {}: {} vs {}",
643 i,
644 original[0][i],
645 restored[0][i]
646 );
647 }
648 }
649
650 #[test]
651 fn test_polygon_multiple_regions() {
652 let polygons = vec![
653 vec![10.0, 10.0, 50.0, 10.0, 50.0, 50.0, 10.0, 50.0],
654 vec![60.0, 60.0, 90.0, 60.0, 90.0, 90.0, 60.0, 90.0],
655 ];
656
657 let mask = coco_polygon_to_mask(&polygons, 100, 100);
658
659 assert_eq!(mask.polygon.len(), 2);
660 assert_eq!(mask.polygon[0].len(), 4);
661 assert_eq!(mask.polygon[1].len(), 4);
662 }
663
664 #[test]
665 fn test_polygon_filters_too_small() {
666 let polygons = vec![
667 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], ];
671
672 let mask = coco_polygon_to_mask(&polygons, 100, 100);
673
674 assert_eq!(mask.polygon.len(), 1);
675 }
676
677 #[test]
682 fn test_decode_rle_simple() {
683 let rle = CocoRle {
690 counts: vec![1, 2, 1, 2], size: [3, 2], };
694
695 let result = decode_rle(&rle);
701 assert!(result.is_ok());
702
703 let (mask, height, width) = result.unwrap();
704 assert_eq!(height, 3);
705 assert_eq!(width, 2);
706 assert_eq!(mask.len(), 6);
707
708 }
713
714 #[test]
715 fn test_decode_rle_all_background() {
716 let rle = CocoRle {
717 counts: vec![100], size: [10, 10],
719 };
720
721 let (mask, _, _) = decode_rle(&rle).unwrap();
722 assert!(mask.iter().all(|&v| v == 0));
723 }
724
725 #[test]
726 fn test_decode_rle_all_foreground() {
727 let rle = CocoRle {
728 counts: vec![0, 100], size: [10, 10],
730 };
731
732 let (mask, _, _) = decode_rle(&rle).unwrap();
733 assert!(mask.iter().all(|&v| v == 1));
734 }
735
736 #[test]
737 fn test_decode_rle_invalid_counts() {
738 let rle = CocoRle {
739 counts: vec![50], size: [10, 10],
741 };
742
743 let result = decode_rle(&rle);
744 assert!(result.is_err());
745 }
746
747 #[test]
752 fn test_shoelace_area_square() {
753 let polygon = vec![0.0, 0.0, 100.0, 0.0, 100.0, 100.0, 0.0, 100.0];
755 let area = shoelace_area(&polygon);
756 assert!((area - 10000.0).abs() < 1e-6);
757 }
758
759 #[test]
760 fn test_shoelace_area_triangle() {
761 let polygon = vec![0.0, 0.0, 100.0, 0.0, 50.0, 100.0];
764 let area = shoelace_area(&polygon);
765 assert!((area - 5000.0).abs() < 1e-6);
766 }
767
768 #[test]
769 fn test_calculate_coco_area_polygon() {
770 let seg =
771 CocoSegmentation::Polygon(vec![vec![0.0, 0.0, 100.0, 0.0, 100.0, 100.0, 0.0, 100.0]]);
772 let area = calculate_coco_area(&seg).unwrap();
773 assert!((area - 10000.0).abs() < 1e-6);
774 }
775}