diff options
Diffstat (limited to 'src/subtitles/extraction')
| -rw-r--r-- | src/subtitles/extraction/embedded.rs | 3 | ||||
| -rw-r--r-- | src/subtitles/extraction/mod.rs | 33 | ||||
| -rw-r--r-- | src/subtitles/extraction/whisper.rs | 3 |
3 files changed, 21 insertions, 18 deletions
diff --git a/src/subtitles/extraction/embedded.rs b/src/subtitles/extraction/embedded.rs index 920f52b..39698cf 100644 --- a/src/subtitles/extraction/embedded.rs +++ b/src/subtitles/extraction/embedded.rs @@ -5,8 +5,7 @@ use anyhow::Context; use crate::{subtitles::SubtitleCue, subtitles::extraction::*}; pub fn extract_embedded_subtitles( - // stream index to use when storing extracted subtitles, this index already - // has to be in TRACKS when this function is called! + // stream index to use when storing extracted subtitles stream_ix: StreamIndex, context: ffmpeg::codec::Context, time_base: ffmpeg::Rational, diff --git a/src/subtitles/extraction/mod.rs b/src/subtitles/extraction/mod.rs index 5070fdb..6495b62 100644 --- a/src/subtitles/extraction/mod.rs +++ b/src/subtitles/extraction/mod.rs @@ -3,22 +3,30 @@ mod embedded; /// Synthesis of subtitles from audio using whisper.cpp mod whisper; -use std::{collections::BTreeMap, sync::mpsc, thread}; +use std::{collections::BTreeMap, fmt::Display, sync::mpsc, thread}; use ffmpeg::Rational; use relm4::{ComponentSender, Worker}; -use crate::subtitles::{SUBTITLE_TRACKS, StreamIndex, SubtitleCue}; +use crate::subtitles::{StreamIndex, SubtitleCue}; pub struct SubtitleExtractor {} +#[derive(Debug, Clone)] +pub struct ExtractionArgs { + pub url: String, + pub whisper_stream_index: Option<usize>, +} + +impl Display for ExtractionArgs { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{} {:?}", self.url, self.whisper_stream_index) + } +} + #[derive(Debug)] pub enum SubtitleExtractorMsg { - ExtractFromUrl { - url: String, - // the index of the audio stream on which to run a whisper transcription - whisper_stream_index: Option<usize>, - }, + Extract(ExtractionArgs), } #[derive(Debug)] @@ -38,10 +46,10 @@ impl Worker for SubtitleExtractor { fn update(&mut self, msg: SubtitleExtractorMsg, sender: ComponentSender<Self>) { match msg { - SubtitleExtractorMsg::ExtractFromUrl { + SubtitleExtractorMsg::Extract(ExtractionArgs { url, whisper_stream_index: whisper_audio_stream_ix, - } => { + }) => { self.handle_extract_from_url(url, whisper_audio_stream_ix, sender); } } @@ -55,12 +63,8 @@ impl SubtitleExtractor { whisper_audio_stream_ix: Option<usize>, sender: ComponentSender<Self>, ) { - // Clear existing tracks - SUBTITLE_TRACKS.write().clear(); - match self.extract_subtitles(&url, whisper_audio_stream_ix, sender.clone()) { Ok(_) => { - log::info!("Subtitle extraction completed successfully"); sender .output(SubtitleExtractorOutput::ExtractionComplete) .unwrap(); @@ -125,7 +129,8 @@ impl SubtitleExtractor { } // wait for extraction to complete - for (_, (_, join_handle)) in subtitle_extractors { + for (packet_tx, join_handle) in subtitle_extractors.into_values() { + drop(packet_tx); join_handle .join() .unwrap() diff --git a/src/subtitles/extraction/whisper.rs b/src/subtitles/extraction/whisper.rs index bd6fba7..be4346a 100644 --- a/src/subtitles/extraction/whisper.rs +++ b/src/subtitles/extraction/whisper.rs @@ -21,8 +21,7 @@ struct WhisperCue { } pub fn generate_whisper_subtitles( - // stream index to use when storing generated subtitles, this index - // already has to be in TRACKS when this function is called! + // stream index to use when storing generated subtitles stream_ix: StreamIndex, context: ffmpeg::codec::Context, time_base: ffmpeg::Rational, |