Skip to main content

fpo_rust/
config.rs

1//! Plate configuration for inference, parsed from a YAML file.
2
3use anyhow::Context;
4use serde::Deserialize;
5use std::path::Path;
6
7/// Interpolation method used when resizing the plate image.
8#[derive(Debug, Clone, PartialEq, Eq, Deserialize, Default)]
9#[serde(rename_all = "lowercase")]
10pub enum ImageInterpolation {
11    Nearest,
12    #[default]
13    Linear,
14    Cubic,
15    Area,
16    Lanczos4,
17}
18
19/// Colour mode of the plate image fed into the model.
20#[derive(Debug, Clone, PartialEq, Eq, Deserialize, Default)]
21#[serde(rename_all = "lowercase")]
22pub enum ImageColorMode {
23    #[default]
24    Grayscale,
25    Rgb,
26}
27
28/// Padding colour used during letter-box resizing.
29///
30/// YAML accepts a single integer (grayscale) or a three-element sequence [r, g, b].
31#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
32#[serde(untagged)]
33pub enum PaddingColor {
34    Gray(u8),
35    Rgb([u8; 3]),
36}
37
38impl Default for PaddingColor {
39    fn default() -> Self {
40        PaddingColor::Rgb([114, 114, 114])
41    }
42}
43
44impl PaddingColor {
45    /// Return the grayscale byte value (uses the first channel for RGB).
46    pub fn as_gray(&self) -> u8 {
47        match self {
48            PaddingColor::Gray(v) => *v,
49            PaddingColor::Rgb([r, _, _]) => *r,
50        }
51    }
52
53    /// Return the RGB byte triple.
54    pub fn as_rgb(&self) -> [u8; 3] {
55        match self {
56            PaddingColor::Gray(v) => [*v, *v, *v],
57            PaddingColor::Rgb(c) => *c,
58        }
59    }
60}
61
62/// Inference configuration for a plate-recognition model, read from a YAML file.
63#[derive(Debug, Clone, Deserialize)]
64pub struct PlateConfig {
65    /// Maximum number of character slots predicted by the model.
66    pub max_plate_slots: usize,
67    /// Full alphabet used by the model (one character per class).
68    pub alphabet: String,
69    /// Padding character appended to plates shorter than `max_plate_slots`.
70    pub pad_char: char,
71    /// Image height the model expects.
72    pub img_height: u32,
73    /// Image width the model expects.
74    pub img_width: u32,
75    /// Whether to preserve aspect ratio via letter-box padding.
76    #[serde(default)]
77    pub keep_aspect_ratio: bool,
78    /// Interpolation method used when resizing.
79    #[serde(default)]
80    pub interpolation: ImageInterpolation,
81    /// Colour mode: grayscale (1-channel) or rgb (3-channel).
82    #[serde(default)]
83    pub image_color_mode: ImageColorMode,
84    /// Padding colour used when `keep_aspect_ratio` is true.
85    #[serde(default)]
86    pub padding_color: PaddingColor,
87    /// Optional list of region / country labels the model can predict.
88    #[serde(default)]
89    pub plate_regions: Option<Vec<String>>,
90}
91
92impl PlateConfig {
93    /// Load a `PlateConfig` from a YAML file.
94    pub fn from_yaml(path: impl AsRef<Path>) -> anyhow::Result<Self> {
95        let text = std::fs::read_to_string(path.as_ref())
96            .with_context(|| format!("Cannot read config: {}", path.as_ref().display()))?;
97        let cfg: PlateConfig = serde_yml::from_str(&text)
98            .with_context(|| format!("Cannot parse config: {}", path.as_ref().display()))?;
99        Ok(cfg)
100    }
101
102    /// Number of colour channels (1 for grayscale, 3 for RGB).
103    pub fn num_channels(&self) -> u32 {
104        match self.image_color_mode {
105            ImageColorMode::Rgb => 3,
106            ImageColorMode::Grayscale => 1,
107        }
108    }
109
110    /// `true` if the model includes a region-recognition head.
111    pub fn has_region_recognition(&self) -> bool {
112        self.plate_regions
113            .as_ref()
114            .map_or(false, |v| !v.is_empty())
115    }
116
117    /// Index of `pad_char` in the alphabet.
118    pub fn pad_idx(&self) -> usize {
119        self.alphabet
120            .chars()
121            .position(|c| c == self.pad_char)
122            .unwrap_or(0)
123    }
124}