Skip to main content

fpm_rs/measurements/
stack.rs

1use serde::{Deserialize, Deserializer, Serialize, de::Error as _};
2use std::path::Path;
3
4use crate::{
5    error::{Error, Result},
6    image_io::{GrayscaleScaling, load_grayscale, load_grayscale_tiff_pages},
7};
8
9use super::{FrameMetadata, ImageSet, MeasurementSpec, PreprocessingConfig};
10
11/// An in-memory stack of image-plane intensity measurements.
12#[derive(Clone, Debug, Serialize)]
13pub struct MeasurementStack {
14    data: Vec<f64>,
15    image_shape: (usize, usize),
16    frames: usize,
17    pub frame_metadata: Vec<FrameMetadata>,
18    pub dark_frame: Option<Vec<f64>>,
19    pub flat_field: Option<Vec<f64>>,
20    pub background: Option<Vec<f64>>,
21    pub masks: Option<Vec<u8>>,
22    pub preprocessing: PreprocessingConfig,
23}
24
25impl<'de> Deserialize<'de> for MeasurementStack {
26    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
27    where
28        D: Deserializer<'de>,
29    {
30        #[derive(Deserialize)]
31        struct Representation {
32            data: Vec<f64>,
33            image_shape: (usize, usize),
34            frames: usize,
35            frame_metadata: Vec<FrameMetadata>,
36            dark_frame: Option<Vec<f64>>,
37            flat_field: Option<Vec<f64>>,
38            background: Option<Vec<f64>>,
39            masks: Option<Vec<u8>>,
40            preprocessing: PreprocessingConfig,
41        }
42
43        let representation = Representation::deserialize(deserializer)?;
44        let stack = Self {
45            data: representation.data,
46            image_shape: representation.image_shape,
47            frames: representation.frames,
48            frame_metadata: representation.frame_metadata,
49            dark_frame: representation.dark_frame,
50            flat_field: representation.flat_field,
51            background: representation.background,
52            masks: representation.masks,
53            preprocessing: representation.preprocessing,
54        };
55        stack.validate().map_err(D::Error::custom)?;
56        Ok(stack)
57    }
58}
59
60impl MeasurementStack {
61    pub fn from_vec(
62        data: Vec<f64>,
63        image_shape: (usize, usize),
64        frame_metadata: Vec<FrameMetadata>,
65    ) -> Result<Self> {
66        let frame_len = image_shape
67            .0
68            .checked_mul(image_shape.1)
69            .ok_or_else(|| Error::InvalidShape("measurement shape overflows".into()))?;
70        if frame_len == 0 {
71            return Err(Error::InvalidShape(
72                "measurement dimensions must be non-zero".into(),
73            ));
74        }
75        if data.is_empty() || !data.len().is_multiple_of(frame_len) {
76            return Err(Error::InvalidMeasurements(format!(
77                "data length {} is not a positive multiple of frame size {frame_len}",
78                data.len()
79            )));
80        }
81        if data.iter().any(|value| !value.is_finite()) {
82            return Err(Error::InvalidMeasurements(
83                "measurements contain non-finite values".into(),
84            ));
85        }
86        let frames = data.len() / frame_len;
87        let frame_metadata = if frame_metadata.is_empty() {
88            (0..frames).map(FrameMetadata::new).collect()
89        } else {
90            if frame_metadata.len() != frames {
91                return Err(Error::InvalidMeasurements(format!(
92                    "{} metadata entries for {frames} frames",
93                    frame_metadata.len()
94                )));
95            }
96            frame_metadata
97        };
98        let stack = Self {
99            data,
100            image_shape,
101            frames,
102            frame_metadata,
103            dark_frame: None,
104            flat_field: None,
105            background: None,
106            masks: None,
107            preprocessing: PreprocessingConfig::default(),
108        };
109        stack.validate()?;
110        Ok(stack)
111    }
112
113    pub fn from_frames(frames: &[Vec<f64>], image_shape: (usize, usize)) -> Result<Self> {
114        let data = frames.iter().flatten().copied().collect();
115        Self::from_vec(data, image_shape, Vec::new())
116    }
117
118    /// Loads an in-memory stack from single-channel PNG or TIFF images while
119    /// preserving native 8-bit or 16-bit detector counts.
120    pub fn from_image_files<P: AsRef<Path>>(
121        paths: &[P],
122        frame_metadata: Vec<FrameMetadata>,
123    ) -> Result<Self> {
124        if paths.is_empty() {
125            return Err(Error::InvalidMeasurements(
126                "at least one image path is required".into(),
127            ));
128        }
129        let mut shape = None;
130        let mut data = Vec::new();
131        for path in paths {
132            let frame = load_grayscale(path, GrayscaleScaling::NativeCounts)?;
133            if let Some(expected) = shape {
134                if frame.shape() != expected {
135                    return Err(Error::InvalidMeasurements(format!(
136                        "image {} has shape {:?}, expected {expected:?}",
137                        path.as_ref().display(),
138                        frame.shape()
139                    )));
140                }
141            } else {
142                shape = Some(frame.shape());
143            }
144            data.extend(frame.into_vec());
145        }
146        let metadata = if frame_metadata.is_empty() {
147            paths
148                .iter()
149                .enumerate()
150                .map(|(index, path)| {
151                    let mut metadata = FrameMetadata::new(index);
152                    metadata.label = Some(path.as_ref().display().to_string());
153                    metadata
154                })
155                .collect()
156        } else {
157            frame_metadata
158        };
159        let shape = shape.ok_or_else(|| {
160            Error::InvalidMeasurements("image paths did not produce a frame".into())
161        })?;
162        Self::from_vec(data, shape, metadata)
163    }
164
165    /// Loads every page of a grayscale 8-bit or 16-bit TIFF as one frame.
166    pub fn from_tiff_stack(
167        path: impl AsRef<Path>,
168        frame_metadata: Vec<FrameMetadata>,
169    ) -> Result<Self> {
170        let path = path.as_ref();
171        let pages = load_grayscale_tiff_pages(path, GrayscaleScaling::NativeCounts)?;
172        let first = pages.first().ok_or_else(|| {
173            Error::InvalidMeasurements("TIFF stack does not contain an image".into())
174        })?;
175        let shape = first.shape();
176        if pages.iter().any(|page| page.shape() != shape) {
177            return Err(Error::InvalidMeasurements(
178                "all TIFF pages must have the same dimensions".into(),
179            ));
180        }
181        let metadata = if frame_metadata.is_empty() {
182            (0..pages.len())
183                .map(|index| {
184                    let mut metadata = FrameMetadata::new(index);
185                    metadata.label = Some(format!("{}#page={index}", path.display()));
186                    metadata
187                })
188                .collect()
189        } else {
190            frame_metadata
191        };
192        Self::from_vec(
193            pages.into_iter().flat_map(|page| page.into_vec()).collect(),
194            shape,
195            metadata,
196        )
197    }
198
199    /// Loads a JSON manifest. Relative image paths are resolved against the
200    /// manifest's parent directory. Declared preprocessing is configured but is
201    /// not applied until [`Self::apply_preprocessing`] is called.
202    pub fn from_manifest(path: impl AsRef<Path>) -> Result<Self> {
203        let path = path.as_ref();
204        let manifest = MeasurementSpec::load(path)?;
205        let base_directory = path.parent().unwrap_or_else(|| Path::new("."));
206        Self::from_manifest_definition(manifest, base_directory)
207    }
208
209    pub fn from_manifest_definition(
210        manifest: MeasurementSpec,
211        base_directory: impl AsRef<Path>,
212    ) -> Result<Self> {
213        let base_directory = base_directory.as_ref();
214        let paths: Vec<_> = manifest
215            .frames
216            .iter()
217            .map(|frame| resolve_path(base_directory, &frame.path))
218            .collect();
219        let metadata = manifest
220            .frames
221            .iter()
222            .enumerate()
223            .map(|(index, frame)| FrameMetadata {
224                frame_index: index,
225                illumination_index: frame.illumination_index,
226                original_frame_index: Some(index),
227                original_illumination_index: frame.illumination_index,
228                exposure_time: frame.exposure_time,
229                weight: frame.weight,
230                label: frame
231                    .label
232                    .clone()
233                    .or_else(|| Some(frame.path.display().to_string())),
234            })
235            .collect();
236        let mut stack = Self::from_image_files(&paths, metadata)?;
237        if let Some(path) = &manifest.dark_frame {
238            stack.dark_frame = Some(load_manifest_image(
239                base_directory,
240                path,
241                stack.image_shape,
242            )?);
243        }
244        if let Some(path) = &manifest.flat_field {
245            stack.flat_field = Some(load_manifest_image(
246                base_directory,
247                path,
248                stack.image_shape,
249            )?);
250        }
251        if let Some(background) = &manifest.background {
252            stack.background = Some(load_manifest_image_set(
253                base_directory,
254                background,
255                stack.image_shape,
256                stack.frames,
257            )?);
258        }
259        if let Some(mask) = &manifest.mask {
260            stack.masks = Some(
261                load_manifest_image_set(base_directory, mask, stack.image_shape, stack.frames)?
262                    .into_iter()
263                    .map(|value| u8::from(value != 0.0))
264                    .collect(),
265            );
266        }
267        stack.preprocessing = manifest.preprocessing;
268        stack.validate()?;
269        Ok(stack)
270    }
271
272    pub fn frame_count(&self) -> usize {
273        self.frames
274    }
275
276    pub fn image_shape(&self) -> (usize, usize) {
277        self.image_shape
278    }
279
280    pub fn frame_len(&self) -> usize {
281        self.image_shape.0 * self.image_shape.1
282    }
283
284    pub fn as_slice(&self) -> &[f64] {
285        &self.data
286    }
287
288    pub fn validate(&self) -> Result<()> {
289        let frame_len = self
290            .image_shape
291            .0
292            .checked_mul(self.image_shape.1)
293            .ok_or_else(|| Error::InvalidShape("measurement shape overflows".into()))?;
294        let expected_len = frame_len.checked_mul(self.frames).ok_or_else(|| {
295            Error::InvalidMeasurements("measurement frame count overflows".into())
296        })?;
297        if frame_len == 0 || self.frames == 0 || self.data.len() != expected_len {
298            return Err(Error::InvalidMeasurements(
299                "stored frame count, image shape, and data length are inconsistent".into(),
300            ));
301        }
302        if self.data.iter().any(|value| !value.is_finite()) {
303            return Err(Error::InvalidMeasurements(
304                "measurements contain non-finite values".into(),
305            ));
306        }
307        if self.frame_metadata.len() != self.frames {
308            return Err(Error::InvalidMeasurements(format!(
309                "{} metadata entries for {} frames",
310                self.frame_metadata.len(),
311                self.frames
312            )));
313        }
314        for (index, metadata) in self.frame_metadata.iter().enumerate() {
315            if metadata.frame_index != index {
316                return Err(Error::InvalidMeasurements(format!(
317                    "metadata entry {index} identifies frame {}",
318                    metadata.frame_index
319                )));
320            }
321            if !metadata.exposure_time.is_finite() || metadata.exposure_time <= 0.0 {
322                return Err(Error::InvalidMeasurements(format!(
323                    "frame {index} has invalid exposure {}",
324                    metadata.exposure_time
325                )));
326            }
327            if !metadata.weight.is_finite() || metadata.weight < 0.0 {
328                return Err(Error::InvalidMeasurements(format!(
329                    "frame {index} has invalid weight {}",
330                    metadata.weight
331                )));
332            }
333        }
334        if let Some(dark) = &self.dark_frame {
335            self.validate_correction("dark frame", dark, false)?;
336        }
337        if let Some(flat) = &self.flat_field {
338            self.validate_correction("flat field", flat, false)?;
339            if flat.iter().any(|&value| value <= 0.0) {
340                return Err(Error::InvalidMeasurements(
341                    "flat-field values must be positive".into(),
342                ));
343            }
344        }
345        if let Some(background) = &self.background {
346            self.validate_correction("background", background, true)?;
347        }
348        if let Some(masks) = &self.masks
349            && masks.len() != frame_len
350            && masks.len() != self.data.len()
351        {
352            return Err(Error::InvalidMeasurements(format!(
353                "mask length {} must be {frame_len} or {}",
354                masks.len(),
355                self.data.len()
356            )));
357        }
358        if self.preprocessing.subtract_dark && self.dark_frame.is_none() {
359            return Err(Error::InvalidMeasurements(
360                "dark subtraction requested without a dark frame".into(),
361            ));
362        }
363        if self.preprocessing.divide_flat_field && self.flat_field.is_none() {
364            return Err(Error::InvalidMeasurements(
365                "flat-field correction requested without a flat field".into(),
366            ));
367        }
368        if self.preprocessing.subtract_background && self.background.is_none() {
369            return Err(Error::InvalidMeasurements(
370                "background subtraction requested without a background".into(),
371            ));
372        }
373        Ok(())
374    }
375
376    pub fn frame(&self, index: usize) -> Result<&[f64]> {
377        if index >= self.frames {
378            return Err(Error::FrameOutOfRange {
379                index,
380                frames: self.frames,
381            });
382        }
383        let start = index * self.frame_len();
384        Ok(&self.data[start..start + self.frame_len()])
385    }
386
387    pub fn frame_mut(&mut self, index: usize) -> Result<&mut [f64]> {
388        if index >= self.frames {
389            return Err(Error::FrameOutOfRange {
390                index,
391                frames: self.frames,
392            });
393        }
394        let frame_len = self.frame_len();
395        let start = index * frame_len;
396        Ok(&mut self.data[start..start + frame_len])
397    }
398
399    pub fn frame_weight(&self, index: usize) -> Result<f64> {
400        self.frame_metadata
401            .get(index)
402            .map(|metadata| metadata.weight)
403            .ok_or(Error::FrameOutOfRange {
404                index,
405                frames: self.frames,
406            })
407    }
408
409    /// Returns the mask for a frame. Zero-valued mask entries are excluded.
410    /// A single-frame mask is broadcast to every measurement frame.
411    pub fn frame_mask(&self, index: usize) -> Result<Option<&[u8]>> {
412        if index >= self.frames {
413            return Err(Error::FrameOutOfRange {
414                index,
415                frames: self.frames,
416            });
417        }
418        let Some(masks) = &self.masks else {
419            return Ok(None);
420        };
421        let frame_len = self.frame_len();
422        if masks.len() == frame_len {
423            Ok(Some(masks))
424        } else {
425            let start = index * frame_len;
426            Ok(Some(&masks[start..start + frame_len]))
427        }
428    }
429
430    pub fn with_dark_frame(mut self, dark: Vec<f64>) -> Result<Self> {
431        self.validate_correction("dark frame", &dark, false)?;
432        self.dark_frame = Some(dark);
433        self.preprocessing.subtract_dark = true;
434        Ok(self)
435    }
436
437    pub fn with_flat_field(mut self, flat: Vec<f64>) -> Result<Self> {
438        self.validate_correction("flat field", &flat, false)?;
439        if flat.iter().any(|&value| value <= 0.0) {
440            return Err(Error::InvalidMeasurements(
441                "flat-field values must be positive".into(),
442            ));
443        }
444        self.flat_field = Some(flat);
445        self.preprocessing.divide_flat_field = true;
446        Ok(self)
447    }
448
449    pub fn with_background(mut self, background: Vec<f64>) -> Result<Self> {
450        self.validate_correction("background", &background, true)?;
451        self.background = Some(background);
452        self.preprocessing.subtract_background = true;
453        Ok(self)
454    }
455
456    pub fn with_masks(mut self, masks: Vec<u8>) -> Result<Self> {
457        let frame_len = self.frame_len();
458        if masks.len() != frame_len && masks.len() != self.data.len() {
459            return Err(Error::InvalidMeasurements(format!(
460                "mask length {} must be {frame_len} or {}",
461                masks.len(),
462                self.data.len()
463            )));
464        }
465        self.masks = Some(masks);
466        Ok(self)
467    }
468
469    pub fn normalize_exposure(mut self) -> Self {
470        self.preprocessing.normalize_exposure = true;
471        self
472    }
473
474    pub fn clamp_negative(mut self) -> Self {
475        self.preprocessing.clamp_negative = true;
476        self
477    }
478
479    pub fn apply_preprocessing(mut self) -> Result<Self> {
480        self.validate()?;
481        let frame_len = self.frame_len();
482        for frame_index in 0..self.frames {
483            let exposure = self.frame_metadata[frame_index].exposure_time;
484            if self.preprocessing.normalize_exposure && (!exposure.is_finite() || exposure <= 0.0) {
485                return Err(Error::InvalidMeasurements(format!(
486                    "frame {frame_index} has invalid exposure {exposure}"
487                )));
488            }
489            for pixel in 0..frame_len {
490                let index = frame_index * frame_len + pixel;
491                let mut value = self.data[index];
492                if let Some(dark) = self
493                    .preprocessing
494                    .subtract_dark
495                    .then_some(self.dark_frame.as_deref())
496                    .flatten()
497                {
498                    value -= dark[pixel];
499                }
500                if let Some(background) = self
501                    .preprocessing
502                    .subtract_background
503                    .then_some(self.background.as_deref())
504                    .flatten()
505                {
506                    value -= background[if background.len() == frame_len {
507                        pixel
508                    } else {
509                        index
510                    }];
511                }
512                if let Some(flat) = self
513                    .preprocessing
514                    .divide_flat_field
515                    .then_some(self.flat_field.as_deref())
516                    .flatten()
517                {
518                    value /= flat[pixel];
519                }
520                if self.preprocessing.normalize_exposure {
521                    value /= exposure;
522                }
523                if self.preprocessing.clamp_negative {
524                    value = value.max(0.0);
525                }
526                self.data[index] = value;
527            }
528        }
529        Ok(self)
530    }
531
532    fn validate_correction(&self, name: &str, values: &[f64], per_frame: bool) -> Result<()> {
533        let frame_len = self.frame_len();
534        let valid_length =
535            values.len() == frame_len || (per_frame && values.len() == self.data.len());
536        if !valid_length || values.iter().any(|value| !value.is_finite()) {
537            return Err(Error::InvalidMeasurements(format!(
538                "{name} must contain finite values and have length {frame_len}{}",
539                if per_frame {
540                    format!(" or {}", self.data.len())
541                } else {
542                    String::new()
543                }
544            )));
545        }
546        Ok(())
547    }
548}
549
550pub(super) fn resolve_path(base_directory: &Path, path: &Path) -> std::path::PathBuf {
551    if path.is_absolute() {
552        path.to_owned()
553    } else {
554        base_directory.join(path)
555    }
556}
557
558pub(super) fn load_manifest_image(
559    base_directory: &Path,
560    path: &Path,
561    expected_shape: (usize, usize),
562) -> Result<Vec<f64>> {
563    let resolved = resolve_path(base_directory, path);
564    let image = load_grayscale(&resolved, GrayscaleScaling::NativeCounts)?;
565    if image.shape() != expected_shape {
566        return Err(Error::InvalidMeasurements(format!(
567            "image {} has shape {:?}, expected {expected_shape:?}",
568            resolved.display(),
569            image.shape()
570        )));
571    }
572    Ok(image.into_vec())
573}
574
575pub(super) fn load_manifest_image_set(
576    base_directory: &Path,
577    images: &ImageSet,
578    expected_shape: (usize, usize),
579    frame_count: usize,
580) -> Result<Vec<f64>> {
581    match images {
582        ImageSet::Single(path) => load_manifest_image(base_directory, path, expected_shape),
583        ImageSet::PerFrame(paths) => {
584            if paths.len() != frame_count {
585                return Err(Error::InvalidMeasurements(format!(
586                    "manifest image set has {} entries for {frame_count} frames",
587                    paths.len()
588                )));
589            }
590            let mut values = Vec::with_capacity(expected_shape.0 * expected_shape.1 * frame_count);
591            for path in paths {
592                values.extend(load_manifest_image(base_directory, path, expected_shape)?);
593            }
594            Ok(values)
595        }
596    }
597}