fpm_rs/datasets/
loader.rs

1use std::{
2    collections::BTreeMap,
3    fs::File,
4    io::BufReader,
5    path::{Path, PathBuf},
6};
7
8use serde::{Deserialize, Serialize};
9
10use crate::{
11    Array2, Complex64, Result,
12    configuration::SimulationConfiguration,
13    error::Error,
14    measurements::{ImageSet, MeasurementSpec, MeasurementStack},
15    reconstruction::ReconstructionProblem,
16};
17
18/// Current version of the language-neutral dataset bundle format.
19pub const DATASET_FORMAT_VERSION: u32 = 1;
20
21/// The `dataset.json` entry point defined by `dataset_spec.md`.
22#[derive(Clone, Debug, Serialize, Deserialize)]
23#[serde(deny_unknown_fields)]
24pub struct DatasetManifest {
25    pub format_version: u32,
26    pub measurement_manifest: PathBuf,
27    pub configuration: PathBuf,
28    /// Optional JSON-serialized `Array2<Complex64>` in reconstruction space.
29    #[serde(default)]
30    pub ground_truth_object: Option<PathBuf>,
31    /// Optional JSON-serialized binary `Array2<u8>` in reconstruction space.
32    #[serde(default)]
33    pub valid_object_mask: Option<PathBuf>,
34    #[serde(default)]
35    pub provenance: BTreeMap<String, String>,
36    #[serde(default)]
37    pub measurement_units: Option<String>,
38}
39
40/// A validated dataset loaded from a bundle conforming to `dataset_spec.md`.
41#[derive(Clone, Debug)]
42pub struct Dataset {
43    source_path: Option<PathBuf>,
44    measurements: MeasurementStack,
45    configuration: SimulationConfiguration,
46    ground_truth_object: Option<Array2<Complex64>>,
47    valid_object_mask: Option<Array2<u8>>,
48    provenance: BTreeMap<String, String>,
49    measurement_units: Option<String>,
50}
51
52impl Dataset {
53    /// Constructs a dataset from already-loaded values.
54    pub fn new(
55        measurements: MeasurementStack,
56        configuration: SimulationConfiguration,
57    ) -> Result<Self> {
58        measurements.validate()?;
59        configuration.validate()?;
60        ReconstructionProblem::new(
61            measurements.clone(),
62            configuration.compiled_models.reconstruction_model.clone(),
63        )?;
64        Ok(Self {
65            source_path: None,
66            measurements,
67            configuration,
68            ground_truth_object: None,
69            valid_object_mask: None,
70            provenance: BTreeMap::new(),
71            measurement_units: None,
72        })
73    }
74
75    pub fn measurements(&self) -> &MeasurementStack {
76        &self.measurements
77    }
78
79    /// Root directory of the loaded bundle, or `None` for programmatic data.
80    pub fn source_path(&self) -> Option<&Path> {
81        self.source_path.as_deref()
82    }
83
84    pub fn configuration(&self) -> &SimulationConfiguration {
85        &self.configuration
86    }
87
88    pub fn ground_truth_object(&self) -> Option<&Array2<Complex64>> {
89        self.ground_truth_object.as_ref()
90    }
91
92    pub fn valid_object_mask(&self) -> Option<&Array2<u8>> {
93        self.valid_object_mask.as_ref()
94    }
95
96    pub fn provenance(&self) -> &BTreeMap<String, String> {
97        &self.provenance
98    }
99
100    pub fn measurement_units(&self) -> Option<&str> {
101        self.measurement_units.as_deref()
102    }
103
104    pub fn reconstruction_problem(&self) -> Result<ReconstructionProblem<MeasurementStack>> {
105        ReconstructionProblem::new(
106            self.measurements.clone(),
107            self.configuration
108                .compiled_models
109                .reconstruction_model
110                .clone(),
111        )
112    }
113
114    pub fn subset(&self) -> super::DatasetSubsetBuilder<'_> {
115        super::DatasetSubsetBuilder::new(self)
116    }
117
118    fn with_metadata(
119        mut self,
120        ground_truth_object: Option<Array2<Complex64>>,
121        valid_object_mask: Option<Array2<u8>>,
122        provenance: BTreeMap<String, String>,
123        measurement_units: Option<String>,
124    ) -> Result<Self> {
125        validate_dataset_metadata(
126            &self.configuration,
127            ground_truth_object.as_ref(),
128            valid_object_mask.as_ref(),
129            &provenance,
130            measurement_units.as_deref(),
131        )?;
132        self.ground_truth_object = ground_truth_object;
133        self.valid_object_mask = valid_object_mask;
134        self.provenance = provenance;
135        self.measurement_units = measurement_units;
136        Ok(self)
137    }
138}
139
140/// Loads an already-converted dataset bundle from a local directory.
141#[derive(Clone, Debug)]
142pub struct DatasetLoader {
143    root: PathBuf,
144}
145
146impl DatasetLoader {
147    pub fn new(root: impl Into<PathBuf>) -> Result<Self> {
148        let root = root.into();
149        if !root.is_dir() {
150            return Err(Error::Dataset(format!(
151                "dataset path does not exist or is not a directory: {}",
152                root.display()
153            )));
154        }
155        Ok(Self { root })
156    }
157
158    pub fn root(&self) -> &Path {
159        &self.root
160    }
161
162    pub fn manifest_path(&self) -> PathBuf {
163        self.root.join("dataset.json")
164    }
165
166    pub fn load(&self) -> Result<Dataset> {
167        let manifest_path = self.manifest_path();
168        let manifest: DatasetManifest =
169            serde_json::from_reader(BufReader::new(File::open(&manifest_path)?))?;
170        if manifest.format_version != DATASET_FORMAT_VERSION {
171            return Err(Error::Dataset(format!(
172                "unsupported dataset format version {} in {}; expected {}",
173                manifest.format_version,
174                manifest_path.display(),
175                DATASET_FORMAT_VERSION
176            )));
177        }
178        validate_dataset_manifest_paths(&manifest)?;
179        let measurement_manifest_path = self.root.join(&manifest.measurement_manifest);
180        let measurement_spec = MeasurementSpec::load(&measurement_manifest_path)?;
181        validate_measurement_paths(&measurement_spec)?;
182        let measurement_base = measurement_manifest_path
183            .parent()
184            .unwrap_or_else(|| Path::new("."));
185        let measurements =
186            MeasurementStack::from_manifest_definition(measurement_spec, measurement_base)?;
187        let configuration = SimulationConfiguration::load(self.root.join(&manifest.configuration))?;
188        let ground_truth_object = manifest
189            .ground_truth_object
190            .as_ref()
191            .map(|path| load_json_array(self.root.join(path), "ground-truth object"))
192            .transpose()?;
193        let valid_object_mask = manifest
194            .valid_object_mask
195            .as_ref()
196            .map(|path| load_json_array(self.root.join(path), "valid-object mask"))
197            .transpose()?;
198        let mut dataset = Dataset::new(measurements, configuration)?;
199        dataset.source_path = Some(self.root.clone());
200        dataset.with_metadata(
201            ground_truth_object,
202            valid_object_mask,
203            manifest.provenance,
204            manifest.measurement_units,
205        )
206    }
207}
208
209fn validate_dataset_manifest_paths(manifest: &DatasetManifest) -> Result<()> {
210    validate_contained_path("measurement manifest", &manifest.measurement_manifest)?;
211    validate_contained_path("configuration", &manifest.configuration)?;
212    if let Some(path) = &manifest.ground_truth_object {
213        validate_contained_path("ground-truth object", path)?;
214    }
215    if let Some(path) = &manifest.valid_object_mask {
216        validate_contained_path("valid-object mask", path)?;
217    }
218    Ok(())
219}
220
221fn validate_measurement_paths(spec: &MeasurementSpec) -> Result<()> {
222    for frame in &spec.frames {
223        validate_contained_path("measurement frame", &frame.path)?;
224    }
225    if let Some(path) = &spec.dark_frame {
226        validate_contained_path("dark frame", path)?;
227    }
228    if let Some(path) = &spec.flat_field {
229        validate_contained_path("flat field", path)?;
230    }
231    if let Some(images) = &spec.background {
232        validate_image_set_paths("background", images)?;
233    }
234    if let Some(images) = &spec.mask {
235        validate_image_set_paths("mask", images)?;
236    }
237    Ok(())
238}
239
240fn validate_image_set_paths(label: &str, images: &ImageSet) -> Result<()> {
241    match images {
242        ImageSet::Single(path) => validate_contained_path(label, path),
243        ImageSet::PerFrame(paths) => paths
244            .iter()
245            .try_for_each(|path| validate_contained_path(label, path)),
246    }
247}
248
249fn validate_contained_path(label: &str, path: &Path) -> Result<()> {
250    if path.as_os_str().is_empty()
251        || path
252            .components()
253            .any(|component| !matches!(component, std::path::Component::Normal(_)))
254    {
255        return Err(Error::Dataset(format!(
256            "{label} must be a safe relative path contained by its manifest directory: {}",
257            path.display()
258        )));
259    }
260    Ok(())
261}
262
263fn load_json_array<T>(path: PathBuf, label: &str) -> Result<Array2<T>>
264where
265    T: for<'de> Deserialize<'de>,
266{
267    let file = File::open(&path).map_err(|error| {
268        Error::Dataset(format!(
269            "failed to open {label} at {}: {error}",
270            path.display()
271        ))
272    })?;
273    serde_json::from_reader(BufReader::new(file)).map_err(|error| {
274        Error::Dataset(format!(
275            "failed to parse {label} at {}: {error}",
276            path.display()
277        ))
278    })
279}
280
281fn validate_dataset_metadata(
282    configuration: &SimulationConfiguration,
283    ground_truth_object: Option<&Array2<Complex64>>,
284    valid_object_mask: Option<&Array2<u8>>,
285    provenance: &BTreeMap<String, String>,
286    measurement_units: Option<&str>,
287) -> Result<()> {
288    if let Some(ground_truth) = ground_truth_object {
289        if ground_truth.shape() != configuration.reconstruction_shape {
290            return Err(Error::Dataset(format!(
291                "ground-truth shape {:?} differs from reconstruction shape {:?}",
292                ground_truth.shape(),
293                configuration.reconstruction_shape
294            )));
295        }
296        if ground_truth
297            .as_slice()
298            .iter()
299            .any(|value| !value.re.is_finite() || !value.im.is_finite())
300        {
301            return Err(Error::Dataset(
302                "ground-truth object contains non-finite values".into(),
303            ));
304        }
305    }
306    if let Some(mask) = valid_object_mask {
307        if ground_truth_object.is_none() {
308            return Err(Error::Dataset(
309                "a valid-object mask requires a ground-truth object".into(),
310            ));
311        }
312        if mask.shape() != configuration.reconstruction_shape {
313            return Err(Error::Dataset(format!(
314                "valid-object mask shape {:?} differs from reconstruction shape {:?}",
315                mask.shape(),
316                configuration.reconstruction_shape
317            )));
318        }
319        if mask.as_slice().iter().any(|&value| value > 1) || !mask.as_slice().contains(&1) {
320            return Err(Error::Dataset(
321                "valid-object mask must contain only zero and one and select at least one pixel"
322                    .into(),
323            ));
324        }
325    }
326    if provenance
327        .iter()
328        .any(|(key, value)| key.trim().is_empty() || value.trim().is_empty())
329    {
330        return Err(Error::Dataset(
331            "dataset provenance keys and values must not be empty".into(),
332        ));
333    }
334    if measurement_units.is_some_and(|units| units.trim().is_empty()) {
335        return Err(Error::Dataset(
336            "measurement units must not be empty when provided".into(),
337        ));
338    }
339    Ok(())
340}
fpm_rs/datasets/loader.rs

fpm_rs/datasets/
loader.rs
