usls/results/
mask.rs

1use aksr::Builder;
2use anyhow::Result;
3use image::GrayImage;
4use rayon::prelude::*;
5
6use crate::{InstanceMeta, MaskStyle, Polygon};
7
8/// Mask: Gray Image.
9#[derive(Builder, Default, Clone)]
10pub struct Mask {
11    /// The grayscale image representing the mask.
12    mask: GrayImage,
13    /// Metadata associated with the mask instance.
14    meta: InstanceMeta,
15    /// Optional styling information for visualization.
16    style: Option<MaskStyle>,
17}
18
19// #[derive(Builder, Default, Clone)]
20// pub struct Masks(Vec<Mask>);
21
22impl std::fmt::Debug for Mask {
23    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
24        let mut f = f.debug_struct("Mask");
25        f.field("dimensions", &self.dimensions());
26        if let Some(id) = &self.meta.id() {
27            f.field("id", id);
28        }
29        if let Some(name) = &self.meta.name() {
30            f.field("name", name);
31        }
32        if let Some(confidence) = &self.meta.confidence() {
33            f.field("confidence", confidence);
34        }
35        if let Some(track_id) = &self.meta.track_id() {
36            f.field("track_id", track_id);
37        }
38        f.finish()
39    }
40}
41
42impl PartialEq for Mask {
43    fn eq(&self, other: &Self) -> bool {
44        self.mask == other.mask
45    }
46}
47
48impl Mask {
49    impl_meta_methods!();
50    pub fn new(u8s: &[u8], width: u32, height: u32) -> Result<Self> {
51        let mask: image::ImageBuffer<image::Luma<_>, Vec<_>> =
52            image::ImageBuffer::from_raw(width, height, u8s.to_vec())
53                .ok_or(anyhow::anyhow!("Failed to build ImageBuffer."))?;
54
55        Ok(Self {
56            mask,
57            ..Default::default()
58        })
59    }
60
61    pub fn to_vec(&self) -> Vec<u8> {
62        self.mask.to_vec()
63    }
64
65    pub fn height(&self) -> u32 {
66        self.mask.height()
67    }
68
69    pub fn width(&self) -> u32 {
70        self.mask.width()
71    }
72
73    pub fn dimensions(&self) -> (u32, u32) {
74        self.mask.dimensions()
75    }
76
77    pub fn polygon(&self) -> Option<Polygon> {
78        let polygons = self.polygons();
79        if polygons.is_empty() {
80            return None;
81        }
82
83        polygons.into_iter().max_by(|x, y| {
84            x.area()
85                .partial_cmp(&y.area())
86                .unwrap_or(std::cmp::Ordering::Equal)
87        })
88    }
89
90    pub fn polygons(&self) -> Vec<Polygon> {
91        let contours: Vec<imageproc::contours::Contour<i32>> =
92            imageproc::contours::find_contours_with_threshold(self.mask(), 0);
93        let polygons: Vec<Polygon> = contours
94            .into_par_iter()
95            .filter_map(|contour| {
96                if contour.border_type == imageproc::contours::BorderType::Hole
97                    && contour.points.len() <= 2
98                {
99                    return None;
100                }
101                let coords: Vec<[f32; 2]> = contour
102                    .points
103                    .iter()
104                    .map(|p| [p.x as f32, p.y as f32])
105                    .collect();
106
107                let mut polygon = match Polygon::try_from(coords) {
108                    Ok(p) => p.verify(),
109                    Err(_) => return None,
110                };
111                if let Some(x) = self.name() {
112                    polygon = polygon.with_name(x);
113                }
114                if let Some(x) = self.id() {
115                    polygon = polygon.with_id(x);
116                }
117                if let Some(x) = self.confidence() {
118                    polygon = polygon.with_confidence(x);
119                }
120                Some(polygon)
121            })
122            .collect();
123
124        polygons
125    }
126}
usls/results/mask.rs

usls/results/
mask.rs
