1use ndarray::{Array2, Array3};
2use std::{
3 collections::{HashMap, VecDeque},
4 mem::size_of,
5 path::{Path, PathBuf},
6 sync::{Arc, Mutex},
7};
8
9use crate::{
10 Result,
11 error::Error,
12 image_io::{
13 GrayscaleScaling, grayscale_image_shape, grayscale_tiff_page_shapes, load_grayscale,
14 load_grayscale_tiff_page,
15 },
16};
17
18use super::{
19 FrameMetadata, MeasurementSpec, MeasurementStack, PreprocessingConfig,
20 stack::{load_manifest_image, load_manifest_image_set, resolve_path},
21};
22
23#[derive(Debug)]
39pub struct LazyMeasurementStack {
40 paths: Vec<PathBuf>,
41 page_indices: Vec<Option<usize>>,
42 image_shape: (usize, usize),
43 frame_metadata: Vec<FrameMetadata>,
44 dark_frame: Option<Vec<f64>>,
45 flat_field: Option<Vec<f64>>,
46 background: Option<Vec<f64>>,
47 masks: Option<Vec<u8>>,
48 preprocessing: PreprocessingConfig,
49 cache_capacity: usize,
50 cache_byte_capacity: Option<usize>,
51 cache: Mutex<FrameCache>,
52}
53
54#[derive(Debug, Default)]
55struct FrameCache {
56 frames: HashMap<usize, Arc<Vec<f64>>>,
57 least_recently_used: VecDeque<usize>,
58 bytes: usize,
59}
60
61impl FrameCache {
62 fn touch(&mut self, index: usize) {
63 self.least_recently_used.retain(|&cached| cached != index);
64 self.least_recently_used.push_back(index);
65 }
66
67 fn evict_least_recently_used(&mut self) -> bool {
68 let Some(index) = self.least_recently_used.pop_front() else {
69 return false;
70 };
71 if let Some(frame) = self.frames.remove(&index) {
72 self.bytes = self
73 .bytes
74 .saturating_sub(frame.len().saturating_mul(size_of::<f64>()));
75 }
76 true
77 }
78}
79
80impl Clone for LazyMeasurementStack {
81 fn clone(&self) -> Self {
82 Self {
83 paths: self.paths.clone(),
84 page_indices: self.page_indices.clone(),
85 image_shape: self.image_shape,
86 frame_metadata: self.frame_metadata.clone(),
87 dark_frame: self.dark_frame.clone(),
88 flat_field: self.flat_field.clone(),
89 background: self.background.clone(),
90 masks: self.masks.clone(),
91 preprocessing: self.preprocessing.clone(),
92 cache_capacity: self.cache_capacity,
93 cache_byte_capacity: self.cache_byte_capacity,
94 cache: Mutex::new(FrameCache::default()),
97 }
98 }
99}
100
101impl LazyMeasurementStack {
102 pub fn from_image_files<P: AsRef<Path>>(
104 paths: &[P],
105 frame_metadata: Vec<FrameMetadata>,
106 ) -> Result<Self> {
107 if paths.is_empty() {
108 return Err(Error::InvalidMeasurements(
109 "at least one image path is required".into(),
110 ));
111 }
112 let paths: Vec<PathBuf> = paths.iter().map(|path| path.as_ref().to_owned()).collect();
113 let image_shape = grayscale_image_shape(&paths[0])?;
114 for path in &paths[1..] {
115 let shape = grayscale_image_shape(path)?;
116 if shape != image_shape {
117 return Err(Error::InvalidMeasurements(format!(
118 "image {} has shape {shape:?}, expected {image_shape:?}",
119 path.display()
120 )));
121 }
122 }
123 let frame_metadata =
124 metadata_or_labels(&paths, frame_metadata, |path, _| path.display().to_string());
125 let frame_count = paths.len();
126 Self::new(paths, vec![None; frame_count], image_shape, frame_metadata)
127 }
128
129 pub fn from_tiff_stack(
132 path: impl AsRef<Path>,
133 frame_metadata: Vec<FrameMetadata>,
134 ) -> Result<Self> {
135 let path = path.as_ref().to_owned();
136 let shapes = grayscale_tiff_page_shapes(&path)?;
137 let image_shape = *shapes.first().ok_or_else(|| {
138 Error::InvalidMeasurements("TIFF stack does not contain an image".into())
139 })?;
140 if shapes.iter().any(|&shape| shape != image_shape) {
141 return Err(Error::InvalidMeasurements(
142 "all TIFF pages must have the same dimensions".into(),
143 ));
144 }
145 let paths = vec![path; shapes.len()];
146 let frame_metadata = metadata_or_labels(&paths, frame_metadata, |path, index| {
147 format!("{}#page={index}", path.display())
148 });
149 let page_indices = (0..shapes.len()).map(Some).collect();
150 Self::new(paths, page_indices, image_shape, frame_metadata)
151 }
152
153 pub fn from_manifest(path: impl AsRef<Path>) -> Result<Self> {
157 let path = path.as_ref();
158 let manifest = MeasurementSpec::load(path)?;
159 let base_directory = path.parent().unwrap_or_else(|| Path::new("."));
160 Self::from_manifest_definition(manifest, base_directory)
161 }
162
163 pub fn from_manifest_definition(
165 manifest: MeasurementSpec,
166 base_directory: impl AsRef<Path>,
167 ) -> Result<Self> {
168 let base_directory = base_directory.as_ref();
169 let paths: Vec<_> = manifest
170 .frames
171 .iter()
172 .map(|frame| resolve_path(base_directory, &frame.path))
173 .collect();
174 let metadata = manifest
175 .frames
176 .iter()
177 .enumerate()
178 .map(|(index, frame)| FrameMetadata {
179 frame_index: index,
180 illumination_index: frame.illumination_index,
181 original_frame_index: Some(index),
182 original_illumination_index: frame.illumination_index,
183 exposure_time: frame.exposure_time,
184 weight: frame.weight,
185 label: frame
186 .label
187 .clone()
188 .or_else(|| Some(frame.path.display().to_string())),
189 })
190 .collect();
191 let mut stack = Self::from_image_files(&paths, metadata)?;
192 if let Some(path) = &manifest.dark_frame {
193 stack.dark_frame = Some(load_manifest_image(
194 base_directory,
195 path,
196 stack.image_shape,
197 )?);
198 }
199 if let Some(path) = &manifest.flat_field {
200 stack.flat_field = Some(load_manifest_image(
201 base_directory,
202 path,
203 stack.image_shape,
204 )?);
205 }
206 if let Some(background) = &manifest.background {
207 stack.background = Some(load_manifest_image_set(
208 base_directory,
209 background,
210 stack.image_shape,
211 stack.frame_count(),
212 )?);
213 }
214 if let Some(mask) = &manifest.mask {
215 stack.masks = Some(
216 load_manifest_image_set(
217 base_directory,
218 mask,
219 stack.image_shape,
220 stack.frame_count(),
221 )?
222 .into_iter()
223 .map(|value| u8::from(value != 0.0))
224 .collect(),
225 );
226 }
227 stack.preprocessing = manifest.preprocessing;
228 stack.validate()?;
229 Ok(stack)
230 }
231
232 fn new(
233 paths: Vec<PathBuf>,
234 page_indices: Vec<Option<usize>>,
235 image_shape: (usize, usize),
236 frame_metadata: Vec<FrameMetadata>,
237 ) -> Result<Self> {
238 let stack = Self {
239 paths,
240 page_indices,
241 image_shape,
242 frame_metadata,
243 dark_frame: None,
244 flat_field: None,
245 background: None,
246 masks: None,
247 preprocessing: PreprocessingConfig::default(),
248 cache_capacity: 1,
249 cache_byte_capacity: None,
250 cache: Mutex::new(FrameCache::default()),
251 };
252 stack.validate()?;
253 Ok(stack)
254 }
255
256 pub fn with_dark_frame(mut self, dark: Vec<f64>) -> Result<Self> {
258 self.validate_correction("dark frame", &dark, false)?;
259 self.dark_frame = Some(dark);
260 self.preprocessing.subtract_dark = true;
261 self.reset_cache();
262 Ok(self)
263 }
264
265 pub fn with_flat_field(mut self, flat: Vec<f64>) -> Result<Self> {
267 self.validate_correction("flat field", &flat, false)?;
268 if flat.iter().any(|&value| value <= 0.0) {
269 return Err(Error::InvalidMeasurements(
270 "flat-field values must be positive".into(),
271 ));
272 }
273 self.flat_field = Some(flat);
274 self.preprocessing.divide_flat_field = true;
275 self.reset_cache();
276 Ok(self)
277 }
278
279 pub fn with_background(mut self, background: Vec<f64>) -> Result<Self> {
281 self.validate_correction("background", &background, true)?;
282 self.background = Some(background);
283 self.preprocessing.subtract_background = true;
284 self.reset_cache();
285 Ok(self)
286 }
287
288 pub fn with_masks(mut self, masks: Vec<u8>) -> Result<Self> {
290 let frame_len = self.frame_len();
291 let stack_len = self.stack_len()?;
292 if masks.len() != frame_len && masks.len() != stack_len {
293 return Err(Error::InvalidMeasurements(format!(
294 "mask length {} must be {frame_len} or {stack_len}",
295 masks.len()
296 )));
297 }
298 self.masks = Some(masks);
299 Ok(self)
300 }
301
302 pub fn normalize_exposure(mut self) -> Self {
304 self.preprocessing.normalize_exposure = true;
305 self.reset_cache();
306 self
307 }
308
309 pub fn clamp_negative(mut self) -> Self {
311 self.preprocessing.clamp_negative = true;
312 self.reset_cache();
313 self
314 }
315
316 pub fn with_preprocessing(mut self, preprocessing: PreprocessingConfig) -> Result<Self> {
319 self.preprocessing = preprocessing;
320 self.reset_cache();
321 self.validate()?;
322 Ok(self)
323 }
324
325 pub fn with_cache_capacity(mut self, capacity: usize) -> Result<Self> {
327 if capacity == 0 {
328 return Err(Error::InvalidParameter {
329 name: "cache_capacity",
330 reason: "must be greater than zero".into(),
331 });
332 }
333 self.cache_capacity = capacity;
334 self.reset_cache();
335 Ok(self)
336 }
337
338 pub fn with_cache_byte_capacity(mut self, capacity: usize) -> Result<Self> {
341 let frame_bytes = self.frame_bytes()?;
342 if capacity < frame_bytes {
343 return Err(Error::InvalidParameter {
344 name: "cache_byte_capacity",
345 reason: format!("must be at least one decoded frame ({frame_bytes} bytes)"),
346 });
347 }
348 self.cache_byte_capacity = Some(capacity);
349 self.reset_cache();
350 Ok(self)
351 }
352
353 pub fn frame_count(&self) -> usize {
355 self.paths.len()
356 }
357
358 pub fn image_shape(&self) -> (usize, usize) {
360 self.image_shape
361 }
362
363 pub fn frame_len(&self) -> usize {
365 self.image_shape.0 * self.image_shape.1
366 }
367
368 pub fn frame_metadata(&self) -> &[FrameMetadata] {
370 &self.frame_metadata
371 }
372
373 pub fn preprocessing(&self) -> &PreprocessingConfig {
375 &self.preprocessing
376 }
377
378 pub fn paths(&self) -> &[PathBuf] {
381 &self.paths
382 }
383
384 pub fn cached_frame_count(&self) -> usize {
386 self.cache.lock().map_or(0, |cache| cache.frames.len())
390 }
391
392 pub fn cached_byte_count(&self) -> usize {
394 self.cache.lock().map_or(0, |cache| cache.bytes)
395 }
396
397 pub fn frame(&self, index: usize) -> Result<Arc<Vec<f64>>> {
399 if index >= self.frame_count() {
400 return Err(Error::FrameOutOfRange {
401 index,
402 frames: self.frame_count(),
403 });
404 }
405 {
406 let mut cache = self.lock_cache()?;
407 if let Some(frame) = cache.frames.get(&index).cloned() {
408 cache.touch(index);
409 return Ok(frame);
410 }
411 }
412
413 let path = &self.paths[index];
416 let mut values = match self.page_indices[index] {
417 Some(page_index) => {
418 load_grayscale_tiff_page(path, page_index, GrayscaleScaling::NativeCounts)?
419 }
420 None => load_grayscale(path, GrayscaleScaling::NativeCounts)?,
421 };
422 if values.dim() != self.image_shape {
423 return Err(Error::InvalidMeasurements(format!(
424 "image {} changed shape to {:?}, expected {:?}",
425 path.display(),
426 values.dim(),
427 self.image_shape
428 )));
429 }
430 if !values.is_standard_layout() {
431 return Err(Error::NonStandardLayout {
432 context: "decoded lazy measurement frame",
433 shape: values.shape().to_vec(),
434 strides: values.strides().to_vec(),
435 });
436 }
437 let decoded_strides = values.strides().to_vec();
438 let values_slice = values
439 .as_slice_mut()
440 .ok_or_else(|| Error::NonStandardLayout {
441 context: "decoded lazy measurement frame",
442 shape: vec![self.image_shape.0, self.image_shape.1],
443 strides: decoded_strides,
444 })?;
445 self.preprocess_frame(index, values_slice)?;
446 let (values, offset) = values.into_raw_vec_and_offset();
447 if offset.unwrap_or(0) != 0 || values.len() != self.frame_len() {
448 return Err(Error::InvalidMeasurements(
449 "decoded lazy frame allocation does not match its logical image".into(),
450 ));
451 }
452 let loaded = Arc::new(values);
453 let loaded_bytes = loaded.len().checked_mul(size_of::<f64>()).ok_or_else(|| {
454 Error::InvalidMeasurements("decoded lazy frame byte length overflows".into())
455 })?;
456
457 let mut cache = self.lock_cache()?;
458 if let Some(existing) = cache.frames.get(&index).cloned() {
459 cache.touch(index);
460 return Ok(existing);
461 }
462 while cache.frames.len() >= self.cache_capacity
463 || self
464 .cache_byte_capacity
465 .is_some_and(|capacity| cache.bytes.saturating_add(loaded_bytes) > capacity)
466 {
467 if !cache.evict_least_recently_used() {
468 break;
469 }
470 }
471 cache.frames.insert(index, loaded.clone());
472 cache.bytes = cache.bytes.checked_add(loaded_bytes).ok_or_else(|| {
473 Error::InvalidMeasurements("lazy cache byte accounting overflows".into())
474 })?;
475 cache.touch(index);
476 Ok(loaded)
477 }
478
479 pub fn frame_weight(&self, index: usize) -> Result<f64> {
481 self.frame_metadata
482 .get(index)
483 .map(|metadata| metadata.weight)
484 .ok_or(Error::FrameOutOfRange {
485 index,
486 frames: self.frame_count(),
487 })
488 }
489
490 pub fn frame_mask(&self, index: usize) -> Result<Option<&[u8]>> {
492 if index >= self.frame_count() {
493 return Err(Error::FrameOutOfRange {
494 index,
495 frames: self.frame_count(),
496 });
497 }
498 let Some(masks) = &self.masks else {
499 return Ok(None);
500 };
501 let frame_len = self.frame_len();
502 if masks.len() == frame_len {
503 Ok(Some(masks))
504 } else {
505 let start = index * frame_len;
506 Ok(Some(&masks[start..start + frame_len]))
507 }
508 }
509
510 pub fn materialize(&self) -> Result<MeasurementStack> {
512 let mut data = Vec::with_capacity(self.stack_len()?);
513 for frame in 0..self.frame_count() {
514 data.extend_from_slice(self.frame(frame)?.as_slice());
515 }
516 let mut stack =
517 MeasurementStack::from_vec(data, self.image_shape, self.frame_metadata.clone())?;
518 if let Some(masks) = &self.masks {
519 if masks.len() == self.frame_len() {
520 stack =
521 stack.with_masks(Array2::from_shape_vec(self.image_shape, masks.clone())?)?;
522 } else {
523 stack = stack.with_per_frame_masks(Array3::from_shape_vec(
524 (self.frame_count(), self.image_shape.0, self.image_shape.1),
525 masks.clone(),
526 )?)?;
527 }
528 }
529 Ok(stack)
530 }
531
532 pub fn validate(&self) -> Result<()> {
535 if self.paths.is_empty()
536 || self.frame_metadata.len() != self.paths.len()
537 || self.page_indices.len() != self.paths.len()
538 || self.cache_capacity == 0
539 {
540 return Err(Error::InvalidMeasurements(
541 "lazy paths, page indices, and metadata must have equal non-zero counts".into(),
542 ));
543 }
544 let frame_bytes = self.frame_bytes()?;
545 for (index, metadata) in self.frame_metadata.iter().enumerate() {
546 if metadata.frame_index != index
547 || !metadata.exposure_time.is_finite()
548 || metadata.exposure_time <= 0.0
549 || !metadata.weight.is_finite()
550 || metadata.weight < 0.0
551 {
552 return Err(Error::InvalidMeasurements(format!(
553 "lazy frame metadata {index} is invalid"
554 )));
555 }
556 }
557 if let Some(dark) = &self.dark_frame {
558 self.validate_correction("dark frame", dark, false)?;
559 }
560 if let Some(flat) = &self.flat_field {
561 self.validate_correction("flat field", flat, false)?;
562 if flat.iter().any(|&value| value <= 0.0) {
563 return Err(Error::InvalidMeasurements(
564 "flat-field values must be positive".into(),
565 ));
566 }
567 }
568 if let Some(background) = &self.background {
569 self.validate_correction("background", background, true)?;
570 }
571 if let Some(masks) = &self.masks {
572 let frame_len = self.frame_len();
573 let stack_len = self.stack_len()?;
574 if masks.len() != frame_len && masks.len() != stack_len {
575 return Err(Error::InvalidMeasurements(
576 "lazy mask dimensions are invalid".into(),
577 ));
578 }
579 }
580 if self.preprocessing.subtract_dark && self.dark_frame.is_none() {
581 return Err(Error::InvalidMeasurements(
582 "dark subtraction requested without a dark frame".into(),
583 ));
584 }
585 if self.preprocessing.divide_flat_field && self.flat_field.is_none() {
586 return Err(Error::InvalidMeasurements(
587 "flat-field correction requested without a flat field".into(),
588 ));
589 }
590 if self.preprocessing.subtract_background && self.background.is_none() {
591 return Err(Error::InvalidMeasurements(
592 "background subtraction requested without a background".into(),
593 ));
594 }
595 if self
596 .cache_byte_capacity
597 .is_some_and(|capacity| capacity < frame_bytes)
598 {
599 return Err(Error::InvalidMeasurements(
600 "lazy byte cache must fit at least one decoded frame".into(),
601 ));
602 }
603 Ok(())
604 }
605
606 fn preprocess_frame(&self, frame_index: usize, values: &mut [f64]) -> Result<()> {
607 let frame_len = self.frame_len();
608 let exposure = self.frame_metadata[frame_index].exposure_time;
609 let dark = if self.preprocessing.subtract_dark {
610 Some(self.dark_frame.as_deref().ok_or_else(|| {
611 Error::InvalidMeasurements("dark subtraction requested without a dark frame".into())
612 })?)
613 } else {
614 None
615 };
616 let background = if self.preprocessing.subtract_background {
617 Some(self.background.as_deref().ok_or_else(|| {
618 Error::InvalidMeasurements(
619 "background subtraction requested without a background".into(),
620 )
621 })?)
622 } else {
623 None
624 };
625 let flat = if self.preprocessing.divide_flat_field {
626 Some(self.flat_field.as_deref().ok_or_else(|| {
627 Error::InvalidMeasurements(
628 "flat-field correction requested without a flat field".into(),
629 )
630 })?)
631 } else {
632 None
633 };
634 for (pixel, value) in values.iter_mut().enumerate() {
635 if let Some(dark) = dark {
636 *value -= dark[pixel];
637 }
638 if let Some(background) = background {
639 *value -= background[if background.len() == frame_len {
640 pixel
641 } else {
642 frame_index * frame_len + pixel
643 }];
644 }
645 if let Some(flat) = flat {
646 *value /= flat[pixel];
647 }
648 if self.preprocessing.normalize_exposure {
649 *value /= exposure;
650 }
651 if self.preprocessing.clamp_negative {
652 *value = value.max(0.0);
653 }
654 if !value.is_finite() {
655 return Err(Error::InvalidMeasurements(format!(
656 "preprocessing frame {frame_index} produced a non-finite value at pixel {pixel}"
657 )));
658 }
659 }
660 Ok(())
661 }
662
663 fn validate_correction(&self, name: &str, values: &[f64], per_frame: bool) -> Result<()> {
664 let frame_len = self.frame_len();
665 let stack_len = self.stack_len()?;
666 let valid_length = values.len() == frame_len || (per_frame && values.len() == stack_len);
667 if !valid_length || values.iter().any(|value| !value.is_finite()) {
668 return Err(Error::InvalidMeasurements(format!(
669 "{name} must contain finite values and have length {frame_len}{}",
670 if per_frame {
671 format!(" or {stack_len}")
672 } else {
673 String::new()
674 }
675 )));
676 }
677 Ok(())
678 }
679
680 fn frame_bytes(&self) -> Result<usize> {
681 self.image_shape
682 .0
683 .checked_mul(self.image_shape.1)
684 .filter(|&length| length > 0)
685 .and_then(|length| length.checked_mul(size_of::<f64>()))
686 .ok_or_else(|| Error::InvalidShape("lazy measurement shape is invalid".into()))
687 }
688
689 fn stack_len(&self) -> Result<usize> {
690 self.frame_len()
691 .checked_mul(self.frame_count())
692 .ok_or_else(|| {
693 Error::InvalidMeasurements("lazy measurement stack length overflows".into())
694 })
695 }
696
697 fn reset_cache(&mut self) {
698 self.cache = Mutex::new(FrameCache::default());
699 }
700
701 fn lock_cache(&self) -> Result<std::sync::MutexGuard<'_, FrameCache>> {
702 self.cache
703 .lock()
704 .map_err(|_| Error::Numerical("lazy measurement cache lock was poisoned".into()))
705 }
706}
707
708fn metadata_or_labels(
709 paths: &[PathBuf],
710 frame_metadata: Vec<FrameMetadata>,
711 label: impl Fn(&Path, usize) -> String,
712) -> Vec<FrameMetadata> {
713 if frame_metadata.is_empty() {
714 paths
715 .iter()
716 .enumerate()
717 .map(|(index, path)| {
718 let mut metadata = FrameMetadata::new(index);
719 metadata.label = Some(label(path, index));
720 metadata
721 })
722 .collect()
723 } else {
724 frame_metadata
725 }
726}
727
728#[cfg(test)]
729mod tests {
730 use std::panic::{AssertUnwindSafe, catch_unwind};
731
732 use super::*;
733
734 #[test]
735 fn poisoned_cache_fails_frame_reads_and_hides_unavailable_statistics() {
736 let stack = LazyMeasurementStack::new(
737 vec![PathBuf::from("unreadable-after-poison.png")],
738 vec![None],
739 (1, 1),
740 vec![FrameMetadata::new(0)],
741 )
742 .unwrap();
743
744 let result = catch_unwind(AssertUnwindSafe(|| {
745 let _guard = stack.cache.lock().unwrap();
746 panic!("intentional cache-lock poison for test");
747 }));
748 assert!(result.is_err());
749
750 assert!(
751 matches!(stack.frame(0), Err(Error::Numerical(message)) if message.contains("poisoned"))
752 );
753 assert_eq!(stack.cached_frame_count(), 0);
754 assert_eq!(stack.cached_byte_count(), 0);
755 }
756}