subx_core/services/vad/
audio_loader.rs

1//! Direct audio loader: Uses Symphonia to directly decode various audio formats and obtain i16 sample data.
2//!
3//! Supports MP4, MKV, OGG, WAV and other formats, returning sample data and audio information.
4use crate::services::vad::detector::AudioInfo;
5use crate::{Result, error::SubXError};
6use log::{debug, trace, warn};
7use std::fs::File;
8use std::path::Path;
9use symphonia::core::audio::SampleBuffer;
10use symphonia::core::codecs::CodecRegistry;
11use symphonia::core::codecs::DecoderOptions;
12use symphonia::core::formats::FormatOptions;
13use symphonia::core::io::MediaSourceStream;
14use symphonia::core::probe::Hint;
15use symphonia::core::probe::Probe;
16use symphonia::default::{get_codecs, get_probe};
17
18/// Direct audio loader using Symphonia to decode and obtain raw sample data.
19pub struct DirectAudioLoader {
20    probe: &'static Probe,
21    codecs: &'static CodecRegistry,
22}
23
24#[cfg(test)]
25mod tests {
26    use super::*;
27
28    #[tokio::test]
29    async fn test_direct_mp4_loading() {
30        // Test direct audio loading using assets/SubX - The Subtitle Revolution.mp4
31        //
32        // The media asset ships in this repository's `assets/` directory
33        // (relocated there when the test suite split landed), and is resolved
34        // from the crate manifest directory so the test is independent of the
35        // working directory.
36        let asset = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR"))
37            .join("assets")
38            .join("SubX - The Subtitle Revolution.mp4");
39        let loader = DirectAudioLoader::new().expect("Failed to initialize DirectAudioLoader");
40        let (samples, info) = loader
41            .load_audio_samples(&asset, 2_147_483_648)
42            .expect("load_audio_samples failed");
43        assert!(!samples.is_empty(), "Sample data should not be empty");
44        assert!(info.sample_rate > 0, "sample_rate should be greater than 0");
45        assert!(
46            info.total_samples > 0,
47            "total_samples should be greater than 0"
48        );
49    }
50}
51
52impl DirectAudioLoader {
53    /// Creates a new audio loader instance.
54    pub fn new() -> Result<Self> {
55        Ok(Self {
56            probe: get_probe(),
57            codecs: get_codecs(),
58        })
59    }
60
61    /// Loads i16 samples and audio information from an audio file path.
62    ///
63    /// `max_audio_bytes` caps the accepted file size as a defense-in-depth
64    /// guard; callers should thread the configured limit from
65    /// `GeneralConfig.max_audio_bytes`.
66    pub fn load_audio_samples<P: AsRef<Path>>(
67        &self,
68        path: P,
69        max_audio_bytes: u64,
70    ) -> Result<(Vec<i16>, AudioInfo)> {
71        let path_ref = path.as_ref();
72        debug!(
73            "[DirectAudioLoader] Start loading audio file: {:?}",
74            path_ref
75        );
76        crate::core::fs_util::check_file_size(path_ref, max_audio_bytes, "Audio")
77            .map_err(|e| SubXError::audio_processing(e.to_string()))?;
78        // Open the media source.
79        let file = File::open(path_ref).map_err(|e| {
80            warn!(
81                "[DirectAudioLoader] Failed to open audio file: {:?}, error: {}",
82                path_ref, e
83            );
84            SubXError::audio_processing(format!("Failed to open audio file: {}", e))
85        })?;
86        debug!(
87            "[DirectAudioLoader] Successfully opened audio file: {:?}",
88            path_ref
89        );
90
91        // Create the media source stream.
92        let mss = MediaSourceStream::new(Box::new(file), Default::default());
93        debug!("[DirectAudioLoader] MediaSourceStream created");
94
95        // Create a hint to help format probing based on file extension.
96        let mut hint = Hint::new();
97        if let Some(ext) = path_ref.extension().and_then(|e| e.to_str()) {
98            debug!(
99                "[DirectAudioLoader] Detected extension: {} (used for format probing)",
100                ext
101            );
102            hint.with_extension(ext);
103        } else {
104            debug!("[DirectAudioLoader] No extension detected, using default format probing");
105        }
106
107        // Probe the media format.
108        let probed = self
109            .probe
110            .format(&hint, mss, &FormatOptions::default(), &Default::default())
111            .map_err(|e| {
112                warn!("[DirectAudioLoader] Format probing failed: {}", e);
113                SubXError::audio_processing(format!("Failed to probe format: {}", e))
114            })?;
115        debug!("[DirectAudioLoader] Format probing succeeded");
116        let mut format = probed.format;
117
118        // List all tracks and their channel info before selecting
119        for (idx, t) in format.tracks().iter().enumerate() {
120            let sr = t
121                .codec_params
122                .sample_rate
123                .map(|v| v.to_string())
124                .unwrap_or("None".to_string());
125            let ch = t
126                .codec_params
127                .channels
128                .map(|c| c.count().to_string())
129                .unwrap_or("None".to_string());
130            debug!(
131                "[DirectAudioLoader] Track[{}]: id={}, sample_rate={}, channels={}",
132                idx, t.id, sr, ch
133            );
134        }
135
136        // Select the first audio track that contains sample_rate as the audio source.
137        let track = format
138            .tracks()
139            .iter()
140            .find(|t| t.codec_params.sample_rate.is_some())
141            .ok_or_else(|| {
142                warn!("[DirectAudioLoader] No audio track with sample_rate found");
143                SubXError::audio_processing("No audio track found".to_string())
144            })?;
145        // Clone necessary info first to avoid borrow conflicts
146        let track_id = track.id;
147        let sample_rate = track.codec_params.sample_rate.ok_or_else(|| {
148            warn!("[DirectAudioLoader] Audio track sample_rate is unknown");
149            SubXError::audio_processing("Sample rate unknown".to_string())
150        })?;
151        let channels = track.codec_params.channels.map(|c| c.count() as u16);
152        let time_base = track.codec_params.time_base;
153        debug!(
154            "[DirectAudioLoader] Selected track: id={}, sample_rate={}, channels={:?}",
155            track_id, sample_rate, channels
156        );
157
158        // Create decoder for the track.
159        let dec_opts = DecoderOptions::default();
160        let mut decoder = self
161            .codecs
162            .make(&track.codec_params, &dec_opts)
163            .map_err(|e| {
164                warn!("[DirectAudioLoader] Failed to create decoder: {}", e);
165                SubXError::audio_processing(format!("Failed to create decoder: {}", e))
166            })?;
167        debug!("[DirectAudioLoader] Decoder created successfully");
168
169        // Decode packets and collect samples.
170        let mut samples = Vec::new();
171        let mut packet_count = 0;
172        let mut last_pts: u64 = 0;
173        while let Ok(packet) = format.next_packet() {
174            if packet.track_id() != track_id {
175                continue;
176            }
177            packet_count += 1;
178            trace!(
179                "[DirectAudioLoader] Decoding packet {} (track_id={})",
180                packet_count, track_id
181            );
182            let decoded = decoder.decode(&packet).map_err(|e| {
183                warn!("[DirectAudioLoader] Failed to decode packet: {}", e);
184                SubXError::audio_processing(format!("Decode error: {}", e))
185            })?;
186            // Create a sample buffer for this packet using its signal spec and capacity.
187            let spec = *decoded.spec();
188            let mut sample_buf = SampleBuffer::<i16>::new(decoded.capacity() as u64, spec);
189            sample_buf.copy_interleaved_ref(decoded);
190            let sample_len = sample_buf.samples().len();
191            trace!(
192                "[DirectAudioLoader] Packet decoded successfully, got {} samples",
193                sample_len
194            );
195            samples.extend_from_slice(sample_buf.samples());
196            // Directly record the timestamp of the last packet
197            last_pts = packet.ts;
198        }
199        debug!(
200            "[DirectAudioLoader] Packet decoding finished, total {} packets, {} samples accumulated",
201            packet_count,
202            samples.len()
203        );
204
205        // Calculate total samples and audio duration
206        let total_samples = samples.len();
207        // Use Timebase to calculate duration_seconds
208        let duration_seconds = if let Some(tb) = time_base {
209            if last_pts > 0 {
210                let (num, den) = (tb.numer, tb.denom);
211                last_pts as f64 * num as f64 / den as f64
212            } else {
213                total_samples as f64 / (sample_rate as f64 * channels.unwrap_or(1) as f64)
214            }
215        } else {
216            total_samples as f64 / (sample_rate as f64 * channels.unwrap_or(1) as f64)
217        };
218        // If channels is None, try to infer channel count from duration_seconds
219        let channels = channels.unwrap_or_else(|| {
220            let ch = if duration_seconds > 0.0 {
221                (total_samples as f64 / (sample_rate as f64 * duration_seconds)).round() as u16
222            } else {
223                1
224            };
225            debug!("[DirectAudioLoader] Inferred channel count: {}", ch);
226            ch
227        });
228        debug!(
229            "[DirectAudioLoader] Audio info: sample_rate={}, channels={}, duration_seconds={:.3}, total_samples={}",
230            sample_rate, channels, duration_seconds, total_samples
231        );
232
233        Ok((
234            samples,
235            AudioInfo {
236                sample_rate,
237                channels,
238                duration_seconds,
239                total_samples,
240            },
241        ))
242    }
243}
subx_core/services/vad/audio_loader.rs

subx_core/services/vad/
audio_loader.rs
