Skip to main content

edgefirst_decoder/model_config/
overrides.rs

1// SPDX-FileCopyrightText: Copyright 2026 Au-Zone Technologies
2// SPDX-License-Identifier: Apache-2.0
3
4//! Caller overrides and their string grammar, shared by the Rust, Python and
5//! C surfaces and by command lines.
6
7use std::fmt;
8use std::str::FromStr;
9
10use edgefirst_tensor::{LetterboxPlacement, Normalization, NormalizationPreset};
11
12use super::{ModelConfigError, ModelTask};
13use crate::schema::DecodeMode;
14
15fn bad(what: &str, value: &str, grammar: &str) -> ModelConfigError {
16    ModelConfigError::Override(format!("{what} `{value}` is not one of {grammar}"))
17}
18
19const NORMALIZE_GRAMMAR: &str = "unsigned | signed | raw | imagenet | R,G,B/R,G,B";
20const RESIZE_GRAMMAR: &str = "stretch | letterbox | letterbox-center | letterbox-topleft";
21const TASK_GRAMMAR: &str = "detect | segment | semantic";
22const DECODE_GRAMMAR: &str = "accurate | fast";
23
24/// A `normalize` override: a preset name or `R,G,B/R,G,B`, a per-channel
25/// mean and standard deviation in the 0–1 domain as Hugging Face and
26/// torchvision write them.
27///
28/// ```
29/// use edgefirst_decoder::model_config::NormalizeOverride;
30/// use edgefirst_tensor::NormalizationPreset;
31///
32/// let n: NormalizeOverride = "0.485,0.456,0.406/0.229,0.224,0.225".parse().unwrap();
33/// assert_eq!(n.to_string(), "0.485,0.456,0.406/0.229,0.224,0.225");
34/// assert_eq!(
35///     "unsigned".parse::<NormalizeOverride>().unwrap(),
36///     NormalizeOverride::Preset(NormalizationPreset::Unsigned)
37/// );
38/// assert!("0,0,0/0,1,1".parse::<NormalizeOverride>().is_err());
39/// ```
40#[derive(Debug, Clone, Copy, PartialEq)]
41#[non_exhaustive]
42pub enum NormalizeOverride {
43    /// A named preset.
44    Preset(NormalizationPreset),
45    /// Explicit mean and standard deviation, 0–1 domain.
46    MeanStd {
47        /// Per-channel mean.
48        mean: [f32; 3],
49        /// Per-channel standard deviation.
50        std: [f32; 3],
51    },
52}
53
54impl NormalizeOverride {
55    /// The normalization on 0–255 values.
56    ///
57    /// # Errors
58    ///
59    /// [`ModelConfigError::Override`] for a non-finite mean or a standard
60    /// deviation that is not finite and positive.
61    pub fn to_normalization(&self) -> Result<Normalization, ModelConfigError> {
62        match self {
63            Self::Preset(p) => Ok(Normalization::preset(*p)),
64            Self::MeanStd { mean, std } => Normalization::unit_mean_std(mean, std, 255.0)
65                .map_err(|e| ModelConfigError::Override(format!("normalize: {e}"))),
66        }
67    }
68}
69
70fn triple(s: &str) -> Option<[f32; 3]> {
71    let mut it = s.split(',').map(|v| v.trim().parse::<f32>().ok());
72    let t = [it.next()??, it.next()??, it.next()??];
73    it.next().is_none().then_some(t)
74}
75
76impl FromStr for NormalizeOverride {
77    type Err = ModelConfigError;
78
79    fn from_str(s: &str) -> Result<Self, Self::Err> {
80        let s = s.trim();
81        if let Some(p) = NormalizationPreset::from_str_code(s) {
82            return Ok(Self::Preset(p));
83        }
84        let (mean, std) = s
85            .split_once('/')
86            .and_then(|(m, d)| Some((triple(m)?, triple(d)?)))
87            .ok_or_else(|| bad("normalize", s, NORMALIZE_GRAMMAR))?;
88        let n = Self::MeanStd { mean, std };
89        n.to_normalization()?;
90        Ok(n)
91    }
92}
93
94impl fmt::Display for NormalizeOverride {
95    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
96        match self {
97            Self::Preset(p) => f.write_str(p.as_str()),
98            Self::MeanStd { mean, std } => write!(
99                f,
100                "{},{},{}/{},{},{}",
101                mean[0], mean[1], mean[2], std[0], std[1], std[2]
102            ),
103        }
104    }
105}
106
107/// A `resize` override: `stretch`, or `letterbox` with an optional
108/// placement (`letterbox-center`, `letterbox-topleft`). A bare `letterbox`
109/// follows the model family's convention; an explicit placement keeps the
110/// family's pad.
111///
112/// ```
113/// use edgefirst_decoder::model_config::ResizeOverride;
114/// use edgefirst_tensor::LetterboxPlacement;
115///
116/// let r: ResizeOverride = "letterbox-topleft".parse().unwrap();
117/// assert_eq!(r, ResizeOverride::Letterbox(Some(LetterboxPlacement::TopLeft)));
118/// assert_eq!(ResizeOverride::Letterbox(None).to_string(), "letterbox");
119/// ```
120#[derive(Debug, Clone, Copy, PartialEq, Eq)]
121#[non_exhaustive]
122pub enum ResizeOverride {
123    /// Scale each axis to the input size.
124    Stretch,
125    /// Letterbox, at the placement given or the family's.
126    Letterbox(Option<LetterboxPlacement>),
127}
128
129impl FromStr for ResizeOverride {
130    type Err = ModelConfigError;
131
132    fn from_str(s: &str) -> Result<Self, Self::Err> {
133        let s = s.trim();
134        match s {
135            "stretch" => Ok(Self::Stretch),
136            "letterbox" => Ok(Self::Letterbox(None)),
137            _ => s
138                .strip_prefix("letterbox-")
139                .and_then(LetterboxPlacement::from_str_code)
140                .map(|p| Self::Letterbox(Some(p)))
141                .ok_or_else(|| bad("resize", s, RESIZE_GRAMMAR)),
142        }
143    }
144}
145
146impl fmt::Display for ResizeOverride {
147    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
148        match self {
149            Self::Stretch => f.write_str("stretch"),
150            Self::Letterbox(None) => f.write_str("letterbox"),
151            Self::Letterbox(Some(p)) => write!(f, "letterbox-{}", p.as_str()),
152        }
153    }
154}
155
156impl FromStr for ModelTask {
157    type Err = ModelConfigError;
158
159    fn from_str(s: &str) -> Result<Self, Self::Err> {
160        Self::from_str_code(s.trim()).ok_or_else(|| bad("task", s, TASK_GRAMMAR))
161    }
162}
163
164impl fmt::Display for ModelTask {
165    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
166        f.write_str(self.as_str())
167    }
168}
169
170fn parse_labels(value: &str) -> Result<Vec<String>, ModelConfigError> {
171    let labels: Vec<String> = value.split(',').map(|l| l.trim().to_owned()).collect();
172    if labels.iter().any(String::is_empty) {
173        return Err(ModelConfigError::Override(format!(
174            "labels `{value}` has an empty name; expected comma-separated class names"
175        )));
176    }
177    Ok(labels)
178}
179
180/// Values that replace what discovery found, applied after every source and
181/// recorded as [`ConfigSource::Override`](super::ConfigSource::Override).
182///
183/// ```
184/// use edgefirst_decoder::model_config::{ModelOverrides, ResizeOverride};
185/// use edgefirst_tensor::LetterboxPlacement;
186///
187/// let mut o = ModelOverrides::default();
188/// o.set("resize", "letterbox-topleft").unwrap();
189/// o.set("labels", "background, person").unwrap();
190/// assert_eq!(o.resize, Some(ResizeOverride::Letterbox(Some(LetterboxPlacement::TopLeft))));
191/// assert_eq!(o.labels.as_deref(), Some(&["background".to_string(), "person".to_string()][..]));
192/// assert!(o.set("resize", "fill").is_err());
193/// ```
194#[derive(Debug, Clone, Default, PartialEq)]
195#[non_exhaustive]
196pub struct ModelOverrides {
197    /// What the outputs decode to.
198    pub task: Option<ModelTask>,
199    /// Input normalization.
200    pub normalize: Option<NormalizeOverride>,
201    /// Resize mode.
202    pub resize: Option<ResizeOverride>,
203    /// Semantic decode mode.
204    pub decode: Option<DecodeMode>,
205    /// Class names, in channel order. Resolution checks their number only
206    /// against a class count the model's outputs state; end-to-end
207    /// detections and class maps state none, so any number is accepted
208    /// there and the caller checks it.
209    pub labels: Option<Vec<String>>,
210}
211
212impl ModelOverrides {
213    /// No overrides.
214    #[must_use]
215    pub fn new() -> Self {
216        Self::default()
217    }
218
219    /// Whether no override is set.
220    #[must_use]
221    pub fn is_empty(&self) -> bool {
222        *self == Self::default()
223    }
224
225    /// Sets the override named `key` (`task`, `normalize`, `resize`,
226    /// `decode` or `labels`) from its string form. `labels` is a
227    /// comma-separated list of non-empty names.
228    ///
229    /// # Errors
230    ///
231    /// [`ModelConfigError::Override`] for an unknown key or a value outside
232    /// the key's grammar.
233    pub fn set(&mut self, key: &str, value: &str) -> Result<(), ModelConfigError> {
234        match key.trim() {
235            "task" => self.task = Some(value.parse()?),
236            "normalize" => self.normalize = Some(value.parse()?),
237            "resize" => self.resize = Some(value.parse()?),
238            "decode" => {
239                let v = value.trim();
240                self.decode = Some(
241                    DecodeMode::from_str_code(v).ok_or_else(|| bad("decode", v, DECODE_GRAMMAR))?,
242                )
243            }
244            "labels" => self.labels = Some(parse_labels(value)?),
245            other => {
246                return Err(ModelConfigError::Override(format!(
247                    "unknown override `{other}`; expected task, normalize, resize, decode or labels"
248                )))
249            }
250        }
251        Ok(())
252    }
253}
254
255#[cfg(test)]
256mod tests {
257    use super::*;
258
259    #[test]
260    fn override_grammar_round_trips() {
261        for s in [
262            "imagenet",
263            "unsigned",
264            "signed",
265            "raw",
266            "0.485,0.456,0.406/0.229,0.224,0.225",
267            "0.5,0.5,0.5/0.5,0.5,0.5",
268        ] {
269            let n: NormalizeOverride = s.parse().unwrap();
270            assert_eq!(n.to_string(), s);
271        }
272        for s in [
273            "stretch",
274            "letterbox",
275            "letterbox-center",
276            "letterbox-topleft",
277        ] {
278            assert_eq!(s.parse::<ResizeOverride>().unwrap().to_string(), s);
279        }
280        for s in ["detect", "segment", "semantic"] {
281            assert_eq!(s.parse::<ModelTask>().unwrap().to_string(), s);
282        }
283        let mut o = ModelOverrides::new();
284        for s in ["accurate", "fast"] {
285            o.set("decode", s).unwrap();
286            assert_eq!(o.decode.unwrap().as_str(), s);
287        }
288    }
289
290    #[test]
291    fn normalize_mean_std_is_scaled_to_255() {
292        let n: NormalizeOverride = "0.5,0.5,0.5/0.5,0.5,0.5".parse().unwrap();
293        assert_eq!(n.to_normalization().unwrap(), Normalization::signed());
294    }
295
296    #[test]
297    fn bad_override_lists_accepted_forms() {
298        let e = "fill".parse::<ResizeOverride>().unwrap_err().to_string();
299        assert!(e.contains("letterbox-topleft"), "{e}");
300        assert!("letterbox-top_left".parse::<ResizeOverride>().is_err());
301        assert!("unit".parse::<NormalizeOverride>().is_err());
302        let e = "0.1,0.2/0.3"
303            .parse::<NormalizeOverride>()
304            .unwrap_err()
305            .to_string();
306        assert!(e.contains("R,G,B/R,G,B"), "{e}");
307        assert!("1,2,3,4/1,2,3".parse::<NormalizeOverride>().is_err());
308        for bad_std in [
309            "0,0,0/0,1,1",
310            "0,0,0/-1,1,1",
311            "nan,0,0/1,1,1",
312            "0,0,0/inf,1,1",
313        ] {
314            assert!(bad_std.parse::<NormalizeOverride>().is_err(), "{bad_std}");
315        }
316        assert!("classify".parse::<ModelTask>().is_err());
317        assert!(ModelOverrides::new().set("decode", "slow").is_err());
318    }
319
320    #[test]
321    fn set_by_key() {
322        let mut o = ModelOverrides::new();
323        assert!(o.is_empty());
324        o.set("task", "semantic").unwrap();
325        o.set("normalize", "imagenet").unwrap();
326        o.set("decode", " fast ").unwrap();
327        o.set("labels", "a, b").unwrap();
328        assert_eq!(o.task, Some(ModelTask::Semantic));
329        assert_eq!(o.decode, Some(DecodeMode::Fast));
330        assert_eq!(
331            o.labels.as_deref(),
332            Some(&["a".to_string(), "b".to_string()][..])
333        );
334        assert!(o.set("colour", "x").is_err());
335        for empty in ["", "a,,b", "a, "] {
336            assert!(o.set("labels", empty).is_err(), "{empty:?}");
337        }
338    }
339}