fpm_rs/callbacks/
save_image.rs1use 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 {
13 frequency: usize,
14 directory: PathBuf,
15}
16
17impl SaveImageEvery {
18 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
71pub struct SavePupilEvery {
73 frequency: usize,
74 directory: PathBuf,
75}
76
77impl SavePupilEvery {
78 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
128pub struct SaveResidualsEvery {
133 frequency: usize,
134 directory: PathBuf,
135}
136
137impl SaveResidualsEvery {
138 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}