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