Skip to main content

fpm_rs/callbacks/
save_image.rs

1use std::{fs, path::PathBuf};
2
3use crate::{
4    Result,
5    diagnostics::DiagnosticRequest,
6    reconstruction::{save_grayscale, save_signed_grayscale},
7};
8
9use super::{Callback, CallbackAction, CallbackHook, StepContext};
10
11pub struct SaveImageEvery {
12    frequency: usize,
13    directory: PathBuf,
14}
15
16impl SaveImageEvery {
17    pub fn new(frequency: usize, directory: impl Into<PathBuf>) -> Self {
18        Self {
19            frequency: frequency.max(1),
20            directory: directory.into(),
21        }
22    }
23}
24
25impl Callback for SaveImageEvery {
26    fn requires(&self) -> Vec<DiagnosticRequest> {
27        vec![
28            DiagnosticRequest::ObjectAmplitude,
29            DiagnosticRequest::ObjectPhase,
30        ]
31    }
32
33    fn requires_for(&self, hook: CallbackHook, iteration: usize) -> Vec<DiagnosticRequest> {
34        if hook == CallbackHook::IterationEnd && iteration.is_multiple_of(self.frequency) {
35            self.requires()
36        } else {
37            Vec::new()
38        }
39    }
40
41    fn on_start(&mut self, _context: &StepContext<'_>) -> Result<CallbackAction> {
42        fs::create_dir_all(&self.directory)?;
43        Ok(CallbackAction::Continue)
44    }
45
46    fn on_iteration_end(&mut self, context: &StepContext<'_>) -> Result<CallbackAction> {
47        if context.iteration.is_multiple_of(self.frequency) {
48            if let Some(amplitude) = &context.diagnostics.object_amplitude {
49                save_grayscale(
50                    amplitude,
51                    self.directory
52                        .join(format!("amplitude_{:05}.png", context.iteration)),
53                    false,
54                )?;
55            }
56            if let Some(phase) = &context.diagnostics.object_phase {
57                save_grayscale(
58                    phase,
59                    self.directory
60                        .join(format!("phase_{:05}.png", context.iteration)),
61                    true,
62                )?;
63            }
64        }
65        Ok(CallbackAction::Continue)
66    }
67}
68
69pub struct SavePupilEvery {
70    frequency: usize,
71    directory: PathBuf,
72}
73
74impl SavePupilEvery {
75    pub fn new(frequency: usize, directory: impl Into<PathBuf>) -> Self {
76        Self {
77            frequency: frequency.max(1),
78            directory: directory.into(),
79        }
80    }
81}
82
83impl Callback for SavePupilEvery {
84    fn requires(&self) -> Vec<DiagnosticRequest> {
85        vec![DiagnosticRequest::Pupil]
86    }
87
88    fn requires_for(&self, hook: CallbackHook, iteration: usize) -> Vec<DiagnosticRequest> {
89        if hook == CallbackHook::IterationEnd && iteration.is_multiple_of(self.frequency) {
90            self.requires()
91        } else {
92            Vec::new()
93        }
94    }
95
96    fn on_start(&mut self, _context: &StepContext<'_>) -> Result<CallbackAction> {
97        fs::create_dir_all(&self.directory)?;
98        Ok(CallbackAction::Continue)
99    }
100
101    fn on_iteration_end(&mut self, context: &StepContext<'_>) -> Result<CallbackAction> {
102        if context.iteration.is_multiple_of(self.frequency) {
103            if let Some(amplitude) = &context.diagnostics.pupil_amplitude {
104                save_grayscale(
105                    amplitude,
106                    self.directory
107                        .join(format!("pupil_amplitude_{:05}.png", context.iteration)),
108                    false,
109                )?;
110            }
111            if let Some(phase) = &context.diagnostics.pupil_phase {
112                save_grayscale(
113                    phase,
114                    self.directory
115                        .join(format!("pupil_phase_{:05}.png", context.iteration)),
116                    true,
117                )?;
118            }
119        }
120        Ok(CallbackAction::Continue)
121    }
122}
123
124/// Periodically saves one signed intensity-residual PNG per measured frame.
125///
126/// Each image is scaled symmetrically around zero so mid-gray represents zero,
127/// dark pixels are negative residuals, and bright pixels are positive residuals.
128pub struct SaveResidualsEvery {
129    frequency: usize,
130    directory: PathBuf,
131}
132
133impl SaveResidualsEvery {
134    pub fn new(frequency: usize, directory: impl Into<PathBuf>) -> Self {
135        Self {
136            frequency: frequency.max(1),
137            directory: directory.into(),
138        }
139    }
140}
141
142impl Callback for SaveResidualsEvery {
143    fn requires(&self) -> Vec<DiagnosticRequest> {
144        vec![DiagnosticRequest::ResidualImages]
145    }
146
147    fn requires_for(&self, hook: CallbackHook, iteration: usize) -> Vec<DiagnosticRequest> {
148        if hook == CallbackHook::IterationEnd && iteration.is_multiple_of(self.frequency) {
149            self.requires()
150        } else {
151            Vec::new()
152        }
153    }
154
155    fn on_start(&mut self, _context: &StepContext<'_>) -> Result<CallbackAction> {
156        fs::create_dir_all(&self.directory)?;
157        Ok(CallbackAction::Continue)
158    }
159
160    fn on_iteration_end(&mut self, context: &StepContext<'_>) -> Result<CallbackAction> {
161        if context.iteration.is_multiple_of(self.frequency)
162            && let Some(images) = &context.diagnostics.residual_images
163        {
164            for (frame, residual) in images.iter().enumerate() {
165                save_signed_grayscale(
166                    residual,
167                    self.directory.join(format!(
168                        "residual_{:05}_frame_{frame:05}.png",
169                        context.iteration
170                    )),
171                )?;
172            }
173        }
174        Ok(CallbackAction::Continue)
175    }
176}