Processes a single spectrum frame through the neural network.
(
&mut self,
spectrum: &[(f32, f32); NUM_FREQ_BINS],
)
| 206 | |
| 207 | /// Processes a single spectrum frame through the neural network. |
| 208 | pub fn process_frame( |
| 209 | &mut self, |
| 210 | spectrum: &[(f32, f32); NUM_FREQ_BINS], |
| 211 | ) -> Result<[(f32, f32); NUM_FREQ_BINS], Box<dyn std::error::Error + Send + Sync>> { |
| 212 | // Fill input buffer (flat layout: [1, 257, 1, 2]) |
| 213 | for (i, &(re, im)) in spectrum.iter().enumerate() { |
| 214 | self.state.input_buf[i * 2] = re; |
| 215 | self.state.input_buf[i * 2 + 1] = im; |
| 216 | } |
| 217 | |
| 218 | // Create tensor references using (shape, slice) tuples |
| 219 | let input_tensor = |
| 220 | TensorRef::from_array_view(([1usize, NUM_FREQ_BINS, 1, 2], &self.state.input_buf[..]))?; |
| 221 | let conv_tensor = |
| 222 | TensorRef::from_array_view(([2usize, 1, 16, 16, 33], &self.state.conv[..]))?; |
| 223 | let tra_tensor = TensorRef::from_array_view(([2usize, 3, 1, 1, 16], &self.state.tra[..]))?; |
| 224 | let inter_tensor = |
| 225 | TensorRef::from_array_view(([2usize, 1, 33, 16], &self.state.inter[..]))?; |
| 226 | |
| 227 | // Run inference |
| 228 | let outputs = self.session.run(ort::inputs![ |
| 229 | input_tensor, |
| 230 | conv_tensor, |
| 231 | tra_tensor, |
| 232 | inter_tensor, |
| 233 | ])?; |
| 234 | |
| 235 | // Extract outputs using try_extract_tensor |
| 236 | let (_, output_enh_data) = outputs[0].try_extract_tensor::<f32>()?; |
| 237 | let (_, output_conv_data) = outputs[1].try_extract_tensor::<f32>()?; |
| 238 | let (_, output_tra_data) = outputs[2].try_extract_tensor::<f32>()?; |
| 239 | let (_, output_inter_data) = outputs[3].try_extract_tensor::<f32>()?; |
| 240 | |
| 241 | // Update state |
| 242 | self.state.conv.copy_from_slice(output_conv_data); |
| 243 | self.state.tra.copy_from_slice(output_tra_data); |
| 244 | self.state.inter.copy_from_slice(output_inter_data); |
| 245 | |
| 246 | // Extract enhanced spectrum |
| 247 | for (i, pair) in self.state.output_buf.iter_mut().enumerate() { |
| 248 | *pair = (output_enh_data[i * 2], output_enh_data[i * 2 + 1]); |
| 249 | } |
| 250 | |
| 251 | Ok(self.state.output_buf) |
| 252 | } |
| 253 | } |