Skip to main content

edgefirst_client/coco/
convert.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright © 2025 Au-Zone Technologies. All Rights Reserved.
3
4//! Coordinate and geometry conversion functions between COCO and EdgeFirst
5//! formats.
6//!
7//! ## Coordinate Systems
8//!
9//! - **COCO**: Pixel coordinates, top-left origin for bboxes
10//! - **EdgeFirst JSON**: Normalized 0-1, top-left origin for Box2d
11//! - **EdgeFirst Arrow**: Normalized 0-1, center-point for box2d column
12
13use super::types::{CocoCompressedRle, CocoRle, CocoSegmentation};
14use crate::{Box2d, Error, MaskData, Polygon};
15
16// =============================================================================
17// Bounding Box Conversion
18// =============================================================================
19
20/// Convert COCO bbox `[x, y, w, h]` (top-left, pixels) to EdgeFirst `Box2d`
21/// (top-left, normalized).
22///
23/// EdgeFirst `Box2d` uses `{x, y, w, h}` where `(x, y)` is the top-left corner,
24/// normalized to the range `[0, 1]`.
25///
26/// # Arguments
27/// * `bbox` - COCO bounding box `[x_min, y_min, width, height]` in pixels
28/// * `image_width` - Image width in pixels
29/// * `image_height` - Image height in pixels
30///
31/// # Returns
32/// EdgeFirst `Box2d` with normalized coordinates
33///
34/// # Example
35/// ```
36/// use edgefirst_client::coco::coco_bbox_to_box2d;
37///
38/// let coco_bbox = [100.0, 50.0, 200.0, 150.0]; // x=100, y=50, w=200, h=150
39/// let box2d = coco_bbox_to_box2d(&coco_bbox, 640, 480);
40///
41/// assert!((box2d.left() - 100.0 / 640.0).abs() < 1e-6);
42/// assert!((box2d.top() - 50.0 / 480.0).abs() < 1e-6);
43/// ```
44pub 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
57/// Convert EdgeFirst `Box2d` to COCO bbox `[x, y, w, h]` (top-left, pixels).
58///
59/// # Arguments
60/// * `box2d` - EdgeFirst `Box2d` (normalized 0-1, top-left origin)
61/// * `image_width` - Image width in pixels
62/// * `image_height` - Image height in pixels
63///
64/// # Returns
65/// COCO bbox `[x_min, y_min, width, height]` in pixels
66pub 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
78/// Validate that a COCO bounding box is within image bounds.
79///
80/// # Arguments
81/// * `bbox` - COCO bounding box `[x, y, w, h]` in pixels
82/// * `image_width` - Image width in pixels
83/// * `image_height` - Image height in pixels
84///
85/// # Returns
86/// `Ok(())` if valid, `Err` with description if invalid
87pub 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    // Allow slight overflow for floating point precision
102    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
120// =============================================================================
121// Polygon Conversion
122// =============================================================================
123
124/// Convert COCO polygon segmentation to EdgeFirst `Polygon` format.
125///
126/// COCO polygons: `[[x1,y1,x2,y2,...], [x3,y3,...]]` (nested, pixel
127/// coordinates) EdgeFirst Polygon: `Vec<Vec<(f32, f32)>>` (nested, normalized
128/// 0-1)
129///
130/// # Arguments
131/// * `polygons` - COCO polygon array (nested Vec of pixel coordinates)
132/// * `image_width` - Image width in pixels
133/// * `image_height` - Image height in pixels
134///
135/// # Returns
136/// EdgeFirst `Polygon` with normalized coordinates
137pub 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) // Need at least 3 points
148        .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) // Still need 3+ points after conversion
161        .collect();
162
163    Polygon::new(converted)
164}
165
166/// Convert EdgeFirst `Polygon` format to COCO polygon segmentation.
167///
168/// # Arguments
169/// * `polygon` - EdgeFirst `Polygon` with normalized coordinates
170/// * `image_width` - Image width in pixels
171/// * `image_height` - Image height in pixels
172///
173/// # Returns
174/// COCO polygon array (nested Vec of pixel coordinates)
175pub 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) // Need at least 3 points
187        .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
195// =============================================================================
196// RLE Decoding
197// =============================================================================
198
199/// Decode uncompressed RLE to binary mask.
200///
201/// **CRITICAL**: RLE uses column-major (Fortran) order, starting with
202/// background.
203///
204/// # Arguments
205/// * `rle` - COCO RLE with counts array
206///
207/// # Returns
208/// Binary mask as `Vec<u8>` in row-major order, plus `(height, width)`
209pub 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    // Validate counts sum
214    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    // Decode to column-major flat array
223    let mut column_major = vec![0u8; total_pixels];
224    let mut pos = 0usize;
225    let mut is_foreground = false; // Starts with background
226
227    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    // Convert column-major to row-major
239    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
253/// Decode LEB128 encoded string to counts array.
254///
255/// Based on pycocotools encoding.
256fn 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            // Decode based on character ranges (pycocotools encoding)
271            let decoded = if (48..96).contains(&byte) {
272                byte - 48 // '0'-'_'
273            } else if byte >= 96 {
274                byte - 96 + 48 // 'a' and above
275            } 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        // Sign extend if needed
288        if shift < 32 && (value & (1 << (shift - 1))) != 0 {
289            value |= (-1i64) << shift;
290        }
291
292        counts.push(value);
293    }
294
295    // Convert from diff encoding to absolute counts
296    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
306/// Decode compressed RLE (LEB128) to binary mask.
307pub 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
318// =============================================================================
319// Contour Extraction
320// =============================================================================
321
322/// Convert binary mask to polygon contours.
323///
324/// Uses a simple boundary tracing algorithm to extract outer contours from
325/// a binary segmentation mask.
326///
327/// # Arguments
328/// * `mask` - Binary mask (0 = background, 1 = foreground) in row-major order
329/// * `width` - Image width
330/// * `height` - Image height
331///
332/// # Returns
333/// Vector of contours, each contour is a vector of `(x, y)` pixel coordinates
334pub 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                // Check if this is a boundary pixel (has at least one neighbor that's 0 or
346                // edge)
347                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
369/// Trace a contour starting from the given point using 8-connectivity.
370fn 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    // Direction vectors for 8-connectivity: E, SE, S, SW, W, NW, N, NE
383    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; // Start going east
387    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        // Find next boundary pixel
398        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; // Turn around and search from there
410                    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; // Safety limit
423        }
424    }
425
426    if contour.len() >= 3 {
427        Some(contour)
428    } else {
429        None
430    }
431}
432
433/// Convert RLE segmentation to EdgeFirst `Polygon`.
434///
435/// Decodes the RLE, extracts contours, and normalizes to `[0, 1]` range.
436pub 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    // Normalize contours to 0-1 range
445    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
463/// Convert any COCO segmentation to EdgeFirst `Polygon`.
464///
465/// Handles all segmentation types: polygon, RLE, and compressed RLE.
466pub 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
487// =============================================================================
488// RLE → MaskData (PNG) Conversion
489// =============================================================================
490
491/// Convert COCO RLE segmentation to PNG-encoded MaskData (1-bit binary).
492pub 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
497/// Convert any COCO segmentation to MaskData for RLE variants.
498///
499/// Returns `None` for polygon segmentation (use `coco_segmentation_to_polygon`
500/// instead).
501pub 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
512// =============================================================================
513// RLE Encoding
514// =============================================================================
515
516/// Encode a binary mask (row-major, 0/1 values) as COCO RLE.
517///
518/// COCO RLE uses column-major (Fortran) order, starting with a background
519/// count. This function transposes the input from row-major to column-major
520/// before run-length encoding.
521///
522/// # Arguments
523/// * `mask` - Binary mask in row-major order (0 = background, nonzero = foreground)
524/// * `width` - Image width in pixels
525/// * `height` - Image height in pixels
526///
527/// # Returns
528/// A `CocoRle` with counts and size fields.
529pub 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    // Convert row-major to column-major
543    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    // Run-length encode (starts with background count)
551    let mut counts = Vec::new();
552    let mut current = 0u8; // start with background
553    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
572// =============================================================================
573// Area Calculation
574// =============================================================================
575
576/// Calculate area from COCO segmentation (in pixels²).
577pub fn calculate_coco_area(segmentation: &CocoSegmentation) -> Result<f64, Error> {
578    match segmentation {
579        CocoSegmentation::Polygon(polygons) => {
580            // Use shoelace formula for polygon area
581            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
600/// Calculate polygon area using the shoelace formula.
601fn 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    // =========================================================================
626    // Bounding Box Tests
627    // =========================================================================
628
629    #[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        // Negative dimensions
680        assert!(validate_coco_bbox(&[10.0, 20.0, -100.0, 80.0], 640, 480).is_err());
681        // Zero dimensions
682        assert!(validate_coco_bbox(&[10.0, 20.0, 0.0, 80.0], 640, 480).is_err());
683        // Out of bounds
684        assert!(validate_coco_bbox(&[600.0, 400.0, 100.0, 100.0], 640, 480).is_err());
685    }
686
687    // =========================================================================
688    // Polygon Tests
689    // =========================================================================
690
691    #[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        // Check normalized coordinates
700        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); // 4 points * 2 coords
717
718        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],                         // Only 1 point - should be filtered
766            vec![10.0, 10.0, 50.0, 50.0],             // Only 2 points - should be filtered
767            vec![10.0, 10.0, 50.0, 10.0, 50.0, 50.0], // 3 points - should be kept
768        ];
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        // Empty segmentation polygon vec![] should be handled gracefully
778        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    // =========================================================================
787    // RLE Tests
788    // =========================================================================
789
790    #[test]
791    fn test_decode_rle_simple() {
792        // 2x3 image with pattern:
793        // 0 1
794        // 1 1
795        // 0 0
796        // Column-major: [0,1,0], [1,1,0] → counts: [1,1,1, 0,2,1] simplified to
797        // [1,2,1,2]
798        let rle = CocoRle {
799            counts: vec![1, 2, 1, 2], /* bg=1, fg=2, bg=1, fg=2 (wait, that's 6 not 6... let me
800                                       * recalc) */
801            size: [3, 2], // height=3, width=2
802        };
803
804        // Total pixels = 6
805        // counts = [1, 2, 1, 2] sums to 6 ✓
806        // Column 0: bg=1 (pixel 0), fg=2 (pixels 1,2)
807        // Column 1: bg=1 (pixel 3), fg=2 (pixels 4,5)
808
809        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        // Row-major layout:
818        // Row 0: mask[0]=col0_row0, mask[1]=col1_row0
819        // Row 1: mask[2]=col0_row1, mask[3]=col1_row1
820        // Row 2: mask[4]=col0_row2, mask[5]=col1_row2
821    }
822
823    #[test]
824    fn test_decode_rle_all_background() {
825        let rle = CocoRle {
826            counts: vec![100], // All background
827            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], // No background, all foreground
838            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], // Only 50 pixels, but image is 100
849            size: [10, 10],
850        };
851
852        let result = decode_rle(&rle);
853        assert!(result.is_err());
854    }
855
856    // =========================================================================
857    // Area Calculation Tests
858    // =========================================================================
859
860    #[test]
861    fn test_shoelace_area_square() {
862        // 100x100 square
863        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        // Triangle with vertices at (0,0), (100,0), (50,100)
871        // Area = 0.5 * base * height = 0.5 * 100 * 100 = 5000
872        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    // =========================================================================
886    // RLE Encoding Tests
887    // =========================================================================
888
889    #[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        // All background → single count of 100
895        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        // All foreground → [0 bg, 100 fg]
904        assert_eq!(rle.counts, vec![0, 100]);
905    }
906
907    #[test]
908    fn test_encode_decode_rle_roundtrip() {
909        // Create a known mask: 2x2 square in top-left of 4x4 image
910        #[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        // Decode back and verify pixel-for-pixel match
924        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        // Only pixel (0,0) is foreground in 3x3 image
933        #[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}