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 {
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
124pub 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}