1use 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#[derive(Debug, Clone, Copy, PartialEq)]
41#[non_exhaustive]
42pub enum NormalizeOverride {
43 Preset(NormalizationPreset),
45 MeanStd {
47 mean: [f32; 3],
49 std: [f32; 3],
51 },
52}
53
54impl NormalizeOverride {
55 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
121#[non_exhaustive]
122pub enum ResizeOverride {
123 Stretch,
125 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#[derive(Debug, Clone, Default, PartialEq)]
195#[non_exhaustive]
196pub struct ModelOverrides {
197 pub task: Option<ModelTask>,
199 pub normalize: Option<NormalizeOverride>,
201 pub resize: Option<ResizeOverride>,
203 pub decode: Option<DecodeMode>,
205 pub labels: Option<Vec<String>>,
210}
211
212impl ModelOverrides {
213 #[must_use]
215 pub fn new() -> Self {
216 Self::default()
217 }
218
219 #[must_use]
221 pub fn is_empty(&self) -> bool {
222 *self == Self::default()
223 }
224
225 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}