iris/segmentation/
mod.rs

1pub mod components;
2pub mod watershed;
3
4pub use components::ComponentStats;
5
6use crate::dnn::{OnnxModel, WeightLoader};
7use crate::error::Result;
8use crate::image::Image;
9use burn::tensor::{Int, Tensor, backend::Backend};
10
11/// Semantic segmentation mask output.
12pub struct SegmentationMask<B: Backend> {
13    /// Mask tensor of shape [H, W] containing class labels or probabilities.
14    pub mask: Tensor<B, 2, Int>,
15}
16
17pub struct Segmenter<B: Backend> {
18    #[allow(dead_code)]
19    model: Option<OnnxModel<B>>,
20}
21
22impl<B: Backend> Segmenter<B> {
23    pub fn new(model: OnnxModel<B>) -> Self {
24        Self { model: Some(model) }
25    }
26
27    /// Loads a Segmenter with default pretrained weights implicitly.
28    pub fn pretrained(device: &B::Device) -> Result<Self> {
29        if let Ok(model) = OnnxModel::load("weights/segmenter.onnx", device) {
30            Ok(Self { model: Some(model) })
31        } else if let Ok(model) = OnnxModel::load("segmenter_mock.onnx", device) {
32            Ok(Self { model: Some(model) })
33        } else {
34            Ok(Self { model: None })
35        }
36    }
37
38    /// Loads a Segmenter from an ONNX model.
39    pub fn from_onnx(path: impl AsRef<std::path::Path>, device: &B::Device) -> Result<Self> {
40        let model = OnnxModel::load(path, device)?;
41        Ok(Self { model: Some(model) })
42    }
43
44    /// Loads a Segmenter from a Safetensors model.
45    pub fn from_safetensors(path: impl AsRef<std::path::Path>, device: &B::Device) -> Result<Self> {
46        let _weights = WeightLoader::load_safetensors::<B>(path, device)?;
47        Ok(Self { model: None })
48    }
49
50    /// Loads a Segmenter from a native Burn model.
51    pub fn from_burn(path: impl AsRef<std::path::Path>, device: &B::Device) -> Result<Self> {
52        let _weights = WeightLoader::load_bin::<B>(path, device, [100, 100])?;
53        Ok(Self { model: None })
54    }
55
56    /// Performs segmentation prediction on the image.
57    pub fn segment(&self, image: &Image<B>) -> Result<SegmentationMask<B>> {
58        if let Some(ref model) = self.model {
59            let input = model.preprocess(image)?;
60            // Shape: [1, NumClasses, H, W]
61            let out: Tensor<B, 4> = model.predict_raw(input)?;
62
63            // Take argmax along the class axis (axis 1) to get class labels per pixel, shape [1, 1, H, W]
64            let class_indices = out.argmax(1);
65            let squeezed = class_indices.squeeze::<2>(); // shape [H, W]
66
67            Ok(SegmentationMask { mask: squeezed })
68        } else {
69            let shape = image.shape();
70            let device = image.tensor.device();
71            let mask = Tensor::<B, 2, Int>::zeros([shape[1], shape[2]], &device);
72            Ok(SegmentationMask { mask })
73        }
74    }
75}
76
77impl<B: Backend> Default for Segmenter<B> {
78    fn default() -> Self {
79        Self { model: None }
80    }
81}
82
83#[cfg(test)]
84mod tests {
85    use super::*;
86    use crate::test_helpers::{TestBackend, test_device};
87    use burn::tensor::TensorData;
88
89    #[test]
90    fn test_segmenter() {
91        let device = test_device();
92        let flat_data = vec![0.5f32; 3 * 8 * 8];
93        let tensor =
94            Tensor::<TestBackend, 3>::from_data(TensorData::new(flat_data, [3, 8, 8]), &device);
95        let img = Image::new(tensor);
96
97        let segmenter = Segmenter::<TestBackend>::default();
98        let mask = segmenter.segment(&img).unwrap();
99        assert_eq!(mask.mask.dims(), [8, 8]);
100    }
101}
iris/segmentation/mod.rs

iris/segmentation/
mod.rs
