data_contracts/
preprocess.rs

1use thiserror::Error;
2
3#[derive(Debug, Clone, Copy)]
4pub struct ImageStats {
5    pub mean: [f32; 3],
6    pub std: [f32; 3],
7    pub aspect: f32,
8}
9
10impl ImageStats {
11    pub fn tiny_input(&self) -> [f32; 4] {
12        [self.mean[0], self.mean[1], self.mean[2], self.aspect]
13    }
14
15    pub fn feature_vector(&self, box_count: f32) -> [f32; 8] {
16        [
17            self.mean[0],
18            self.mean[1],
19            self.mean[2],
20            self.std[0],
21            self.std[1],
22            self.std[2],
23            self.aspect,
24            box_count,
25        ]
26    }
27}
28
29#[derive(Debug, Error)]
30pub enum ImageStatsError {
31    #[error("expected {expected} values, got {actual}")]
32    InvalidLength { expected: usize, actual: usize },
33    #[error("image dimensions must be non-zero")]
34    ZeroDimensions,
35}
36
37pub fn stats_from_rgb_u8(
38    width: u32,
39    height: u32,
40    rgb: &[u8],
41) -> Result<ImageStats, ImageStatsError> {
42    if width == 0 || height == 0 {
43        return Err(ImageStatsError::ZeroDimensions);
44    }
45    let expected = width as usize * height as usize * 3;
46    if rgb.len() != expected {
47        return Err(ImageStatsError::InvalidLength {
48            expected,
49            actual: rgb.len(),
50        });
51    }
52
53    let mut sum = [0f32; 3];
54    let mut sumsq = [0f32; 3];
55    for chunk in rgb.chunks_exact(3) {
56        let r = chunk[0] as f32 / 255.0;
57        let g = chunk[1] as f32 / 255.0;
58        let b = chunk[2] as f32 / 255.0;
59        sum[0] += r;
60        sum[1] += g;
61        sum[2] += b;
62        sumsq[0] += r * r;
63        sumsq[1] += g * g;
64        sumsq[2] += b * b;
65    }
66
67    stats_from_sums(width as usize, height as usize, sum, sumsq)
68}
69
70pub fn stats_from_rgba_u8(
71    width: u32,
72    height: u32,
73    rgba: &[u8],
74) -> Result<ImageStats, ImageStatsError> {
75    if width == 0 || height == 0 {
76        return Err(ImageStatsError::ZeroDimensions);
77    }
78    let expected = width as usize * height as usize * 4;
79    if rgba.len() != expected {
80        return Err(ImageStatsError::InvalidLength {
81            expected,
82            actual: rgba.len(),
83        });
84    }
85
86    let mut sum = [0f32; 3];
87    let mut sumsq = [0f32; 3];
88    for chunk in rgba.chunks_exact(4) {
89        let r = chunk[0] as f32 / 255.0;
90        let g = chunk[1] as f32 / 255.0;
91        let b = chunk[2] as f32 / 255.0;
92        sum[0] += r;
93        sum[1] += g;
94        sum[2] += b;
95        sumsq[0] += r * r;
96        sumsq[1] += g * g;
97        sumsq[2] += b * b;
98    }
99
100    stats_from_sums(width as usize, height as usize, sum, sumsq)
101}
102
103pub fn stats_from_chw_f32(
104    width: usize,
105    height: usize,
106    chw: &[f32],
107) -> Result<ImageStats, ImageStatsError> {
108    if width == 0 || height == 0 {
109        return Err(ImageStatsError::ZeroDimensions);
110    }
111    let expected = width * height * 3;
112    if chw.len() != expected {
113        return Err(ImageStatsError::InvalidLength {
114            expected,
115            actual: chw.len(),
116        });
117    }
118
119    let mut sum = [0f32; 3];
120    let mut sumsq = [0f32; 3];
121    let pixels_per_channel = width * height;
122    for c in 0..3 {
123        let start = c * pixels_per_channel;
124        let slice = &chw[start..start + pixels_per_channel];
125        for v in slice {
126            sum[c] += *v;
127            sumsq[c] += *v * *v;
128        }
129    }
130
131    stats_from_sums(width, height, sum, sumsq)
132}
133
134fn stats_from_sums(
135    width: usize,
136    height: usize,
137    sum: [f32; 3],
138    sumsq: [f32; 3],
139) -> Result<ImageStats, ImageStatsError> {
140    let pix_count = (width * height) as f32;
141    let mean = [sum[0] / pix_count, sum[1] / pix_count, sum[2] / pix_count];
142    let std = [
143        (sumsq[0] / pix_count - mean[0] * mean[0]).max(0.0).sqrt(),
144        (sumsq[1] / pix_count - mean[1] * mean[1]).max(0.0).sqrt(),
145        (sumsq[2] / pix_count - mean[2] * mean[2]).max(0.0).sqrt(),
146    ];
147    Ok(ImageStats {
148        mean,
149        std,
150        aspect: width as f32 / height as f32,
151    })
152}
data_contracts/preprocess.rs

data_contracts/
preprocess.rs
